Hyper Params & loss function for ViT-L/16

[google-research]

Hyper Params & loss function for ViT-L/16

[akolesnikoff]

Hi @owmohamm and sorry for the late reply.

AFAIK weight decay does not have significant regularization effect in the modern networks that have normalization layers (e.g. LayerNorm) or standardized weights. I can recommend to look into this two references to understand this effect:

1. L2 Regularization versus Batch and Weight Normalization by Laarhoven: https://arxiv.org/abs/1706.05350.
2. An Exponential Learning Rate Schedule for Deep Learning by Li and Arora - https://arxiv.org/abs/1910.07454.

Hope it helps.

[andsteing]

Hello!

1. The checkpoints published in this repository were pretrained for 90 epochs.
2. Same data augmentation and model initialization
3. We use sigmoid cross entropy (see e.g. tf.nn.softmax_cross_entropy_with_logits) - you could implement this as follows:

def sigmoid_xent(*, logits, labels):
  log_p = jax.nn.log_sigmoid(logits)
  log_not_p = jax.nn.log_sigmoid(-logits)
  nll = -jnp.sum(labels * log_p + (1. - labels) * log_not_p, axis=-1)
  return jnp.mean(nll)

Best!

[owmohamm]

Thanks for the info!!

Just to confirm, you were referring to tf.nn.sigmoid_cross_entropy_with_logits and not tf.nn.softmax_cross_entropy_with_logits, right?

In regards to the weight decay, what value was used and was it applied to full model or were the bias and norm terms excluded? When I try to pretrain the model the bias terms tend to grow much larger than those in the released weights

[andsteing]

You're 100% right, that was an unfortunate typo on my side!

As for weight decay, we used weight_decay=0.1 and we apply the following function to all kernel weights:

def decay_weight_fn(w, lr, weight_decay):
  return (1.0 - lr * decay) * w if decay else w

As for bias terms, we initialized the head biases to -10 :

params["head"]["bias"] = jnp.full_like(params["head"]["bias"], -10.)

HTH

[owmohamm]

Thanks for your help!!, that resolves any questions I had.

[andsteing]

Correction : Actually, for the relevant pretraining run, we used flax.optim.Adam(weight_decay=0.1) instead of the manual version that I described above (so this would indeed apply to the weight decay to all the weights).

[owmohamm]

That's interesting, was this done to combat overfitting?
In my experiments (not this architecture but in general) I've noticed that adding weight decay for the bias terms usually gives an inferior model, but in this case i.e. pretraining ViT-L/16 on imagenet-21k the model tends to overfit quite easily.

[pengzhiliang]

Hello, everyone, @owmohamm @andsteing @akolesnikoff . I meet some same questions when pretrain the ViT-B/16 architecture on imagenet-21k:

1. I also find that the bias in our weight in larger than the released weights.
   As you says, do I need to apply the weight decay to all weight items including bias and norm? Then, should all the biases be initialized to -10 or is it just the bias of the final classifier?

params["head"]["bias"] = jnp.full_like(params["head"]["bias"], -10.)

2. About the multilabel.
   Where should I download multilabel information? Or does it mean that the same image in the 21k dataset of 14M already belongs to different classes? I'm not sure.
3. When I use softmax cross entropy loss, the value of loss after pre-training is about 1.75. However, when using sigmoid cross-entropy loss, the 10th epoch is only about 0.0003. I am not sure if something is wrong.

I have been troubled by these for a long time and hope to get help, thank u very mch!

[andsteing]

To your questions:

1. We used flax.optim.Adam(weight_decay=0.1) during pre-training for these models. Only final classifier head biases were initialized.
2. As for ImageNet21k dataset multi-labeledness, see e.g. ImageNet-21K Pretraining for the Masses. Note that JFT is also multi-label, but ImageNet-1k is not (note that we also use sigmoid loss for ImageNet-1k though, see #62)
3. With the implementation above you should get loss curves that resemble what I posted in #62

[pengzhiliang]

Ok, thank u very much!