LongShortAlignment

This repository is the official implementation for ICML 2025 paper Long-Short Alignment for Effective Long-Context Modeling in LLMs.

Introduction

Large language models (LLMs) have exhibited impressive performance, yet their effectiveness remains limited by the fixed context window of the Transformer architecture. Among the challenges of long-context modeling, 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 (e.g., positional encodings, data structures) to the output distribution of the model. Specifically, we highlight the critical role of long-short alignment---the consistency of output distributions across sequences of varying lengths---and introduce a metric called Long-Short Misalignment to quantify this phenomenon. Building on these findings, we develop a regularization term that promotes long-short alignment during training.

Our experiments demonstrate that the long-short misalignment metric strongly correlates with long-context benchmark performance, significantly outperforming traditional training loss in predictive accuracy. Besides, incorporating the metric as a regularization term attains consistent improvements in a plug-and-play manner over standard cross-entropy fine-tuning.

Requirements

Python 3.10 + PyTorch 2.3 + Transformers 4.37

pip install -r requirements.txt

Reproduce the paper

Long-Short Alignment fine-tuning

To conduct long-context fine-tuning with our proposed regularization term, run accelerate config and enable DeepSpeed acceleration. deepspeed/zero3.json is the configuration file used for training.

cd finetune
sh train.sh

The training data is randomly sampled from RedPajama-Data-1T-Sample inside finetune.py. We use Llama-2-7B as the base model and CLEX as the RoPE adjustment method. To run models with EABF, please downgrade the version of transformers to 4.37.0.

Evaluation on Long-context Benchmark

In the paper, we evaluate models on LongBench and BABILong. Please refer to the respective code repositories.

Citation

If you use our code, please cite

@inproceedings{du2025long,
      title={Long-Short Alignment for Effective Long-Context Modeling in LLMs},
      author={Tianqi Du and Haotian Huang and Yifei Wang and Yisen Wang},
      booktitle={International Conference on Machine Learning (ICML)},
      year={2025}
}