In-context classification: Calibration and noise in linear attention with softmax readout
Authors: Arman Rysmakhanov, Mary Letey, Yue M. Lu, Jacob Zavatone-Veth, Cengiz Pehlevan
Organizations: Applied Mathematics, Harvard University · Williams College · Kempner Institute, Harvard University · Center for Brain Science and Society of Fellows, Harvard University
We study in-context binary classification in a simple attention model trained to minimise a logistic loss. Each prompt has a latent logistic teacher, and the learner must infer its direction from Bernoulli-labeled context examples. We derive an asymptotic prediction for the test logit in a high-dimensional regime with context length proportional to token dimension, and total training prompts proportional to dimension squared. Our formula cleanly separates calibration, test-time context noise, and non-isotropic parameter-estimation noise as interpretable effects on in-context performance. Our analysis also identifies the ridgeless separability regime and addresses the effects of context length, pretraining-set size, regularisation, and sampling temperature on in-context classification performance. The theory's predictions closely agree with finite-dimensional simulations and fully trained linear-attention models, while experiments with standard softmax attention show qualitatively related behaviour beyond the regime assumed by our analysis.
Figures & tables
Figure 1: The distribution of empirical htest⊤γ∗ matches our derived distribution R . The ridgeless theory predicts recovery of the true logit distribution as αtrain,τ→∞ . Increasing τ (right-hand plots) reduces variance, but alignment with Gtrue requires both variance reduction and a1Cβ/β→1. For all, β=1,λ=10−6,d=32,αtest=αtrain.
Figure 2: Empirical simulation refers to directly sampling n sequences of length ℓ , computing γ , and evaluating this parameter matrix on freshly sampled sequences. This is compared to the theoretical prediction in Result 3.3 . The critical ratio is given by eq. ( 16 ). Here d=64 , αtrain=4=αtest , β=1 . In (a), the empirical simulations take λ=10−8 .
Figure 3: Simulation versus theory for classification error. We can compare to random guessing error 1/2, and Bayes error 1−E[σ(β\absolutevalueGq)] . For short contexts, the model effectively guesses, but nears the Bayes floor as context length increases. d=32,αtrain=αtest,τ=10,β=1,λ=10−8 .
Figure 4: Effects of data parameters αtrain and τ on anisotropy (a) and calibration (b),(c). We compute a1,a2 (and correspondingly, m,v ) using Result 3.2 . In (b),(c), R is taken simply as a real-valued variable and the function p(R) is computed by numerically taking the expectation in ( 22 ). In all plots we take αtest=108 (effectively infinite) to isolate the effects of αtrain and τ ; all labels α refer to αtrain . For all, λ=0 and β=1 . For (b) αtrain=4 . For (c) τ=20 .
Figure 5: Heatmap of critical τc against α and β .
Figure 6: Similarly to Figure 1 , we compare the theory to full-parameter linear attention, namely, we train the full K,Q,V matrices. Here we plot a quantile-quantile comparison of the model’s output logit on new prompt matrices Z against the theory prediction R and the true latent βGq . In all plots, d=16,λ=10−6,β=1 . Regularisation is applied to ∥K∥F2+∥Q∥F2+∥V∥F2 .
Figure 7: An analogous comparison to Figure 1 , but the logits here are from a trained transformer-style softmax attention block. In particular, we use standard scaled dot-product softmax attention with fully trainable query Q , key K , value V , and output O weights. We initialise these at identity to improve training dynamics. The model is trained for 100 epochs with AdamW with weight decay 0.1 and early stopping. We set d=32 and β=4 .
Linear attention is a tractable model for understanding the mechanisms governing in-context learning in transformers. For linear regression tasks, recent asymptotic analyses have characterised its learning and generalisation behaviour. We extend this theory to nonlinear single-index targets, y=f(x⊤w)+ε. Our main result establishes a nonlinearity-noise equivalence: linear attention extracts only the linear Hermite component of f, while the remaining nonlinear structure contributes to the generalisation error as effective noise. This reduction allows results from the corresponding linear theory to be transferred to nonlinear tasks. We illustrate its implications for finite pretraining data and for the transition from task memorisation to task generalisation as task diversity increases. These results identify a limitation of the reduced linear-attention model and provide a tractable starting point for studying nonlinear in-context learning.
Mary Letey, Arman Rysmakhanov, Yue M. Lu +2
Applied Mathematics, Harvard University Kempner Institute, Harvard University · Williams College · Center for Brain Science, Harvard University Society of Fellows, Harvard University
Transformer-based large language models exhibit in-context learning, enabling adaptation to downstream tasks via few-shot prompting with demonstrations. In practice, such models are often fine-tuned to improve zero-shot performance on downstream tasks, allowing them to solve tasks without examples and thereby reducing inference costs. However, fine-tuning can degrade in-context learning, limiting the performance of fine-tuned models on tasks not seen during fine-tuning. Using linear attention models, we provide a theoretical analysis that characterizes how fine-tuning objectives modify attention parameters and identifies conditions under which this leads to degraded few-shot performance. We show that fine-tuning all attention parameters can harm in-context learning, whereas restricting updates to the value matrix improves zero-shot performance while preserving in-context learning. We further show that incorporating an auxiliary few-shot loss enhances in-context learning primarily on the target task, at the expense of degraded in-context learning ability on tasks not seen during fine-tuning. We provide empirical evidence from synthetic and real-world datasets consistent with the qualitative predictions of our theory.
Chungpa Lee, Jy-yong Sohn, Kangwook Lee
Yonsei University, Korea · University of Wisconsin–Madison, USA · KRAFTON, Korea +1
Recent work has demonstrated that transformers and linear attention models can perform in-context learning (ICL) on simple function classes, such as linear regression. In this paper, we empirically study how these two attention mechanisms differ in their ICL behavior on the canonical linear-regression task of Garg et al. We evaluate learning quality (MSE), convergence, and generalization behavior of each architecture. We also analyze how increasing model depth affects ICL performance. Our results illustrate both the similarities and limitations of linear attention relative to quadratic attention in this setting.