Train-Test disparity

[DRSY]

Hi, thanks for your inspiring work and open-sourced codebase.

I noticed that in the "skim attention" module, the 1/0 mask is directly multiplied by the attention probability matrix to prevent skipped tokens from being attended by other tokens. In the paper, it is said that "By doing so, the remaining tokens will have the
identical computational value as actually pruning.". However, during training, the attention probability of remaining tokens does not sum up to 1 because some probability masses are zeroed-out. On the contrary, during inference, the attention probability of remaining tokens always equals 1. This causes train-test disparity and contradicts what is described in the paper.

Looking forward to your reply.
Thanks

[ChandlerGuan]

Hello, thanks for following our work.

This is actually a bug in our design. We have managed to fix it by multiplying the mask by the attention probability matrix before the Softmax function in attention calculation. That is, Eq.(4) in the paper should be revised from $SkimAttn(H^i)=Attn(H^i)*M^i=Softmax((Q(H^i)K(H^i))*M^i$ to $SkimAttn(H^i)=MakedSoftmax((Q(H^i)K(H^i), M^i)=\frac{exp((Q(H^i)K(H^i))*M^i+\epsilon}{\sum{exp((Q(H^i)K(H^i))*M^i+\epsilon}}$. It is fixed with dd44976.