100サンプルでLLMのコンテキストウィンドウを拡張する

Extending LLMs' Context Window with 100 Samples

January 13, 2024

著者: Yikai Zhang, Junlong Li, Pengfei Liu

cs.AI

要旨

大規模言語モデル（LLMs）は、事前学習されたコンテキストウィンドウを超える外挿能力が限られていることが知られており、長い入力が必要な下流タスクへの適用が制約されています。最近の研究では、LLaMA、PaLM、GPT-NeoXなどの著名なLLMで採用されている人気の位置符号化手法であるロータリーポジション埋め込み（RoPE）を修正することで、LLMのコンテキストウィンドウを拡張しようと試みています。しかし、Position Interpolation（PI）やYaRNなどの先行研究はリソース集約的であり、その適用性を評価するための比較実験が不足しています。本研究では、LLMのアテンションエントロピー（すなわち、アテンションスコアの情報エントロピー）が安定性を維持するための内在的な必要性を特定し、RoPEの基本周波数を調整し、アテンションロジットをスケーリングすることで、LLMが効率的に大きなコンテキストウィンドウに適応できるようにするRoPEの新たな拡張を提案します。我々は、様々なコンテキスト要求タスクにおいて、異なるコンテキストウィンドウサイズでのファインチューニング性能とロバスト性の優位性を検証しました。特に、我々の手法は、LLaMA-2-7B-Chatのコンテキストウィンドウを16,384に拡張し、わずか100サンプルと6トレーニングステップで、驚異的な効率性を示しました。最後に、特定の下流タスクにおけるコンテキストウィンドウ拡張にデータ構成とトレーニングカリキュラムがどのように影響するかを探り、長い会話でLLMをファインチューニングすることが良い出発点であることを示唆しています。我々は、コードとSFTデータをhttps://github.com/GAIR-NLP/Entropy-ABFで公開しています。

English

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.