Transformer通过镜面下降法学习上下文中的隐变量混合模型。
Transformers Learn Latent Mixture Models In-Context via Mirror Descent
- 将注意力机制建模为基于混合转移分布的上下文学习问题。
- 三层Transformer可精确实现一步镜面下降，逼近贝叶斯最优预测。
- 训练后模型的注意力与预测分布符合理论构造，深层模型表现接近多步优化。
序列建模需判断上下文中哪些历史标记具有因果相关性及其重要性，这本质上由Transformer中的注意力层处理，但其内在学习机制仍不清晰。本文将估计标记重要性形式化为上下文学习问题，提出基于混合转移分布的框架：一个隐变量决定过去标记对下一个标记的影响，该隐变量的分布由未观察到的混合权重参数化，而Transformer需在上下文中学习这些权重。我们证明Transformer可通过镜面下降学习这些权重。具体地，我们构建了一个三层Transformer，可精确实现一步镜面下降，并证明所得估计器是贝叶斯最优预测的一阶近似。实验验证了该构造的可学习性：从零训练的Transformer展现出与理论一致的预测分布、注意力模式及学习的转移矩阵；更深模型的表现接近多步镜面下降。
原文摘要 · Abstract (English)
Sequence modelling requires determining which past tokens are causally relevant from the context and their importance: a process inherent to the attention layers in transformers, yet whose underlying learned mechanisms remain poorly understood. In this work, we formalize the task of estimating token importance as an in-context learning problem by introducing a framework based on Mixture of Transition Distributions, where a latent variable determines the influence of past tokens on the next. The distribution over this latent variable is parameterized by unobserved mixture weights that transformers must learn in-context. We demonstrate that transformers can implement Mirror Descent to learn these weights from the context. Specifically, we give an explicit construction of a three-layer transformer that exactly implements one step of Mirror Descent and prove that the resulting estimator is a first-order approximation of the Bayes-optimal predictor. Corroborating our construction and its learnability via gradient descent, we empirically show that transformers trained from scratch learn solutions consistent with our theory: their predictive distributions, attention patterns, and learned transition matrix closely match the construction, while deeper models achieve performance comparable to multi-step Mirror Descent.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。