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

Recent years have seen transformers dominate NLP and multi-modal tasks, yet their internal algorithms remain elusive. Early work linked transformer layers to optimization steps like gradient descent [von Oswald et al., 2023], but these often focused on linear tasks or simplified attention mechanisms. In kernel methods, attention has been associated with kernel functions or gradient flows [Cheng et al., 2024], but without end-to-end convergence guarantees. Nonlinear regression, especially Gaussian kernel ridge regression (KRR), offers a well-understood mathematical solution, making it an ideal testbed for mechanistic analysis. Prior studies have explored kernelized attention and functional gradient descent, but the precise algorithmic role of softmax attention in implementing iterative solvers for nonlinear problems has remained unclear. This paper advances the field by explicitly connecting the standard softmax attention mechanism to classical iterative algorithms, bridging the gap between deep learning architectures and numerical linear algebra.

The core challenge is whether a standard softmax-attention transformer can implement a convergent solver with prediction guarantees for nonlinear kernel ridge regression. Existing theories mainly address linear tasks or kernelized attention, leaving the role of softmax normalization in nonlinear settings ambiguous. The difficulty lies in interpreting the layered structure as an iterative algorithm, understanding how attention normalization corresponds to preconditioning, and how MLP layers approximate scalar operations. Clarifying this would not only deepen theoretical understanding but also enable the design of more interpretable and reliable models that can perform controlled numerical computations within the deep learning framework.

This work’s key innovations include: 1) The mechanistic interpretation of softmax attention as implementing row-normalized kernel matrix-vector products, serving as a preconditioner in Richardson iteration. 2) The construction of a transformer architecture with explicit depth and width bounds that approximates the KRR solution with prediction error ε. 3) The theoretical proof linking layer-wise error trajectories to classical iterative solver steps, validated through empirical experiments. 4) Demonstrating that the combination of softmax normalization and Gaussian kernels is essential for the mechanism’s success, supported by ablation studies. These contributions collectively provide a new understanding of how deep models can perform controlled, convergent numerical algorithms.

• �� Formulate the Gaussian kernel ridge regression as a dual linear system (K + λI)w = y.
• �� Design a single-head transformer with input preprocessing, including special tokens and feature augmentation.
• �� Use softmax attention to produce row-normalized kernel matrices, acting as a preconditioner.
• �� Implement iterative Richardson updates via transformer layers, with MLPs approximating scalar multiplications such as D−1w and D−1y.
• �� Construct the network with depth O(log(1/ε)) to control convergence, and width O(pN/ε) to ensure accurate scalar approximations.
• �� Theoretically analyze the error propagation, showing the network’s output converges to the KRR predictor within ε, with explicit bounds on depth and width.

The experiments involve training GPT-2 style models on Gaussian process regression datasets with varying input distributions and prompt lengths. The models are evaluated by comparing layer-wise predictions with classical iterative solvers—preconditioned Richardson, conjugate gradient, gradient descent, and Nesterov methods. The primary focus is on whether the layer-wise error trajectories match those of Richardson iteration, which the results confirm. Ablation studies test the importance of softmax normalization and Gaussian kernels, showing that replacing these components weakens the mechanism. The experiments demonstrate the model’s ability to approximate the KRR solution with high accuracy, validating the theoretical claims across diverse data settings.

The trained transformer achieves ε-accuracy with depth O(log(1/ε)) and width O(pN/ε), with error trajectories mirroring Richardson iteration steps. Layer-wise errors decrease smoothly, outperforming gradient-based methods in convergence speed. Ablation experiments reveal that removing softmax normalization or changing the kernel significantly impairs performance. Across multiple distributions, the model’s predictions closely follow classical iterative solver trajectories, confirming the mechanistic interpretation. These results establish that standard softmax attention can realize a convergent numerical solver for nonlinear regression.

The mechanism enables designing interpretable, provably convergent transformers for scientific computing, large-scale regression, and control tasks. It suggests that deep models can perform controlled numerical algorithms, opening new avenues in AI for physics simulations, inverse problems, and data-driven scientific discovery. The approach could be extended to multi-head architectures and other kernels, broadening its applicability in industry and research.

The analysis assumes bounded data and specific kernel functions, limiting immediate applicability to more complex, real-world data distributions. Extending the mechanism to high-dimensional, non-Gaussian kernels requires further theoretical development. Computational costs grow with desired accuracy, and the current construction relies on smooth scalar functions approximable by ReLU networks. Future work should address these limitations, explore robustness, and generalize to broader settings.