Attention with Trained Embeddings Provably Selects Important Tokens

2025-05-22

AI Summary

This work addresses the theoretical gap in understanding how token importance is implicitly encoded in language model embeddings. Focusing on binary classification, it analyzes the co-adaptation between learnable token embeddings and softmax attention under gradient descent training. Through gradient flow analysis and logistic loss optimization, the paper rigorously establishes—within the optimization dynamics framework—that (i) after a single training step, embeddings align with output vectors proportionally to token frequency; and (ii) at convergence, attention concentrates precisely on discriminative tokens while the <cls> embedding maximizes the classification margin. Empirical validation on IMDB and Yelp confirms strong alignment between theoretical predictions and observed embedding behavior and attention patterns. The key contribution is the first rigorous demonstration that standard training inherently performs implicit importance-based token selection, thereby providing the first formal theoretical foundation for interpreting Transformer embeddings.

Technology Category

• Natural Language Processing: (Large) Language Models
• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Large Vision Models

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

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., $ exttt{Softmax}( p^ op E_X^ op ) E_X v = frac{ sum_{i=1}^T exp(p^ op E_{x_i}) E_{x_i}^ op v}{sum_{j=1}^T exp(p^ op E_{x_{j}}) }$, where $E_X = [ E_{x_1} , dots, E_{x_T} ]^ op$ contains the embeddings of the input sequence, $p$ is the embedding of the $mathrm{langle cls angle}$ token and $v$ the output vector. First, we show that, already after a single step of gradient training with the logistic loss, the embeddings $E_X$ 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 $mathrm{langle cls angle}$ 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.

Problem

• Understanding theoretical properties of token embeddings in language models
• Analyzing how gradient descent shapes embedding structure for classification
• Proving attention mechanisms select important tokens based on label relevance

Innovation

• Trained embeddings align with token importance
• Softmax attention selects predictive tokens
• Gradient flow maximizes margin for selection