拡散言語モデルの推論を広げる頻度認識リマスキング

拡散言語モデル（DLM）は、通常の自己回帰モデルのように左から右へ一度に文章を生成するのではない。複数回のデノイジングを通じて、各位置のトークンを確定したり、再びマスクして予測し直したりする。この「いつ確定し、いつやり直すか」という判断は、最終的な回答の品質だけでなく、推論に使われる計算量にも大きく関わる。

今回の研究は、DLMのテスト時スケーリングをどのように設計するかを扱う。既存の確信度認識型の手法は、ある時点での確信度やエントロピーを使ってトークンを確定する。しかし、数学のような難しい推論問題では、局所的に高い確信度が正しい推論経路を意味するとは限らない。モデルが早い段階で一見もっともらしい誤った経路に固定されると、追加サンプリングを行っても同じ経路を繰り返す可能性がある。

手法の要点

• 2段階の計算配分：最初は高速な基礎サンプラーですべての問題を処理し、未解決の難しい問題にだけ追加の計算を割り当てる。
• 予測の変化回数を追跡：各デノイジング段階で、各トークンの予測がどれほど頻繁に変化したかを記録する。これは単一時点の確信度とは異なる、推論過程に基づく信号である。
• 頻度認識リマスキング：予測が頻繁に変わるトークンに対して、リマスキングの判断を変える。モデルが最も好む一つの経路へ早く収束しすぎるのを防ぎ、別の推論経路を探索しやすくすることが狙いだ。
• 難しい例に限定して適用：すべての入力の推論コストを一律に増やすのではなく、基礎戦略で解けなかった例に複雑なデコードを集中させる。

実験にはLLaDA-8B-InstructとDream-7B-Instructが使われた。MATH-500では、強力な適応的計算配分ベースラインに対して、絶対正解率がLLaDAで最大3.73ポイント、Dreamで最大7.33ポイント向上した。HumanEvalでは、絶対カバレッジの向上幅がそれぞれ最大11.39ポイントと10.98ポイントに達した。また、基礎となる確信度認識戦略で100回サンプリングしても解けなかった例について、頻度認識リマスキングが3.75%から18.75%を追加で解決できたと報告されている。

意味と今後の論点

この研究のポイントは、単純にサンプリング回数を増やすことではない。デノイジング中の予測変化の履歴を、瞬間的な確信度を補う動的な信号として利用し、追加計算を必要な問題に振り向けている。反復的に生成するDLMでは、途中の判断を見直せるため、このような戦略はモデルを再学習せずに探索能力を引き出す方法になり得る。

一方、結果は対象となった2つのモデルと、論文で示された数学・コードベンチマークに基づく。ほかのタスクでも変化頻度が安定した信号になるのか、追加リマスキングが実運用の遅延や計算コストに見合うのかは、今後の検証が必要だ。それでも、簡単な問題では素早く確定し、難しい問題では推論経路を柔軟に保つというDLM推論の設計方針を示している。

出典：Hugging Face Daily Papers