$S^3$: 拡散言語モデルにおける推論時（テスト時）向け層化スケーリング探索
arXiv cs.LG / 2026/4/9
💬 オピニオンSignals & Early TrendsIdeas & Deep AnalysisModels & Research
要点
- 本論文では、拡散言語モデルに対するテスト時スケーリング手法としてS^3（Stratified Scaling Search）を提案し、推論計算を復元（denoising）の過程で再配分することで、末尾でbest-of-Kのみを使う場合に比べて出力品質を向上させる。
- 各ステップで複数の候補となる復元軌道を展開し、それらを軽量な参照不要（reference-free）検証器でスコアリングする。そして多様性を探索フロンティアに保ちながら、有望な候補を選択的に再サンプリングする。
- 本手法は、報酬に傾いたサンプリング分布を近似することで、高品質な出力が得られる確率を高めつつ、元のモデルの事前分布（prior）に基づいた制約を維持する。
- LLaDA-8B-Instructを用いたMATH-500、GSM8K、ARC-Challenge、TruthfulQAでの実験により、基盤モデルやデコードスケジュールを変更することなく一貫した改善が確認される。特に数学的推論タスクで最大の向上が見られる。
- これらの結果は、復元軌道に対する古典的な検証器ガイド付き探索が、拡散言語モデルにおけるテスト時スケーリングの実用的な仕組みであることを示唆している。
