Abstract
We study the approximation capabilities and on-convergence behaviors of one-layer transformers on the noiseless and noisy in-context reasoning of next-token prediction. Existing theoretical results focus on understanding the in-context reasoning behaviors for either the first gradient step or when the number of samples is infinite. Furthermore, no convergence rates nor generalization abilities were known. Our work addresses these gaps by showing that there exists a class of one-layer transformers that are provably Bayes-optimal with both linear and ReLU attention. When being trained with gradient descent, we show via a finite-sample analysis that the expected loss of these transformers converges at linear rate to the Bayes risk. Moreover, we prove that the trained models generalize to unseen samples as well as exhibit learning behaviors that were empirically observed in previous works. Our theoretical findings are further supported by extensive empirical validations.
Abstract (translated)
我们研究了一层变压器在无噪声和有噪声的上下文推理中的下-token预测时的近似能力和在线收敛行为。现有的理论结果主要集中在理解要么第一次梯度步骤，要么样本数量无限时的上下文推理行为上。此外，没有已知的收敛速率或泛化能力。我们的工作通过展示一类一层变压器，在使用线性注意力和ReLU注意力时都是贝叶斯最优的，填补了这些空白。当用梯度下降法训练时，我们通过对有限样本分析显示，这些变压器的预期损失以线性的速度收敛到贝叶斯风险。此外，我们还证明了经过训练的模型能够泛化至未见过的数据，并且表现出与之前工作中观察到的学习行为一致的行为模式。我们的理论发现进一步通过广泛的实证验证得到了支持。
URL
https://arxiv.org/abs/2505.15009
