PaperlayerView exact revision on arXiv

Source-linked AI summary

YaRN: Efficient Context Window Extension of Large Language Models

Bowen Peng, Jeffrey Quesnelle, Honglu Fan, Enrico Shippole

arXiv:2309.00071v3cs.CLcs.AIcs.LG

TL;DR

Pretrained language models struggle to generalize beyond their training context windows. YaRN extends RoPE-based context windows through compute-efficient interpolation and Dynamic Scaling, achieving state-of-the-art extension with limited fine-tuning and extrapolating beyond fine-tuning lengths.

  • Problem

    Pretrained language models have limited ability to generalize beyond the context window seen during training.

  • Method

    YaRN combines RoPE interpolation techniques to extend context windows, while Dynamic Scaling enables inference-time extension without modifying the attention mechanism.

  • Results

    YaRN reaches state-of-the-art context-window-extension performance after fine-tuning on less than ∼0.1% of original pre-training data, while Dynamic-YaRN enables more than 2x extension without fine-tuning.

  • Takeaways & Limitations

    YaRN supports efficient extrapolation to longer contexts, including testing beyond the context length represented in a fine-tuning dataset.

  • Takeaways & Limitations

    The optimal base for NTK-aware interpolation is difficult to determine and usually must be found empirically, increasing fine-tuning difficulty and cost.

Abstract

from arXiv · show

Rotary Position Embeddings (RoPE) have been shown to effectively encode positional information in transformer-based language models. However, these models fail to generalize past the sequence length they were trained on. We present YaRN (Yet another RoPE extensioN method), a compute-efficient method to extend the context window of such models, requiring 10x less tokens and 2.5x less training steps than previous methods. Using YaRN, we show that LLaMA models can effectively utilize and extrapolate to context lengths much longer than their original pre-training would allow, while also surpassing previous the state-of-the-art at context window extension. In addition, we demonstrate that YaRN exhibits the capability to extrapolate beyond the limited context of a fine-tuning dataset. Code is available at https://github.com/jquesnelle/yarn

1 INTRODUCTION

Pretrained language models struggle to generalize beyond their training context windows, motivating RoPE-based interpolation methods and YaRN. YaRN combines prior interpolation ideas to extend context efficiently, achieving state-of-the-art performance with minimal fine-tuning and enabling further extension through Dynamic Scaling.

  • Motivation: Context-window length is a major pretrained-LLM limitation, and positional encodings generally cannot generalize to sequences significantly longer than training lengths.ALiBi provides only limited generalization according to the cited discussion.
  • Prior work: Position Interpolation and related methods extend RoPE context lengths by modifying positional representations and fine-tuning on small amounts of data.The related methods include NTK-aware, Dynamic NTK, and NTK-by-parts interpolation.
  • Contribution: YaRN presents an improved RoPE context-extension method for LLaMA, GPT-NeoX, and PaLM model families.The paper also provides an account of prior unpublished NTK-aware, Dynamic NTK, and NTK-by-parts work.
  • Results: Less than ∼0.1% of the original pre-training data is sufficient for YaRN to reach state-of-the-art context-window-extension performance after fine-tuning.The cited result reports the data fraction, while the surrounding contribution statement identifies the performance as state-of-the-art.
  • Results: More than 2x context-window extension is possible with Dynamic-YaRN without fine-tuning.Dynamic-YaRN combines YaRN with the inference-time Dynamic Scaling technique.

2 BACKGROUND AND RELATED WORK

RoPE encodes relative position through frequency-specific rotations, and context extension can be expressed as transformations of positions and frequencies. The background compares interpolation approaches, including PI, ReRoPE, and LM-Infinite, while defining the notation used for extending a pretrained context length.

  • Rotary Position Embeddings: RoPE associates hidden dimensions with separate frequencies so attention depends on relative token distance.The attention layer first converts hidden vectors into query and key vectors before applying the rotary representation.
  • Interpolation framework: Interpolation methods modify RoPE through a position map g(m) and a frequency transformation h(θ).The modified function is expressed as f′W(xm,m,θ)=fW(xm,g(m),h(θ)).
  • Context extension notation: The target context length L′ exceeds pretrained length L, with scale factor s defined by their ratio.The paper also introduces wavelength λd as the token length required for a full rotary rotation at dimension d.
  • Related work: Position Interpolation uniformly scales RoPE dimensions, whereas ReRoPE and LM-Infinite modify attention mechanisms rather than using purely embedding interpolation.ReRoPE additionally requires two attention passes and is incompatible with Flash Attention 2 in the cited discussion; LM-Infinite is likewise not immediately compatible.

3 METHODOLOGY

YaRN addresses RoPE’s limited length generalization by combining frequency-aware interpolation, attention-logit scaling, and dynamic scaling. The resulting method preserves relative local information while extending context efficiently across fine-tuned and non-fine-tuned settings.

  • 3.1 LOSS OF HIGH FREQUENCY INFORMATION - "NTK-AWARE" INTERPOLATION: PI degrades as scaling grows because uniformly stretching RoPE removes high-frequency components, with prior fine-tunes reaching roughly s = 8 before outputs degrade.NTK-aware interpolation instead scales high frequencies less and low frequencies more to distribute interpolation pressure.
  • 3.1 LOSS OF HIGH FREQUENCY INFORMATION - "NTK-AWARE" INTERPOLATION: NTK-aware interpolation retains the NTK motivation but requires empirically finding an optimal base for each intended extension factor, increasing fine-tuning difficulty and cost.NTK-by-parts preserves the underlying observation while using wavelength-dependent interpolation strategies.
  • 3.2 LOSS OF RELATIVE LOCAL DISTANCES - "NTK-BY-PARTS" INTERPOLATION: NTK-by-parts leaves short-wavelength dimensions unchanged, interpolates dimensions with wavelengths at least as large as L, and mixes both strategies for intermediate dimensions.The method uses α and β to define boundaries between these interpolation regimes.
  • 3.3 YARN: YaRN combines NTK-by-parts interpolation with attention scaling by 1/t, altering attention without modifying its code and adding zero inference or training overhead.Rotary embeddings are generated in advance and reused across forward passes.
  • 3.3 YARN: YaRN surpasses previous methods in both fine-tuned and non-fine-tuned scenarios while remaining directly compatible with Flash Attention 2.The method is presented as an improved RoPE extension for model families including LLaMA, GPT-NeoX, and PaLM.
  • 3.4 DYNAMIC SCALING - "DYNAMIC NTK" INTERPOLATION: Dynamic Scaling updates the scale factor for each forward pass, allowing graceful degradation across changing sequence lengths instead of abrupt failure beyond the trained limit.Dynamic NTK combines this inference-time procedure with NTK-aware interpolation and works exceptionally well without fine-tuning when L′ = L.

4 EXPERIMENTS

Experiments evaluate YaRN’s training behavior, long-context language modeling, retrieval, benchmark preservation, and computational efficiency across extended LLaMA models. YaRN extrapolates to unseen context lengths, outperforms competing interpolation methods under matched budgets, preserves short-context abilities, and converges fastest.

  • 4.1 TRAINING: YaRN converges faster and maintains lower training loss than other interpolation techniques for extending LLaMA 7B to 32k context.The comparison uses 400 training steps on 32k-chunked PG19 data.
  • 4.2 LONG SEQUENCE LANGUAGE MODELING: YaRN extrapolates Llama 2 models to 128k context, including a model trained with only 64k-context data and 200 additional steps from a 64k checkpoint.The evaluation uses sliding-window perplexity on ten Proof-pile documents of at least 128k tokens.
  • 4.2 LONG SEQUENCE LANGUAGE MODELING: YaRN consistently achieves lower perplexity than the other interpolation methods in both non-fine-tuned and fine-tuned scenarios at the same number of training steps.The comparison is reported as an ablation across four interpolation methods.
  • 4.2 LONG SEQUENCE LANGUAGE MODELING: YaRN achieves higher passkey retrieval scores than other interpolation methods for similarly trained 32k LLaMA 7B models.Passkeys are five-digit numbers placed at random locations in otherwise meaningless text, with accuracy evaluated across prompt lengths.
  • 4.4 STANDARDIZED BENCHMARKS: YaRN models show minimal performance degradation relative to their respective Llama 2 baselines, while extending context windows.The iterative extension from 64k to 128k produces an average 0.49% score drop between the s = 16 and s = 32 models.
  • 4.5 COMPUTATIONAL EFFICIENCY: Modifying cached rotary position embeddings adds no computational or memory cost, and YaRN converges fastest among the compared methods.Training-time comparisons are reported in A100-hours.

5 CONCLUSION

The conclusion presents YaRN as an efficient RoPE interpolation method that extends context windows while preserving original benchmark abilities. It also supports extrapolation from shorter fine-tuning data and faster convergence through transfer learning.

  • 5 CONCLUSION: YaRN improves upon existing RoPE interpolation methods and serves as a drop-in replacement for PI with minimal implementation effort.The authors report no downsides in their conclusion.
  • 5 CONCLUSION: Fine-tuned YaRN models preserve original abilities on multiple benchmarks while attending to much larger context sizes.The conclusion links benchmark preservation with large-context capability.
  • 5 CONCLUSION: YaRN supports efficient extrapolation with shorter fine-tuning datasets and transfer learning for faster convergence under compute-constrained scenarios.The conclusion characterizes this as enabling “train short, and test long.”

6 REPRODUCIBILITY

This section documents the provided training and evaluation code, reviews Position Interpolation and NTK-aware interpolation, and explains the base-change objective behind NTK-aware scaling.

  • Reproducibility: The supplementary material includes the complete code used to train YaRN models and reproduce several evaluations, plus implementations of referenced extension methods.YaRN training uses publicly available PG19 tokenized into contiguous 64k-token chunks.
  • Position Interpolation: Position Interpolation modifies RoPE position indices so sequences beyond the pretrained limit can be handled after a small amount of fine-tuning.The new context window is L′ > L, where L is the pretrained limit.
  • NTK-aware interpolation: The NTK-aware interpolation is defined as a RoPE modification using functions parameterized by the scale factor s.The paper introduces this method through a formal definition following its RoPE notation.
  • NTK-aware interpolation: NTK-aware interpolation improves extension for non-fine-tuned models but can underperform Position Interpolation after fine-tuning because some dimensions extrapolate out of bounds.Its theoretical scale factor also does not accurately describe the true context extension scale in practice.
  • Base-change objective: The base-change design spreads interpolation pressure across hidden dimensions, scaling the lowest frequency like linear positional scaling while keeping the highest frequency constant.A new base b′ is chosen so the final dimension matches the wavelength of linear interpolation with scale factor s.

A.3 THE IMPACT OF PRE-SOFTMAX SCALING OF YARN ON PERPLEXITY

The experiments examine how YaRN’s pre-softmax scaling factor affects perplexity across long documents and token positions. They find that a suitable factor improves perplexity broadly, with the best value generally stable and close to the recommended formula.

  • Experimental setup: The study evaluates 896 16k-token RedPajama documents, dividing the 16k-token range into eight position segments and testing different pre-softmax scalings.RedPajama was selected as the open-source dataset closest to LLaMA’s training data as far as the authors were aware.
  • Analysis design: The evaluation compares mean perplexity changes and sample-level minima across token-position segments rather than only reporting aggregate document scores.Documents are split into 2048-token chunks for the position-dependent analysis.
  • Perplexity impact: For a suitable scaling factor, samples can achieve better perplexity across the extended context window.The analysis compares perplexity against the t = 1 setting and uses sliding-window perplexity summaries.
  • Perplexity impact: The best scaling value is mostly consistent across different samples and token positions.Figure 6 counts samples attaining minimal perplexity at each tested scaling across position segments.
  • Recommended scaling: The finding remains consistent across different scale factors s, and the best t closely follows YaRN’s recommended formula.The experiments use s = 8 for the illustrated comparisons and report variation using one-standard-deviation shading.

B.2 TRAINING EFFICIENCY OF YARN

YaRN achieves results comparable to Position Interpolation for extending Llama 2 7B from 4096 to 8192 tokens while using substantially less training.

  • Training efficiency: 2.5x less training steps and data: YaRN uses the same fine-tuning methodology as the PI comparison while reducing the training budget.YaRN was fine-tuned for 400 steps versus PI’s 1000 steps.
  • Training efficiency: Comparable results: YaRN matches PI closely when extending Llama 2 7B from 4096 to 8192 context length.The comparison is based on sliding-window perplexity over ten 128k-token Proof-pile documents.

B.3 COMPARING THE PERPLEXITY OF VARIOUS METHODS OVER A SLIDING WINDOW

The sliding-window evaluations compare YaRN with other long-context extensions on Llama 2 across larger target context sizes. YaRN shows strong performance throughout the targeted context and good long-sequence performance on GovReport.

  • Model comparisons: The evaluation compares Llama 2 models fine-tuned with YaRN at s = 16 and s = 32 against Together.ai and NTK-aware Code Llama.Results are summarized in Table 7 and visualized in Figure 7.
  • Model comparisons: Strong performance across the entire targeted context size makes YaRN the first method reported here to successfully extend Llama 2’s effective context size.The claim concerns the targeted context evaluated in the comparison.
  • GovReport evaluation: Fine-tuning with YaRN achieves good performance on long sequences in the 32k maximal-context GovReport evaluation without Dynamic Scaling.This result uses 50 untruncated GovReport documents with at least 16k tokens per sample.
  • Evaluation protocol: The Proof-pile visualization reports sliding-window perplexity for documents truncated to each model’s evaluation context-window size.The sliding window uses S = 256.

B.4 GOVREPORT EVALUATIONS

The GovReport evaluation is described in Section 4.2, with results presented as sliding-window perplexity for 50 long documents using a 32k context window.

  • Section 4.2 evaluates models on GovReport documents.
  • Table 8 reports sliding-window perplexity with S = 256.
  • The evaluation covers 50 long GovReport documents at a fixed 32k context window.

B.5 PASSKEY RETRIEVAL

Passkey retrieval tests evaluate whether extended-context models can retrieve randomly placed passkeys across context lengths up to 128k. The results also compare perplexity-based and retrieval-based indicators of long-context capability, including Dynamic-YaRN without fine-tuning.

  • Passkey retrieval: > 99% passkey retrieval accuracy was achieved by both 7b and 13b YaRN models fine-tuned at 128k across the entire 8k–128k evaluation range.
  • Passkey retrieval: The passkey task placed keys uniformly at random across evaluation windows ranging from 8k to 128k over 10 iterations.
  • Perplexity versus retrieval: Code Llama 13b retrieved the passkey accurately at 128k despite increasing perplexity above 100k.
  • Perplexity versus retrieval: YaRN with s = 32 had higher passkey accuracy and similar perplexity to s = 16 after 200 additional training steps.
  • Dynamic Scaling: Dynamic Scaling updates interpolation factor s at inference time for methods including PI, NTK-by-parts, and YaRN.
  • Dynamic Scaling: Dynamic Scaling prevents perplexity blow-up beyond Llama 2’s pretrained context window, while Dynamic-YaRN outperforms Dynamic-PI without fine-tuning.
Loading 2309.00071v3…