MARS: Enabling Autoregressive Models Multi-Token Generation
Abstractの概要
MARS（Mask AutoRegression）は、指示チューニング済みの自己回帰型言語モデルに対し、標準的な左から右へのデコード動作を維持しつつ、1回のフォワードパスで複数トークンを予測可能にする軽量なファインチューニング手法である。本手法は追加パラメータやアーキテクチャの変更を必要とせず、既存の教師ありファインチューニングデータを用いて、マスク予測目的関数と自己回帰損失を組み合わせて学習する。著者らは自己回帰とブロックマスク予測の間に4つのギャップを特定し、そのうち3つ（アテンションパターン、ロジットの整列、生成順序）は排除可能な設計上の選択であり、トークンマスキングのみが本質的な差異であると主張している。Qwen2.5-0.5Bおよび7Bモデルを用いた6つのベンチマークでの実験により、同一チェックポイントが1トークンモードではベースラインと同等またはやや改善された品質で動作し、マルチトークンモードでは信頼度閾値によって制御されるより高いスループットを実現できることが示された。また、ブロックレベルのKVキャッシュ戦略が導入され、アルゴリズム的なトークン/フォワードパスの改善をバッチ推論における実時間の高速化に変換する。
新規性
主な新規性は、投機的ドラフトモデル、追加のデコードヘッド、アーキテクチャの変更を一切必要とせず、自己回帰とブロックマスク予測の間の3つの特定されたギャップを解消することで、ファインチューニングのみでマルチトークン生成を自己回帰モデルに追加できることを実証した点にある。また、より大きなブロックサイズにおいて自己回帰の性能を維持するための重要なメカニズムとして、クリーン入力ストリームに対する補助的なSFT損失を導入し、実用的なバッチ推論のためのブロックレベルKVキャッシュ方式を提案している。
成果
1トークンモードでは、MARSは0.5Bスケール（平均+1.7）および7Bスケール（平均+1.5）の両方で6つのベンチマークにおいてAR SFTベースラインと同等またはそれ以上の性能を達成した。マルチトークンモードではτ=0.95において、MARS-7Bは平均わずか1.3ポイントの低下で1フォワードパスあたり1.68トークンを生成し、ブロックレベルKVキャッシュを用いることでバッチサイズ4のQwen2.5-7Bにおいて最大1.71倍の実時間高速化を達成した。アブレーション実験により、補助的なSFT損失を除去するとブロックサイズの増加に伴い品質が悪化する（0.5Bでブロックサイズが4から16に増加すると平均が28.4から22.2に低下）一方、これを含めることでブロックサイズ全体にわたり性能が安定化することが確認された。
論文の注目点
- MARSは因果アテンション、右シフトロジット、左から右へのトークン受理を維持することで厳密な自己回帰互換性を保持し、自己回帰とブロックマスク予測の間の4つの特定されたギャップのうち3つを解消し、トークンマスキングのみを本質的な差異として残している。
- マスク予測とクリーンストリーム自己回帰SFT損失の組み合わせは、より大きなブロックサイズでの品質維持に経験的に不可欠である：これがない場合、0.5Bでの平均精度はブロックサイズが4から16に増加すると6.2ポイント低下するが、ある場合はわずか0.7ポイントの低下にとどまる。
- 推論は信頼度閾値による制御可能な速度・品質トレードオフをサポートし、ブロックレベルKVキャッシュはアルゴリズム的なトークン/フォワードパスの改善を実時間のバッチ高速化に変換するために必要であり、Qwen2.5-7Bのバッチサイズ4で最大1.71倍の高速化を達成した。
参考リンク
- arXiv: https://arxiv.org/abs/2604.07023v1
- Fugu-MT: https://fugumt.com/fugumt/paper_check/2604.07023v1
- Hugging Face Papers: https://huggingface.co/papers/2604.07023
- GitHub: https://github.com/Xalp/MARS