Optimal Attention Temperature Enhances In-Context Learning under Distribution Shift

• URL: http://arxiv.org/abs/2511.01292v1
• Date: Mon, 03 Nov 2025 07:18:27 GMT
• Title: Optimal Attention Temperature Enhances In-Context Learning under Distribution Shift
• Abstract summary: This paper provides the first theoretical and empirical study of attention temperature for ICL under distribution shift.<n>We prove that shifts in input covariance or label noise substantially impair ICL, but that an optimal attention temperature exists which minimizes this error.<n>Our results establish attention temperature as a principled and powerful mechanism for improving the robustness of ICL in pretrained Transformers.
• Score: 8.135786025034626
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Pretrained Transformers excel at in-context learning (ICL), inferring new tasks from only a handful of examples. Yet, their ICL performance can degrade sharply under distribution shift between pretraining and test data, a regime increasingly common in real-world deployments. While recent empirical work hints that adjusting the attention temperature in the softmax can enhance Transformer performance, the attention temperature's role in ICL under distribution shift remains unexplored. This paper provides the first theoretical and empirical study of attention temperature for ICL under distribution shift. Using a simplified but expressive "linearized softmax" framework, we derive closed-form generalization error expressions and prove that shifts in input covariance or label noise substantially impair ICL, but that an optimal attention temperature exists which minimizes this error. We then validate our predictions through extensive simulations on linear regression tasks and large-scale experiments with GPT-2 and LLaMA2-7B on question-answering benchmarks. Our results establish attention temperature as a principled and powerful mechanism for improving the robustness of ICL in pretrained Transformers, advancing theoretical understanding and providing actionable guidance for selecting attention temperature in practice.

Related papers

• Beyond Post-Hoc Temperature Scaling: Bilevel Optimization for LLM Calibration [8.906593271577266]
  Preference alignment often makes large language models (LLMs) overconfident and poorly calibrated.<n>We propose maximizing the entropy of predictive distributions as the calibration objective.<n>Our method yields well-calibrated LLMs with particular advantages in out-of-domain generalization.
  arXiv  Detail & Related papers  (2026-08-07T17:05:10Z)
• Mitigating Error Amplification in Fast Adversarial Training [58.74042726356826]
  Fast Adversarial Training (FAT) has proven effective in enhancing model robustness by encouraging networks to learn perturbation-invariant representations.<n>FAT often suffers from catastrophic overfitting (CO), where the model overfits to the training attack and fails to generalize to unseen ones.<n>We propose a Distribution-aware Dynamic Guidance (DDG) strategy that dynamically adjusts both the perturbation budget and supervision signal.
  arXiv  Detail & Related papers  (2026-04-27T11:23:18Z)
• Learning to Adapt: In-Context Learning Beyond Stationarity [23.324726233034614]
  Transformer models have become foundational across a wide range of scientific and engineering domains.<n>Key capability underlying their success is in-context learning (ICL)<n>Recent theoretical efforts have begun to uncover the mechanisms behind this phenomenon.
  arXiv  Detail & Related papers  (2026-04-13T03:37:08Z)
• Understanding the Generalization of Bilevel Programming in Hyperparameter Optimization: A Tale of Bias-Variance Decomposition [53.68517860700599]
  We propose an ensemble hypergradient strategy to reduce the variance in HPO algorithms effectively.<n>We establish a connection between excess error and hypergradient estimation, offering some understanding of empirical observations.
  arXiv  Detail & Related papers  (2026-02-20T02:52:13Z)
• Data Distribution as a Lever for Guiding Optimizers Toward Superior Generalization in LLMs [60.68927774057402]
  We show, for the first time, that a lower simplicity bias induces a better generalization.<n>Motivated by this insight, we demonstrate that the training data distribution by upsampling or augmenting examples learned later in training similarly reduces SB and leads to improved generalization.<n>Our strategy improves the performance of multiple language models including Phi2-2.7B, Llama3.2-1B, Gemma3-1B-PT, Qwen3-0.6B-Base-achieving relative accuracy gains up to 18% when fine-tuned with AdamW and Muon.
  arXiv  Detail & Related papers  (2026-01-31T07:40:36Z)
• Unlabeled Data Can Provably Enhance In-Context Learning of Transformers [6.129841115228296]
  Large language models (LLMs) exhibit impressive in-context learning (ICL) capabilities, yet the quality of their predictions is fundamentally limited.<n>We propose a novel augmented ICL framework, in which the prompt includes a small set of labeled examples alongside a block of unlabeled inputs.<n>This enables the transformer to implicitly extract useful information from both labeled and unlabeled data, leading to provable improvements in ICL accuracy.
  arXiv  Detail & Related papers  (2026-01-15T04:23:32Z)
• From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers [67.02076505996284]
  We study how the choice of pretraining data distribution steers a shallow transformer toward one behavior or the other.<n>Our results shed light on the algorithmic biases of pretrained transformers and offer conceptual guidelines for data-driven control of their learned behaviors.
  arXiv  Detail & Related papers  (2025-12-21T08:10:26Z)
• CurES: From Gradient Analysis to Efficient Curriculum Learning for Reasoning LLMs [53.749193998004166]
  Curriculum learning plays a crucial role in enhancing the training efficiency of large language models.<n>We propose CurES, an efficient training method that accelerates convergence and employs Bayesian posterior estimation to minimize computational overhead.
  arXiv  Detail & Related papers  (2025-10-01T15:41:27Z)
• Exploring the Impact of Temperature Scaling in Softmax for Classification and Adversarial Robustness [8.934328206473456]
  This study delves into the often-overlooked parameter within the softmax function, known as "temperature"<n>Our empirical studies, adopting convolutional neural networks and transformers, reveal that moderate temperatures generally introduce better overall performance.<n>For the first time, we discover a surprising benefit of elevated temperatures: enhanced model robustness against common corruption, natural perturbation, and non-targeted adversarial attacks like Projected Gradient Descent.
  arXiv  Detail & Related papers  (2025-02-28T00:07:45Z)
• Technical Debt in In-Context Learning: Diminishing Efficiency in Long Context [24.905102026459428]
  Transformers have demonstrated remarkable in-context learning capabilities, adapting to new tasks by simply conditioning on demonstrations without parameter updates.<n>It remains unclear to what extent the transformers optimally learn in-context compared to principled learning algorithms.
  arXiv  Detail & Related papers  (2025-02-07T00:26:45Z)
• Non-asymptotic Convergence of Training Transformers for Next-token Prediction [48.9399496805422]
  Transformers have achieved extraordinary success in modern machine learning due to their excellent ability to handle sequential data. This paper provides a fine-grained non-asymptotic analysis of the training dynamics of a one-layer transformer. We show that the trained transformer presents non-token prediction ability with dataset shift.
  arXiv  Detail & Related papers  (2024-09-25T20:22:06Z)
• How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? [82.51626700527837]
  Transformer-based large language models displayed impressive in-context learning capabilities, where a pre-trained model can handle new tasks without fine-tuning. We analyze how the mechanics of how Transformer to achieve ICL contribute to the technical challenges of the training problems in Transformers.
  arXiv  Detail & Related papers  (2024-02-23T21:07:20Z)
• On the Limitations of Temperature Scaling for Distributions with Overlaps [8.486166869140929]
  We show that for empirical risk minimizers for a general set of distributions, the performance of temperature scaling degrades with the amount of overlap between classes. We prove that optimizing a modified form of the empirical risk induced by the Mixup data augmentation technique can in fact lead to reasonably good calibration performance.
  arXiv  Detail & Related papers  (2023-06-01T14:35:28Z)
• How Can Increased Randomness in Stochastic Gradient Descent Improve Generalization? [0.0]
  We study the role of the SGD learning rate and batch size in generalization. We show that increasing SGD temperature encourages the selection of local minima with lower curvature.
  arXiv  Detail & Related papers  (2021-08-21T13:18:49Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.