PC-Sampler: Position-Aware Calibration of Decoding Bias in Masked Diffusion Models
This addresses a key bottleneck in non-autoregressive sequence generation for tasks like logical reasoning and planning, offering an incremental but significant improvement over existing MDM decoding methods.
The paper tackled the sensitivity of masked diffusion models (MDMs) to decoding strategies, which suffer from poor global trajectory control and bias toward trivial tokens, by introducing PC-Sampler, a decoding method that improved performance by over 10% on average across seven benchmarks.
Recent advances in masked diffusion models (MDMs) have established them as powerful non-autoregressive alternatives for sequence generation. Nevertheless, our preliminary experiments reveal that the generation quality of MDMs is still highly sensitive to the choice of decoding strategy. In particular, widely adopted uncertainty-based samplers suffer from two key limitations: a lack of global trajectory control and a pronounced bias toward trivial tokens in the early stages of decoding. These shortcomings restrict the full potential of MDMs. In this work, we introduce Position-Aware Confidence-Calibrated Sampling (PC-Sampler), a novel decoding strategy that unifies global trajectory planning with content-aware informativeness maximization. PC-Sampler incorporates a position-aware weighting mechanism to regulate the decoding path and a calibrated confidence score to suppress the premature selection of trivial tokens. Extensive experiments on three advanced MDMs across seven challenging benchmarks-including logical reasoning and planning tasks-demonstrate that PC-Sampler consistently outperforms existing MDM decoding strategies by more than 10% on average, significantly narrowing the performance gap with state-of-the-art autoregressive models. All codes are available at https://github.com/NEUIR/PC-Sampler.