Attention with Trained Embeddings Provably Selects Important Tokens

Token embeddings play a crucial role in language modeling but, despite this practical relevance, their theoretical understanding remains limited. Our paper addresses the gap by characterizing the structure of embeddings obtained via gradient descent. Specifically, we consider a one-layer softmax attention model with a linear head for binary classification, i.e., textttSoftmax( p^top EX^top ) EX v = frac sumi=1T exp(p^top Exi) Exi^top vsumj=1T exp(p^top E_xj) , where EX = [ Ex1 , dots, ExT ]^top contains the embeddings of the input sequence, p is the embedding of the langle cls rangle token and v the output vector. First, we show that, already after a single step of gradient training with the logistic loss, the embeddings EX capture the importance of tokens in the dataset by aligning with the output vector v proportionally to the frequency with which the corresponding tokens appear in the dataset. Then, after training p via gradient flow until convergence, the softmax selects the important tokens in the sentence (i.e., those that are predictive of the label), and the resulting langle cls rangle embedding maximizes the margin for such a selection. Experiments on real-world datasets (IMDB, Yelp) exhibit a phenomenology close to that unveiled by our theory.