How to Train Long-Context Language Models (Effectively)

要約

我々は、ロングコンテクスト情報を効果的に利用するための言語モデル(LM)の継続的な学習と教師付き微調整(SFT)について研究する。パープレキシティや単純なNIAH（Needle-in-a-Haystack）テストの代わりに、幅広いロングコンテクスト課題を使用し、ロングコンテクスト能力をより良く明らかにするため、SFT後のモデルを命令データで評価する。また、SFT後のモデルを命令データで評価することで、ロングコンテキストの能力をより明確にすることができる。その結果、(1)コードリポジトリや書籍は、ロングデータの優れたソースであるが、高品質のショートデータと組み合わせることが重要であること、(2)評価長を超えるシーケンス長でのトレーニングは、ロングコンテキストの性能を向上させること、(3)SFTの場合、ショートインストラクションデータセットのみを使用することで、ロングコンテキストのタスクで強力な性能が得られることがわかった。Llama-3から初期化され、40Bのトークンで学習された我々の最終モデルProLong-8Bは、128Kの長さで同サイズのモデルの中で最先端のロングコンテクスト性能を示す。ProLongは、ロングコンテキストの学習中に見たトークンの数がLlama-3.18B-Instructの5%しかないにもかかわらず、ロングコンテキストのタスクの大部分でLlama-3.18B-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 tasks, and we evaluate models after SFT with instruction data 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. 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 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.18B-Instruct on the majority of long-context tasks despite having seen 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