Itai Gat Heli Ben-Hamu Marton Havasi Daniel Haziza Jeremy Reizenstein Gabriel Synnaeve David Lopez-Paz Brian Karrer Yaron Lipman

要約

自己回帰型の次トークン予測言語モデルは強力な能力を有するが、特にデコード段階において推論にかかる計算コストおよびメモリコストが高いため、実用的な展開において大きな課題に直面している。本研究では、標準的な次トークン予測（NTP）とマスク付きトークン予測（MATP）を単一のアーキテクチャ内で統合することで、生成を高速化するシンプルかつ柔軟なパラダイム「Set Block Decoding（SBD）」を提案する。SBDは、モデルが複数の未来トークン（連続的でなくてもよい）を並列でサンプリング可能となる点が、従来の高速化手法との主要な差異である。この柔軟性により、離散拡散文脈から得られる高度なソルバーを活用でき、精度を損なうことなく大幅な高速化を実現する。SBDはアーキテクチャの変更や追加の学習ハイパーパラメータを必要とせず、正確なKVキャッシュとの互換性を維持しており、既存の次トークン予測モデルに対するファインチューニングによって実装可能である。Llama-3.1 8BおよびQwen-3 8Bをファインチューニングした結果、SBDを用いることで、同等のNTP学習と比較して生成に必要な前方伝搬（forward pass）回数を3〜5倍削減しつつ、同等の性能を達成できることを実証した。