Learning Linear Regression with Low-Rank Tasks in-Context
High-dimensional analysis of linear attention models reveals prediction decomposition, implicit regularization, and phase transitions in low-rank tasks.
Key Findings
Methodology
This paper employs high-dimensional asymptotic analysis of a linear attention model trained on low-rank regression tasks. The prediction distribution is derived, revealing a decomposition into an algorithmic component and noise terms—memorization and structural noise. The model learns an efficient low-rank algorithm exploiting task structure, with finite-sample fluctuations inducing implicit regularization. The analysis uncovers a phase transition governed by task complexity, delineating regimes of under- and over-utilization of task structure, with implications for generalization and robustness.
Key Results
- The model’s prediction in the high-dimensional limit decomposes into a core algorithmic signal plus two independent noise components: memorization noise, which diminishes with task familiarity, and structural noise, which dominates under structural mismatch. Empirical results confirm that increasing the context length α reduces algorithmic error but can increase memorization noise, especially in out-of-distribution tasks. The phase transition occurs at a critical task rank ratio ρ, where the model shifts from under- to over-utilization of task structure, affecting generalization performance.
- Finite-sample fluctuations during pretraining induce an implicit regularization effect, stabilizing low-rank learning. The effective regularization strength scales inversely with the sample ratio α, preventing ill-posed solutions in the idealized infinite-data limit. Numerical experiments validate the theoretical predictions, showing that models trained with finite data outperform those trained with perfect information in certain regimes, highlighting the importance of statistical noise in learning stability.
Significance
This work advances the theoretical understanding of in-context learning by providing a mechanistic explanation of how transformers implicitly regularize and exploit task structure in high dimensions. It bridges the gap between empirical observations and formal theory, offering insights into the phase transition phenomena and the role of statistical fluctuations. The findings inform the design of more robust models capable of leveraging structural priors, with broad implications for multi-task learning, transfer learning, and the development of general-purpose AI systems.
Technical Contribution
The paper introduces a rigorous high-dimensional framework for analyzing linear attention models on structured tasks, deriving explicit formulas for prediction components. It demonstrates how statistical fluctuations induce an implicit regularization, preventing ill-posed solutions. The phase transition analysis quantifies the trade-off between task exploitation and robustness, providing a theoretical basis for observed empirical phenomena. The approach can be extended to more complex models and real-world tasks, offering a new lens for understanding deep learning mechanisms.
Novelty
This is the first comprehensive high-dimensional theoretical analysis connecting the prediction decomposition, implicit regularization, and phase transitions in transformer models trained on low-rank tasks. Unlike prior work focusing solely on performance metrics, this study elucidates the underlying mechanisms, revealing how statistical fluctuations serve as a natural regularizer and how task structure induces a phase transition. It bridges the gap between toy models and practical deep learning systems, offering a novel mechanistic perspective.
Limitations
- The analysis assumes a linear attention mechanism and high-dimensional asymptotics, which may not fully capture the complexities of deep, nonlinear transformers in real applications.
- Experiments are primarily on synthetic data with controlled low-rank structures; real-world data may exhibit more complex, non-linear dependencies.
- The low-rank assumption limits applicability to tasks with more intricate or high-rank structures; extending the theory to non-linear and deep models remains a challenge.
Future Work
Future research will extend the framework to multi-layer, nonlinear transformers, exploring how depth and non-linearity influence the mechanisms uncovered here. Investigating the interplay between explicit regularization techniques and implicit regularization from finite data could yield more robust training strategies. Additionally, applying this theory to real-world datasets will test its practical relevance and guide the development of models that better leverage task structure for improved generalization.
AI Executive Summary
Transformers have revolutionized AI, yet their internal mechanisms—especially how they learn and generalize across tasks—remain elusive. This study tackles this challenge by analyzing a simplified linear attention model trained on low-rank regression tasks in the high-dimensional limit. The core insight is that predictions decompose into an algorithmic signal plus two types of noise: memorization and structural. The model learns an efficient low-rank algorithm that exploits task structure, with finite-sample fluctuations during training inducing an implicit regularization effect that stabilizes learning.
A key discovery is the phase transition phenomenon: when the task's intrinsic rank exceeds a critical threshold, the model shifts from under- to over-utilizing task structure, impacting its ability to generalize. The analysis shows that statistical fluctuations in finite data act as a natural regularizer, preventing ill-posed solutions that appear in the idealized infinite-data limit. This explains why models trained with finite samples often outperform those with perfect information, highlighting the constructive role of data noise.
These findings deepen our understanding of in-context learning, revealing how models implicitly balance task exploitation and robustness. The theoretical framework offers a pathway to designing more stable and generalizable AI systems, with implications spanning multi-task learning, transfer learning, and beyond. Future work aims to extend these insights to deeper, nonlinear architectures and real-world datasets, promising a richer understanding of deep learning's inner workings.
Deep Dive
Abstract
In-context learning (ICL) is a key building block of modern large language models, yet its theoretical mechanisms remain poorly understood. It is particularly mysterious how ICL operates in real-world applications where tasks have a common structure. In this work, we address this problem by analyzing a linear attention model trained on low-rank regression tasks. Within this setting, we precisely characterize the distribution of predictions and the generalization error in the high-dimensional limit. Moreover, we find that statistical fluctuations in finite pre-training data induce an implicit regularization. Finally, we identify a sharp phase transition of the generalization error governed by task structure. These results provide a framework for understanding how transformers learn to learn the task structure.