マルチトークン残差予測（Multi-token Residual Prediction, MRP）は、Modal ResearchとNYU上海校のHeavyBall Researchによる共同研究の成果であり、拡散言語モデル（DLM）の推論効率を最適化する革新的な手法です。本記事では、その設計原理、実験結果、および2つの推論シナリオにおける応用について詳述します。
MRPの核心は、マルチトークン予測（MTP）の改良にあります。自己回帰モデルでは、MTPは軽量ヘッドを用いて主幹ネットワークの隠れ状態から複数の後続トークンを予測し、投機的検証と組み合わせて高速化を実現します。しかし、この手法を拡散言語モデルに直接適用すると、複数ステップの予測時に精度が急激に低下する問題が発生しました。例えば、SDAR-4Bモデルでは、直接的なMTPが1ステップ目でGSM8K精度84.8%を達成したものの、4ステップ目では1.9%にまで落ち込みました。
研究チームは、完全な分布予測の難しさが原因であると特定し、代わりに隣接するデノイジングステップ間の残差（差分）を予測することを提案しました。拡散過程のマルコフ性により、各ステップで変更される位置は限られており、残差信号は本質的に低複雑性です。この洞察に基づき、MRPモジュールは3層の小型Transformerとして設計され、凍結されたDLM主幹に付加されます。訓練時には、トークンセットの開示前後で主幹を2回実行し、未開示位置におけるKLダイバージェンスを最小化することで残差を学習します。
実験では、MRPが静的および動的の両シナリオで優れた性能を示しました。静的モード（各ステップで固定数のトークンを開示）では、MRPを投機的デコードのドラフトモデルとして使用することで、無損失の高速化が可能です。SGLang実装において、SDAR-8BモデルはGSM8Kで90.4%の精度と1.40倍のスループット向上を達成しました。直接デコード（検証スキップ）では、速度向上は1.89倍に達し、精度低下はわずかです。動的モード（信頼度に基づくバッチ開示）では、低閾値による過剰開示で失われた精度を回復できます。例えば、閾値τ=0.5において、SDAR-1.7BモデルのGSM8K精度は41.6%から59.1%（+17.5ポイント）に改善しました。
MRPの柔軟性により、ユーザーはアプリケーション要件に応じて、無損失高速化と有損高速化の間で自由に選択できます。この制御は、独自の推論スタックを持つ開発者にとって特に重要です。研究チームはSDAR-1.7B、4B、8Bモデルで広範な評価を行い、GSM8K、MATH500、HumanEval、MBPPなどのベンチマークで顕著な改善を確認しました。コードとSGLang実装は公開されています。