Scale-invariant Attention

TL;DR

This paper introduces a scale-invariant attention mechanism for large language models that improves zero-shot generalization from short to long contexts and enhances long-context retrieval performance.

Contribution

It proposes two scale-invariance conditions for attention mechanisms and demonstrates a simple transformation that achieves these, improving long-context generalization.

Findings

Significant reduction in validation loss when generalizing to longer contexts

Effective at long-context retrieval tasks

Improves zero-shot transfer from short to long contexts

Abstract

One persistent challenge in LLM research is the development of attention mechanisms that are able to generalise from training on shorter contexts to inference on longer contexts. We propose two conditions that we expect all effective long context attention mechanisms to have: scale-invariant total attention, and scale-invariant attention sparsity. Under a Gaussian assumption, we show that a simple position-dependent transformation of the attention logits is sufficient for these conditions to hold. Experimentally we find that the resulting scale-invariant attention scheme gives considerable benefits in terms of validation loss when zero-shot generalising from training on short contexts to validation on longer contexts, and is effective at long-context retrieval.