arXiv:2605.08475cs.LGcs.AI2026-05

Transformer通过预条件理查森迭代实现高斯核回归,误差可控。

Transformers Can Implement Preconditioned Richardson Iteration for In-Context Gaussian Kernel Regression

论文配图:Transformers Can Implement Preconditioned Richardson Iteration for In-Context Gaussian Kernel Regression
图 1 · 摘自论文原文
  • 用softmax注意力模拟预条件理查森迭代求解核线性系统。
  • 在数据有界下,用O(log(1/ε))层和O(√(N/ε))宽的MLP达ε精度。
  • 实验验证了其与经典求解器的误差轨迹高度一致,适合理解非线性ICL机制。

对上下文学习(ICL)的机制研究发现,线性回归等任务可通过迭代算法实现,常使用线性或ReLU注意力变体。对于非线性ICL,先前工作将softmax和核注意力关联到函数梯度型动态,但尚不清楚标准Transformer能否实现具有端到端预测误差保证的收敛求解器。本文研究高斯核下的上下文核岭回归(KRR),证明标准softmax注意力Transformer可在前向传播中通过执行预条件理查森迭代近似KRR预测器。在数据有界假设下,构造出单头Transformer,仅需O(log(1/ε))层与O(√(N/ε))宽度的MLP,即可对长度为N的提示实现ε精度预测。该结构揭示了Transformer内部的功能分解:softmax注意力生成行归一化的高斯核算子以支持跨标记交互,而ReLU MLP层局部近似更新所需的标记内标量运算。我们还训练了类GPT-2的Transformer在高斯过程回归任务上,通过线性探测比较各层输出与经典KRR求解器的步骤输出,发现其误差曲线最符合预条件理查森迭代。消融实验进一步支持此解释。理论与实证共同表明,预条件理查森迭代是softmax注意力Transformer实现非线性上下文高斯核回归的具体机制。

原文摘要 · Abstract (English)

Mechanistic accounts of in-context learning (ICL) have identified iterative algorithms for linear regression and related linear prediction tasks, often using linear or ReLU attention variants. For nonlinear ICL, prior work has related softmax and kernelized attention to functional-gradient-type dynamics, but it remains unclear whether a standard transformer with softmax attention can implement a convergent solver with an end-to-end prediction-error guarantee. In this paper, we study in-context kernel ridge regression (KRR) with Gaussian kernels and show that a standard softmax-attention transformer can approximate the KRR predictor during its forward pass by implementing preconditioned Richardson iteration on the associated kernel linear system. Under bounded-data assumptions, we construct a single-head transformer with $O(\log(1/ε))$ blocks and MLP width $O(\sqrt{N/ε})$ that achieves $ε$-accurate prediction for prompts of length $N$. Our construction reveals a functional decomposition within the transformer architecture: softmax attention produces a row-normalized Gaussian-kernel operator needed for cross-token interactions, while ReLU MLP layers act locally to approximate the intra-token scalar arithmetic required by the update. Empirically, we train GPT-2-style transformers on Gaussian-process regression tasks to further test the preconditioned Richardson interpretation. Through linear probing, we compare the transformer's layer-wise predictions with the step-wise outputs of classical KRR solvers and find that its error profiles align most consistently with preconditioned Richardson iteration. Ablation studies further support this interpretation. Together, our theory and experiments identify preconditioned Richardson iteration as a concrete mechanism that softmax-attention transformers can realize for nonlinear in-context Gaussian-kernel regression.

Transformer核回归迭代算法ICL机制

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