Resonance R o PE : Improving Context Length Generalization of Large Language Models (2024)

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.

## introduction recent advancements in large language models (llms) have demonstrated their potential across a wide spectrum of natural language processing tasks, showcasing their ability to handle complex interactions, document analyses, professional writing, and advanced reasoning with a unified approach @xcite @xcite . as these models are increasingly adapted for complex applications, challenges arise in scenarios requiring the comprehension or generation 1 https://github.com/sheryc/resonance_rope . recent efforts to enhance tstl performance have focused on llms equipped with rotary position embedding (rope) @xcite , such as llama @xcite and mistral @xcite , owing to their exceptional capabilities and widespread adoption. these initiatives aim to refine the test-time computation of rope position embedding by introducing a scaling factor to either the position index of each token @xcite or rope's base value @xcite @xcite . these methods ensure that the position embeddings for out-of-distribution (ood) positions remain within the range experienced during pre-training. this minimizes the need for the model to adapt to new position embedding value ranges, a task that is inherently difficult. in this paper, we introduce resonance rope, a novel technique designed to further narrow the generalization gap on position embeddings in tstl scenarios. recognizing that rope's position embedding is governed by a complex, non-linear function, we posit that minimizing extrapolation on ood positions, while crucial, is insufficient. we argue that it is equally vital to address the interpolation of rope features at the ood positions. by implementing resonance rope, we slightly scale each rope feature to correspond to an integer wavelength. this adjustment aligns each rope feature's wavelength with a specific token span length, enabling it to "resonate" with a particular local context length. this simple modification effectively reduces the generalization gap for over half of the position embedding features in llama and llama2 under tstl scenarios. furthermore, our approach is compatible with rope and any rope-based scaling techniques, enhancing their performance in tstl situations without the need for additional computational resources during training or inference. additionally, to facilitate further research on position embeddings, we present a new synthetic benchmark tailored for tstl scenarios, named posgen. improving position embeddings for tstl requires a detailed analysis of the cause of failures in handling longer contexts. however, current benchmarks, such as those measuring perplexity in long context @xcite @xcite and most synthetic tstl tasks @xcite face a common issue: the difficulty of generating the next token increases with context length. this makes it difficult to determine whether a model's failure is due to its inability to generate more complex tokens or its failure to recognize out-of-distribution (ood) positions. posgen addresses this limitation by standardizing the difficulty level of token generation across all positions. this ensures that any observed shortcomings are directly related to the model's inability to identify and handle new token positions effectively. our contributions in this study are threefold: ## proposed method: resonance rope in this section, we introduce resonance rope, a universal improvement for rope and rope-based scaling methods to (further) improve their length extrapolation performance. suppose we abstract rope's equation 4, 5: for any in a tstl scenario where we generalize an llm from length l to length l ′ , let us denote a scaled rope function by f . to perform well on ood positions it should reduce the feature gap h( f ) between token features seen during training and token features after scaling that we can define for each i-th feature as: ## evaluating position embeddings with posgen in this section, we propose our new position embedding evaluation suite: posgen, based on an analysis of common failure patterns on existing position embedding evaluation methods. we consider a next token prediction task, where we expect the model to generate the token x l given the input sequence {x 0 , • • • , x l-1 }. in tstl scenarios, when a model succeeds in correctly generating a token up to position l but fails systematically afterwards, we observe two failure patterns: currently, neither perplexity-based evaluations @xcite @xcite nor synthetic tstl evaluations @xcite can effectively distinguish these two failure patterns, since the token generation difficulty tends to increase with respect to the sequence length in these tasks. to facilitate research on better position representations, we design posgen, which controls the difficulty in generating tokens throughout the sequence to be identical, which effectively distinguishes the two types of tstl failures. failures in this benchmark are only due to the inability to recognize new token positions in tstl scenarios. our posgen framework comprises three subtasks, with each extracting the general token dependency pattern of a different type of reasoning task. suppose that we define a fixed function h : v j+@xmath0, where v is the model's vocabulary and j, k are predefined constants controlling the task's difficulty. the three subtasks of pos-gen are as follows: ## experiments we evaluate resonance rope on three different tstl tasks: a small-scale evaluation on our proposed posgen task, and llm-scale evaluations with llama2-chat @xcite on both language modeling perplexity and real-world long context applications. ## conclusion we introduce resonance rope, a novel enhancement of rope that focuses on minimizing the interpolation of rope features for ood positions, thereby reducing the generalization gap and improving llm's performance on train-short-testlong (tstl) scenarios. additionally, we present a novel synthetic benchmark, posgen, which provides a fine-grained analysis of the model's tstl performance regarding various token dependency patterns. extensive experiments on our proposed posgen and two llm-based evaluations demonstrate resonance rope's efficacy in identifying ood positions and its compatibility with current rope scaling strategies. future work includes exploring resonance rope's performance on other foundational models, and the identification of more optimal wavelength combinations for rope features.