Attention dropout causing problem in attention score distribution

[RicRicci22]

Attention dropout causing problem in attention score distribution

[ArthurZucker]

hey! sorry not sure to understand, the dropout should not work with model.eval() no?

[RicRicci22]

Hi @ArthurZucker!
Yes, this is correct, the thing that I was pointing out is that during training the dropout acting on the attention weights is also modifying them using the inverse of the dropout probability.
On a standard layer, I understand that this makes it so that the output is similar in magnitude during training and evaluation.

However, in an attention layer, this causes the attention weights to not sum to 1.
In turn, during inference, since the dropout is ineffective, the attention weights do sum to 1, and thus there is this discrepancy between train and test that I think can cause some troubles.

It is like the network is always making inferences on slightly out-of-distribution samples.

Not sure if I explained it better now!

[ArthurZucker]

That's for sure, but all models are trained that way 😄
I never thought about this, but dropout in general would be bad for inference, feel free to do some benchmarks I am curious!

[dhruvbird]

Yes, I have also noticed this problem and put together a notebook to demonstrate what is happening. https://colab.research.google.com/drive/10f5pqC4XO5grmP1soT-Yh12-JOFg_i3w?usp=sharing

Due to dropout's behaviour in training, it will scale up the softmax outputs. This causes probabilities to be less than or greater than 1.0 and not exactly 1.0 during training, whereas at test/inference time, this behaviour of dropout is not seen because dropout becomes a no-op during inference, and all probabilities add up to 1.0. I think this might be a problem, but I haven't seen it addressed systematically anywhere. This is not new and has been discussed before on the PyTorch forums: pytorch/pytorch#42929

The way I'd solve this is apply dropout before running softmax so that after softmax, the probabilities add up to 1.0.

[ArthurZucker]

TBH if you want you can open a PR to see if this improves performances of let's say Llama3 on MMLU for example! That would be relevant to say wether or not this has potential impact!

[dhruvbird]

TBH if you want you can open a PR to see if this improves performances of let's say Llama3 on MMLU for example! That would be relevant to say wether or not this has potential impact!

Just changing the code and running inference probably won't help and will most likely make things worse since the model was trained in a specific way and inference should try to keep that the same. In my mind the only? way to actually test this theory is to train 2 models and compare them on specific benchmarks. I lack the GPU resources to do so though. @ArthurZucker if you have some resources I'm happy to send a PR and you could help me validate this?