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

arXiv: https://arxiv.org/abs/2511.01292 · HTML: https://mlsanity.com/paper/2511.01292/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 39 references that are in the map.
- First posted on arXiv 3 Nov 2025.

## Builds on (most cited first, 30 of 33)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Emergent Abilities of Large Language Models](https://mlsanity.com/paper/2206.07682/) (2022)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [What Makes Good In-Context Examples for GPT-3?](https://mlsanity.com/paper/2101.06804/) (2021)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Crowdsourcing Multiple Choice Science Questions](https://mlsanity.com/paper/1707.06209/) (2017)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [YaRN: Efficient Context Window Extension of Large Language Models](https://mlsanity.com/paper/2309.00071/) (2023)
- [Are Emergent Abilities of Large Language Models a Mirage?](https://mlsanity.com/paper/2304.15004/) (2023)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [cosFormer: Rethinking Softmax in Attention](https://mlsanity.com/paper/2202.08791/) (2022)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [The Effect of Sampling Temperature on Problem Solving in Large Language Models](https://mlsanity.com/paper/2402.05201/) (2024)
- [Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression](https://mlsanity.com/paper/2306.15063/) (2023)
- [Vision Transformer for Small-Size Datasets](https://mlsanity.com/paper/2112.13492/) (2021)
- [How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?](https://mlsanity.com/paper/2310.08391/) (2023)
- [Competition Dynamics Shape Algorithmic Phases of In-Context Learning](https://mlsanity.com/paper/2412.01003/) (2024)
- [OpenICL: An Open-Source Framework for In-context Learning](https://mlsanity.com/paper/2303.02913/) (2023)
- [Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression](https://mlsanity.com/paper/2310.17086/) (2023)
- [Bridging the Divide: Reconsidering Softmax and Linear Attention](https://mlsanity.com/paper/2412.06590/) (2024)
- [Softmax is not Enough (for Sharp Size Generalisation)](https://mlsanity.com/paper/2410.01104/) (2024)
- [Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond](https://mlsanity.com/paper/2407.10005/) (2024)
- [On the Noise Robustness of In-Context Learning for Text Generation](https://mlsanity.com/paper/2405.17264/) (2024)

## Built on by (most cited first, 1 of 1)

- [Symmetry Reveals Layerwise Dynamics: How Transformers Perform In-Context Classification](https://mlsanity.com/paper/2604.11613/) (2026)

## Abstract (arXiv)

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 on pretrained LLMs (GPT-2 and Llama2-7B) 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.