[T5] lm_head weights initialization: set variance to reciprocal of hidden dim
Author: Birch-san
Pull request: yes
Labels: name: WIP

# What does this PR do?

![image](https://github.com/huggingface/transformers/assets/6141784/bddb088b-5bae-4bbb-bbb8-feff450c814e)

before this PR: lm_head weights were initialized with variance of 1, and it output activations with variance ~= hidden_dim. this is a very high variance for logits, and resulted in initial cross-entropy loss of ~110, which is Very High.

after this PR: lm_head weights initialized with variance of reciprocal of hidden_dim. this outputs activations with variance ~= 1. this is results in initial cross-entropy loss of ~11, which is high, but in line with what we'd expect.

Fixes https://github.com/huggingface/transformers/issues/16749 (again)

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

## Who can review?

@ArthurZucker @younesbelkada

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_26441). All of your documentation changes will be reflected on that endpoint.

I would like to argue for setting the initialization of both input and output embedding layers to `normal(0, 0.05)` as discussed in this comment: https://github.com/huggingface/transformers/issues/16749#issuecomment-1101607487. I am not sure why this was actually not done (cc @patrickvonplaten). 

To provide more context: neither the default initialization (`normal(0, 1)`) nor what's proposed in this PR works in my experiment when I train from random weights. `normal(0, 0.05)` works and should probably be the default to align with the original implementation. 

I wasn't sure what to make of Patrick's `stddev=0.05`, as no explanation was given.

I surmise that the reasoning was:

- `VocabEmbedding#__init__` initializes embedding [via the `mtf.layers.embedding_weights` helper](https://github.com/tensorflow/mesh/blob/a32810e32709e0eaad3b241475d3be0957409adc/mesh_tensorflow/transformer/transformer.py#L2155).  
  It **might** default to `_scale_variable_like_classifier_weights=False` (I didn't see any `.gin` file overriding this, but then again I never found the `.gin` file used to train T5), which would make `initializer=None`.
  - side-note: I didn't see any capability to support an untied lm_head. maybe that means we should default to using a tied lm_head? though for vocab as small as 32k, it's [probably better to leave it untied](https://twitter.com/Birchlabs/status/1684178632871120898).
- when `initializer is None`, `mtf.layers.embedding_weights` initializes the embedding [via  `tf.random_normal_initializer()`](https://github.com/tensorflow/mesh/blob/a32810e32709e0eaad3b241475d3be0957409adc/mesh_tensorflow/layers.py#L2096).
- `tf.random_normal_initializer()` [defaults to std=0.05](https://www.tensorflow.org/api_docs/python/tf/random_normal_initializer).

That said, when `hidden_dim==512` (as in t5-small): std=0.05 is _awfully_ close to `std=hidden_dim**-.5=0.044`.

I tried initializing both the embedding and the lm_head with std=0.05. It performed moreorless the same (well, fractionally worse).

<img width="1097" alt="image" src="https://github.com/huggingface/transformers/assets/6141784/31305e4e-bf3b-4bba-afdb-f0baf3768a57">

Thanks for the quick turnaround. In my experiment:
- Current default `std=1`: Starts at a high loss and converges to high value.
- `lm_head` with `std=hidden_dim**-.5`: Starts at a lower loss than before but diverges after a few hundred steps.
- `initializer_factor=0.05`: Starts at an even lower loss and converges nicely. 

how big is your hidden_dim? for t5-small, `512**-.5=0.044` is a very similar number to 0.05.

I am also running my quick tests on `flan-t5-small`.

okay yeah certainly I agree that we should modify _both_ the embedding **and** the lm_head weights.

I think we don't _actually_ have any info from the mesh tensorflow repository on how to initialize an _untied_ `lm_head`. I don't think they support that.

I tried a model with a bigger hidden dim (t5 base, 768). so that there would be a bigger difference between 0.05 and `hidden_dim**-.5`.

<img width="1083" alt="image" src="https://github.com/huggingface/transformers/assets/6141784/2d340649-9b8e-4fb3-934e-6ab92b2a4fbe">

no discernable difference by 600 steps.

so I don't know which is more _correct_ out of `dim**-.5` or `0.05`, but I think empirically there's not a big difference.

certainly I think initializing the _embedding_ to 0.05 matches the MTF repository.

but I don't think the MTF repository supports an untied head at all, so I think that's up to us? unless there's some precedent for what to do there. there's two schools of thought:

- initialize the lm_head as though it were tied to the embedding (0.05)
- initialize the lm_head like you'd initialize most Linear layers (`dim**-.5`)

I kinda think treating the lm_head as untied (`dim**-.5`) makes more sense.

numerically and empirically I think there's not much difference. I think the more important thing is don't initialize either the embedding or the lm_head with `std=1`.

oh, there _is_ a reference to tied vs untied lm_head:  
https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586-L594

I think they're saying "if the embedding is tied, we anticipate that the hidden states will have too much variance, and so we correct it". so whatever std they use to initialize the embedding, is higher than the std an lm_head would want? this feels a bit weird. it kinda sounds like the embedding was initialized with something much larger than 0.05, and they're trying to compensate. is this a clue as to why HF initializes embedding to std=1?

whereas if it's _untied_, they treat it as a `dense` layer with `kernel_initializer=None`:  
https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L88-L90  
so I think they initialize untied lm_head as `std=hidden_dim**-.5`?

maybe we can learn something about how to initialize the weights, based on how they ended up distributed after Google finished training the model?

I took an off-the-shelf model:

```python
from transformers import T5ForConditionalGeneration
model: T5ForConditionalGeneration = T5ForConditionalGeneration.from_pretrained('google/t5-v1_1-small')
```

and I [analysed how all of its layers were distributed](https://gist.github.com/Birch-san/7709286ec73c6795666050a4f2786309).

embedding std is `11.6375`  
lm_head std is `1.1614`

considering encoder self-attn, decoder self-attn, decoder cross-attn…  
attn q is `~0.05` or `~hidden_dim**-.5` (hard to tell the difference)  
attn k is `0.3~0.4`  
attn v is `0.6~0.8`  
attn o is `0.6~1.0`

considering encoder FFN, decoder FFN…  
FFN wi_0 is `0.32~0.36`  
FFN wi_1 is `0.85~1.17`  
FFN wo is `0.53~0.73`

encoder per-block layernorms `~0.02`  
encoder final layernorm `~0.05` or `~hidden_dim**-.5` (hard to tell the difference)  

decoder per-block layernorms `0.07~0.10`  
decoder final layernorm `~0.7`



t5-small pretraining `.gin` config apparently [lives here](https://storage.cloud.google.com/t5-data/pretrained_models/small/operative_config.gin?_ga=2.225153114.-1337093425.1644544368). I hadn't been able to find it in the MTF repository.  
_from https://github.com/allenai/unifiedqa/issues/40#issuecomment-1049151987._

The T5X repository has become somewhat of a reference implementation to pretrain T5X. Here they use stddev=1.0 as an initialization: https://github.com/google-research/t5x/blob/b051e46075fdcb02fcdd4dc648dd9560243bfdb2/t5x/examples/t5/network.py#L289

That aligns with what we currently have in Transformers. Should we maybe open an issue in T5X asking about the differences because it indeed seems like in the original implementation they used 0.05 as init.

@patrickvonplaten  
> Here they use stddev=1.0 as an initialization

do they though? it looks to me like they use `scale=1.0`, not `std=1.0`.

- since embedding is [untied](https://github.com/google-research/t5x/blob/b051e46075fdcb02fcdd4dc648dd9560243bfdb2/t5x/examples/t5/t5_1_1/base.gin#L55),
- [lm_head is a DenseGeneral](https://github.com/google-research/t5x/blob/b051e46075fdcb02fcdd4dc648dd9560243bfdb2/t5x/examples/t5/network.py#L269) with default initialization
- DenseGeneral's default initialization [is `variance_scaling(1.0, 'fan_in', 'truncated_normal')`](https://github.com/google-research/t5x/blob/b051e46075fdcb02fcdd4dc648dd9560243bfdb2/t5x/examples/t5/layers.py#L351)
- `1.0` [is the `scale` parameter](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.initializers.variance_scaling.html) in `std=(scale/n)**.5`
  - when `mode='fan_in'`, `n` refers to "the number of input units in the weights tensor", which I guess is equivalent to PyTorch's `Linear(in_features=…)`
  - so `n=hidden_dim`
  - in other words: `std=hidden_dim**-.5`

@patrickvonplaten  
> Should we maybe open an issue in T5X asking about the differences because it indeed seems like in the original implementation they used 0.05 as init.

what did you think of my investigation that suggested that MTF [initialized untied lm_head using `std=hidden_dim**-.5`](https://github.com/huggingface/transformers/pull/26441#issuecomment-1741877125)?

> The T5X repository has become somewhat of a reference implementation to pretrain T5X. Here they use stddev=1.0 as an initialization: https://github.com/google-research/t5x/blob/b051e46075fdcb02fcdd4dc648dd9560243bfdb2/t5x/examples/t5/network.py#L289
> 
> That aligns with what we currently have in Transformers. Should we maybe open an issue in T5X asking about the differences because it indeed seems like in the original implementation they used 0.05 as init.

Agree that there is a difference! Let's maybe open an issue in T5X and link it here? Think T5X is still pretty active

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

https://github.com/google-research/t5x/issues/1477

Thanks a lot for posting on T5X @Birch-san - :crossed_fingers: someone from the T5 team looks into it

in the meanwhile, maybe worth un-staling/re-opening this PR?

Haha it appears the stale bot is over-aggressive 

I ran into exactly this issue today, and the proposed fixed made the grad_norm during training misbehave less (at least so far. Here are the charts for reference. 

<img width="765" alt="image" src="https://github.com/user-attachments/assets/44754bdd-6dca-4180-93c4-da25bd17d130">

Purple is on master, and blue is with the fix in this MR. 

It would be super if we can merge this PR into master so that anyone trying to train a T5 model from scratch (i.e. no fine-tuning using pre-trained weights) can benefit from this fix. 


@dhruvbird Thanks for sharing your results! It appears this wasn't merged in as it was pending confirmation from the T5 team. Although this would technically be a breaking change, it does appear to be useful for training the model. WDYT about adding @ArthurZucker ?

Yeah sounds good in general, the issue is that it could affect the entire field of diffusion that relies on t5 cc @sayakpaul

In the diffusion community, T5 is not generally trained from scratch, though. Just a heads up. 

I also note that attention is not being scaled by the usual `head_dim**-.5`:  
https://github.com/huggingface/transformers/blob/1bd9d1c89998d7cf7f1ce0290e8944e9299cf407/src/transformers/models/t5/modeling_t5.py#L525

instead, this scale factor is fused into the weight init of the Q projection:  
https://github.com/huggingface/transformers/blob/1bd9d1c89998d7cf7f1ce0290e8944e9299cf407/src/transformers/models/t5/modeling_t5.py#L864

it's likely Google did this for performance reasons (I expect only marginal, and certainly outmoded by flash attention).  
they made this decision knowing they were using Adafactor optimizer, which exhibits equivalent training dynamics either way.

under AdamW, this optimization is harmful. you can only do this trick if your optimizer scales lr by the RMS of the params (e.g. Adafactor, [LARS+AdamW](https://github.com/kakaobrain/torchlars/blob/3b3d7e9c7bd35a31b2c2fa6f213beb0bf6881892/torchlars/lars.py#L133), or [AdamWScaled](https://github.com/PiotrNawrot/nanoT5/blob/1c82d67bf8dea635be68a3b2a68a43b68b665193/nanoT5/utils/copied_utils.py#L339C17-L339C26)).

so I also recommend to change the Q weight init to be the same as the K and V weight init (`d_model**-.5`), and to scale `Q•K.T` by `head_dim**-.5` (i.e. typical attention formulation). upgrading to torch SDPA would of course be nice too.

#31167 is coming, I gotta review it, scaling would be breaking as it is not done currently