# Attention with Trained Embeddings Provably Selects Important Tokens

arXiv: https://arxiv.org/abs/2505.17282 · HTML: https://mlsanity.com/paper/2505.17282/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 39 of its 48 references that are in the map.
- First posted on arXiv 22 May 2025.

## Builds on (most cited first, 30 of 39)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Parameter-Efficient Transfer Learning for NLP](https://mlsanity.com/paper/1902.00751/) (2019)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [One Fits All:Power General Time Series Analysis by Pretrained LM](https://mlsanity.com/paper/2302.11939/) (2023)
- [Implicit Regularization in Deep Matrix Factorization](https://mlsanity.com/paper/1905.13655/) (2019)
- [In Search of the Real Inductive Bias: On the Role of Implicit Regularization in Deep Learning](https://mlsanity.com/paper/1412.6614/) (2014)
- [Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss](https://mlsanity.com/paper/2002.04486/) (2020)
- [A mathematical perspective on Transformers](https://mlsanity.com/paper/2312.10794/) (2023)
- [The emergence of clusters in self-attention dynamics](https://mlsanity.com/paper/2305.05465/) (2023)
- [Representational Strengths and Limitations of Transformers](https://mlsanity.com/paper/2306.02896/) (2023)
- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer](https://mlsanity.com/paper/2305.16380/) (2023)
- [Transformers as Support Vector Machines](https://mlsanity.com/paper/2308.16898/) (2023)
- [Max-Margin Token Selection in Attention Mechanism](https://mlsanity.com/paper/2306.13596/) (2023)
- [On the Role of Attention in Prompt-tuning](https://mlsanity.com/paper/2306.03435/) (2023)
- [A Mathematical Theory of Attention](https://mlsanity.com/paper/2007.02876/) (2020)
- [Dynamic metastability in the self-attention model](https://mlsanity.com/paper/2410.06833/) (2024)
- [How Far Can Transformers Reason? The Globality Barrier and Inductive Scratchpad](https://mlsanity.com/paper/2406.06467/) (2024)
- [Scaling Laws for Associative Memories](https://mlsanity.com/paper/2310.02984/) (2023)
- [From Self-Attention to Markov Models: Unveiling the Dynamics of Generative Transformers](https://mlsanity.com/paper/2402.13512/) (2024)
- [Mechanics of Next Token Prediction with Self-Attention](https://mlsanity.com/paper/2403.08081/) (2024)
- [Implicit Bias and Fast Convergence Rates for Self-attention](https://mlsanity.com/paper/2402.05738/) (2024)
- [Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot](https://mlsanity.com/paper/2406.06893/) (2024)
- [Implicit Regularization of Gradient Flow on One-Layer Softmax Attention](https://mlsanity.com/paper/2403.08699/) (2024)
- [Attention layers provably solve single-location regression](https://mlsanity.com/paper/2410.01537/) (2024)

## Built on by (most cited first, 1 of 1)

- [Focus and Dilution: The Multi-stage Learning Process of Attention](https://mlsanity.com/paper/2605.01199/) (2026)

## Abstract (arXiv)

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.