Abstract
The attention mechanism has been widely used in deep neural networks as amodel component. By now, it has become a critical building block in manystate-of-the-art natural language models. Despite its great success establishedempirically, the working mechanism of attention has not been investigated at asufficient theoretical depth to date. In this paper, we set up a simple textclassification task and study the dynamics of training a simple attention-basedclassification model using gradient descent. In this setting, we show that, forthe discriminative words that the model should attend to, a persisting identityexists relating its embedding and the inner product of its key and the query.This allows us to prove that training must converge to attending to thediscriminative words when the attention output is classified by a linearclassifier. Experiments are performed, which validates our theoretical analysisand provides further insights.