PaperlayerView exact revision on arXiv

Source-linked AI summary

Training-Free Long-Context Scaling of Large Language Models

Chenxin An, Fei Huang, Jun Zhang, Shansan Gong, Xipeng Qiu, Chang Zhou, Lingpeng Kong

arXiv:2402.17463v2cs.CL

TL;DR

Long-context performance weakens beyond LLMs’ pretraining lengths, while long-context finetuning is costly and training data can be difficult to access. DCA uses chunk-based attention without continual training; it extends Llama2 70B beyond 100k tokens and performs comparably to or better than continually trained models on long-context tasks.

  • Problem

    LLMs can degrade on sequences longer than their pretraining length, while long-context finetuning is costly and its training corpora are not always accessible.

    Source: §1 · p. 1, §2.2 · p. 3

  • Method

    DCA reuses pretrained position indices and decomposes long-sequence attention into intra-chunk, inter-chunk, and successive-chunk components.

    Source: §1 · pp. 1–2

  • Results

    Llama2 70B with DCA handles more than 100k tokens, and training-free models perform comparably to or better than continually trained models on long-context understanding benchmarks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    DCA extrapolates more than 8x the training length without further training.

    Source: §5 · p. 9

Abstract

from arXiv · show

The ability of Large Language Models (LLMs) to process and generate coherent text is markedly weakened when the number of input tokens exceeds their pretraining length. Given the expensive overhead of finetuning large-scale models with longer sequences, we propose Dual Chunk Attention (DCA), which enables Llama2 70B to support context windows of more than 100k tokens without continual training. By decomposing the attention computation for long sequences into chunk-based modules, DCA manages to effectively capture the relative positional information of tokens within the same chunk (Intra-Chunk) and across distinct chunks (Inter-Chunk), as well as integrates seamlessly with Flash Attention. In addition to its impressive extrapolation capability, DCA achieves performance on practical long-context tasks that is comparable to or even better than that of finetuned models. When compared with proprietary models, our training-free 70B model attains 94% of the performance of gpt-3.5-16k, indicating it is a viable open-source alternative. All code and data used in this work are released at \url{https://github.com/HKUNLP/ChunkLlama}.

1. Introduction

DCA addresses the cost and data barriers of long-context finetuning with a training-free, chunk-based attention framework. It extends context lengths substantially and performs comparably to or better than continually trained models on long-context understanding benchmarks.

1. Introduction

  • Long-context finetuning is costly and depends on training corpora with limited accessibility, while training-free local-information methods can lose long-range dependencies.Scaled positional encodings can also increase PPL in training-free settings, especially beyond twice the training length.Source: §1 · p. 1
  • DCA reuses pretrained position indices and divides long-sequence attention into intra-chunk, inter-chunk, and successive-chunk components.The chunk-based attention calculation integrates with Flash Attention 2.Source: §1 · pp. 1–2
  • Llama2 70B with DCA handles more than 100k tokens, while 4k-context LLMs extend beyond 32k with negligible PPL increase.The paper also reports that previous training-free methods typically falter beyond 8k tokens.Source: §1 · p. 2
  • DCA can extend existing 32k-context models to 192k while maintaining high passkey retrieval accuracy and low perplexity.The paper describes DCA as orthogonal to PI and NTK scaled positional encodings.Source: §1 · p. 2
  • Training-free models achieve performance comparable to or better than costly continually trained models on zero-shot and few-shot long-context understanding benchmarks.The evaluation covers language modeling, passkey retrieval, question answering, and summarization, including 70B models.Source: §1 · p. 2

2. Background

The background explains how positional information is represented in attention and why extending RoPE beyond its training range is challenging.

2.1. Positional Encoding

  • Standard positional embeddings add a position-dependent feature vector to each input token.The embedding function maps absolute position indices into a d-dimensional feature space.Source: §2.1 · p. 2
  • RoPE adds positional information in the attention layer, where query-key inner products encode relative position offsets.The relative position matrix is Toeplitz, with each entry representing the query position minus the key position.Source: §2.1 · p. 2

2.2. Extrapolation of RoPE

  • RoPE can degrade on sequences longer than those seen in pretraining because they introduce unseen relative positions.In the example, training uses a 6-token window and inference uses 12 tokens; PI halves position offsets, reducing positional resolution.Source: §2.2 · pp. 2–3, Figure 1 · p. 3

3. Method

DCA partitions long-sequence attention into chunks and assigns position indices according to within-chunk and cross-chunk relations, with a special rule for neighboring chunks.

3.1. Intra-Chunk Attention

  • Within each chunk, DCA keeps query and key position indices within the chunk-size range and computes their relative-position difference.Source: §3.1 · p. 3

3.2. Inter-Chunk Attention

  • For cross-chunk attention, DCA assigns inter-chunk queries the pretraining maximum position index c −1 while key indices repeat cyclically within the chunk.Source: §3.2 · p. 4, Figure 2 · p. 4

3.3. Successive-Chunk Attention

  • Successive-chunk attention adjusts positions for the w neighboring tokens, addressing the boundary case where adjacent tokens can otherwise receive a relative distance of 4.Source: §3.3 · p. 4

3. Method

  • DCA combines intra-chunk, inter-chunk, and successive-chunk attention to construct relative positions and compute attention scores.Source: §3.3 · pp. 4–5

3.4. Normalization

  • DCA applies softmax to normalize its inner products and integrates with Flash Attention 2 at comparable GPU memory usage and inference speed to Llama self-attention.Source: §3.4 · p. 5

4. Experiments

The experiments evaluate DCA across Llama2 variants and long-context models on language modeling, practical tasks, and attention analyses. They report context extrapolation, task performance, and efficiency and ablation findings.

4.1. Experimental Setup

  • DCA evaluation covers Llama2 7B, 13B, and 70B variants, plus Together-32k and CodeLlama; language modeling spans PG19 contexts from 4k to 192k.Few-shot evaluation includes NarrativeQA, QMSum, QuALITY, and Qasper.Source: §4 · p. 5, §4.1 · p. 5

4.2. Long-Sequence Language Modeling

  • On PG19, CHUNKLLAMA2 extends context beyond 32k with only a 0.02 PPL increase and surpasses finetuned models at 16k.The 70B variant's PPL rises from 5.18 to 5.59 across context lengths.Source: §4.2 · p. 6, Table 1 · p. 6
  • DCA extends Llama2 70B to 96k with a 0.56 PPL increase over its 4k-context result; PI- and NTK-based models reach 192k.The enhanced models retain 90% passkey retrieval accuracy at up to 192k tokens.Source: §4.2 · p. 6, Table 2 · p. 7

4.3. Practical Tasks

  • 70B DCA matches Longlora at 37.8 vs. 37.2 while requiring no training steps.Longlora uses LoRA-based tuning and supports a 32k context window.Source: §4.3 · pp. 6–7, Table 3 · p. 7
  • The training-free Llama2 70B chat model reaches 94% of gpt-3.5-turbo-16k performance at a 16k context window.On zero-shot L-Eval tasks, the 7B/13B models also show comparable performance to open-source models with further training.Source: §4.3 · pp. 7–8, Table 4 · p. 8

4.4. Analysis

  • On a single A100 GPU, DCA maintains similar inference speed and GPU memory use to Flash Attention without considerable overhead.Figure 3 compares original PyTorch self-attention, Flash Attention, and DCA across prompt lengths.Source: §4.4 · p. 8, Figure 3 · p. 8
  • Combining intra-chunk, inter-chunk, and successive-chunk attention yields both low PPL and high passkey retrieval accuracy.Intra-chunk attention alone maintains low PPL but hinders retrieval across chunks; adding inter-chunk attention improves retrieval at 12k, while losing locality raises PPL.Source: §4.4 · pp. 8–9, Figure 4 · p. 9

5. Conclusion

DCA is an efficient, training-free approach to extending LLM context windows beyond their training limits. It reuses existing position indices and a multi-faceted attention mechanism to extrapolate beyond eight times the training length.

5. Conclusion

  • DCA uses existing position indices and multi-faceted attention to extrapolate more than 8x the training length without further training.The conclusion characterizes the approach as novel and efficient for overcoming LLM context-length limitations.Source: §5 · p. 9

Impact Statement

DCA modifies inference code rather than requiring extensive retraining and is compatible with Flash Attention. The paper states that it preserves performance within the training length and benefits performance beyond it.

Impact Statement

  • DCA requires inference-code modifications, is compatible with Flash Attention, and preserves model performance within the training length while benefiting it beyond that range.The authors describe it as a cost-effective solution for long-context scenarios and compatible with models already finetuned for long contexts.Source: Impact Statement · p. 9

A.1. Passkey retrieval

The passkey tests compare DCA with positional-extension baselines across context lengths and document depths, showing strong retrieval at extended lengths but depth-specific failure patterns.

A.1. Passkey retrieval

  • 100% passkey retrieval accuracy across tested depths held through 18k tokens for CHUNKLLAMA2.The Llama2 13B model had a 4k pretraining context; tests used 20 passkeys per depth.Source: §A.1 · p. 13
  • At extended lengths, CHUNKLLAMA2 retained high retrieval accuracy through 32k, while DCA offers a learning-free path to 192k with existing long-context models.The 192k comparison includes Together-32k 7B and its DCA-enhanced counterpart.Source: §A.1 · p. 13, Figure 7 · p. 14
  • NTK-Aware accuracy fell to 40%–80% when passkeys were in middle sections, while CHUNKLLAMA2 first dropped at the beginning as context expanded.PI failure cases appeared largely unrelated to document depth.Source: §A.1 · p. 13
  • The 24k retrieval tests used models with 4k pretraining contexts and no further training, with 20 cases at each document depth.Figure 5 plots input context length against passkey depth.Source: Figure 5 · p. 13, §A.1 · p. 13

A.2. More Examples

The example illustrates how DCA segments a sequence longer than pretraining length and assigns attention calculations to tokens within and across chunks.

A.2. More Examples

  • For a 12-token example with pre-training length 8, when query[11] and query[0] perform an inner product, their relative position is 11, exceeding the pre-training size.Source: §A.2 · p. 14
  • For a 12-token sequence with pretraining length 8, DCA uses chunk size 4 to divide the input into three chunks.The chunk size is a hyperparameter smaller than the pretraining length.Source: §A.2 · p. 14
  • Intra-chunk attention handles tokens in the same chunk, with a maximum relative position of 3 in the example.The resulting position matrix is shown in Appendix Figure 8(a).Source: §A.2 · p. 14
  • Inter-chunk attention handles tokens in different chunks; with pretraining length 8, the maximum position index is 7.Its resulting position matrix is shown in Appendix Figure 8(b); successive-chunk attention handles tokens in successive chunks.Source: §A.2 · pp. 14–15

A.3. Flash Attention

DCA implements its three attention components as separate Flash Attention calculations for each query, then combines their outputs using shared normalization.

A.3. Flash Attention

  • For query i, DCA computes attention over keys j ≤ i using three Flash Attention operations with costs O(i − n*s), O(s), and O(s*(n − 1)), where n = floor(i/s).These correspond to intra-chunk, successive-chunk, and inter-chunk attention, respectively.Source: §A.3 · p. 15
  • The operations select intra-chunk keys from s*n to i, successive-chunk keys from s*(n−1) to s*n, and inter-chunk keys before s*(n−1).The pseudocode applies rotary position embeddings to the keys and to each attention-specific query.Source: §A.3 · p. 15
  • DCA sums the three attention-map normalizers, then weights each attention output by its map sum before dividing by the total.The resulting output has shape [1, d].Source: §A.3 · p. 15

A.4. In-Context Examples Selection

The study compares random and retrieval-based in-context example selection, finding that closest examples can perform poorly on question answering.

A.4. In-Context Examples Selection

  • The selection comparison used random training-set examples or retrieval-based examples ranked by query relevance, including BM25.The authors call the highest-scoring examples EXAMPLE BEST and the lowest-scoring examples EXAMPLE WORST.Source: §A.4 · pp. 15–16
  • On both question-answering datasets, the closest retrieved examples produced the poorest outcomes, while random and lowest-scoring examples performed relatively similarly.The authors suggest highly similar examples may prompt copying of an answer that is wrong for the current question.Source: §A.4 · p. 16, Table 5 · p. 16
  • QMSum summarization performance was generally less influenced by prompt selection than the two question-answering datasets.The comparison used CHUNKLLAMA2 7B/13B.Source: §A.4 · p. 16, Table 5 · p. 16

A.5. Performance on Unseen Data

Using the paper itself as a 19,388-token test document, the authors report strong elementary and promising difficult-query performance for ChunkLlama 70B, but remaining difficulty with global method understanding.

A.5. Performance on Unseen Data

  • The evaluation used the paper’s own content as an unseen 19,388-token input, omitting its title, abstract, and conclusion, with straightforward and method-comprehension questions.Source: §A.5 · p. 16
  • CHUNKLLAMA2 outperformed NTK in instruction interpretation and answer accuracy, although the 13B model still gave imprecise answers, including a wrong finetuning-corpus response.The correct corpus was ShareGPT and AlpacaGPT4, not the original Llama2 corpus.Source: §A.5 · pp. 16–17
  • ChunkLlama 70B could explain the inter-chunk and successive-chunk strategies but still struggled with questions requiring global understanding of DCA.Source: §A.5 · p. 16
  • ChunkLlama 70B performed highly on elementary and promisingly on harder paper-content questions, while Dynamic-NTK failed every listed easy and hard test case.The easy questions mostly required no deep understanding; the hard questions typically required understanding the method.Source: §A.5 · p. 16, Table 6 · p. 17, Table 7 · p. 18
  • ChunkLlama 70B correctly named four L-Eval tasks—TOFEL, QuALITY, Coursera, and SFiction—across input lengths of 3k to 27k.The 13B response instead listed tasks not identified as L-Eval tasks in the passage.Source: §A.5 · p. 17
Loading 2402.17463v2…