Large differences between T5 weight initialization in TF and torch
- `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.
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.