Ordinary Least Squares is a Special Case of Transformer

TL;DR

Proves that a single-layer linear Transformer is mathematically equivalent to Ordinary Least Squares (OLS) regression via spectral decomposition, enabling one-pass statistical inference.

cs.LG 🔴 Advanced 2026-04-15 68 views
Xiaojun Tan Yuchen Zhao
Transformer statistical inference spectral analysis linear regression neural networks

Key Findings

Methodology

The authors employ spectral decomposition of the empirical covariance matrix to establish a parameter configuration where the linear attention mechanism replicates the OLS closed-form solution. By decomposing Ω = VΛ V→, they construct matrices WQ, WK, WV as L=VΛ^{-1/2}, aligning attention computations with the inverse covariance operation. This configuration allows the Transformer to perform OLS in a single forward pass. Empirical validation involves training on synthetic data, demonstrating convergence of parameters and outputs to the theoretical OLS solution, confirming the algebraic equivalence and stability of the approach.

Key Results

  • In a 1D linear regression task with 500 samples, the trained linear Transformer’s parameter L converged to the theoretical L* after 5000 epochs, with output errors approaching noise level, confirming exact replication of OLS. Spectral analysis verified that attention computes the inverse covariance, establishing the structural isomorphism. The model’s internal parameters and output closely matched the analytical solution, demonstrating the method’s effectiveness.
  • The training dynamics show rapid alignment within 1000 epochs, with the model’s output stabilizing at the OLS solution, and parameters recovering the inverse covariance structure exactly. This confirms that the Transformer can implicitly learn the statistical operator corresponding to OLS, providing a new perspective on its inference capabilities.
  • Analysis of slow and fast memory mechanisms reveals that the fixed matrix L encodes long-term statistical patterns, while attention scores dynamically adapt to real-time input, enabling context-sensitive inference. This decoupling explains the Transformer’s ability to handle long sequences and adapt to distribution shifts.

Significance

This work fundamentally reinterprets Transformer as a statistical inference engine capable of executing classical estimators in a single forward pass. It bridges the gap between deep learning and traditional statistics, offering a mathematically grounded explanation for the model’s remarkable ability to perform in-context learning and memory. The spectral approach provides a clear algebraic framework, enabling future research to design more interpretable and theoretically justified architectures. Moreover, the connection to Hopfield networks and associative memory capacity expansion highlights the broader implications for neural memory systems, potentially leading to more efficient and transparent AI models.

Technical Contribution

The paper introduces a rigorous spectral decomposition framework that maps the parameters of a single-layer linear Transformer to the inverse covariance matrix, enabling exact realization of the OLS estimator. It demonstrates that attention mechanisms, when configured appropriately, perform the inverse covariance operation, effectively executing statistical projection in one step. Additionally, the study uncovers a decoupled memory architecture—long-term (slow) memory encoded in learned weights and short-term (fast) memory via attention scores—providing a new understanding of the internal dynamics. This theoretical foundation paves the way for designing more interpretable, efficient, and statistically grounded neural architectures, extending classical associative memory models into deep learning contexts.

Novelty

This is the first work to rigorously establish the algebraic equivalence between a single-layer linear Transformer and the OLS solution, using spectral decomposition to explicitly construct parameters. Unlike prior views that treat Transformers as iterative approximators, this research reveals their inherent statistical operator nature, capable of executing classical estimators directly. The integration of spectral analysis with neural architecture design offers a novel perspective, bridging statistical inference and deep learning in a unified framework. The insights into memory mechanisms and the evolution toward exponential capacity via Softmax attention further distinguish this work as a foundational contribution.

Limitations

  • The current analysis is restricted to single-layer linear Transformers; extending to multi-layer nonlinear models remains an open challenge, limiting direct applicability to real-world architectures.
  • Sensitivity to data distribution shifts implies that performance may degrade under covariate shift, necessitating robustness enhancements.
  • Spectral decomposition becomes computationally expensive for high-dimensional data, posing scalability issues in large-scale applications.

Future Work

Future research should explore multi-layer nonlinear extensions, aiming to generalize the algebraic equivalence to more complex architectures. Developing scalable spectral algorithms and robustness techniques will be crucial for practical deployment. Additionally, integrating nonlinear activations and attention mechanisms with the spectral framework could unlock richer modeling capabilities, bridging the gap between theoretical insights and real-world AI systems.

AI Executive Summary

Transformers have revolutionized sequence modeling, yet their mathematical essence remains elusive. This study offers a groundbreaking perspective by demonstrating that a single-layer linear Transformer can exactly implement the classical statistical method—Ordinary Least Squares (OLS)—through spectral decomposition of the data covariance matrix. By carefully configuring the attention mechanism’s parameters, the authors show that the Transformer’s forward pass performs the OLS projection in a single step, eliminating the need for iterative optimization.

The core of this discovery lies in spectral analysis: decomposing the empirical covariance matrix into eigenvalues and eigenvectors, then constructing attention weights that invert this covariance. This algebraic construction aligns the Transformer’s linear attention with the inverse covariance operation fundamental to OLS. Empirical experiments on synthetic data confirm that, through gradient descent, the model parameters converge to the theoretical configuration, and the output matches the classical OLS solution with high precision.

Beyond this, the authors reveal a decoupled memory mechanism within the Transformer: the learned weights encode long-term statistical patterns (slow memory), while attention scores dynamically adapt to current inputs (fast memory). This insight explains the model’s ability to handle long sequences and adapt to distribution shifts, linking it to associative memory models like Hopfield networks. The evolution from linear prototypes to full nonlinear Transformers, especially with Softmax attention, is discussed as a transition from quadratic to exponential memory capacity.

This work fundamentally redefines the understanding of Transformer architectures, positioning them as statistical operators capable of executing classical inference tasks directly. It bridges deep learning and statistical theory, paving the way for more interpretable, efficient, and theoretically grounded models. Despite current limitations to linear, single-layer models, the framework sets a foundation for future extensions to complex, multi-layer architectures, promising a new era of explainable AI.

Deep Dive

Abstract

The statistical essence of the Transformer architecture has long remained elusive: Is it a universal approximator, or a neural network version of known computational algorithms? Through rigorous algebraic proof, we show that the latter better describes Transformer's basic nature: Ordinary Least Squares (OLS) is a special case of the single-layer Linear Transformer. Using the spectral decomposition of the empirical covariance matrix, we construct a specific parameter setting where the attention mechanism's forward pass becomes mathematically equivalent to the OLS closed-form projection. This means attention can solve the problem in one forward pass, not by iterating. Building upon this prototypical case, we further uncover a decoupled slow and fast memory mechanism within Transformers. Finally, the evolution from our established linear prototype to standard Transformers is discussed. This progression facilitates the transition of the Hopfield energy function from linear to exponential memory capacity, thereby establishing a clear continuity between modern deep architectures and classical statistical inference.

cs.LG cs.AI math.ST stat.ML