Contextual dropout module for attention layers:
Dropout has been widely used in attention layers [1]}, [2]}, [3]}. For example, it can be applied to multi-head attention weights after the softmax operation (see illustrations in Figure REF ). The weights are of dimension \([H, N_K, N_Q]\) , where \(H\)  is the number of heads, \(N_K\)  the number of keys, and \(N_Q\)  the number of queries. In this case, we find that setting \(\alpha ^l\in \mathbb {R}^{H}\)  gives good performance. Intuitively, this coincides with the choice of channel dimension for convolutional layers, as heads in attention could be analogized as channels in convolution.
