Large differences between T5 weight initialization in TF and torch
Author: jorgemcgomes
Pull request: no

- `transformers` version: 4.18.0, master branch

### Who can help
@patrickvonplaten

I found some significant differences in weight init between the PT and TF implementations of T5.

The **embeddings** (model.shared):

- In PT, according to `T5PreTrainedModel._init_weights`, they are initialized with random normal with std=1.0:
`module.shared.weight.data.normal_(mean=0.0, std=factor * 1.0)`

- In TF (TFT5Model), the embeddings are initialized as such:
`self.shared = TFSharedEmbeddings(config.vocab_size, config.d_model, name="shared")`
Since initializer_range is not being provided, it is using the default, which is `hidden_size**-0.5` (see TFSharedEmbeddings).

This means that in the base model (d=768), the weights in PT are being initialized with **stdev=1.0**, and in TF they are being initialized with **stdev=0.036**.

The **LM head** (model.lm_head):

- In PT, the initializer is not specified, meaning it is being initialized with a uniform distribution in [-sqrt(1/d_model), sqrt(1/d_model)] (https://pytorch.org/docs/stable/generated/torch.nn.Linear.html). The weights don't seem to be initialized in _init_weights either.
`lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)`

- In TF, the initializer is explicitly provided (TFT5ForConditionalGeneration):
`lm_head_initializer = tf.keras.initializers.RandomNormal(mean=0, stddev=config.initializer_factor)`

So, in the base model, the weights in PT are initialized with a uniform distribution of **[-0.036, 0.036]**, and in TF they are initialized with a random normal with **stdev=1.0**.

I'm not entirely sure about the actual implications of this in model training. But at least the lm_head weights will have a huge impact in loss values initially.

Based on other transformer models I've seen, the "correct" answer seems to be that both weights should be initialised with stdev=1.0. But none of the implementations actually does this.

Comments:
cc @gante @Rocketknight1 

Thanks a lot for the issue @jorgemcgomes! I think for PyTorch and Tensorflow we actually never really made sure that the init is correct because we mostly focused on fine-tuning. But we should correct this! 

I think I made sure that the init is correct in Flax's T5 implementation so we could/should use this as a gold-standard. So let's look there at the Embedding and lm_head:
- https://github.com/huggingface/transformers/blob/048443db863214aef9c8341517b427edced63c81/src/transformers/models/t5/modeling_flax_t5.py#L1366
- https://github.com/huggingface/transformers/blob/048443db863214aef9c8341517b427edced63c81/src/transformers/models/t5/modeling_flax_t5.py#L1384

So it looks like for both the word embeddings and the lm_head the init should be:

`random_normal(mean=0, stddev=config.initializer_factor)`

Guess PyTorch got one right and TF the other one. 

@craffel can you confirm this, that a gaussian normal distribution is used as an init for T5's word embeddings and language model head (in case it's not tied to the word embeddings)

Yep, MTF initializes embeddings as a standard Gaussian. https://github.com/tensorflow/mesh/blob/a32810e32709e0eaad3b241475d3be0957409adc/mesh_tensorflow/layers.py#L2096

@jorgemcgomes Thanks for spotting this! Would you be willing to make a PR to bring the TF/PT implementations in line with the JAX one?

Sure. I can do that in the coming days. But there might be more to this.

Based on the experiences I was doing (my problem/data is very specific, and I'm running some modifications in T5, so take this with a grain of salt), "fixing" the lm_head init (from σ=0.03 to σ=1.0) caused huge initial train/valid losses, even causing instability with the same LR.

There's this interesting bit:

https://github.com/huggingface/transformers/blob/de8b06f9bf908ef1e6317ecb1f74a02313eee72e/src/transformers/models/t5/modeling_t5.py#L1662-L1667

* with tie_word_embeddings=True, the input to the final layer is scaled down by d^-0.5 and multiplied with standard gaussian weights (the embeddings weights).
* with tie_word_embeddings=False, the input to the final layer is **not** scaled down, and **if the proposed fix is introduced** it is also multiplied with standard gaussian weights (the lm_head weights). This doesn't sound right, and can explain the large loss values and instability I mentioned.

And it might also explain why the current PT implementation of T5v1.1 appears to be working fine: the sequence input is not scaled down, but it is being multiplied with small weights instead (initialised with σ=d^-0.5). Two wrongs that cancel each other?
This would mean that the current PT implementation is "fine", but TF and Flax are broken.

That explanation makes sense to me. Just to confirm, is training stable in the current version with the small TF/Flax init?

To summarise, based on my experiments with a non-tied LM head (T5v1.1):

- small embeddings init, small lm_head init --> stable
- small embeddings init, large lm_head init [as in TF] --> unstable
- large embeddings init, small lm_head init [as in PT] --> stable
- large embeddings init, large lm_head init [as in Flax] --> unstable

The init of the embeddings doesn't seem to matter that much at all. Maybe layer norm takes care of that?

And large lm_head inits (as found in the current TF and Flax implementations) are always unstable.

Training

> Yep, MTF initializes embeddings as a standard Gaussian. https://github.com/tensorflow/mesh/blob/a32810e32709e0eaad3b241475d3be0957409adc/mesh_tensorflow/layers.py#L2096

Thanks for looking this up. So I think both embeddings should then be initialized as: tf.random_normal_initializer(
    mean=0.0, stddev=0.05, seed=None
)

meaning `self.config.initializer_factor` should be set to 0.05. 

The most important thing is to match the original code-base here. Don't think we need to run different pretrainings to find the best init scheme since stability is always data-dependent. 

=> Seems like the Flax init methods were good to me so I'd suggest to just apply this to PT and TF as well 

@jorgemcgomes, 

Would you like to open a PR to fix the initialization for T5 here as described in the comment above? Otherwise happy to take over the issue!

Please take over the issue @patrickvonplaten . This got pretty muddy and I'm not sure what is the right approach here.

I have also run into this issue, and I commented on #26441 requesting it be merged since properly initializing the lm_head weights to have mean=0.0 and stdev=1.0/sqrt(d_model) seems to fix the issue with exploding gradients and loss going up that I am seeing.