プライマリパスとデコード処理を分離する Dual-Flow Transformers──推論効率を維持しつつ継続予測の計算量を拡大

論文概要

Large Language Models (LLM) の推論コスト削減を目的として、プロンプト処理（prefill）と自己回帰デコード（decode）の計算特性を分離する新しいアーキテクチャ「Dual-Flow Transformer」が提案された。従来の Transformer は深さや幅の拡張によって両方のフェーズの計算量が同時に増加するが、本手法ではプライマリフローがプロンプトを処理して単一の永続的 KV キャッシュを書き込み、補助フローは最終プロンプト位置以降でのみ有効化される設計を採用している。これにより、KV キャッシュを追加で生成することなく継続予測の計算量を増加させ、複数のアーキテクチャおよびデータ構成においてバリデーションロスの低減を実現している。

Figure 1: 標準的なTransformerとDual-Flowのバリデーションロスの比較。

新規性と貢献

本研究の最大の貢献は、プリフィルパスとデコード追加計算を分離（decouple）するという設計思想にある。具体的には以下の点が新しい。

プライマリフローが生成する単一の永続的 KV キャッシュを共有しつつ、補助フローによって継続予測の演算のみを選択的に追加する構造。
2つのフロー間で主要な attention、MLP、出力マトリクスを共有し、トークン埋め込みと軽量な結合のみを分けることで、グループ実行時の重みやキャッシュの再利用を可能にした点。
MoE モデルにおいて、プライマリと補助のエキスパート数を独立したコントロールとして扱い、プリフィルコストと継続コスト、予測品質のトレードオフを柔軟に調整できる点。

提案手法の詳細

Dual-Flow Transformer は、完全な因果言語モデルとして動作する「プライマリフロー」と、プロンプト処理時にはスキップされ最終位置から起動する「補助フロー」の2つで構成される。

なぜその設計にしたか: プリフィルは演算器がボトルネックになりやすく、デコードはメモリ帯域がボトルネックになるというハードウェア的特性の差異を活かすためである。デコード時にのみ追加の演算器を活用し、KV キャッシュのサイズを増やさないことでメモリ帯域への負荷を抑えながら表現力を高める狙いがある。
重みの共有と軽量カップリング: アテンション機構や MLP のマトリクスを共有しつつ、トークン埋め込み層を分離することで、ハードウェアの効率的なメモリローディングと演算のグループ化を両立している。

Figure 2: コンポーネントアブレーションにおける各構成のバリデーションロス比較。

評価・考察

論文では、一致させたトークン数（matched-token）の比較において、Dual-Flow が従来の構造と比較して一貫して低いバリデーションロスを達成することを示している。

MoE レジームにおける評価: プリフィルのエキスパート計算量を固定した状態でデコード計算量を増加させる実験や、固定されたデコードエキスパート予算を2つのフロー間で再配分する実験を実施。
トレードオフの確認: プリフィル・デコード・品質の間のトレードオフを詳細に分析し、フェーズ固有のエキスパート割り当てが予測性能の向上に寄与することを実証している。

Figure 3: 固定プリフィルスライスにおける補助エキスパート数とバリデーションロスの関係。

応用例と今後の展望

本手法は、API サービングなどで大量のリクエストを処理する大規模言語モデルの推論インフラにおいて実務的なインパクトを持つ。特に、コンテキスト長が長く、デコードのレイテンシとスループットが厳しく問われるクラウド上の LLM ホスティング事業者（数千〜数万リクエスト/秒規模のインフラ運用企業）にとって、KV キャッシュの肥大化を防ぎながらモデルの表現力を引き上げる有効な選択肢となる。今後は、より多様な MoE 構成や実ハードウェアにおけるレイテンシの実測評価が期待される。

結論

Dual-Flow Transformer は、プリフィルとデコードのフェーズ特性の違いを利用して計算量を効果的に分離・再配分する新しいアーキテクチャである。KV キャッシュの増大を招かずに継続予測の品質を高められる点を示し、今後の効率的な推論基盤設計に新しい方向性を提示した。

注釈

Prefill（プリフィル）: 入力プロンプトを並列処理して KV キャッシュを生成するフェーズ。
Decode（デコード）: 過去の KV キャッシュを参照しながらトークンを1つずつ自己回帰的に生成するフェーズ。