Optimal Attention Temperature Enhances In-Context Learning under Distribution Shift
This addresses robustness issues in real-world deployments of in-context learning for users of pretrained Transformers, offering a principled mechanism to enhance performance under distribution shift.
The paper tackles the problem of in-context learning performance degradation in pretrained Transformers under distribution shift, showing that an optimal attention temperature minimizes generalization error and improves robustness in experiments with models like GPT-2 and LLaMA2-7B.
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.