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., $\texttt{Softmax}( p^\top E_X^\top ) E_X v = \frac{ \sum_{i=1}^T \exp(p^\top E_{x_i}) E_{x_i}^\top v}{\sum_{j=1}^T \exp(p^\top E_{x_{j}}) }$, where $E_X = [ E_{x_1} , \dots, E_{x_T} ]^\top$ contains the embeddings of the input sequence, $p$ is the embedding of the $\mathrm{\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 $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 \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.
Abstract (translated)
令牌嵌入在语言建模中扮演着关键角色，尽管具有实际重要性，但对其理论理解仍然有限。我们的论文通过表征经梯度下降获得的嵌入结构来填补这一空白。具体来说，我们考虑了一个单层softmax注意力模型，并带有用于二分类的线性头部，即 $\texttt{Softmax}( p^\top E_X^\top ) E_X v = \frac{\sum_{i=1}^T \exp(p^\top E_{x_i}) E_{x_i}^\top v}{\sum_{j=1}^T \exp(p^\top E_{x_{j}})}$，其中 $E_X = [ E_{x_1}, \dots, E_{x_T}]^\top$ 包含输入序列的嵌入，$p$ 是 $\mathrm{\langle cls \rangle}$ 令牌的嵌入，而 $v$ 是输出向量。首先，我们证明即使经过一次梯度训练（使用逻辑损失），嵌入 $E_X$ 已经通过与输出向量 $v$ 的比例对齐来捕捉数据集中令牌的重要性，这种对齐方式与其在数据集中的出现频率成正比。然后，在通过梯度流训练 $p$ 直到收敛之后，softmax 会选择句子中重要的令牌（即那些能预测标签的令牌），并且最终产生的 $\mathrm{\langle cls \rangle}$ 嵌入会最大化这样的选择所带来的间隔。在现实世界数据集（IMDB、Yelp）上的实验展示了与我们理论揭示的现象学非常接近的结果。
URL
https://arxiv.org/abs/2505.17282
