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 .
Applied Mathematics, Harvard University Kempner Institute, Harvard University · Williams College · Center for Brain Science, Harvard University Society of Fellows, Harvard University