Diffusion Language Model Inference with Monte Carlo Tree Search

2025-12-12

AI Summary

Diffusion language models (DLMs) face a combinatorial search challenge in parallel denoising inference—jointly optimizing mask position selection and token commitment—where existing heuristic or auxiliary-training-based approaches struggle to balance efficiency and generation quality. This paper proposes MEDAL, the first framework to integrate Monte Carlo Tree Search (MCTS) into the DLM initialization phase. MEDAL introduces a confidence-guided action space pruning mechanism and a residual-mask-bit prioritization strategy, enabling training-free, principled, search-based inference. Evaluated across multiple benchmarks, MEDAL achieves up to a 22.0% improvement over state-of-the-art methods, significantly enhancing both local token accuracy and global sequence coherence. By unifying search-based reasoning with diffusion-based language modeling, MEDAL establishes a novel paradigm for DLM inference that is theoretically grounded, empirically effective, and computationally efficient.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models

Application Category

• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• User Modeling, Personalization and Recommendation: ML for personalized search and recommendations
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Diffusion language models (DLMs) have recently emerged as a compelling alternative to autoregressive generation, offering parallel generation and improved global coherence. During inference, DLMs generate text by iteratively denoising masked sequences in parallel; however, determining which positions to unmask and which tokens to commit forms a large combinatorial search problem. Existing inference methods approximate this search using heuristics, which often yield suboptimal decoding paths; other approaches instead rely on additional training to guide token selection. To introduce a principled search mechanism for DLMs inference, we introduce MEDAL, a framework that integrates Monte Carlo Tree SEarch initialization for Diffusion LAnguage Model inference. We employ Monte Carlo Tree Search at the initialization stage to explore promising unmasking trajectories, providing a robust starting point for subsequent refinement. This integration is enabled by restricting the search space to high-confidence actions and prioritizing token choices that improve model confidence over remaining masked positions. Across multiple benchmarks, MEDAL achieves up to 22.0% improvement over existing inference strategies, establishing a new paradigm for search-based inference in diffusion language models.

Problem

• Optimizes combinatorial search in diffusion language model inference
• Improves token selection and unmasking trajectories without extra training
• Enhances decoding paths and model confidence for better text generation

Innovation

• Monte Carlo Tree Search initializes diffusion model inference
• Restricts search space to high-confidence unmasking actions
• Prioritizes token choices improving model confidence globally