arXiv:2505.17282cs.LGcs.CL2025-05NeurIPS

训练后注意力能自动选出重要词元,理论揭示其机制。

Attention with Trained Embeddings Provably Selects Important Tokens

  • 通过梯度下降使词元嵌入与标签相关性对齐。
  • 训练收敛后,注意力精准选出预测标签的关键词元。
  • 适合研究注意力机制原理或模型可解释性的读者。

词元嵌入在语言建模中至关重要,但其理论理解仍有限。本文以单层Softmax注意力模型配合线性头进行二分类,分析梯度下降下的嵌入结构。首先,仅经过一次梯度更新,嵌入即与输出向量对齐,比例于词元在数据集中出现的频率。随后,在梯度流下训练⟨cls⟩嵌入直至收敛,此时Softmax注意力能准确选择句子中对标签具有预测性的关键词元,且⟨cls⟩嵌入最大化该选择的边界。真实数据集(IMDB、Yelp)上的实验结果与理论预测高度吻合。

原文摘要 · Abstract (English)

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.

注意力机制理论分析嵌入对齐可解释性

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。