Long-Short Alignment for Effective Long-Context Modeling in LLMs

TL;DR

This paper introduces a novel approach to improve long-context modeling in large language models by focusing on aligning output distributions across different sequence lengths, leading to better length generalization.

Contribution

It proposes the concept of long-short alignment, introduces a metric called Long-Short Misalignment, and develops a regularization method to enhance length generalization in LLMs.

Findings

Long-Short Misalignment correlates with length generalization performance.

Regularization promoting long-short alignment improves model performance on long sequences.

The approach offers a new perspective beyond positional encodings for long-context modeling.

Abstract

Large language models (LLMs) have exhibited impressive performance and surprising emergent properties. However, their effectiveness remains limited by the fixed context window of the transformer architecture, posing challenges for long-context modeling. Among these challenges, length generalization -- the ability to generalize to sequences longer than those seen during training -- is a classical and fundamental problem. In this work, we propose a fresh perspective on length generalization, shifting the focus from the conventional emphasis on input features such as positional encodings or data structures to the output distribution of the model. Specifically, through case studies on synthetic tasks, we highlight the critical role of \textbf{long-short alignment} -- the consistency of output distributions across sequences of varying lengths. Extending this insight to natural language tasks, we propose a metric called Long-Short Misalignment to quantify this phenomenon, uncovering a strong correlation between the metric and length generalization performance. Building on these findings, we develop a regularization term that promotes long-short alignment during training. Extensive experiments validate the effectiveness of our approach, offering new insights for achieving more effective long-context modeling in LLMs. Code is available at https://github.com/PKU-ML/LongShortAlignment.