How do Transformers perform In-Context Autoregressive Learning?

TL;DR

By analyzing a linear-attention Transformer trained on an autoregressive process $s_{t+1}=Ws_t$, the paper investigates how context enables in-context autoregressive learning (ICL). It shows that, under token augmentation and commuting context matrices, the in-context mapping $\Gamma_{\theta^*}$ implements a one-step gradient update on the inner objective, while in non-augmented settings, the mapping becomes a geometry-driven operation enabled by orthogonal heads and learnable positional encoding. The authors derive exact minimizers for unitary and orthogonal contexts, reveal head-orthogonality, and study the role of softmax versus linear attention, validating the theory with experiments that extend to non-commuting $W$. These insights shed light on the mechanisms behind Transformers' context-sensitive computation and offer theoretical grounding for in-context learning in linear attention architectures.

Abstract

Transformers have achieved state-of-the-art performance in language modeling tasks. However, the reasons behind their tremendous success are still unclear. In this paper, towards a better understanding, we train a Transformer model on a simple next token prediction task, where sequences are generated as a first-order autoregressive process $s_{t+1} = W s_t$. We show how a trained Transformer predicts the next token by first learning $W$ in-context, then applying a prediction mapping. We call the resulting procedure in-context autoregressive learning. More precisely, focusing on commuting orthogonal matrices $W$, we first show that a trained one-layer linear Transformer implements one step of gradient descent for the minimization of an inner objective function, when considering augmented tokens. When the tokens are not augmented, we characterize the global minima of a one-layer diagonal linear multi-head Transformer. Importantly, we exhibit orthogonality between heads and show that positional encoding captures trigonometric relations in the data. On the experimental side, we consider the general case of non-commuting orthogonal matrices and generalize our theoretical findings.

Figures (10)

• Figure 1: Illustration of the autoregressive process in $\mathbb{R}^3$. Dots and crosses correspond to two different orthogonal matrices $W$.
• Figure 2: Trigonometric formula implemented by the Transformer in-context. The minima of the training loss correspond to implementing, up to multiplying factors: $2 \cos{\theta}R_{\theta} - I_2 = R_{2\theta}.$
• Figure 3: Matrices $\mathtt{A}$, $\mathtt{B}$, $\mathtt{B}^\top \mathtt{A}$ and $P$ after training model \ref{['eq:att']} on loss \ref{['eq:optim']} with random initialization. We take $d=10$ and $T=15$. Left: Unitary context case with $H=10.$Right: Orthogonal context case, with $H=8 < d$, which leads to low rank $\mathtt{B}^\top \mathtt{A}$. In both cases, we obtain arbitrarily small final loss. We recover parameters corresponding to our Propositions \ref{['prop:optimality_unitary']} and \ref{['prop:optimality_ortho']}.
• Figure 4: Histograms of the mean squared errors (MSE) when fitting an AR process to sequences in $D$ (original, in blue) or $D_{\mathrm{shuffle}}$ (shuffled, in orange). We only display MSEs bigger than a threshold of $10^{-12}$.
• Figure 5: Evolution of the MSE with depth $L$. We compare with $L$ steps of gradient descent on the inner loss \ref{['eq:inner']}. At initialization, the MSE is between $1$ and $2$.

Theorems & Definitions (26)

• Remark 1
• Definition 1: In-context autoregressive learning
• Remark 2
• Lemma 1
• Proposition 1: vonoswald2023uncovering
• Proposition 2: In-context autoregressive learning with gradient-descent
• Lemma 2
• Proposition 3: Unitary optimal in-context mapping
• Proposition 4: Quadratic loss
• Proposition 5