具有训练后嵌入的注意力机制可证明地选择重要词元
Attention with Trained Embeddings Provably Selects Important Tokens
AI总结：
本文研究了单层softmax注意力模型中梯度下降训练后的词元嵌入结构，证明其能按词频对齐输出向量并选择重要词元以最大化分类间隔，理论结果在IMDB和Yelp数据集上得到验证。
摘要：
词元嵌入在语言建模中发挥着至关重要的作用，但尽管具有这种实际相关性，对其的理论理解仍然有限。我们的论文通过刻画通过梯度下降获得的嵌入结构来填补这一空白。具体而言，我们考虑了一个带有线性头用于二分类的单层 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）上的实验展示了与我们的理论所揭示的现象接近的特性。
英文摘要：
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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
