Fix embeddings resizing in TF models
Author: jplu
Pull request: yes

# What does this PR do?

Currenlty when the embeddings are resized the biases are not resized in same time. In TF there is no explicit link between the decoder weights and biases in a dense layer contrarily than in PT. This PR fixes this issue by resizing in same time the biases, even thought I don't know if this is the best solution. @LysandreJik @sgugger what do you think?

Comments:
Thanks @sgugger for your useful comments. I was thinking the same about `get_output_embeddings` but I didn't want to change to much things in same time.

I like very much the solution you proposed and I'm totally fine with it!

@sgugger I have reworked the resizing for the bias and applied it on BERT at first for testing. Are you agree with this new way to do? If yes, I will do the same for the other models.

@sgugger @patrickvonplaten @LysandreJik This PR takes care of resizing all the bias, and if we start to change how the embeddings are resized + modify the generation, I think it would be a bit too much and out of the scope of this PR. Then, what I propose is to keep how it was at the beginning in `generation_tf_utils.py` and the `self.get_output_embeddings` methods and move this discussion on another PR. In this another PR I would like as well to fully review how the resizing is done, because the number of line of codes can be largely reduced and simplified. What do you think?

It would be awesome if we can keep the `get_output_embeddings()` method and leave `generate()` as it is and only focus on the resizing problem here. I'm 100% on board with fixing the resizing problem and it'd be awesome to do this orthogonally to `get_output_embeddings()`. 

A couple of reasons why I would like to keep `get_output_embeddings()` (I can copy this to the new PR as well):
1) Consistency with PyTorch. In PyTorch `get_output_embeddings()` is even more integrated with other functionalities (like weight tying) and I think we should stay consistent in TF and PT
2) `get_output_embeddings()` is an important function IMO to quickly get the correct logit matrix. Without this function it's not at all always obvious how to get the output embeddings for some models (especially EncoderDecoder, RAG, ...). A unified API for all models is of great help here IMO and I use it a lot actually
3) Don't want to tie the capability of a model to `generate()` with the `MODEL_FOR_....` classes - this is inconsistent with PyTorch and unnecessarily creates a dependency IMO.



Thanks a lot @patrickvonplaten for sharing this! I think we should move this talk to a more suited place, and meanwhile I will revert that part of the changes.

I disagree with you on this @patrickvonplaten 

> 1. Consistency with PyTorch. In PyTorch get_output_embeddings() is even more integrated with other functionalities (like weight tying) and I think we should stay consistent in TF and PT

The weight tying cannot be done the same way in TF (and honestly the resizing on the PyTorch side is a bit hacky and very hard to understand it kind of goes against our principle of no magic code), so this alone is not an argument for keeping the `get_output_embeddings` method

> 2. get_output_embeddings() is an important function IMO to quickly get the correct logit matrix. Without this function it's not at all always obvious how to get the output embeddings for some models (especially EncoderDecoder, RAG, ...). A unified API for all models is of great help here IMO and I use it a lot actually

The problem is that this function is always implemented to return the input embeddings, so the function as it is does not do anything more than `get_input_embeddings` while giving the user a false sense of what it returns. (Note that there is no model in TF apart from mobileBERT that has the capability of having different weights for the embeddings and the decoder, the weights are **always** tied).

> 3. Don't want to tie the capability of a model to `generate()` with the `MODEL_FOR_....` classes - this is inconsistent with PyTorch and unnecessarily creates a dependency IMO.

The PyTorch side has no assert, so in that case, the consistent thing is to remove the assert entirely. 

I could be convinced to leave the `get_output_embeddings` method for mobileBERT only since it's the only model where it returns something useful, but it's dangerous to have it otherwise (unless we had a way to untie the weights, but that's for another PR!)

Ok we debriefed a bit with @patrickvonplaten to avoid spamming the PR. I had missed that some models are already using an output embeddings that is different from the input embeddings (most models are tied), like T5 or mT5. So those, like `mobileBERT`, will definitely need the `get_output_embeddings` method. Right now though, the resizing does not work for those models.

In the end, we both agree on keeping that method, add the `get_output_bias` method and the `resize_embeddings` should use the outputs of those two methods as well as `get_input_embeddings` in all the things it has to resize. To check if the input embeddings and output_embeddings are the same (and not resize them twice) we could use the `._handle_name` attribute of their weights (or something else if you have a better idea).

Does that all make sense?

Ok, I'm totally fine with this 👍 ! Nevertheless, there are still few things I don't get.

> Right now though, the resizing does not work for those models.

What do you mean by the resizing does not work? Which one? Do you have a more specific example?

> To check if the input embeddings and output_embeddings are the same (and not resize them twice) we could use the ._handle_name attribute of their weights (or something else if you have a better idea).

I don't understand this sentence, do you have an example? What do we have to check if the input/output embeddings are different if we get them with two separate methods (namely get_input_embeddings and get_output_embeddings).

The new T5 and mT5 models have an output embedding layer that is sometimes tied to the input embeddings (so same weights like BERT) and sometimes different. When it's different, it is not resized.

> I don't understand this sentence, do you have an example? What do we have to check if the input/output embeddings are different if we get them with two separate methods (namely get_input_embeddings and get_output_embeddings).

The output embeddings are, very often, the same as the input embeddings (BERT situation) so in most instances `get_output_embeddings` will return the same thing as `get_input_embeddings` (which is why we initially decided to remove `get_output_embeddings` when discussing together). However, in some cases, it returns something different (mT5 and T5 as mentioned above, or mobileBERT) which is (with avoiding a breaking change) the main argument to keep this `get_output_embeddings` method. However, when taking its result in `resize_embeddings`, we should check if we get a different result from `get_input_embeddings`. Is that clearer?

Crystal clear!!! Thanks a lot for the details! I will proceed to the changes once the sprint is finished 👍 

@patrickvonplaten I have put back the `get_output_embeddings`, does it seems ok for you now, or did I forget something?

Did I miss anything else?

@LysandreJik any idea why all the tests are failing with a timeout?

Yes, I can see why. Seeing with how to fix it, you'll probably have to rebase.

@jplu, please kindly rebase your branch - this is yet another edge case I haven't expected - fixed in master. Thank you!

Thanks @stas00 and @LysandreJik for having fixed the issue.

@LysandreJik @sgugger @patrickvonplaten is there anything else I have to do in this PR? or it looks ok for you?

@patrickvonplaten any other comments? Or you are fine with the current version?

Not sure why it says @mfuntowicz force-pushed on this PR, but sadly it seems like the history was messed up a bit. Maybe it can be resolved by just `git reset --hard` to the commit before Morgan's force-push? 

Oh no @mfuntowicz killed all my open PRs 😄 ok trying to fix this :)

Ok should be good now!

LGTM! @LysandreJik just missing your approval. The Flax tests do not pass and I don't know why :(