Llama 3モデル群：大規模言語モデルの新たな展開

論文解説

問題設定と動機

現代の人工知能システムは大規模な基盤モデルに大きく依存していますが、最も高性能なモデル（GPT-4など）は依然としてクローズドのままです。一般に利用可能なフロンティア級のモデルを構築するには、データキュレーション、計算規模、複雑性の管理という3つの柱にわたる極めて大きなエンジニアリング上のボトルネックを克服する必要があります。

著者らは、最大128Kのコンテキストウィンドウ、ゼロショットおよびマルチステップのツール利用、高度な推論、コード生成をサポートする、8B、70B、405Bパラメータからなる多言語対応モデル群であるLlama 3を開発しました。

主要なアイデア

• 高密度（Dense）アーキテクチャによる複雑性の管理: Mixture-of-Experts（MoE）アーキテクチャに切り替えるのではなく、事前学習の安定性を最大化するため、標準的な高密度Transformer（405Bモデルでは126層、次元数16,384、128アテンションヘッド）を採用し、Grouped-Query Attention（GQA）および基底周波数 $\theta =500,000$ のRoPEを使用しています。
• データスケーリングと厳格なフィルタリング: 事前学習データは、Llama 2の1.8Tトークンから、多様なドメインにわたる約15.6Tトークン（一般知識約50%、数学/推論25%、コード17%、多言語8%）へと大幅に拡大されました。
• 予測可能なスケーリング則: 2段階のスケーリング則フレームワークにより、トレーニング開始前に4桁にわたる計算量とダウンストリームタスク性能（ARC Challengeなど）を関連付けています。
• 反復的な事後学習（Post-Training）: 教師あり微調整（SFT）、リジェクションサンプリング（RS）、Direct Preference Optimization（DPO）を組み合わせた6ラウンドのアライメントを実施し、ターゲットを絞った合成データ生成も活用しています。
• 4D並列処理とインフラ: テンソル並列（TP）、コンテキスト並列（CP）、パイプライン並列（PP）、完全シャードデータ並列（FSDP/DP）を組み合わせて、最大16K基のH100 GPU上で405Bモデルを事前学習しました。

手法（主要な数式を含む）

1. スケーリング則の定式化

計算量に最適化されたトークン割り当ては、計算予算 $C$ のべき乗則関数としてモデル化されます：

$$N^{\star }(C)=A{C}^{\alpha }$$

$6\times 10^{18}$ から $10^{22}$ FLOPsの間の計算予算に対する経験的フィッティングにより、$(\alpha ,A)=(0.53,0.29)$ が特定されました。目標予算である $3.8\times 10^{25}$ FLOPsへと外挿することで、16.55Tトークンで402Bパラメータモデルを学習することが示唆され、最終的な405Bパラメータというアーキテクチャの決定指針となりました。

2. コンテキストウィンドウの拡張とアテンションマスキング

8Kから128Kトークンへのコンテキスト拡張は、約800Bトークンにわたる6つの継続事前学習ステージで実施されます。同一の長いコンテキストシーケンス内にパックされた複数ドキュメント間のクロスコンタミネーションを防ぐため、カスタムのドキュメントアテンションマスクが適用されます。

Context Parallelism（CP）は入力シーケンスを $2\times \text{CP}$ 個のチャンクに分割し、キー（$K$）およびバリュー（$V$）の表現に対してAll-Gather通信パターンを実行します。GQAによってクエリ（$Q$）に対して $K$ と $V$ のサイズが縮小されるため、All-Gatherの通信オーバーヘッド $O(S)$ は、アテンションの計算量 $O(S^2)$ と比較して無視できるほど小さくなります。

3. Direct Preference Optimization（DPO）の適応

選好最適化において、反復的な末尾の生成や時期尚早な停止を防ぐため、フォーマットトークン（ヘッダーおよび終了マーカー）はDPO損失から明示的にマスクされます。さらに、選択されたシーケンス（chosen）に対して補助的な負の対数尤度（NLL）正則化項が追加されます：

$${\mathcal{L}}_{\text{reg}}=0.2\cdot {\mathcal{L}}_{\text{NLL}}(\text{chosen})$$

結果（主要な数値を含む）

• フラッグシップ405Bモデルのベンチマーク:
  • MMLU（5-shot）: 87.3%（0-shot CoT: 88.6%）
  • GSM8K（8-shot, CoT）: 96.8%
  • MATH（0-shot, CoT）: 73.8%
  • HumanEval（0-shot）: 89.0%
  • ARC Challenge（0-shot）: 96.9%
  • IFEval: 88.6%
• 小型モデル（8Bおよび70B）: 両モデルとも、同規模のオープンソースの競合モデル（Gemma 2 9B、Mistral 7B、Mixtral 8x22Bなど）を上回る性能を示しました。例えば、Llama 3 8Bは0-shot CoT MMLUで73.0%、GSM8Kで84.5%を達成しています。
• ロングコンテキスト: 最大128Kトークンの単一Needle-In-A-Haystackで100%の完全な検索性能を達成し、405BモデルのMulti-needle再現率は98.1%を記録しました。
• 事前学習の効率性: 16,384基のH100 GPU上の構成全体でBF16 Model FLOPs Utilization（MFU）は38〜43%に達し、466回の中断が発生したにもかかわらず、54日間のスナップショット期間において90%以上の有効学習時間を維持しました。

限界点と未解決の課題

• ハードウェアの信頼性と電力変動: 16K GPU規模では、ハードウェアの中断が1日に平均数回発生しました（予期せぬ中断の78%はハードウェア/GPUの問題に起因）。チェックポイント保存時や集団通信ステップ時における数十メガワット規模の同時電力変動は、データセンターの電力網に大きな負荷を与えます。
• 敵対的頑健性のギャップ: 言い換えベンチマーク（PAWS）では耐性を示したものの、数学的推論およびQAの敵対的ベンチマーク（GSM-Plus、Adv SQuAD）では、非敵対的ベンチマークと比較して大幅な性能低下が見られました。
• マルチモーダルの提供状況: ビジョンおよび音声アダプタは研究環境において競合力のある能力を示しましたが、現在も活発に開発中であるため、初期の一般公開からは除外されました。