Wavefront Decoding、ループ型言語モデルのデコードを並列化

導入：ループ構造が生むデコード遅延

ループ型言語モデルは、同じ重み共有ブロックを何度も適用することで、パラメータ数を大きく増やさずに実効的な計算深度を高めます。しかし生成時には、1トークンを出力するたびに再帰ブロックを複数回、順番に呼び出す必要があります。モデル内部の重み共有はパラメータ効率に貢献する一方、推論の遅延という別の課題を生みます。

Wavefront Decoding（WFD）は、この課題に対してモデルの再学習や独立したドラフターモデルの追加を行わず、デコードのスケジュールを組み替えます。中間的な再帰状態を活用し、複数のトークン位置と深度を同じ計算の流れに載せる手法です。

核心：ドラフトと検証を同時に進める

一般的な投機的デコードでは、まずドラフターが候補を作り、その後ターゲットモデルがまとめて検証します。WFDは、ループ型モデルの次の性質を使ってこの段階分離を崩します。

• 中間状態をドラフトに利用する。 全再帰を終えていない浅い状態でも次トークン予測を生成できるため、新しい位置の候補作成に使えます。
• 重み共有をバッチ化に活用する。 位置や再帰深度が異なる状態でも、同じブロックを通るため、1回のバッチ呼び出しで処理できます。
• 対角線状の波面を作る。 新しい位置は浅い深度から開始し、先行する位置はより深い再帰へ進め、全深度の検証に到達させます。
• 拒否された候補を修正する。 浅いドラフトが検証を通らない場合は、全深度の予測で結果を補正します。

つまりWFDは、単に多くの状態を同時に計算するだけではありません。系列位置と再帰深度にまたがる依存関係を波面として再配置し、浅いドラフトと深い検証を連続する再帰呼び出しの中で重ね合わせます。

結果とKV共有

論文の概要によると、6種類のSpec-Benchタスクで評価した結果、自回帰デコードに対してOuro-2.6Bでは2.42倍、Huginn-3.5Bでは3.54倍の高速化を達成しました。また、ドラフトを作ってから検証する従来型スケジュールを一貫して上回ったとされています。

さらに、波面処理で発生するKVトラフィックを減らすため、異なる再帰段階をまたぐKV共有も導入されました。この最適化により、Huginn-3.5Bで報告された高速化は最大4.81倍に達します。ただし、これらは対象モデル、ベンチマーク、実装条件に基づく結果であり、すべてのモデルやハードウェアで同じ効果が得られることを意味しません。コードは公開されており、再現検証が可能です。

意義

WFDは、ループ型モデルの反復計算を単なる逐次コストではなく、スケジュール可能な推論資源として扱います。中間状態の予測力、重み共有による混在バッチ処理、そしてKV管理を組み合わせることで、別のドラフターモデルを学習せずにデコードを高速化できる可能性を示しました。

一方で、効果は中間予測の品質、波面スケジューラの実装、KVキャッシュの帯域効率に左右されます。それでも、モデルの深度方向をどのように推論システムの並列性へ変換するかという、ループ型アーキテクチャに固有の重要な設計課題を明確にしています。

出典：Hugging Face Daily Papers