E^2-LLM: Efficient and Extreme Length Extension of Large Language Models

Jiaheng Liu,Zhiqi Bai,Yuanxing Zhang,Chenchen Zhang,Yu Zhang,Ge Zhang,Jiakai Wang,Haoran Que,Yukang Chen,Wenbo Su,Tiezheng Ge,Jie Fu,Wenhu Chen,Bo Zheng

2024-02-22

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 Large Language Models, 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.

Computation and Language,Artificial Intelligence

What problem does this paper attempt to address?

The paper aims to address the high computational cost faced by large language models (LLMs) when handling long contexts. Specifically, existing long context extension methods typically require additional training processes to support longer context windows and demand substantial GPU resources. To tackle these issues, the authors propose an efficient and extreme length extension method called E2-LLM. This method supports arbitrary length context windows with only a single training process, significantly reducing computational costs and eliminating the need to collect long context data. Additionally, E2-LLM is based on Rotary Position Embedding (RoPE) and introduces two different enhancement strategies—enhancement of scale parameters and position index parameters—to improve the model's robustness to different relative position differences. Experimental results show that E2-LLM performs excellently in long context tasks across multiple benchmark datasets.