AIWatch.

#ICML2026 に採択された私たちの論文 ”UnMaskFork: Test-Time Scaling for Masked Diffusion via Deterministic Action Branching” は、複数の拡散言語モデルを協調させることで、コーディングや数学の能力を向上できることを示しました。

一つの問題をLLMに長考させたり、解答の生成や改善を何度も繰り返させたりと、推論時のリソースを活用してLLMの性能を引き上げる「推論時スケーリング」が重要な手法として盛んに研究されています。この技術により、追加学習に頼らず推論時の計算を増やすことでLLMの性能を高められるようになり、用途に応じて計算コストと性能のバランスを柔軟に選ぶことができます。

テキストを前から順に生成する通常のLLMと異なり、近年注目を集めるマスク型拡散言語モデル（MDLM）は、全体がマスクされた状態から徐々にマスクを外しテキストを埋めていくことで文章を完成させます。MDLMは複数の箇所を並列に生成できるため高速化が期待できるほか、文章全体を見渡しながら柔軟に生成できるという特長があり、新しいパラダイムの言語モデルとして盛んに研究されています。我々は、通常のLLMで標準的な「温度を上げてランダム性を高め、多様な解答を何度も生成させる」手法はDream-CoderのようなMDLMではうまく機能しないことを実験により明らかにしました。

私たちの提案するUnMaskFork（UMF）は、このランダム性の代わりに「モデルの切り替え」で多様性を生み出します。複数のMDLMが一つの解答のマスク解除を分担し、どのモデルがどの段階を担当するかの有望な順番をモンテカルロ木探索で探索します。。各モデルは他のモデルの途中経過を引き継ぎ、最も自信のある箇所を埋めていきます。この協調アプローチにより、生成品質を保ったまま多様な解答を探索でき、コーディングのベンチマークで既存の推論時スケーリング手法を一貫して上回り、数学でも計算量に応じた着実な性能向上を示しました。

MDLMの発展においても推論時スケーリングは重要な鍵であり、本研究はMDLM特有の困難をUMFで回避できることを示しました。UMFは追加学習やモデル内部の改変を一切必要とせず、学習済みのモデルを推論時に組み合わせるだけで機能します。そのため、異なるデータや方法で学習された拡散言語モデルの多様性を活用することで、そのまま性能向上へ繋げることができます。今後多様なMDLMが登場するほど、UMFの価値はさらに広がっていくと考えています。

本研究は、複数のLLMを協調させるAB-MCTSやSakana Fuguなど、「AIの集合知」を追求する我々の研究の一環です。今後もモデルの多様性を力に変える研究を進めていきます。UMFの詳細なアルゴリズムや、実際にどのように協調が行われるかなどの具体例はブログや論文をご覧ください。

ブログ: https://pub.sakana.ai/umf 論文: https://arxiv.org/abs/2602.04344 🐟