Research map / Architectures and efficiency / Attention and transformers
Attention with Trained Embeddings Provably Selects Important Tokens
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Analyzes how embeddings trained by gradient descent in a one-layer softmax attention classifier encode token importance, showing one gradient step aligns token embeddings with the output vector by label frequency and gradient flow on the CLS embedding yields max-margin selection of important tokens.
Problem
Theoretical understanding of what token embeddings learn under gradient training in attention models, and how this relates to token statistics and selection of important tokens for binary classification.
Approach
Shows that a single gradient step aligns token embeddings with the output vector proportional to a signed frequency statistic, and that gradient flow on the CLS embedding converges in direction to a max-margin solution selecting all completely positive/negative tokens.
Key results
- After one gradient step, token embeddings align with the output vector v proportionally to the average signed frequency alpha_s (Lemma on one-step GD, error bounded by 11 eta_0 d^{-1/4}).
- Under the single-relevant-token assumption, gradient flow on the CLS embedding has norm diverging and its limiting direction selects all completely positive/negative tokens and is a max-margin direction (Theorem 1).
- Experiments on IMDB and Yelp show trained embeddings reflect token posterior differences, qualitatively matching theory; the two-layer model reaches loss around 1e-5 versus about 1e-1 for the one-layer model.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train all parameters of the one-layer attention model to convergence | One-layer attention · Synthetic data | token-importance alignment | Helped |
| Train all parameters of the one-layer attention model to convergence | One-layer attention · IMDB and Yelp | token-importance alignment | Helped |
| Use the two-layer model with a skip connection and layer normalization instead of the one-layer model | Attention model · IMDB and Yelp | loss: of the order of $10^{-1}$ -> of the order of $10^{-5}$ | Helped |
| Use a two-layer model with a skip connection and layer normalization | Attention model · IMDB and Yelp | token-importance alignment | No clear effect |
Surprising findings
- The two-layer model with skip connection and layer norm reaches much smaller training loss (~1e-5) than the one-layer model (~1e-1), yet shows similar embedding structure.
Limitations
- Analysis covers only first-order token statistics and does not capture causal structure between tokens.
- Theory assumes a one-layer model and sequences containing a single relevant token (Assumption on data).
- Sufficient condition for selecting only important tokens is stated under an additional margin assumption.
Experimental setup
Datasets: IMDB Yelp reviews synthetic K-level token data
Benchmarks: IMDB Yelp
Metrics: dot-product of embeddings with CLS embedding and output vect training loss
Builds on: implicit bias of gradient descent on logistic loss max-margin attention SVM equivalence
Closest papers by meaning
- SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking (NeurIPS 2026 · similarity 0.88)
- Attention Sinks: A 'Catch, Tag, Release' Mechanism for Embeddings (NeurIPS 2025 · similarity 0.88)
- Selective Attention: Enhancing Transformer through Principled Context Control (NeurIPS 2024 · similarity 0.87)
- Implicit Optimization Bias of Next-token Prediction in Linear Models (NeurIPS 2024 · similarity 0.87)
- Towards Understanding the Word Sensitivity of Attention Layers: A Study via Random Features (ICML 2024 · similarity 0.87)
- NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.