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

• URL: http://arxiv.org/abs/2605.08475v1
• Date: Fri, 08 May 2026 20:50:58 GMT
• Title: Transformers Can Implement Preconditioned Richardson Iteration for In-Context Gaussian Kernel Regression
• Abstract summary: We study in-context kernel ridge regression (KRR) with Gaussian kernels.<n>We show that a standard softmax-attention transformer can approximate the KRR predictor during its forward pass.<n>We find that its error profiles align most consistently with preconditioned Richardson iteration.
• Score: 13.818160005611752
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Cubit: Token Mixer with Kernel Ridge Regression [30.306738326416223]
  We show that the attention module in Transformers can be interpreted as performing Nadaraya-Watson regression.<n>We propose Cubit, a potential next-generation architecture that leverages Kernel Ridge Regression (KRR)<n>We argue that Cubit, as a KRR-based architecture, provides a stronger mathematical foundation than the vanilla Transformer.
  arXiv  Detail & Related papers  (2026-05-07T16:18:55Z)
• LaplacianFormer:Rethinking Linear Attention with Laplacian Kernel [27.87296519831803]
  quadratic complexity of softmax attention presents a major obstacle for scaling Transformers to high-resolution vision tasks.<n>We propose LaplacianFormer, a Transformer variant that employs a Laplacian kernel as a principled alternative to softmax.<n>Experiments on ImageNet show LaplacianFormer achieves strong performance-efficiency trade-offs while improving attention.
  arXiv  Detail & Related papers  (2026-04-22T09:04:54Z)
• PRISM: Parallel Residual Iterative Sequence Model [52.26239951489612]
  We propose PRISM (Parallel Residual Iterative Sequence Model) to resolve this tension.<n>PRISM introduces a solver-inspired inductive bias that captures key structural properties of multi-step refinement in a parallelizable form.<n>We prove that this formulation achieves Rank-$L$ accumulation, structurally expanding the update manifold beyond the single-step Rank-$1$ bottleneck.
  arXiv  Detail & Related papers  (2026-02-11T12:39:41Z)
• In-context Learning for Mixture of Linear Regressions: Existence, Generalization and Training Dynamics [34.458004744956334]
  We prove that there exists a transformer capable of achieving a prediction error of order $mathcalO(sqrtd/n)$ with high probability.<n>We also analyze the training dynamics of transformers with single linear self-attention layers, demonstrating that, with appropriately parameters, gradient flow optimization over the population mean square loss converges to a global optimum.
  arXiv  Detail & Related papers  (2024-10-18T05:28:47Z)
• Highly Adaptive Ridge [84.38107748875144]
  We propose a regression method that achieves a $n-2/3$ dimension-free L2 convergence rate in the class of right-continuous functions with square-integrable sectional derivatives. Har is exactly kernel ridge regression with a specific data-adaptive kernel based on a saturated zero-order tensor-product spline basis expansion. We demonstrate empirical performance better than state-of-the-art algorithms for small datasets in particular.
  arXiv  Detail & Related papers  (2024-10-03T17:06:06Z)
• On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability [34.43255978863601]
  Several suggest that transformers learn a mesa-optimizer during autorere training. We show that a stronger assumption related to the moments of data is the sufficient necessary condition that the learned mesa-optimizer can perform.
  arXiv  Detail & Related papers  (2024-05-27T05:41:06Z)
• Stable Nonconvex-Nonconcave Training via Linear Interpolation [51.668052890249726]
  This paper presents a theoretical analysis of linearahead as a principled method for stabilizing (large-scale) neural network training. We argue that instabilities in the optimization process are often caused by the nonmonotonicity of the loss landscape and show how linear can help by leveraging the theory of nonexpansive operators.
  arXiv  Detail & Related papers  (2023-10-20T12:45:12Z)
• Transformers as Support Vector Machines [54.642793677472724]
  We establish a formal equivalence between the optimization geometry of self-attention and a hard-margin SVM problem. We characterize the implicit bias of 1-layer transformers optimized with gradient descent. We believe these findings inspire the interpretation of transformers as a hierarchy of SVMs that separates and selects optimal tokens.
  arXiv  Detail & Related papers  (2023-08-31T17:57:50Z)
• Trained Transformers Learn Linear Models In-Context [39.56636898650966]
  Attention-based neural networks as transformers have demonstrated a remarkable ability to exhibit inattention learning (ICL) We show that when transformer training over random instances of linear regression problems, these models' predictions mimic nonlinear of ordinary squares.
  arXiv  Detail & Related papers  (2023-06-16T15:50:03Z)
• What learning algorithm is in-context learning? Investigations with linear models [87.91612418166464]
  We investigate the hypothesis that transformer-based in-context learners implement standard learning algorithms implicitly. We show that trained in-context learners closely match the predictors computed by gradient descent, ridge regression, and exact least-squares regression. Preliminary evidence that in-context learners share algorithmic features with these predictors.
  arXiv  Detail & Related papers  (2022-11-28T18:59:51Z)
• Softmax-free Linear Transformers [90.83157268265654]
  Vision transformers (ViTs) have pushed the state-of-the-art for visual perception tasks. Existing methods are either theoretically flawed or empirically ineffective for visual recognition. We propose a family of Softmax-Free Transformers (SOFT)
  arXiv  Detail & Related papers  (2022-07-05T03:08:27Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.