Transformers Don't In-Context Learn Least Squares Regression

2025-07-12

AI Summary

This work investigates the mechanisms underlying Transformer failure in in-context learning (ICL), particularly on linear regression tasks under out-of-distribution (OOD) prompts. Contrary to prevailing hypotheses—such as implicit least-squares execution—the study finds that ICL does not rely on explicit algorithmic implementation but instead critically depends on spectral properties of the input distribution. Within-distribution, the top two singular vectors of the residual stream remain stable and strongly correlate with low prediction loss; under OOD inputs, however, this spectral structure collapses, precipitating sharp generalization degradation. Through systematic OOD generalization experiments and spectral analysis of residual stream representations, the paper establishes, for the first time, a causal link among input distribution, singular vector structure, and ICL performance. These findings provide a novel perspective on the intrinsic limitations and transferability of ICL, advancing our understanding of when and why contextual inference fails.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: Safety and Robustness
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

In-context learning (ICL) has emerged as a powerful capability of large pretrained transformers, enabling them to solve new tasks implicit in example input-output pairs without any gradient updates. Despite its practical success, the mechanisms underlying ICL remain largely mysterious. In this work we study synthetic linear regression to probe how transformers implement learning at inference time. Previous works have demonstrated that transformers match the performance of learning rules such as Ordinary Least Squares (OLS) regression or gradient descent and have suggested ICL is facilitated in transformers through the learned implementation of one of these techniques. In this work, we demonstrate through a suite of out-of-distribution generalization experiments that transformers trained for ICL fail to generalize after shifts in the prompt distribution, a behaviour that is inconsistent with the notion of transformers implementing algorithms such as OLS. Finally, we highlight the role of the pretraining corpus in shaping ICL behaviour through a spectral analysis of the learned representations in the residual stream. Inputs from the same distribution as the training data produce representations with a unique spectral signature: inputs from this distribution tend to have the same top two singular vectors. This spectral signature is not shared by out-of-distribution inputs, and a metric characterizing the presence of this signature is highly correlated with low loss.

Problem

• Study transformers' in-context learning mechanisms for linear regression
• Examine out-of-distribution generalization failures in ICL transformers
• Analyze pretraining corpus influence on ICL via spectral signatures

Innovation

• Transformers fail OLS generalization tests
• Spectral analysis reveals unique input signatures
• Pretraining corpus shapes ICL behavior