AIチップ大手のNVIDIAの研究者らが最近、「FFN Fusion（FFN融合）」という革新的なアーキテクチャ最適化技術を発表しました。この技術は、Transformerアーキテクチャに固有の逐次計算のボトルネックを解消することで、大規模言語モデル（LLM）の推論効率を大幅に向上させ、高性能AIアプリケーションのより広範な展開を促進することを目的としています。
近年、大規模言語モデルは自然言語処理、科学研究、対話型エージェントなどの分野で強力な能力を示しています。しかし、モデルの規模と複雑さが増すにつれて、推論プロセスに必要な計算資源も大幅に増加し、効率のボトルネックが生じています。TransformerアーキテクチャはLLMの基礎であり、交互の注意機構とフィードフォワードネットワーク（FFN）層は順序どおりに入力を処理する必要があります。この固有の逐次構造は、モデル規模の拡大に伴い、計算とGPU間の通信コストを大幅に増加させ、効率を低下させ、展開コストを高めます。特に、複数のトークンを迅速に生成する必要があるシナリオ（リアルタイムAIアシスタントなど）では、この問題はさらに顕著になります。
この課題に対処するため、NVIDIAの研究者らはFFN融合技術を提案しました。この手法の中核となる考え方は、モデル内で連続し、相互依存性が低いFFN層を、より幅広い単一のFFNにマージすることです。研究者らは、注意機構層を除去した後、LLMには通常、長い連続したFFNシーケンスが存在することに気づきました。これらのシーケンスを分析することで、これらのFFN層間の依存性が低いことがわかり、並列実行が可能であることがわかりました。
FFN融合の数学的基礎は、複数の直列接続されたFFNの重みを連結して、並列計算可能な等価な単一モジュールを作成することです。例えば、3つのFFNが順次積み重ねられ、各FFNの出力が次のFFNの入力となる場合、FFN融合はこの依存関係を解消し、3つのFFNが同時に同じ入力を処理し、その出力を集約できるようにします。理論的な分析によると、融合後のFFNは元のFFNと同じ表現能力を維持します。
Ultra-253B-Base：性能と効率の両面での向上
NVIDIAの研究者らは、FFN融合技術をMetaのLlama-3.1-405B-Instructモデルに適用し、剪定と再構成によってUltra-253B-Baseという新しいモデルを作成しました。実験結果によると、Ultra-253B-Baseは推論速度と資源効率の面で顕著な向上を示しました。具体的には、バッチサイズが32の場合、推論遅延が1.71倍減少、トークンあたりの計算コストが35倍減少しました。
さらに印象的なのは、効率の向上がモデルの能力を犠牲にすることなく達成されたことです。Ultra-253B-Baseは、MMLU85.17％、MMLU-Pro72.25％、HumanEval86.58％、Arena Hard84.92％、MT-Bench9.19％など、複数の権威ある評価ベンチマークで優れた成績を収めました。これらの結果は、通常、元の4050億パラメーターのモデルと同等かそれ以上であり、Ultra-253B-Baseは2530億パラメーターしか含まれていません。さらに、このモデルのメモリ使用量も半分に削減されました。これはkv-cacheの最適化によるものです。
研究者らは、FFN層間の出力をコサイン類似度で分析して、相互依存性が低い領域を特定し、これらの領域を融合の最適な対象としています。FFN融合技術は、490億、700億、2530億パラメーターを含むさまざまな規模のモデルで検証されており、優れた汎用性があることを示しています。
この研究は、詳細な分析と巧妙なアーキテクチャ設計によって、LLMの効率を大幅に向上させることができることを示しています。FFN融合は、より並列化され、ハードウェアに適したLLMを設計するための基盤を築きました。完全なTransformerモジュールの並列化は、より強い層間の依存性のために多くの課題に直面していますが、FFN融合の成功は、将来のLLMの効率最適化のための重要な方向性を示しています。
論文:https://arxiv.org/abs/2503.18908
