デュアルフローTransformer：主プリフィル経路と追加デコード計算の分離

大規模言語モデル（LLM）が多くのリクエストを処理するにつれ、累積的な推論コストは一度きりの学習コストと比べて重要性を増している。論文は、推論の2つのフェーズがハードウェアに異なる負荷をかけると指摘する。プロンプトのプリフィル（prefill）は並列処理が可能で、一般に計算リソースに制約される。一方、自己回帰デコードは逐次処理であり、多くの場合メモリ帯域幅に制約される。従来の幅や深さのスケーリングは、追加した層が両方のフェーズで評価されるため、両方のコストを同時に増加させてしまう。

Liming Liu氏、Mingze Wang氏、Tuo Zhao氏は、この問題に対してDual-Flow Transformerを提案した。このアーキテクチャは、プライマリフロー（主経路）とオーグジリアリフロー（補助経路）で構成される。プライマリフローは完全な因果言語モデルであり、プロンプト全体を処理してKVキャッシュを書き込む。補助フローはプロンプト処理中はオフになり、最後のプロンプト位置以降でのみ活性化される。これにより、永続的な状態を書き込むことも、プライマリフローに影響を与えることもなく、続きの予測計算を追加できる。

2つのフローは主要なアテンション、MLP、出力行列を共有する一方、別々のトークン埋め込みと軽量な結合を持つ。この設計は、グループ実行時に読み込んだ重みや、プライマリキャッシュのキーと値を再利用する機会も生む。論文では、トークン数を揃えた比較で、Dual-Flowがさまざまなアーキテクチャとデータ構成において検証損失を低減できることを示した。

MoE（混合専門家）モデルでは、この分離により、プライマリと補助のエキスパートファンアウトを、プリフィルコスト・続きの生成コスト・予測品質の独立した制御手段にできる。実験では、プリフィルのエキスパート計算を固定してデコード計算を増やす設定と、固定のデコードエキスパート予算を2つのフローに再配分する設定を検討し、プリフィルとデコードの品質トレードオフを明らかにして、フェーズ固有のエキスパート配分の可能性を示した。

この研究は、LLMの推論コストが重要視される中で、単にすべてのフェーズに計算を均等に追加するのではなく、続きの生成時にだけ追加の学習済み計算を投入するという新たな設計方向を示している。長いコンテキストや高並列なサービス、エッジデバイス上での生成タスクなどにおいて、このようなフェーズ分離の考え方はモデル選択や推論最適化に影響を与える可能性がある。