Supportive Token Revealing for Fast Diffusion Language Model Decoding

2026-06-02

AI Summary

Discrete diffusion language models face a trade-off between generation quality and latency in parallel decoding: aggressive token revelation often introduces errors due to premature commitments, while conservative strategies require excessive denoising steps. To address this, this work proposes AXON, a plug-in module that optimizes the mask-revealing order by dynamically identifying the most supportive confident anchor tokens for subsequent denoising, without modifying the underlying decoder. AXON innovatively shifts the revelation criterion from “safest” to “most supportive,” integrating attention weights, confidence scores, and uncertainty signals into a training-agnostic intervention mechanism. Experiments demonstrate that AXON significantly improves the quality–latency trade-off across multiple reasoning and code generation benchmarks, reducing the number of function evaluations while maintaining or even enhancing generation accuracy.

Abstract

Discrete diffusion language models can generate text efficiently by updating multiple masked positions in parallel, but this parallelism introduces a quality-latency trade-off. Aggressive decoding may commit mutually dependent tokens too early, while conservative decoding requires many denoising steps. Existing methods address this tension by deciding which tokens are safe to reveal using confidence or dependency criteria. However, avoiding unsafe commits does not necessarily make the remaining masked sequence easy to decode, since uncertain tokens may depend on masked tokens, creating a bottleneck for denoising steps. We propose AXON, a training-free module that can be added on top of existing parallel decoding strategies for diffusion language models. Rather than replacing the base decoder, AXON monitors the remaining uncertain masked tokens and intervenes only when their current state suggests that additional context is needed. It then shifts the criterion from which tokens are safest to reveal to which confident reveals would best support later denoising. AXON selects anchors, confident masked tokens that uncertain positions attend to, using attention, uncertainty, and confidence signals. Experiments on reasoning and code-generation benchmarks across multiple diffusion language models show that AXON improves the quality-latency trade-off of existing parallel decoders, often reducing the number of function evaluations while maintaining or improving accuracy.

Problem

• diffusion language models
• parallel decoding
• quality-latency trade-off
• token dependency
• denoising bottleneck

Innovation

• diffusion language models
• parallel decoding
• token revealing
• quality-latency trade-off
• training-free intervention