[論文レビュー] Data Engineering for Scaling Language Models to 128K Context

Yao Fu, Rameswar Panda+5 more|arXiv (Cornell University)|Feb 15, 2024

Network Packet Processing and Optimization被引用数 4

ひとこと要約

本稿では、1–50億トークンの各ソースごとの長さアップサンプリングデータを継続的微調整することで、LLaMA-2言語モデルを128Kのコンテキスト長にスケーリングするデータ工学のレシピを提案する。適切なデータミックス設計—ドメイン比をバランスさせつつ、各ドメイン内で長いシーケンスを重み付けすること—により、ニードルインハリューレトリーバルタスクでGPT-4 128Kに近い性能を達成し、短いコンテキスト性能を維持していることが示された。

ABSTRACT

We study the continual pretraining recipe for scaling language models' context lengths to 128K, with a focus on data engineering. We hypothesize that long context modeling, in particular extit{the ability to utilize information at arbitrary input locations}, is a capability that is mostly already acquired through large-scale pretraining, and that this capability can be readily extended to contexts substantially longer than seen during training~(e.g., 4K to 128K) through lightweight continual pretraining on appropriate data mixture. We investigate the extit{quantity} and extit{quality} of the data for continual pretraining: (1) for quantity, we show that 500 million to 5 billion tokens are enough to enable the model to retrieve information anywhere within the 128K context; (2) for quality, our results equally emphasize extit{domain balance} and extit{length upsampling}. Concretely, we find that naively upsampling longer data on certain domains like books, a common practice of existing work, gives suboptimal performance, and that a balanced domain mixture is important. We demonstrate that continual pretraining of the full model on 1B-5B tokens of such data is an effective and affordable strategy for scaling the context length of language models to 128K. Our recipe outperforms strong open-source long-context models and closes the gap to frontier models like GPT-4 128K.

研究の動機と目的

• アーキテクチャの変更なしに、言語モデルを128Kコンテキスト長に効果的にスケーリングするためのデータ工学の手法を調査すること。
• 特にモデルの元の事前学習長を超える長コンテキスト内での任意の位置に正確な情報検索を可能にするという課題に対処すること。
• 大規模事前学習中にすでに部分的に習得済みの長コンテキスト能力を、洗練されたデータで最小限の継続的微調整を行うことで解き放てるかどうかを特定すること。
• 長コンテキスト適応における、データ量、ドメインバランス、シーケンス長分布のトレードオフを評価すること。
• 今後のインstructチューニングを可能にする、実用的で安価かつ効果的なオープンソースモデルの128Kコンテキスト延長のためのレシピを提供すること。

提案手法

• 64K–80Kのコンテキスト窓を用いた完全注意機構を採用し、1–50億トークンの洗練済みデータ上でLLaMA-2 7Bおよび13Bモデルを継続的微調整する。
• アーキテクチャの変更なしに、より長いコンテキスト長をサポートするためのRoPE位置埋め込みベースの調整を適用する。
• 各ドメイン内で長いシーケンスの割合を増やす一方で、元のドメインミックス比を保持する、各ソースごとの長さアップサンプリング戦略を設計する。
• FlashAttentionの活用によりメモリ使用量を削減し、学術用ハードウェアでも長コンテキスト学習を可能にする。
• 4Kから128Kのコンテキスト長にわたり、正確な検索性能を測定するために、ニードルインハリューベンチマークを用いて評価する。
• データミックス設計の影響を明確にするために、YaRN-Mistral 128K や LongLoRA 100K といった強力なベースラインと比較する。

Figure 1: Needle-in-a-Haystack test (Kamradt, 2023 ) performance comparison. The x-axis denotes the length of the document (the “haystack”); the y-axis indicates the position that the “needle” (a short sentence) is located within the document, from 1K to 128K tokens. For example, 50% indicates that

実験結果

リサーチクエスチョン

• RQ1 4Kコンテキストで事前学習されたモデルに対して、わずかな量の洗練済みデータでの継続的微調整で128Kコンテキスト能力が解き放てるか？
• RQ2 長さの長いドメイン（例：書籍）のシーケンスをそのままアップサンプリングすると、コードやウェブコンテンツなどの他のドメインで性能が低下するか？
• RQ3 各ドメイン内で長いシーケンスをアップサンプリングしながら元のドメインミックス比を維持する戦略は、ドメインベースのアップサンプリングよりも効果的か？
• RQ4 データミックス設計が、長コンテキスト性能と短コンテキスト性能のトレードオフにどのように影響するか？
• RQ5 データ中心の戦略が、長コンテキストモデリングにおける効率性と性能の面で、大規模事前学習レシピを上回る可能性があるか？

主な発見

• 各ソースごとの長さアップサンプリングデータを1–50億トークン分だけ継続的微調整することで、LLaMA-2 7Bおよび13Bモデルが128Kコンテキスト長までニードルインハリューテストで強力な性能を発揮することが可能になった。
• ドメイン比を保持しつつ各ドメイン内で長いシーケンスをアップサンプリングする「各ソースごとの長さアップサンプリング」戦略は、ドメイン全体にわたる性能向上が最もバランスが良く、ナーブなドメインベースのアップサンプリングを上回る。
• ナーブなドメインベースのアップサンプリング（例：書籍の過剰代表）は、コードやウェブコンテンツなどの他のドメインで顕著な性能低下を引き起こし、最適でないデータミックス設計であることが示された。
• 本研究で提示されたレシピで訓練されたモデルは、長コンテキスト検索においてGPT-4 128Kと同等の性能を達成しており、YaRN-Mistral 128K や LongLoRA 100K といった強力なオープンソースモデルを顕著に上回った。
• 検証損失だけでは性能の差を検出できない。ニードルインハリューテストにより、損失関数では捉えきれない長コンテキスト検索能力の顕著な差が明らかになった。
• 本研究では、データ工学が、最先端の長コンテキスト性能を達成するためには、モデルアーキテクチャやアルゴリズム的イノベーションと同等に重要であることが示された。

Figure 2: Length and domain distributions of the various data mixture strategies. We use SlimPajama (Soboleva et al., 2023 ) , an open-source reproduction of LLaMA (Touvron et al., 2023a ) pretraining data mixture, as the source dataset. The original data mixture has about 30% of documents that are