Thinking Inside the Mask: In-Place Prompting in Diffusion LLMs

TL;DR

ICE addresses the limitations of prefix-only prompting by leveraging diffusion LLMs' bidirectional context to support In-Place Chain-of-Thought Prompting. It introduces two-phase decoding with a confidence-based early exit, embedding reasoning templates within masked tokens and transitioning to parallel answer generation once average answer-token confidence reaches a threshold. Empirically, ICE delivers up to +17.29% accuracy improvements on GSM8K with up to 4.12× speedups and up to 276.67× acceleration on MMLU (plus gains on GPQA), while remaining compatible with caching-based acceleration. This work highlights that architectural alignment between structured reasoning and the diffusion generation process can substantially reduce compute while enhancing reasoning quality across diverse tasks.

Abstract

Despite large language models (LLMs) have achieved remarkable success, their prefix-only prompting paradigm and sequential generation process offer limited flexibility for bidirectional information. Diffusion large language models (dLLMs) present new opportunities through their bidirectional attention mechanisms and iterative refinement processes, enabling more flexible in-place prompting strategies. We introduce ICE (In-Place Chain-of-Thought Prompting with Early Exit), a novel framework that transforms prefix-only prompting into in-place prompting specifically designed for dLLMs. ICE integrates in-place prompts directly within masked token positions during iterative refinement and employs a confidence-aware early exit mechanism to significantly reduce computational overhead. Extensive experiments demonstrate ICE's effectiveness, achieving up to 17.29% accuracy improvement with 4.12$\times$ speedup on GSM8K, and up to 276.67$\times$ acceleration on MMLU while maintaining competitive performance.

Figures (7)

• Figure 1: Prompt construction in (a) Autoregressive LLMs vs. (b) Diffusion LLMs. Autoregressive LLMs employ unidirectional context with prefix-only prompting, while dLLMs leverage bidirectional context modeling, enabling in-place prompting and concurrent answer accessibility.
• Figure 2: Average confidence of answer section on GSM8K and MATH during generation. The model's confidence in the answer section rapidly converges to a high level and remains stable throughout subsequent iterations, indicating that the model internally determines the correct answer early in the process while continuing to refine the reasoning trace.
• Figure 3: Overview of the ICE framework. ICE integrates two key components: (1) In-Place Chain-of-Thought Prompting, which embeds structured step-by-step reasoning templates directly into the prompt, and (2) Two-Phase Decoding with Confidence-Aware Early Exit Mechanism. During the Reasoning Phase, the model iteratively decodes the thinking section while monitoring the average confidence of the masked answer section. Once the confidence threshold is reached, the framework transitions to the Answer Generation Phase, decoding all answer tokens in parallel to produce the final response.
• Figure 4: Latency-accuracy trade-off comparison between ICE and vanilla baselines. ICE demonstrates superior Pareto frontiers on both GSM8K and MATH datasets.
• Figure 5: Ablation study on reasoning steps ($N_t$).