L ong A lign: A Recipe for Long Context Alignment of Large Language Models (2024)

Abstract: Extending large language models to effectively handle long contexts requires instruction fine-tuning on input sequences of similar length. To address this, we present LongAlign—a recipe of the instruction data, training, and evaluation for long context alignment. First, we construct a long instruction-following dataset using Self-Instruct. To ensure the data diversity, it covers a broad range of tasks from various long context sources. Second, we adopt the packing and sorted batching strategies to speed up supervised fine-tuning on data with varied length distributions. Additionally, we develop a loss weighting method to balance the contribution to the loss across different sequences during packing training. Third, we introduce the LongBench-Chat benchmark for evaluating instruction-following capabilities on queries of 10k-100k in length. Experiments show that LongAlign outperforms existing recipes for LLMs in long context tasks by up to 30%, while also maintaining their proficiency in handling short, generic tasks.

## introduction large language models (llms) with large context windows facilitate tasks such as summarization, question answering on long text and code @xcite . importantly, they may form the foundational support for life-long conversations and complex agent scenarios @xcite . existing works to build long-context llms predominantly focus on context extension @xcite @xcite , that is, position encoding extension and continual training on long text. in this work, we instead focus on the perspective of long context alignment, i.e., instruction † corresponding authors figure 1 : test results on longbench-chat, which contains real-world queries of 10k-100k in length foot_0 . to address them, we present the longalign recipe, covering data, efficient training, and evaluation, respectively. data-wise, to construct a diverse long instruction-following dataset, we collect long sequences from nine sources and use self-instruct @xcite to generate 10k instruction data of 8k-64k length. training-wise, to address the inefficiency under uneven batching, we adopt the packing strat-egy @xcite that packs sequences together up to the maximum length before dispatching them to gpus. however, we identified a bias in loss averaging during this packing training, as packs containing different numbers of sequences are assigned equal weight in the final loss calculation. to mitigate this bias, we propose a loss weighting strategy to balance contributions to the loss across different sequences. in addition, we introduce sorted batching that groups sequences of similar lengths to reduce the intra-batch idle time. evaluation-wise, we develop longbench-chat, a benchmark compromising open-ended questions of 10k-100k length annotated by ph.d. students. it covers diverse aspects of instruction-following abilities such as reasoning, coding, summarization, and multilingual translation over long contexts. @xcite is employed to score the machine-generated responses based on our annotated groundtruths and few-shot scoring examples. extensive experiments show that longalign effectively aligns models to handle contexts of up to 64k tokens in length while maintaining their performance on general tasks without degradation. in addition, we have the following findings: ## longalign in this section, we discuss the methodology in lon-galign, involving the data construction process, training method, and evaluation benchmark. ## experiments in this section, we aim to answer the following research questions through a series of experiments: rq1. during sft, how does the quantity and diversity of the long instruction data influence the model's performance in downstream tasks. rq2. whether incorporating long instruction data during training affects the model's general capabilities and instruction-following / conversational abilities in short context scenarios. rq3. the impact that the packing and sorted batching training methods have on the training efficiency and the final performance of the models. we also incorporate discussions on the scalability of longalign on model size and context length, and the learning curve in long context alignment. ## conclusion this paper aims to find the best practice for long context alignment in the scope of data, training method, and evaluation. our proposed solution, namely longalign, uses self-instruct to construct diverse long instruction data, and efficiently finetune the model with packing combined with loss weighting or sorted batching. moreover, we introduce longbench-chat to facilitate reliable assessment of llm's instruction-following ability on practical long context interactions. through controlled experiments, we find that the amount, diversity of data, as well as the correct training method, are crucial to the final performance.