Extending LLMs' Context Window with 100 Samples
Authors: name: Yikai Zhang; name: Junlong Li; name: Pengfei Liu
Published: 2024-01-13 07:57:01

Abstract: Large Language Models (LLMs) are known to have limited extrapolation ability
beyond their pre-trained context window, constraining their application in
downstream tasks with lengthy inputs. Recent studies have sought to extend
LLMs' context window by modifying rotary position embedding (RoPE), a popular
position encoding method adopted by well-known LLMs such as LLaMA, PaLM, and
GPT-NeoX. However, prior works like Position Interpolation (PI) and YaRN are
resource-intensive and lack comparative experiments to assess their
applicability. In this work, we identify the inherent need for LLMs' attention
entropy (i.e. the information entropy of attention scores) to maintain
stability and introduce a novel extension to RoPE which combines adjusting
RoPE's base frequency and scaling the attention logits to help LLMs efficiently
adapt to a larger context window. We validate the superiority of our method in
both fine-tuning performance and robustness across different context window
sizes on various context-demanding tasks. Notably, our method extends the
context window of LLaMA-2-7B-Chat to 16,384 with only 100 samples and 6
training steps, showcasing extraordinary efficiency. Finally, we also explore
how data compositions and training curricula affect context window extension
for specific downstream tasks, suggesting fine-tuning LLMs with lengthy
conversations as a good starting point. We release our code and SFT data at
https://github.com/GAIR-NLP/Entropy-ABF.