Uno、離散拡散で大規模言語モデルを無損失高速化

導入

大規模言語モデルの多くは、次のトークンを順番に予測する自回帰方式で動作する。この方式は学習と品質の面で扱いやすい一方、生成時には前のトークンが確定するまで次の計算に進めない。そのため、GPUの並列計算能力を十分に活かしにくい。今回紹介するUnoは、離散拡散の考え方を取り入れ、この逐次的なボトルネックを緩和しようとする手法だ。

Unoの仕組み

Unoは、通常の拡散言語モデルのようにARモデルを全面的に置き換えるのではない。自回帰モデルを基盤として残し、並列生成を担当する軽量な拡散パラメータを追加する。主な流れは次の通りだ。

• AR重みを維持：基盤モデルは標準的な次トークン予測目的で学習され、最終的な出力分布と検証を担う。
• 拡散経路を追加：Diffusion Distillationの段階で、複数トークンを同時に提案できるよう追加重みを学習する。
• 下書きと検証を分離：拡散経路が候補列を並列に作り、AR経路がその候補を検証する。受け入れられた候補が多ければ、逐次ステップを減らせる。
• 別のドラフトモデルが不要：推測デコードのように、専用の小型モデルを別途配備する必要がない。

論文はさらにΨ-Specというサンプラー群を提案している。これはコンテキスト長を固定したまま、無損失の高速化や推論時のスケーリングを可能にすることを狙ったものだ。要約によれば、Unoは評価された各バッチサイズで比較対象の推測デコード手法を上回るスループットを示し、基盤ARモデルに対して最大約3倍の高速化を達成したという。また、8BのUnoが、比較対象となったオープンな拡散モデルDiffusionGemmaを上回ったとも説明されている。ただし、提供された要約は後半が途切れているため、比較条件の詳細は論文本文で確認する必要がある。

「無損失」の意味

ここでいう無損失とは、拡散経路の予測が常に正しいという意味ではない。最終的な基準はあくまで元のARモデルであり、拡散経路は複数候補を効率よく作る役割を担う。AR側が候補を検証することで、基盤モデルの出力分布を維持しながら、1回の検証で複数位置を進められる可能性が生まれる。実際の速度向上は、候補の受理率、検証コスト、バッチサイズ、ハードウェアの並列効率に左右される。

意義と論点

Unoは、推測デコードと完全な拡散型LLMの中間に位置する。ARモデルの品質管理を残しつつ、拡散的な並列候補生成を利用するため、品質と速度を別々の経路に割り当てる設計だ。既存のオープンウェイトARモデルへの追加、またはゼロからの学習に対応すると論文は説明している。

一方、Hugging Face上では新規性をめぐる議論も起きている。コメントでは、先行するOrthrusも、ARバックボーンを固定し、学習可能な拡散経路を追加し、KVキャッシュを再利用し、並列に下書きしてARで検証するという大枠を採用していると指摘された。Uno側は、Unoが元の注意機構を維持し、LoRAなどの軽量な追加を使うのに対し、Orthrusは拡散用の注意ヘッドを追加し、拡散ブロックで双方向注意を使う点が異なると説明している。

この優先権の議論とは別に、UnoはLLM推論の現実的な方向性を示す。ARか拡散かを選ぶのではなく、ARに品質保証を任せ、補助的な拡散経路に並列化を任せる発想だ。今後、異なるモデルや負荷条件で報告された効果が再現されれば、混合型デコーディングは推論基盤の重要な選択肢になり得る。

出典：Hugging Face Daily Papers