現代AIの計算上の綱渡り
大規模言語モデル（LLMs）は、現代の人工知能の柱として存在し、産業や科学的発見を再構築する驚くべき能力を示しています。人間のようなテキストを生成し、洗練された対話エージェントを動かし、複雑な研究タスクさえ支援するその熟練度は、それらを不可欠なツールにしています。これらの強力なモデルの中心には、Transformerアーキテクチャが鼓動しています。これは、交互に配置された層によって特徴づけられる設計です。トークンに分割された入力データは、異なるトークンの重要性を重み付けする一連のアテンションメカニズムを流れ、次に、収集された情報を処理するフィードフォワードネットワーク（FFNs）を通過します。この階層的で逐次的な処理は、Transformerが学習し、出力を生成する方法の基本です。
しかし、このアーキテクチャ自体は効果的である一方で、モデルのサイズと複雑さが増大するにつれて、増大する課題を提示しています。逐次的な性質は、各層が一般的に前の層が計算を完了するのを待たなければ、自身の計算を開始できないことを意味します。この段階的な処理は、特に推論フェーズ（訓練済みモデルが実際に予測やテキストを生成するために使用される段階）において、固有のボトルネックを生み出します。高度なAIアシスタントを動かすようなモデルが、数百億、あるいは数兆のパラメータを組み込むようになると、推論に必要な計算リソースと時間は劇的に増加します。この需要の増大は、著しいレイテンシ（応答の遅延）、スループット（時間あたりに処理されるリクエスト数）の低下、および運用コストの増加につながり、最も強力なLLMの広範な展開とリアルタイムアプリケーションを妨げています。その結果、推論効率の向上はAI研究コミュニティ内で最重要課題となり、これらのモデルが提供する驚くべき性能を損なうことなく計算を合理化できる革新的な戦略の探求を促しています。中心的な課題は、特に計算が複数のGPUにまたがり、処理時間に通信オーバーヘッドが加わる分散環境において、逐次実行によって課せられる制約を緩和することにあります。
最適化ランドスケープのナビゲーション：既存ツールとその限界
LLMをよりスリムで高速にするための継続的な取り組みの中で、研究者たちは最適化技術のツールキットを開発してきました。それぞれが効率化への道を提供しますが、しばしば独自の妥協点を伴い、単一の方法が普遍的な解決策となることを妨げています。これらのトレードオフを理解することは、FFN Fusionのような新しいアプローチの必要性を理解する上で重要です。
著名な技術の一つは**量子化（quantization）**です。これは、モデルの重みと活性化を表現するために使用される数値精度を低減することを含みます。標準的な32ビット浮動小数点数を使用する代わりに、モデルは16ビット、8ビット、あるいはそれ以下のビット表現を使用するかもしれません。これにより、モデルのメモリフットプリントが直接縮小され、低精度数値での演算は通常より高速でエネルギー消費も少ないため、計算が大幅に高速化される可能性があります。しかし、量子化にはリスクがないわけではありません。精度を下げると情報が失われ、モデルの精度が低下する可能性があります。このリスクは、非常に低いビット幅でより顕著になり、精度の低下を緩和するために慎重な実装と、時には再トレーニングが必要になります。課題は、性能低下を許容範囲内に抑えながら効率向上を最大化するスイートスポットを見つけることにあります。
もう一つの一般的な戦略は**枝刈り（pruning）**です。この技術は、大規模なニューラルネットワーク内の多くのパラメータが冗長であるか、最終出力への寄与が最小限である可能性があるという原則に基づいて動作します。枝刈りアルゴリズムは、これらの重要度の低い接続やニューロンを特定して削除し、結果としてより小さく、よりスパースなモデルになります。量子化と同様に、枝刈りはメモリ要件と計算負荷を削減します。しかし、どのパラメータを「安全に」削除できるかを正確に特定することは複雑です。積極的な枝刈りは、意図せずに重要なコンポーネントを削除し、大幅な精度低下につながる可能性があります。枝刈り後にモデルをファインチューニングすることは、性能を回復するためにしばしば必要であり、ワークフローに複雑さを加えます。枝刈りされたモデルが効果的であり続けることを保証するために、慎重な調整が不可欠です。
よりアーキテクチャ的に異なるアプローチは、**Mixture-of-Experts（MoE）**モデルです。すべての入力をネットワーク全体で処理する代わりに、MoEモデルは複数の「エキスパート」サブネットワーク（通常はFFN）で構成されます。各入力トークンに対して、ゲーティングメカニズムがこれらのエキスパートの小さなサブセットを動的に選択して計算を実行します。この条件付き計算は、特定の入力に対してモデルの総パラメータのごく一部しかアクティブにならないことを意味し、特に非常に大規模なモデルのトレーニングと推論において、大幅な計算節約につながります。MoEモデルは、合理的な計算コストを維持しながら、数兆のパラメータにスケールアップできます。しかし、その効率はワークロードに大きく依存します。選択的なアクティベーションパターンが良好なハードウェア利用率につながる非常に大きなバッチサイズを処理するのに優れています。より小さいまたは中間のバッチサイズでは、MoEモデルは計算リソースの利用不足に悩まされる可能性があります。なぜなら、並列ハードウェアがスパースにアクティブ化されたエキスパートによって一貫してビジー状態に保たれない可能性があるためです。さらに、MoEモデルの実装と負荷分散は、標準的な「密な」アーキテクチャの展開よりも複雑になる可能性があります。
量子化、枝刈り、およびMoEモデルはLLM最適化における価値ある進歩を表していますが、それらの固有の限界は、代替または補完的な戦略の必要性を浮き彫りにしています。特に、トレーニングと展開の相対的な単純さから依然として人気のある密なモデルアーキテクチャに対して、精度や実装の複雑さへの妥協が少なく、さまざまなシナリオで広範な効率改善を提供できる方法の探求は続いています。
FFN Fusion：Transformerにおける並列処理の再考
この最適化技術のランドスケープの中で、NVIDIAの研究者たちはFFN Fusionと呼ばれる魅力的な新しいアプローチを導入しました。この技術は、パラメータを変更したり、部分を選択的にアクティブ化したりするのではなく、計算シーケンスをどのように並列化できるかを根本的に再考することによって、Transformerアーキテクチャ固有の逐次ボトルネックに直接立ち向かいます。このイノベーションは、ディープTransformerモデル内のFFN層の振る舞いに関する重要な観察から生まれました。
Puzzleという診断ツールを使用して、研究者たちは大規模モデルの内部動作を分析しました。実験的にアテンション層を削除したところ、モデルが驚くほど長い連続したFFN層のシーケンスを保持することが多いことに気づきました。さらに重要なことに、分析により、これらの隣接するFFNによって実行される計算が、しばしば最小限の相互依存性を示すことが明らかになりました。本質的に、シーケンス内のあるFFNの出力は、直後のFFNが必要とする方向パスやコア情報を大幅に変更しないことが多かったのです。これは、従来は次々に実行されていたこれらのFFNが、モデルの全体的な機能を大幅に妨げることなく、同時並行で実行される可能性を秘めていることを示唆していました。
この洞察がFFN Fusionの基盤を形成しました。中心的なアイデアは、エレガントにシンプルでありながら強力です。低い計算依存性を持つ連続したFFN層のシーケンスを特定し、それらを単一の、より幅の広いFFN層に統合して、同等の計算を並列に実行します。Input -> FFN1 -> FFN2 -> FFN3 -> Outputのような連鎖の代わりに、融合された構造はInput -> Fused_FFN (FFN1+FFN2+FFN3と並列に同等) -> Outputとなります。このアーキテクチャ変換は、ネットワークの逐次的な深さを効果的に短縮し、複数のステップを単一の、より広範な計算ステップに置き換えます。これらの低依存性のFFNシーケンスをターゲットにすることで、FFN Fusionは、モデルの表現力と精度を維持しながら、レイテンシと計算コストを削減することを目指しています。Llama-3.1-405B-InstructからUltra-253B-Baseを開発したことは、この技術の可能性を示す主要なデモンストレーションとなりました。
アーキテクチャの錬金術：FFN Fusionの仕組み
FFN Fusionの背後にある魔法は、フィードフォワードネットワークの基礎となる数学的構造の巧妙な操作にあります。それは単に既存のレイヤーを並べて実行することではありません。元のシーケンスの集合的な振る舞いを複製するが、それを同時に行う新しい、統一されたレイヤーを作成することを含みます。
k個の連続したFFN層のシーケンスを考えてみましょう。標準的なTransformerでは、入力xはFFN1を通過し、その出力がFFN2の入力となり、FFNkまで続きます。各ステップは、前のステップの完了に明示的に依存しています。FFN Fusionはこの依存関係の連鎖を断ち切ります。数学的には、FFNは通常、間に非線形活性化関数（GeLUやSwiGLUなど）を持つ2つの線形変換を含みます：FFN(x) = W_out * Activation(W_in * x)。FFN Fusionは、線形変換がしばしば組み合わせられるという事実を活用します。
融合プロセスは、個々のFFN層の重みを連結することによって機能します。具体的には、連続するFFNの入力重み行列（W_in）が、融合層のための単一の、より大きな入力重み行列に（例えば、ブロック対角的に）結合されます。同様に、出力重み行列（W_out）が連結されて、単一の、より幅の広い出力重み行列を形成します。活性化関数は、このより大きな構造内で要素ごとに適用されます。この構造により、融合されたFFNは、元のFFNに対応する並列パス全体で、元の入力xに対して同時に動作することが保証されます。これらの並列パスからの出力は、連結された出力重みの構造によって暗黙的に集約されます。
理論的な裏付けは、元の層間の依存関係が実際に低かった場合、この融合構造が元のFFNシーケンスと同じ表現能力を維持できることを確認しています。鍵となるのは、どのシーケンスが融合に適しているかを特定することです。これを体系的に行うために、NVIDIAの研究者たちは依存性分析技術を採用しました。彼らは、代表的な入力トークンのセットに対して、連続するFFN層の出力隠れ状態間のコサイン距離を測定しました。コサイン距離が小さいことは、あるFFNの出力ベクトルが、シーケンス内の次のFFNの出力ベクトルと非常に類似した方向を指していることを示します。この類似性は、低い機能的依存性を示唆しています。つまり、2番目のFFNは、最初のFFNによって確立された情報表現を大幅に変更していません。層全体で一貫して低いコサイン距離を示すFFNのシーケンスは、融合の主要な候補として特定されました。なぜなら、それらをマージしても、モデルの学習済み表現や全体的なパフォーマンスを妨げる可能性が低いからです。このデータ駆動型アプローチにより、FFN Fusionを、最も効果的で最も破壊的でないモデルの部分にターゲットを絞って適用することが可能になります。
巨大モデルからスプリンターへ：Ultra-253B-Baseの変革
FFN Fusionの実用的な力は、当時公開されていた最大級のモデルの一つであるLlama-3.1-405B-Instructへの適用を通じて鮮やかに実証されました。4050億のパラメータを誇るこのモデルは、推論において重大な計算上の取り組みを必要としました。研究者たちは、FFN Fusionと戦略的な枝刈りを組み合わせてアーキテクチャを洗練するプロセスに着手し、Ultra-253B-Baseと名付けられた新しい、より効率的なモデルを作成しました。
変革プロセスにはいくつかのステップが含まれていました：
- 分析： 依存性分析ツール（コサイン距離の測定）を使用して、研究者たちはLlama-405Bアーキテクチャ内で層間の依存性が低い連続したFFN層のシーケンスを特定しました。
- 融合： 特定されたこれらのFFNシーケンスは、前述のように（重みを連結して）単一の、より幅の広いFFN層に融合されました。これにより、ネットワーク内の逐次ステップ数が直接削減されました。
- 枝刈り： 同時またはその後に、重要度が低いと見なされたパラメータ（標準的な枝刈り技術を通じて、または融合プロセスによって情報提供されて特定された可能性がある）がモデルから削除されました。
この組み合わせアプローチにより、2530億パラメータを持つモデルであるUltra-253B-Baseが生まれました。これは、元の405Bモデルよりも37%以上少ないパラメータという大幅な削減を表しています。融合によって達成されたアーキテクチャの変更は、性能を維持することを目指しながら、このような大幅なサイズ削減を可能にする鍵でした。目標は、単に小さいモデルではなく、FFN Fusionによって解放された並列性の向上のおかげで、根本的により高速で計算コストが低いモデルでした。このケーススタディは、大規模モデルが効率のために大幅に再構築できることを示す重要な概念実証として機能しました。
利益の測定：パフォーマンス、速度、リソース節約
あらゆる最適化技術の真のテストは、その測定可能な影響にあります。Ultra-253B-Baseの場合、Llama-405BベースにFFN Fusionと枝刈りを適用した結果は説得力があり、能力に大きな妥協をすることなく、複数の次元にわたって大幅な改善を示しました。
推論速度とコスト： 最も顕著な利益は推論効率で観察されました。元の405Bパラメータモデルと比較して、Ultra-253B-Baseは以下を達成しました：
- 推論レイテンシが1.71倍改善。これは、モデルが応答を大幅に高速に生成できることを意味し、リアルタイムアプリケーションにとって重要です。
- バッチサイズ32で測定した場合、トークンあたりの計算コストが35倍削減。トークンあたりの計算操作（FLOPs）のこの劇的な減少は、モデルを提供するためのエネルギー消費の削減とハードウェア要件の削減に直接つながります。
モデルパフォーマンスベンチマーク： 重要なことに、これらの効率改善は、モデルの知能や能力を犠牲にして達成されたわけではありません。Ultra-253B-Baseは、標準的なLLMベンチマークのスイートで厳密に評価され、元の、はるかに大きなモデルに非常に匹敵し、場合によってはそれを上回るスコアを達成しました：
- MMLU (Massive Multitask Language Understanding): 85.17%
- MMLU-Pro (より挑戦的なバージョン): 72.25%
- Arena Hard (困難なプロンプトに対する人間の嗜好評価): 84.92%
- HumanEval (コード生成能力): 86.58%
- MT-Bench (マルチターン会話品質): 9.19
これらのスコアは、融合され枝刈りされたモデルが、わずか2530億のパラメータしか持たないにもかかわらず、その405Bパラメータの祖先と同等の、非常に高いレベルの理解、推論、コーディング能力、および会話品質を保持していることを示しています。
メモリ効率： 計算速度とコストを超えて、FFN Fusionはメモリ節約にも貢献しました。アーキテクチャの変更は、おそらく融合によって可能になった他の最適化と組み合わされて、推論中に必要なキー・バリュー（KV）キャッシュのサイズを2倍削減することにつながりました。KVキャッシュは中間的な活性化（アテンションのキーとバリュー）を格納し、特に長い入力シーケンスの場合、かなりのGPUメモリを消費する可能性があります。この要件を半減させることで、メモリ集約度の低いハードウェアでモデルを実行したり、同じメモリ制約内でより長いコンテキストを処理したりすることが可能になります。
これらの定量化可能な結果は、FFN Fusionの有効性を強調しています。それは、より小さいだけでなく、速度、計算操作、メモリ使用量の点で根本的により効率的なモデルの作成を可能にし、同時に挑戦的なベンチマークでトップクラスのパフォーマンスを維持しました。
知識の保存：トレーニングとファインチューニングの重要な役割
FFN Fusionや枝刈りのような技術を通じて、Llama-405Bのような巨大な事前訓練済み言語モデルをアーキテクチャ的に変更することは、必然的にその学習済みパラメータの微妙なバランスを崩します。数学的な等価性は局所的に機能を保存することを目指しますが、ネットワークの全体的な振る舞いは変化する可能性があります。結果として得られるUltra-253B-Baseモデルが、より効率的になるだけでなく、その高いパフォーマンスレベルを維持することを保証するために、慎重に計画された修正後のトレーニングプロセスが不可欠でした。
このプロセスには、主に2つのフェーズが含まれていました：
知識蒸留（Knowledge Distillation）： 最初のステップは、元の、より大きなモデル（または適切な教師モデル）からの知識を、変更されたアーキテクチャに転送することでした。これは蒸留を通じて達成され、Ultra-253B-Baseモデルは教師モデルの出力または内部表現を模倣するように訓練されました。このフェーズでは、具体的には540億トークンという実質的なデータセットが、8kコンテキストウィンドウで処理されました。蒸留は、融合され枝刈りされたモデルが、アーキテクチャの変更中にわずかに乱された可能性のあるニュアンスや能力を取り戻すのに役立ちます。
段階的ファインチューニング（Staged Fine-Tuning）： 蒸留に続いて、モデルは、徐々に長いコンテキスト長を処理するように特別に設計された一連のファインチューニング段階を経ました。これは、しばしば広範な入力に基づいてテキストを処理および生成することが期待される現代のLLMにとって重要です。ファインチューニングは段階的に進められました：
- 16kコンテキストウィンドウでのファインチューニング。
- 32kコンテキストウィンドウでのさらなるファインチューニング。
- 128kコンテキストウィンドウでの最終ファインチューニング段階。
この段階的なアプローチにより、モデルは、新しく形成された融合FFN層や最適化されたKVキャッシュメカニズムを含むパラメータを、非常に長いシーケンスにわたる依存関係と情報フローを効果的に管理するように徐々に適応させることができます。各段階は前の段階に基づいて構築され、異なるコンテキストサイズ全体で安定性と堅牢なパフォーマンスを保証します。
大規模な蒸留と段階的なロングコンテキストファインチューニングを組み合わせたこの綿密なトレーニングレジメンは、アーキテクチャの効率と高忠実度のパフォーマンスとの間のギャップを埋める上で役立ちました。これにより、FFN Fusionによってもたらされた速度、コスト、メモリの利点が、要求の厳しいベンチマークでのモデルの精度と能力を損なわないことが保証されました。
より広い視野：一般化可能性と将来の方向性
Llama-405BからUltra-253B-Baseへの成功した変革は、FFN Fusionの可能性に対する強力な証拠を提供しますが、その真の価値は、より広範な適用可能性と、将来のLLM設計に対する洞察にあります。この研究は、これが巨大なモデルにのみ適用可能な一回限りのトリックではないことを示しました。
スケールを超えた検証： NVIDIAの研究者たちは、FFN Fusionの方法論をさまざまなサイズのモデルで明示的にテストしました。彼らはこの技術を70Bパラメータモデルに成功裏に適用し、元の対応物と比較して同様の効率向上を達成しました。彼らはまた、49Bスケールでの検証も報告しており、FFNの独立性と融合の可能性が最大規模のモデルだけの排他的な特性ではなく、Transformerアーキテクチャのより一般的な特性である可能性があり、より深いFFNシーケンスが自然に発生する大規模スケールでより顕著になる可能性があるという考えをさらに強化しています。これは、FFN FusionがLLM最適化の武器庫における標準的なツールとなり、さまざまなモデルサイズに適用可能であることを示唆しています。
FFN対フルブロック融合： この研究はまた、Transformerブロック内のアテンション層と比較して、FFN層の特定の役割に光を当てました。連続するFFN層はしばしば低い依存性を示し、融合に理想的でしたが、Transformerブロック全体（アテンション層とFFN層の両方を含む）を並列化しようとする試みは、より困難であることが証明されました。分析は、アテンションメカニズムを含むより強い相互依存性を示しました。ブロック全体を同時に融合すると、より重大なパフォーマンス低下が生じ、アテンション層がトークン間で情報を統合する上で、より重要で逐次的に依存する役割を果たしていることを示唆しています。この発見は、効果的な並列化の境界を明確にするのに役立ちます。FFNシーケンスは肥沃な土壌ですが、アテンションメカニズムは異なる最適化戦略を必要とするかもしれません。
LLMアーキテクチャへの影響： FFN Fusionは、単なる事後的な最適化技術以上のものを提供します。将来のLLMを設計するための貴重な洞察を提供します。FFNのシーケンスがしばしば並列化可能なユニットとして扱えるという発見は、Transformer設計の基礎となることが多い厳密な逐次仮定に挑戦します。これは、最初から本質的により並列フレンドリーな新しいアーキテクチャを刺激する可能性があります。将来のモデルは、融合または並列実行を明示的に意図したFFN構造で設計される可能性があり、GPUアーキテクチャがこの種の並列処理を活用するためにさらに最適化されるハードウェア・ソフトウェア協調設計につながる可能性があります。コサイン距離を使用して層間の依存関係を定量化する体系的な方法は、ニューラルネットワーク構造を理解し、再設計するための貴重な分析ツールも提供します。既存のコンポーネントを並列化することに焦点を当てた思慮深いアーキテクチャの再設計を通じて、大幅な効率向上が可能であることを示すことにより、FFN Fusionは、強力でありながら計算的により持続可能なLLMを開発するための道を開きます。それは、最先端AIの増大するリソース需要を緩和するための道筋を強調しています。