How to Train Long-Context Language Models (Effectively)

要約

我々は、ロングコンテクスト情報を効果的に利用するための言語モデル(LM)の継続的な学習と教師付き微調整(SFT)について研究する。まず、モデル開発の指針となる信頼性の高い評価プロトコルを確立する。パープレキシティや単純なNIAH（needle-in-a-haystack）テストの代わりに、ロングコンテキストのダウンストリームタスクの幅広いセットを使用し、SFT後にモデルを評価する。ロバストな評価に支えられ、我々は、継続的な事前学習のためのデータミックス、インストラクションチューニングデータセット、および位置外挿のような他の多くの設計の選択を決定するための徹底的な実験を実行する。その結果、(1)コードリポジトリや書籍はロングデータの優れたソースであるが、高品質のショートコンテキストデータと組み合わせることが極めて重要であること、(2)評価長を超えるシーケンス長で学習することで、ロングコンテキストの性能が向上すること、(3)SFTでは、ショートインストラクションデータセットのみを使用することで、ロングコンテキストのタスクで強力な性能が得られることがわかった。Llama-3から初期化され、40Bのトークンで学習された最終モデルProLong-8Bは、128Kの長さで同サイズのモデルの中で最先端のロングコンテクスト性能を示す。ProLongは、ロングコンテキストの学習で使用するトークンの数がわずか5%であるにもかかわらず、ほとんどのロングコンテキストのタスクでLlama-3.1-8B-Instructを上回る。さらに、ProLongは512Kまでのトークンを効果的に処理することができ、これは公開されているLMの中で最も長いコンテキストウィンドウの一つである。

要約(オリジナル)

We study continued training and supervised fine-tuning (SFT) of a language model (LM) to make effective use of long-context information. We first establish a reliable evaluation protocol to guide model development — instead of perplexity or simple needle-in-a-haystack (NIAH) tests, we use a broad set of long-context downstream tasks, and we evaluate models after SFT as this better reveals long-context abilities. Supported by our robust evaluations, we run thorough experiments to decide the data mix for continued pre-training, the instruction tuning dataset, and many other design choices such as position extrapolation. We find that (1) code repositories and books are excellent sources of long data, but it is crucial to combine them with high-quality short-context data; (2) training with a sequence length beyond the evaluation length boosts long-context performance; (3) for SFT, using only short instruction datasets yields strong performance on long-context tasks. Our final model, ProLong-8B, which is initialized from Llama-3 and trained on 40B tokens, demonstrates state-of-the-art long-context performance among similarly sized models at a length of 128K. ProLong outperforms Llama-3.1-8B-Instruct on the majority of long-context tasks despite using only 5% as many tokens during long-context training. Additionally, ProLong can effectively process up to 512K tokens, one of the longest context windows of publicly available LMs.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL