E^2-LLM: Efficient and Extreme Length Extension of Large Language Models
(2401.06951)Abstract
Typically, training LLMs with long context sizes is computationally expensive, requiring extensive training hours and GPU resources. Existing long-context extension methods usually need additional training procedures to support corresponding long-context windows, where the long-context training data (e.g., 32k) is needed, and high GPU training costs are assumed. To address the aforementioned issues, we propose an Efficient and Extreme length extension method for LLMs, called E 2 -LLM, with only one training procedure and dramatically reduced computation cost, which also removes the need to collect long-context data. Concretely, first, the training data of our E 2 -LLM only requires a short length (e.g., 4k), which reduces the tuning cost greatly. Second, the training procedure on the short training context window is performed only once time, and we can support different evaluation context windows at inference. Third, in E 2 - LLM, based on RoPE position embeddings, we introduce two different augmentation methods on the scale and position index parameters for different samples in training. It aims to make the model more robust to the different relative differences when directly interpolating the arbitrary context length at inference. Comprehensive experimental results on multiple benchmark datasets demonstrate the effectiveness of our E 2 -LLM on challenging long-context tasks.
Overview
-
LLMs are capable AI systems but face limitations in handling long context lengths, hindering their application in tasks requiring lengthy inputs.
-
Traditional solutions for extending context length are computationally costly and require extensive fine-tuning over large datasets, which is impractical.
-
E²-LLM presents an efficient extension method for LLMs that uses minimal additional training on short contexts, yet supports longer inputs without further fine-tuning.
-
The method employs a two-pronged augmentation strategy utilizing Rotary Position Embedding (RoPE) to improve long-context handling adaptively during inference.
-
E²-LLM demonstrated successful results on benchmarks with notably lower GPU memory costs and matching or outperforming other models trained on longer sequences.
Introduction
Large language models have transformed the nature of tasks that AI systems can perform. However, they typically have a limitation in the context length they can handle, posing a challenge in applications such as document summarization, long conversations, and lengthy reasoning tasks. Most models have a preset limit on the number of tokens they can consider, and increasing this limit traditionally implies a massive computational burden and the need for fine-tuning on extensive datasets. This paper proposes E²-LLM, a method that streamlines the length extension of LLMs in an efficient manner, using shorter context lengths for training and supporting evaluation on longer inputs without additional fine-tuning or computational intensiveness.
Methodology
The basis of E²-LLM lies in its two-pronged augmentation strategy, which leverages the Rotary Position Embedding (RoPE) to extend the effective context length with minimal additional training. The first strategy varies the scale parameter of the position embeddings, effectively changing the facing of position indices so that the model learns to deal with varied densities of positions. Secondly, an augmentation on the position index parameters is introduced to allow for offsets, making the model more versatile to different positional ranges. This is crucial as it teaches the LLM to generalize across different lengths and relative differences, a capability that is activated during inference time depending on the given context window.
Experimental Findings
E²-LLM was put to the test on several benchmark datasets designed to challenge the model’s long-context abilities. It was found to perform effectively across these tasks, often matching or outperforming existing LLMs that had been trained extensively for longer context windows. Notably, E²-LLM managed to achieve these results with significantly lower GPU memory costs as it required only once-off training using shorter data sequences (e.g., 4k tokens), yet it was successful in handling much longer contexts (e.g., 32k tokens) effectively.
Implications and Future Work
The ingenuity of E²-LLM opens new doors for efficient utilization of powerful language models without the prohibitive costs associated with training them on long contexts. As future work, the authors intend to apply this methodology to even larger models and examine its performance on more varied datasets and tasks. Furthermore, they plan to explore the method's adaptability to other types of positional encodings and language models. As the computational landscape becomes more demanding, E²-LLM stands out as a promising approach to pushing the boundaries of what LLMs can do without breaking the computational bank.