PaperlayerView exact revision on arXiv

Source-linked AI summary

LongAlign: A Recipe for Long Context Alignment of Large Language Models

Yushi Bai, Xin Lv, Jiajie Zhang, Yuze He, Ji Qi, Lei Hou, Jie Tang, Yuxiao Dong, Juanzi Li

arXiv:2401.18058v1cs.CLcs.LG

TL;DR

Long-context instruction tuning faces gaps in suitable training data, efficient training, and evaluation of real-world instruction following. LongAlign combines diverse long-instruction data, efficient training strategies, and LongBench-Chat; packing and sorted batching accelerate training by over 100% without performance compromise, while loss weighting improves long-context performance by 10%.

  • Problem

    Long-context alignment lacks long instruction-following data and construction methods, efficient training for varied sequence lengths, and robust evaluation on real-world queries.

    Source: §1 · p. 1, §3.4 · p. 4

  • Method

    LongAlign constructs diverse long-instruction data with Self-Instruct, uses packing and sorted batching with loss weighting for training, and introduces LongBench-Chat for evaluation.

    Source: §1 · pp. 1–2

  • Results

    Over 100% faster training without performance compromise; loss weighting improves long-context performance by 10%.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Controlled experiments indicate that data quantity, data diversity, and the training method are crucial to final performance.

    Source: §5 · p. 9

    The constructed data primarily covers long-context QA, summarization, and reasoning; available data for multi-turn dialogue and long-history agent tasks is challenging to collect.

    Source: E Limitations · p. 19

Abstract

from arXiv · show

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. The code, data, and long-aligned models are open-sourced at https://github.com/THUDM/LongAlign.

1 Introduction

Long-context alignment requires instruction data, efficient training, and evaluation suited to long user prompts. LongAlign addresses these needs with diverse long-instruction data, training strategies, and LongBench-Chat.

1 Introduction

  • Self-Instruct generates 10k instructions of 8k–64k length from long sequences collected across nine sources.Source: §1 · p. 1
  • Packing combines sequences up to the maximum length before GPU dispatch, while sorted batching groups similar-length sequences to reduce idle time.Source: §1 · pp. 1–2
  • LongBench-Chat evaluates instruction following on open-ended queries of 10k–100k length across reasoning, coding, summarization, and multilingual translation.Source: §1 · p. 2
  • LongAlign aligns models to contexts up to 64k tokens while maintaining general-task performance without degradation.Source: §1 · p. 2
  • Long-instruction data quantity and diversity affect final performance by up to 30%.Source: §1 · p. 2
  • More long-instruction data improves long-context performance without harming short-context handling.Source: §1 · p. 2
  • Over 100% faster training without performance compromise; loss weighting improves long-context performance by 10%.Source: §1 · p. 2

2 Related Work

Prior work extends context windows through training or inference-time techniques, while instruction tuning has largely focused on short contexts. LongAlign frames long-context alignment as a distinct challenge spanning data, training, and evaluation.

2 Related Work

  • Long-context scaling uses fine-tuning or continual training on longer sequences, or inference-time methods such as sliding-window attention and neighboring-token compression.Source: §2 · p. 2
  • Instruction tuning is established for short contexts, but long sequences create distinct challenges in alignment data, training methods, and evaluation.Source: §2 · p. 2
  • Xiong et al. (2023) generate long instructions by concatenating short ones, but their data and model are not open-sourced; LongAlpaca-12k is available but lacks in-depth comparisons of data and training methods.Source: §2 · p. 2
  • LongAlign targets full-parameter supervised fine-tuning with full attention by tuning data, training methods, and evaluation across tasks.Source: §2 · p. 2

3 LongAlign

LongAlign covers instruction-data construction, efficient long-context training, and evaluation. Its methods address varied sequence lengths while its benchmark targets long-context instruction following.

3.1 Preliminary

  • Supervised fine-tuning concatenates instruction and response, but computes loss only over response tokens to maximize the response probability conditioned on the instruction.Source: §3.1 · pp. 2–3

3.2 Dataset Construction

  • The dataset uses Claude 2.1 to create tasks and answers for 10k long texts from nine varied sources, producing 10k supervised instances ranging from 8k to 64k tokens.Ten percent of the data is in Chinese; prompts cover summaries, information extraction, reasoning, and other task types.Source: §3.2 · p. 3

3.3 Efficient Long-Context Training

  • Packing combines variable-length sequences into near-maximum-length inputs, reducing batch idle time while block-diagonal attention prevents cross-sequence contamination.The approach uses FlashAttention 2 and requires less computation and IO time than a traditional 2D attention mask.Source: §3.3 · p. 3
  • Equal weighting of packed batches biases loss toward longer sequences and those with more target tokens; scaling sequence losses by K/(NiM) balances their contributions.The data’s long-tailed length distribution motivates packing or sorted batching to reduce idle time; sorted batching accelerates training without a noticeable performance impact, and loss weighting improves downstream tasks by 10%.Source: §3.3 · pp. 3–4

3.4 LongBench-Chat

  • LongBench-Chat evaluates 50 real-world queries of 10k–100k length across document QA, summarization, and coding, with 40 English and 10 Chinese tasks.The benchmark avoids likely memorized texts and questions answerable without reading the long text; GPT-4 scores responses using short queries, reference answers, and few-shot examples.Source: §3.4 · p. 4
  • GPT-4 few-shot scoring aligned with and surpassed human annotator agreement, while models with context lengths of 32k or less tended to underperform on LongBench-Chat.GPT-4 few-shot average scores differed from human experts by 0.1 or less.Source: §3.4 · p. 5

4 Experiments

The experiments examine how long-instruction data and training choices affect long- and short-context performance. They report that adding long-instruction data improves long-task performance without compromising short-task performance.

4 Experiments

  • The experiments probe data quantity and diversity, effects on general and short-context abilities, and the efficiency and performance of packing and sorted batching.Source: §4 · p. 5

4.1 Experimental Setup

  • LongAlign-0k, LongAlign-5k, and LongAlign-10k vary the amount of LongAlign data, while LongAlpaca-12k provides a less diverse comparison set based on papers and books.Source: §4.1 · p. 5
  • Evaluation combines LongBench-Chat and LongBench for long-context tasks with MT-Bench and Open LLM Leaderboard tasks for short-context and general capabilities.Source: §4.1 · p. 6

4.2 Influence of Data

  • LongAlign-0k performed worse on long tasks, while increasing long-instruction data improved long-task performance without compromising short-task performance.The authors also conclude that context extension alone was insufficient for strong downstream long-task performance.Source: §4.2 · p. 7

2. Diversity of long instruction data is beneficial for the model’s instruction-following abil-

The reported findings address data diversity, training efficiency, scaling, learning trends, and out-of-distribution generalization. LongAlign's outcomes vary across training methods and model scales, with gains reported on long-context tasks.

2. Diversity of long instruction data is beneficial for the model’s instruction-following abil-

  • LongAlign-10k performs better than LongAlpaca-12k on LongBench-Chat and MTBench, while LongAlpaca-12k slightly leads on LongBench.The authors attribute LongAlpaca-12k's LongBench advantage primarily to 2WikiMQA and NarrativeQA, whose sources resemble its academic-paper and book data.Source: §2 · p. 7

4.3 Impact of Training Methods

  • Packing and sorted batching each require less than half the training time of naïve batching while producing comparable long- and short-task performance.Source: Figure 5 · p. 7, §4.3 · p. 7

4.4 Discussion

  • LongAlign-13B-64k improves LongBench-Chat by 10% over the 7B-scale model, setting a new record among open-sourced models.Source: Table 4 · p. 8, §4.4 · p. 8

4.3 Impact of Training Methods

  • Loss weighting improves LongBench-Chat performance by about 5%–10% in packing training, with minimal and variable effects on other tasks.Source: §4.3 · p. 8

4.4 Discussion

  • Human-annotated data up to 128k tokens is used to align ChatGLM3-6B at a 128k context window.Source: §4.4 · p. 8
  • Long- and short-task learning curves both improve rapidly from 0–500 steps, rise more slowly afterward, and stabilize after 1,000 steps.Source: Figure 6 · p. 8, §4.4 · p. 8
  • LongAlign-trained models generalize to out-of-distribution long-context queries, and larger models show stronger generalization.The examples include writing a review of a research paper.Source: §4.4 · p. 8

5 Conclusion

LongAlign combines long-instruction data construction, efficient fine-tuning, and evaluation for practical long-context interactions. Controlled experiments identify data amount, data diversity, and training method as important to final performance.

  • LongAlign uses Self-Instruct to construct diverse long-instruction data and fine-tunes models with packing plus loss weighting or sorted batching.Source: §5 · p. 8
  • LongBench-Chat is introduced to assess LLM instruction-following ability in practical long-context interactions.Source: §5 · p. 8
  • Controlled experiments find that data amount, data diversity, and training method are crucial to final performance.Source: §5 · p. 9

A Dataset Construction Details

The dataset is built from documents sampled from multiple sources and used to generate long-instruction examples. Claude produces questions and answers from those documents, with translated prompts used for Chinese texts.

B Training Method Details

The packing implementation keeps attention within each original sequence, while its loss weighting normalizes contributions across sequences in a pack.

B Training Method Details

  • A one-dimensional start-index mask lets FlashAttention compute attention only within each packed sequence.Adjacent mask indices delimit each sequence; the first index is 0 and the last is batch_size × seq_len.Source: B Training Method Details · p. 13
  • Target-token positions receive weight 1/N, where N is the number of target tokens in that sequence.Other positions in the weighted one-dimensional mask receive weight 0.Source: B Training Method Details · p. 13
  • The batch loss sums token-level cross-entropy scaled by K/MN, with M sequences and K packs set according to the batch configuration.Source: B Training Method Details · p. 13

C.1 LongBench-Chat

LongBench-Chat combines long-context questions with expert-verified reference answers and GPT-4 scoring; the section also describes LongBench evaluation costs and the Needle Test setup.

C.1 LongBench-Chat

  • LongBench-Chat combines 30 author-proposed questions with 20 LooGLE long-dependency questions, using expert-written answers verified by at least two experts.The author-proposed set contains 20 English and 10 Chinese questions; the 50 questions cover comprehension and reasoning, information retrieval, timeline reorder, and computation.Source: §C.1 · p. 13
  • The evaluation asks an impartial judge to compare answers with reference answers, assess correctness, helpfulness, accuracy, and relevance, and assign an overall integer rating from 1 to 10.Two human experts each scored 200 responses from a pool of 300 on a 1-to-10 scale using the instruction and referenced answer.Source: §C.1 · pp. 13–14

C.2 LongBench

  • A LongBench evaluation run costs about $24 using GPT-4, averaging approximately 800,000 tokens across 12 datasets.Source: §C.2 · p. 15

C.3 Needle Test

  • The Needle Test places a specific San Francisco recommendation sentence in Paul Graham’s essays and asks for it with a prompt based on Claude 2.17.Source: §C.3 · p. 15

D Case Studies on OOD Queries

Case studies probe out-of-distribution queries using the paper as long input and report model responses to review and rap requests.

D Case Studies on OOD Queries

  • The case study used the paper as long input, omitted its Abstract, Conclusion, and Limitations, and posed review and rap requests absent from long-context SFT data.Source: D Case Studies on OOD Queries · p. 15
  • LongAlign-13B-64k produced what the authors describe as a relatively good paper review, with only a few flaws, on an unseen query type.The authors present this as evidence of OOD generalization to a review request absent from long-context SFT data.Source: D Case Studies on OOD Queries · pp. 15, 17
  • Only LongAlign-13B-64k followed the request to summarize the paper in rap, and its response did not rhyme.Source: D Case Studies on OOD Queries · p. 18

E Limitations

The authors identify gaps in both the range of long-instruction data and the scale of training experiments.

E Limitations

  • The data mainly covers long-context QA, summarization, and reasoning, leaving multi-turn dialogue, long-term role-playing, and long-history agent tasks for future work.The authors say collecting such data is challenging because current LLM performance on these tasks does not meet human needs, making Self-Instruct-style construction difficult.Source: E Limitations · p. 19
  • DeepSpeed and GPU-resource limits restricted SFT to 10B-level models and 64k maximum length, so the authors did not run extensive experiments on larger models or longer data.Source: E Limitations · p. 19
Loading 2401.18058v1…