In-Context Learning Operates as Concept Subspace Learning

TL;DR

This paper proposes a concept subspace framework for understanding in-context learning, showing that task information is concentrated in low-dimensional, task-aligned activation subspaces, enabling effective prediction and concept manipulation.

Contribution

It introduces a mechanistic, subspace-based view of in-context learning, demonstrating how structured demonstrations induce low-dimensional, task-specific concept subspaces in neural activations.

Findings

Prediction decomposes into concept-coordinate regression and off-subspace leakage.

A low-dimensional subspace (68-73 dimensions) captures most task information in Llama-3-8B.

Concept swaps redirect predictions, confirming the role of concept subspaces.

Abstract

Regression and Bayesian accounts of in-context learning (ICL) explain how demonstrations can induce predictors, while mechanistic analyses often identify compact activation directions that steer prompted behavior. However, it remains unclear whether structured demonstrations induce low-dimensional concept inference. We study this question through a concept-subspace view of ICL, in which tasks vary only along intrinsic concept coordinates, although inputs are observed in a high-dimensional ambient space. For ridge and least-squares ICL proxies, prediction decomposes exactly into concept-coordinate regression and off-subspace leakage. Under block-diagonal or near-block-diagonal covariance assumptions, the leading estimation and nuisance-sensitivity terms scale with the dimension of the concept subspace, while residual effects are controlled by cross-subspace coupling. This separation gives a mechanistic prediction: recoverable task information should concentrate in a low-dimensional, task-aligned activation subspace. On CounterFact-derived multi-relation prompts with Llama-3-8B, a 68--73-dimensional subspace of the 4096-dimensional residual stream restores 78.8% of the clean--corrupted accuracy gap, whereas patching the complementary subspace restores 0%. Concept swaps redirect predictions toward injected relations, while random and cross-task matched-rank controls are largely ineffective. Additional experiments on Qwen2.5-7B and a controlled cross-lingual rule task show the same qualitative pattern. These results support concept subspaces as compact, task-aligned mediators of recoverable ICL behavior in structured task families, without implying full-circuit recovery.