Optimal Attention Temperature Improves the Robustness of In-Context Learning under Distribution Shift in High Dimensions

Pretrained Transformers can perform in-context learning (ICL) from a few demonstrations, but this ability can fail sharply when the test distribution differs from pretraining—a common deployment setting. We study attention temperature as a simple inference-time control for improving ICL robustness under such shifts. In a high-dimensional linear-regression framework, we analyze a Transformer with "approximate softmax" attention, which preserves softmax's normalization and temperature-dependent selectivity while remaining tractable. We derive a closed-form expression for the ICL generalization error under distribution shift, and show that it is minimized by an explicit optimal attention temperature. This characterization yields interpretable guidance by linking the best temperature to moments of the pre-softmax attention scores, and predicts when temperature adjustment can recover near Bayes-optimal performance. We validate the theory with extensive simulations, and further demonstrate gains for pretrained LLMs on question-answering benchmarks under distribution shift induced by noisy in-context demonstrations. Overall, attention temperature emerges as a principled, lightweight knob for improving the robustness of ICL in pretrained Transformers.

Modern language models can often solve a new task by looking at a few examples in the prompt, without being retrained. But this ability can become unreliable when the examples are different from what the model is used to, for example, when they are noisier or come from a shifted data source. This paper studies a simple way to make this kind of learning more robust. We adjust the attention temperature, a setting inside the Transformer model that controls how strongly the model focuses on particular examples in the prompt. A lower or higher temperature changes whether the model concentrates sharply on a few examples or spreads its attention more broadly. We first analyze this idea in a simplified mathematical setting where the best choice can be derived explicitly. We then test the prediction through simulations and experiments with pretrained language models. The results show that choosing the attention temperature carefully can reduce errors under distribution shift. This suggests a lightweight way to make in-context learning more reliable without retraining the model.