In-context classification: Calibration and noise in linear attention with softmax readout
Arman Rysmakhanov, Mary Letey, Yue M. Lu, Jacob Zavatone-Veth, Cengiz Pehlevan
Abstract
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.