Confidence anchored decoding improves masked diffusion language models
Cord: Confidence-Anchored Order Residual Decoding for Masked Diffusion Language Models
Computation and LanguageArtificial Intelligence
Summary
Deciding the order in which parts of a sentence get filled in matters a lot for language models that generate text by masking and revealing words. The standard way picks the most confident guesses first, but that doesn’t always help later steps. The authors created Cord, a small system that smartly adjusts this order based on the prompt and what's already been written, learning a kind of plan for revealing words all at once. This method improves how well these models solve tasks like reasoning or coding, without needing much extra computing power.
What this means in practice
- •For natural language engineers: Improve output quality of masked diffusion language models on challenging reasoning and coding tasks by refining token reveal order.
- •For chatbot developers: Enhance response accuracy in conversational assistants that use masked diffusion models with better partial response completion.
Authors
Chengwei Zhou, Zhaoyan Jia, Massoud Pedram, Gourav Datta
Abstract
The decoding order of masked diffusion language models (dLLMs) can substantially affect generation quality, yet common decoders typically apply the same reveal rule throughout generation. Confidence provides a strong default, but it measures how decisively a position can be predicted now, not how useful revealing that position will be for subsequent decoding. We propose Cord, a lightweight prompt-conditioned policy that retains confidence decoding as a stable reference and learns when and how to depart from it through geometric residuals over generation direction, proximity to revealed context, and two-sided infilling. The residuals adapt to both the prompt and the evolving partial response. Rather than sampling a new policy action at every decoding step, Cord samples a single low-dimensional coefficient matrix per rollout, which induces an entire state-dependent reveal trajectory. This yields a contextual-bandit formulation trained directly from sequence-level rewards, without step-level credit assignment or backbone updates. Across three dLLMs and five reasoning, code, and constraint-satisfaction benchmarks, Cord improves average task success from 32.0% under confidence decoding to 48.2%, and further improves over the best per-domain fixed residual. Behavioral audits show that the learned residuals induce substantial, task-dependent changes in reveal locality and two-sided infilling. Cord adds only $0.53$M trainable parameters, requires no additional backbone forward passes, and incurs only 5.39% inference latency overhead.