UnMaskFork: 掩码扩散模型的推理时缩放

拡散言語モデルの協調による推論時スケーリングの実現 #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://t.co/FZ25e6XCws 論文: https://t.co/4JC9SYdTyX 🐟 > **引用原帖 Sakana AI (@SakanaAILabs):** > Can test-time scaling work for diffusion language models? > In our #ICML2026 paper "UnMaskFork," we show that having multiple masked diffusion language models collaborate on a single answer improves performance on coding and math tasks. > Blog: https://t.co/FZ25e6XCws > Test-time scaling is an actively researched technique that boosts LLM performance by using inference-time compute, for example, by having a model think longer or repeatedly refine its answers. This allows us to enhance performance simply by increasing computation during inference without relying on additional training, giving us the flexibility to balance compute costs and performance based on the specific use case. > Unlike standard LLMs that generate text left-to-right, masked diffusion language models (MDLMs) generate text by gradually filling in a fully masked sequence. MDLMs can generate multiple parts of a sequence in parallel, offering potential speed-ups, and they can generate flexibly while seeing the entire sequence at once. This makes them an actively studied new paradigm in language modeling. We found that the standard LLM approach of "raising the temperature to increase randomness and generate diverse answers" does not work well for MDLMs like Dream-Coder. > Instead of relying on this randomness, our proposed method, UnMaskFork (UMF), creates diversity through "model switching." Multiple MDLMs share the task of unmasking a single answer, and we use Monte Carlo Tree Search to search for a promising sequence in which different models handle different stages. Each model picks up where the others left off, filling in the parts it is most confident about. This collaborative approach allows us to explore diverse answers while maintaining generation quality, consistently outperforming existing test-time scaling methods on coding benchmarks and scaling effectively on math as well. > Test-time scaling is also crucial for advancing MDLMs, and our work shows that UMF can sidestep the difficulties specific to them. UMF requires no additional training or changes to the models; it works simply by combining pre-trained models at inference time. This allows us to leverage the diversity of diffusion language models trained on different data and with different methods to improve performance. We believe the value of UMF will only grow as more diverse MDLMs emerge. > This work is part of our broader research into "collective intelligence of AI," alongside methods like AB-MCTS and Sakana Fugu that have multiple LLMs collaborate. We'll continue pursuing research that turns model diversity into a source of strength. > For details of the algorithm and illustrative examples showing how this collaboration works, please see our blog and paper. > Paper: https://t.co/4JC9SYdTyX 🐟 > https://x.com/SakanaAILabs/status/2079567069096693872