Information Entropy Invariance: Enhancing Length Extrapolation in Attention Mechanisms

TL;DR

This paper introduces entropy-based scaling methods to improve length extrapolation in attention mechanisms, enabling models to handle much longer contexts effectively without additional training.

Contribution

It proposes two novel entropy-invariance based scaling techniques, InfoScale and CosScale, with theoretical analysis and state-of-the-art experimental results on extended context lengths.

Findings

InfoScale maintains focus on original tokens during length extension.

Combining InfoScale and CosScale achieves state-of-the-art performance on long-context tasks.

Increasing CosScale approximates windowed attention, addressing long-range context challenges.

Abstract

Since the emergence of research on improving the length extrapolation capabilities of large language models in 2021, some studies have made modifications to the scaling factor in the scaled dot-product attention mechanism as part of their proposed methods without rigorous theoretical justifications. To fill this gap, we propose two new scaled temperatures based on information entropy invariance to enhance length extrapolation. First, a training-free method InfoScale is designed for dotproduct attention, and preserves focus on original tokens during length extrapolation by ensuring consistent entropy. Second, we theoretically analyze the impact of scaling (CosScale) on cosine attention. Experimental data demonstrates that combining InfoScale and CosScale achieves state-ofthe-art performance on the GAU-{\alpha} model with a context window extended to 64 times the training length, and outperforms seven existing methods. Our analysis reveals that significantly increasing CosScale approximates the Windowed Attention, and highlights the significance of attention score dilution as a key challenge in long-range context handling. The code and data are available at https://github.com/HT-NEKO/ Information-Entropy-Invariance.