NeurIPS 2026Multimodal diffusion language models

StepRefine

Rethinking in Spikes: Mitigating Hallucinations in MDLMs with Step-Aware Decoding

Key finding: Localized entropy spikes during denoising are closely associated with hallucinated content in multimodal diffusion language models.

Step-aware triggering: A sliding-window entropy criterion identifies uncertain steps and selectively activates refinement.

Continuous refinement: Probability-weighted token embeddings preserve predictions until consecutive distributions stabilize, before returning to discrete decoding.

Empirical gains: StepRefine improves hallucination, general reasoning, and math/science benchmarks across four multimodal diffusion backbones without training.

Code BibTeX

Zhongxing Xu^1,*Zhonghua Wang^1,*Zhe Qian¹Shiyan Su¹Ming Hu¹Xiaocheng Zou⁴

Wei Feng¹Mingquan Lin⁵Yuyin Zhou³Yifan Peng²Hamid Rezatofighi¹Zongyuan Ge¹

¹ Monash University² Cornell University³ UC Santa Cruz⁴ Northeastern University⁵ University of Minnesota

* Equal contribution | Contact: zhongxing.xu@monash.edu · zhonghua.wang@monash.edu

Overview

Multimodal diffusion models generate text through iterative denoising. Although average mask entropy generally decreases, a small number of steps show sharp rises. These entropy spikes are closely associated with hallucinated content in the final response.

StepRefine identifies anomalous steps relative to a local temporal window. Instead of immediately committing to a token, it lets candidate semantics interact with the context through probability-weighted embeddings.

Entropy spikes and hallucinated content — Mean mask entropy decreases overall, but localized spikes coincide with hallucinated content.

Entropy Spike Analysis

Spikes concentrate in early-to-middle denoising, when context is still limited. Perturbing token selection at spike steps changes performance little, while the same perturbation at ordinary steps causes clear degradation.

Entropy spike characterization across decoding steps — Hallucination association, temporal distribution, and sensitivity to random remasking.

Method

StepRefine detects entropy spikes, refines uncertain predictions with soft token embeddings, and resumes discrete decoding when the distribution stabilizes.

Detect

Trigger refinement when the entropy change exceeds the sliding-window mean plus its standard deviation.

\[\Delta H_t > \mu_{\mathrm{window}} + \sigma_{\mathrm{window}}\]

Refine

Feed probability-weighted token embeddings back into the model to preserve competing predictions.

\[\widetilde{\mathbf e}_i = \sum_{v\in\mathcal V} p_i(v)\,\mathbf W_{\mathrm{emb}}(v)\]

Stop

Stop when consecutive distributions stabilize, then commit the refined prediction.

\[\frac{1}{|\mathcal M_t|}\sum_{i\in\mathcal M_t} D_{\mathrm{KL}}\!\left(p_i^{(r)}\,\middle\|\,p_i^{(r-1)}\right) < \varepsilon,\qquad \varepsilon = 0.01\]

Qualitative Visualization

The example below shows how predictions evolve during refinement at an entropy spike. Rather than immediately committing to a discrete token, StepRefine feeds probability-weighted embeddings back into the model, allowing competing candidates to be reconsidered in context. Across refinement rounds, the distribution shifts toward a context-consistent answer before decoding continues.

Predictions across refinement rounds — Qualitative visualization of the predicted token distribution across refinement rounds.

Ablation Studies

A moderate calibration window

A window of 16 steps balances noisy local estimates and over-smoothed trajectories in the reported experiments.

Accuracy under different calibration window sizes — Effect of the calibration window on four benchmarks.

Adaptive triggering and termination

Entropy-aware triggering outperforms the tested fixed schedules. Monitoring distributional stability also improves over a fixed number of refinement rounds.

Reference

@inproceedings{xu2026rethinking,
  title = {Rethinking in Spikes: Mitigating Hallucinations in MDLMs with Step-Aware Decoding},
  author = {Xu, Zhongxing and Wang, Zhonghua and Qian, Zhe and Su, Shiyan and Hu, Ming and Zou, Xiaocheng and Feng, Wei and Lin, Mingquan and Zhou, Yuyin and Peng, Yifan and Rezatofighi, Hamid and Ge, Zongyuan},
  booktitle = {Advances in Neural Information Processing Systems},
  year = {2026}
}