T5 - Flax - Decreasing performance on pretraining
## Environment info
- `transformers` version: 4.9.0.dev0
- Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29
- Python version: 3.8.10
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (tpu)
- Jax version: 0.2.17
- JaxLib version: 0.68

### Who can help
@patrickvonplaten 

### Script
Using a slightly modified version of run_t5_mlm_flax.py that also support streaming large datasets.

## Information
I am posting this as a bug report, since the behaviour is counter intuitive. I am not sure if this is a bug with jax/T5, or if it is actually a behaviour that should be expected from T5.

We are training T5-base (v1.1) on a large, cleaned 250GB Norwegian dataset. We are training from 1M steps, which should equal roughly two complete epochs. With a lr=8e-3, bs=32, seq_length=512, adafactor, we are experiencing a steady decay in loss:
![image](https://user-images.githubusercontent.com/9079808/131355355-22e4542f-6501-450b-b2fb-4c3b967de0aa.png)
 
The image above shows the first 250k steps. We needed to restart here, so I have not patched the event-files together. But the final loss after 1M steps ends on 1.349. Eval accuracy is also increasing.

The weird thing is that the final checkpoint has really terrible performance on our downstream task!

Looking into this issue, we evaluated multiple pre-training steps, by finetuning each of them 60k steps on a task of translating between two Norwegian dialects. 

![image](https://user-images.githubusercontent.com/9079808/131357344-d482bfb8-659e-468f-a9d6-1378d3b4eccd.png)

The red and blue dots are two models done before and after the t5 optimisation submitted by @patrickvonplaten.

The tendency here is very clear. After roughly 200k steps the model starts to suddenly perform worse on the downstream task, even if the loss is decreasing and the eval accuracy of the pretrained model in improving. The detonation happens before 1 epoch of the pretrain dataset, and though it looks like over-fitting, we find this extremely unlikely.

We have more experience with BERT-like models, and here performance on downstream tasks always improves as long as MLM-accuracy is improving. Is this expected behaviour of T5?
 













To me this looks like the model is somehow overfitting on the pretraining data. 

Think it's very hard to know exactly what is happening here though as the model is trained on specific Norwegian data...one thing it would try is to take the fully pretrained model and to check if it has a tendency to generate the same output tokens. Maybe some words are overly represented in the training data and the model overfits to those words?!  
Maybe using dropout and/or weight_decay (the classic methods for regularization) could help here? 

Gently pinging @craffel (hope that's fine) here in case he has seen something similar, has good ideas for analyzing the fully pretrained model and/or other ideas of what could be going on :-)<|||||>@peregilk - also could you link the model repo on the hub here so that I can take a look into the config? <|||||>Thanks. Here is the link to the repo: https://huggingface.co/pere/norwegian-t5-base-NCC-fast

I absolutely agree that this really looks a lot like overfitting. However, I really can not see why this could be happening. The corpus is described [here](https://github.com/NBAiLab/notram/tree/master/corpus) and [here](https://github.com/NBAiLab/notram/blob/master/corpus/official_NCC2.md). It is huge (250GB), and heavily deduplicated and cleaned. It is also a collection from multiple sources. There should really be no repetitive parts here to overfit on.

Please note that English words or phrases often are used in Norwegian today, and that most Norwegian speaks/understands English. We have therefore added roughly 15GB of English text to the corpus. At least in theory, the model should have some basic understanding of English as well. 

Ill do some more tests and report the results.
  

<|||||>Interesting. I will just confirm that the only time we have seen this behavior (train loss going down, MLM accuracy going up, but downstream task performance going down) is when the model is overfitting to the training dataset. When you say "Eval accuracy is also increasing", do you mean you are computing MLM accuracy on a held-out validation set? Or is it on the train set?

Other random things to think about -
1. How many steps are you training on for the downstream task? How are you doing checkpoint selection?
2. What vocabulary are you using?
3. Are you regularizing on the downstream task?<|||||>Thanks both @patrickvonplaten and @craffel for your insightful comments. Highly appreciated.

I have dived into this, and I am now even more confused about what is going on. As Patrick suggested, I looked at how the model doing MLM-like-tasks at various pretraining steps. I am not detecting repeated tokens, or other oddities. I was only able to figure out how to give me the first prediction, so it is a bit hard comparing it directly to BERT, but it looks to be on the same level. Most predictions give grammatical sense. Also for English (<10% of total training corpus).

```
from transformers import AutoTokenizer, T5ForConditionalGeneration, FlaxT5ForConditionalGeneration
model = T5ForConditionalGeneration.from_pretrained('pere/norwegian-t5-base-NCC-fast', from_flax=True)
# Load 250k checkpoint instead 
#model_250k = T5ForConditionalGeneration.from_pretrained('pere/norwegian-t5-base-NCC-fast', from_flax=True, revision='49d7631d423fc64770a8c8e0d55216792031c97d')
tokenizer = AutoTokenizer.from_pretrained('pere/norwegian-t5-base-NCC-fast', use_fast=True)
input_ids = tokenizer.encode('This is a small <extra_id_0> explaining how to <extra_id_1> a language model.', return_tensors='pt')
print(tokenizer.batch_decode(model.generate(input_ids)))
#output -> "section" and "write"
``` 
Since it seems to perform reasonable good on English, I tried running the run_summarization_flax.py example script on the xsum dataset. I modified the script to support t5, and am using the default settings (lr 5e-5, epochs 6). I adjust the batch size to maximum based on vocab size (without adjusting lr), and are getting these ROUGHE2-results.
BART (default example): 16.99
t5-base: 13.66
mt5-base: 10.36
Norwegian-t5: 9.16 (after 100k) - 7.10 (after 1.000k). See graph below:
![image](https://user-images.githubusercontent.com/9079808/132212886-ad3d7e7a-7c99-484d-b6b5-1631314696c8.png)

I think this points in the direction of the error being on the pretrained model (even if the qualitative mlm-test does not detect this).

@craffel: We are using a community script for the streaming of the dataset here. I read though it, and spotted a few inaccuracy. For instance it seem like they are drawing a new eval set at the start of each epoch. This is of course not ideal, but the model is not running more than 2-3 epochs, so I doubt that this is the reason for the eval-accuracy to be improving.

 1) The first figure is from running 1 epoch on a 100.000 example parallell corpus. Max performance here is after nearly 10 epochs, I did not have time to do this on all checkpoints. However, I made a few tests, and checkpoints that are doing bad after 1 epoch, seem to be bad at 10 epochs as well. 

2) A custom built 50k cased vocab file - Norwegian. I have verified that the tokenisation is reasonable.

3) We are using the default dropout=0.1. I do not think weight decay/L2 is set here.     

My main concern here is the added streaming code in the [training script](https://huggingface.co/pere/norwegian-t5-base-NCC-fast/blob/main/run_t5_mlm_flax_streaming.py). If it for some strange reason kept feeding the same examples to the trainer, this is maybe the result that should be expected (?). However, I have been reading through the code carefully, and this does not seem to be happening.

Any ideas? Are there verified T5 pretraining scrips with streaming support that I can test?




 <|||||>Is it maybe possible that you can try the official (non-streaming) t5 MLM script: https://github.com/huggingface/transformers/blob/master/examples/flax/language-modeling/run_t5_mlm_flax.py? This one is more tested. <|||||>@patrickvonplaten Yes. I have tried that - on a smaller dataset. I am not seeing issues like this.

However, I was unable to get it to run on the 250GB set. I created a TPU VM with an external disk. Unfortunately, I am getting memory issues on the VM. <|||||>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/master/CONTRIBUTING.md) are likely to be ignored.<|||||>There is one small issue with this script. But not sure if that is the reason for the problem.

1. [Here](https://huggingface.co/pere/norwegian-t5-base-NCC-fast/blob/main/run_t5_mlm_flax_streaming.py#L514) : Text stream has been tokenized with an additional `</s>` . I believe we should add an additional argument `add_special_tokens=False`. Here adding `eos` token refers document separator not sample separator. 
2. In the [data collator](https://huggingface.co/pere/norwegian-t5-base-NCC-fast/blob/main/run_t5_mlm_flax_streaming.py#L234), at the end of each `src_input`, and `target_input` we should manually add an `eos_token`

@patrickvonplaten  @peregilk 
<|||||>Just to update this issue as well. We have looked closer at this, and have been running a lot of experiments. I will post the results from this when they are done. The short summary is however that we need to increase regularisation (ie weight_decay) when finetuning our T5s that have been pretrained a long time. I have tested both with a large streaming dataset and smaller non-streaming datasets. It seems like we have the same issue in both.

In the examples reported in the graph above, we are able to get decent results when finetuning a T5 that has been pretrained for 200k steps, right "out of the box". There are optimal learning rates, but all in all it is not very sensitive to this. The T5 that has been pretrained for 1M steps is only possible to finetune if you add a significant amount of regularisation. Adding "weight_decay" seem to be the trick. The optimal dropout seem to be close to 0.1 for all models.

I will post our results when they are completed in a few days. So far, we have no explanation/theory as to why this is happening. Might be related to what @sbmaruf is reporting. I simply do not know.<|||||>I have never used any kind of regularization on T5 except for dropout; dropout=0.1 often helps during fine-tuning (have never tried any other value).<|||||>Thanks @sbmaruf,  @craffel and @patrickvonplaten for your comments. Before continuing this thread, I wanted to run more experiments and making sure that we get consistent results. We have now completed these tests.

First let me sum up: We are training a T5 on a large Norwegian corpus from scratch. We are using the Flax example code with the hyperparameters above, and using the T5 v1.1 training regime. We are seeing a good decline in loss and increase in accuracy. This seem to be the case also for the evaluation-set! We have confirmed "the error" on two runs, both on a huge streaming 150GB corpus, and on a smaller 30GB corpus. We are pretty certain that this is a high quality corpus, and it is intensely deduplicated (roughly 20% of the corpus is Norwegian MC4. the major part is born digital public reports). We are pretty sure we not overfitting. We are mainly testing on a simple translation task (Bokmål vs Nynorsk), but have seen the same issue on other tasks, even if we have not run intensively tests here since the datasets are smaller.

Here is what we are observing: It gets incrementally harder to finetune various checkpoints from the pretraining! We are able to finetune late checkpoints but only after adding weight decay. We have experimented with a lot of different learning rates, and the tendency is the same all over.

Here you can see an example from the 100k pretrain checkpoint:
![image](https://user-images.githubusercontent.com/9079808/139430715-a373deab-d0b3-4393-961a-196dd9169ecb.png)

You can see that it finetunes easily. Adding weight decay just makes the training a bit slower.

Here is from the 500k pretrain checkpoint:
![image](https://user-images.githubusercontent.com/9079808/139430954-1bdd4b1b-6bdf-4a41-bbd8-5c988f4e126a.png)

Here we need weight decay to get stable finetuning. 

At the 1M pretrain checkpoint, finetuning is also notably slower. Only with weight decay are we able to complete it:
![image](https://user-images.githubusercontent.com/9079808/139431256-5e8e1b25-96b6-4bc1-8f92-bd98eb0ea675.png)

As you see all the yellow lines (no weight decay) are off the chart here. It simply does not converge at all.

Here are the complete results from [W&B](https://wandb.ai/nbailab/norwegian-t5-base-checkpoints/reports/T5-Norwegian-Evaluation--VmlldzoxMTY4NzI4?accessToken=7jsxdmw82kf29vnqzxevkii6y9reuiak8wv591e256oe5ryf904h125g1h202d0x)

For some really strange reason we are able to end up with a decent result also on the late checkpoints, but only after very careful tuning of the hyperparameters. The model seem to be very unstable and a nightmare to finetune.

Do any of you think the bug reported by @sbmaruf above could cause any of this?

 <|||||>I don't think that adding or not adding the EOS token makes a difference - so I highly doubt that is the reason for your observations...BTW here is the original preprocessing code for pretraining: https://github.com/google-research/text-to-text-transfer-transformer/blob/main/t5/data/preprocessors.py#L1864<|||||>Thanks @patrickvonplaten for the response. 

We have a new version of training corpus available soon. From the figure above it seems like we are able to detact instability after less than a week on a v3-8, maybe even before we reach one epoch. We will see if we are able to get a pytorch version of the T5 pretraining running. Maybe we can train pytorch and flax in parallell and see if they differ.

If you have any idea about what is causing this and/or other tests we can do, please let us know.
<|||||>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/master/CONTRIBUTING.md) are likely to be ignored.