Table of Contents
Fetching ...

Resonance RoPE: Improving Context Length Generalization of Large Language Models

Suyuchen Wang, Ivan Kobyzev, Peng Lu, Mehdi Rezagholizadeh, Bang Liu

TL;DR

Resonance RoPE is introduced, a novel approach designed to narrow the generalization gap in TSTL scenarios by refining the interpolation of RoPE features for OOD positions, significantly improving the model performance without additional online computational costs.

Abstract

This paper addresses the challenge of train-short-test-long (TSTL) scenarios in Large Language Models (LLMs) equipped with Rotary Position Embedding (RoPE), where models pre-trained on shorter sequences face difficulty with out-of-distribution (OOD) token positions in longer sequences. We introduce Resonance RoPE, a novel approach designed to narrow the generalization gap in TSTL scenarios by refining the interpolation of RoPE features for OOD positions, significantly improving the model performance without additional online computational costs. Furthermore, we present PosGen, a new synthetic benchmark specifically designed for fine-grained behavior analysis in TSTL scenarios, aiming to isolate the constantly increasing difficulty of token generation on long contexts from the challenges of recognizing new token positions. Our experiments on synthetic tasks show that after applying Resonance RoPE, Transformers recognize OOD position better and more robustly. Our extensive LLM experiments also show superior performance after applying Resonance RoPE to the current state-of-the-art RoPE scaling method, YaRN, on both upstream language modeling tasks and a variety of downstream long-text applications.

Resonance RoPE: Improving Context Length Generalization of Large Language Models

TL;DR

Resonance RoPE is introduced, a novel approach designed to narrow the generalization gap in TSTL scenarios by refining the interpolation of RoPE features for OOD positions, significantly improving the model performance without additional online computational costs.

Abstract

This paper addresses the challenge of train-short-test-long (TSTL) scenarios in Large Language Models (LLMs) equipped with Rotary Position Embedding (RoPE), where models pre-trained on shorter sequences face difficulty with out-of-distribution (OOD) token positions in longer sequences. We introduce Resonance RoPE, a novel approach designed to narrow the generalization gap in TSTL scenarios by refining the interpolation of RoPE features for OOD positions, significantly improving the model performance without additional online computational costs. Furthermore, we present PosGen, a new synthetic benchmark specifically designed for fine-grained behavior analysis in TSTL scenarios, aiming to isolate the constantly increasing difficulty of token generation on long contexts from the challenges of recognizing new token positions. Our experiments on synthetic tasks show that after applying Resonance RoPE, Transformers recognize OOD position better and more robustly. Our extensive LLM experiments also show superior performance after applying Resonance RoPE to the current state-of-the-art RoPE scaling method, YaRN, on both upstream language modeling tasks and a variety of downstream long-text applications.
Paper Structure (24 sections, 1 theorem, 16 equations, 4 figures, 2 tables, 1 algorithm)

This paper contains 24 sections, 1 theorem, 16 equations, 4 figures, 2 tables, 1 algorithm.

Key Result

Theorem 1

For a RoPE-equipped model with context window $L$, Resonance RoPE$\tilde{f}$ reduces the feature gap on pre-critical dimensions to $0$. Specifically, $\forall {\bm{x}}\in{\mathbb{X}}$, $\forall n\in{\mathbb{N}}\backslash\{0,\cdots,L-1\}$, we have: for all $i = 0,\dots,2c-1$.

Figures (4)

  • Figure 1: An illustration of RoPE's rotation angles $m\theta_6$ and Resonance RoPE's rotation angles $m\tilde{\theta}_6$ in Eqn. \ref{['eqn:rope-subfeature']} in a TSTL scenario with training max length $64$ and testing max length $128$. RoPE's non-integer feature wavelengths create a feature gap between the RoPE features of the training and OOD testing positions, while Resonance RoPE reduces this gap to 0.
  • Figure 2: An example of the three subtasks of PosGen. This figure shows the process of generating the $12$th token shown in the red boxes for each subtask. In this example, $h$ is a modular addition task with the modulus $m=7$ and the difficulty-controlling parameters $j=1, k=3$. The output token depends on: (1) only the local $j+k$ tokens in the recursive task; (2) $k$ local tokens and the beginning $j$ tokens in the CoT task; and (3) $k$ local tokens and $j$ tokens with a varied dependency distance in the semi-recursive task.
  • Figure 3: The validation loss curves of Transformers using RoPE and YaRN PEs with and without our Resonance scaling on the three subtasks of PosGen.
  • Figure 4: The perplexity of LLaMA-Chat 7B with different position embeddings on GovReport and Proofpile.

Theorems & Definitions (2)

  • Theorem 1
  • proof