UNCODE: Empirical Analysis of Decoding Biases in Masked Diffusion Models

June 26, 2026 ยท View on GitHub

UNCODE: Empirical Analysis of Decoding Biases in Masked Diffusion Models

A training-free decoding-calibration framework that fixes two systematic biases
in Masked Diffusion Models โ€” improving reasoning & planning by 7%+ across 3 MDMs and 7 benchmarks.

๐Ÿ“– Introduction โ€ข ๐ŸŽ‰ News โ€ข โš™๏ธ Setup โ€ข ๐Ÿ“ƒ Evaluation โ€ข ๐Ÿ“ˆ Trajectory โ€ข ๐Ÿ’ป Algorithm โ€ข ๐Ÿ“Œ Citation โ€ข ๐Ÿ“ง Contact


๐Ÿ“– Introduction

UNCODE (UNmasking Calibration for DecOding DEbiasing) is a novel, training-free decoding strategy for Masked Diffusion Models (MDMs) that unifies global trajectory planning with content-aware informativeness maximization.

Uncertainty-based samplers, when applied to MDMs, suffer from two systematic decoding biases:

  • ๐Ÿ”ด Rigid Boundary Bias โ€” boundary tokens (BOS/EOS, sentence edges) are decoded first, collapsing decoding into a fixed U-shaped trajectory and committing to an answer before the reasoning is built.
  • ๐ŸŸก Trivial Token Bias โ€” high-frequency, low-information tokens (punctuation, spaces, fillers) get over-prioritized, spending the decoding budget on surface structure instead of reasoning content.

UNCODE fixes both with a position-aware weighting mechanism and a calibrated, frequency-aware confidence score, guiding the decoding path and suppressing premature selection of unimportant tokens โ€” with no fine-tuning and no architecture change.

UNCODE teaser: rigid boundary bias, trivial token bias, and the ideal decoding trajectory.

๐Ÿ“„ Paper: arXiv:2508.13021 ยท ๐ŸŒ Project page: passionate11.github.io/Uncode-project-page

๐ŸŽ‰ News

  • 2026-04-07: Our paper has been accepted to ACL 2026 (Main Conference)! ๐ŸŽ‰
  • 2025-09-12: Release adds enhanced LLaDA decoding support, integrating recent semi- and non-autoregressive sampling strategies: ReMDM, Fast-dLLM, Semi-AR, Margin-, Entropy- and Confidence-based samplers.
  • 2025-08-19: Released our paper on arXiv and code on GitHub.

โœจ Highlights

๐Ÿš€ >7% average gain over the strongest decoding baseline๐Ÿงฉ 3 ร— 7 MDM backbones ร— reasoning & planning benchmarks
โš–๏ธ 44.7 โ‰ˆ 45.3 โ€” LLaDA-1.5 + UNCODE rivals autoregressive Qwen-2.5-7B๐Ÿ”Œ 0 extra training โ€” plug-and-play, decoding-side only

โš™๏ธ Setup

git clone https://github.com/NEUIR/Uncode.git
cd Uncode
conda create --name uncode python==3.10
conda activate uncode
pip install -r requirements.txt

๐Ÿ“ƒ Evaluation

UNCODE and all baseline methods can be evaluated across mathematical reasoning, code generation, and question-answering datasets: HumanEval, MBPP, GSM8K, MATH-500, GPQA, Countdown, and Sudoku. Results are saved to the results/ folder.

Example โ€” HumanEval with UNCODE

Change --task and --mode to evaluate on other datasets / decoding methods.

cd scripts
python eval.py \
    --task 'humaneval' \
    --model_name 'GSAI-ML/LLaDA-8B-Instruct' \
    --device 'cuda:0' \
    --gen_length 256 \
    --steps 256 \
    --block_length 256 \
    --mode pc_sampler \
    --lambd 0.25 \
    --alpha 10 \
    --data_path ../data/humaneval.jsonl \
    --result_path results/humaneval_pc_sampler

Baseline decoding methods

Decoding MethodCommandDecoding MethodCommand
Semi-Autoregressivebash eval_semi_ar.shEntropybash eval_entropy.sh
EB-Samplerbash eval_eb_sampler.shFast-dLLMbash eval_fast_dllm.sh
Marginbash eval_margin.shPC-Samplerbash eval_pc_sampler.sh
ReMDMbash eval_remdm.shLinear-Positionbash eval_linear_position.sh

All scripts live in scripts/. Run them from inside that folder (cd scripts).

Evaluation tools & consistency

  • GSM8K and GPQA are evaluated with lm-eval; the remaining datasets use scripts/eval.py.
  • All methods share the same evaluation scripts to ensure consistent, comparable assessment.

Painting heatmaps

Generate decoding-trajectory heatmaps for different methods:

cd scripts
bash heatmap.sh

Heatmap outputs are saved to the heatmap_results/ folder.

๐Ÿ“ˆ Decoding Trajectory

The decoding strategy strongly shapes the generation order of MDMs. Existing uncertainty-based methods exhibit a U-shaped trajectory (the rigid boundary bias): boundary tokens (BOS/EOS) are unmasked early because the attention mechanism's local positional bias inflates their confidence, after which decoding converges inward.

UNCODE instead introduces explicit trajectory control via position-aware weighting, yielding an adaptive generation order tailored to each task. Trajectories on GSM8K for four representative samplers:

Confidence-basedEntropy-basedMargin-basedUNCODE

๐Ÿ”‘ Key Observations

  • Rigid Boundary Bias โ€” confidence/entropy/margin samplers consistently show the U-shaped pattern, decoding both sequence boundaries first. This limits their ability to capture the global dependencies needed for complex reasoning.
  • Trivial Token Bias โ€” uncertainty-based samplers over-prioritize semantically trivial, high-frequency tokens (newlines, spaces, the, ., !), leading to suboptimal reasoning paths.
  • Debiasing with UNCODE โ€” exponential positional weighting removes the U-shape, producing a natural progression aligned with the logical flow of reasoning.

This adaptive trajectory control directly drives UNCODE's strong 82.2% GSM8K accuracy, well above uncertainty-based alternatives.

๐Ÿ’ป Algorithm

UNCODE addresses the limitations of uncertainty-based sampling through two core components:

  1. Position-Aware Weighting โ€” an exponential decay over position regulates the decoding path, giving flexible control over generation order to match task structure.
  2. Calibrated Confidence Score โ€” a frequency-based adjustment from a reference corpus suppresses premature selection of trivial tokens, promoting semantically rich content.

Across seven benchmarks, UNCODE consistently outperforms existing MDM decoding strategies, narrowing the gap to state-of-the-art autoregressive models.

Workflow

Require: Predictor pฮธp_\theta, prompt p0p_0, answer length LL, steps TT, hyperparameters ฮป,ฮฑ\lambda, \alpha; reference corpus Dโ€ฒ\mathcal{D}'

  1. pDโ€ฒโ†FreqDist(Dโ€ฒ)p_{\mathcal{D}'} \gets \text{FreqDist}(\mathcal{D}')
  2. xโ†Concat(p0,[MASK]ร—L)x \gets \text{Concat}(p_0, \text{[MASK]} \times L)
  3. for t=1t = 1 to TT do
    • Mtโ†{iโˆฃxi=[MASK]}\mathcal{M}_t \gets \{i \mid x^i = \text{[MASK]}\} ย // mask indices
    • if Mt=โˆ…\mathcal{M}_t = \emptyset then break
    • x^0,p^iโ†pฮธ(โ‹…โˆฃx)\hat{x}_0, \hat{p}^i \gets p_{\theta}(\cdot \mid x)
    • for each position iโˆˆMti \in \mathcal{M}_t do
      • C(i)โ†p^iโ‹…logโกpDโ€ฒ(xi)\mathcal{C}^{(i)} \gets \hat{p}^i \cdot \log p_{\mathcal{D}'}(x^i)
      • C(i)โ†minโก(C(i),ฮฑ)\mathcal{C}^{(i)} \gets \min(\mathcal{C}^{(i)}, \alpha) ย // clip salience
      • w(i)โ†eโˆ’ฮปโ‹…(iโˆ’โˆฃp0โˆฃ)w^{(i)} \gets e^{-\lambda \cdot (i - |p_0|)}
      • score(i)โ†w(i)โ‹…C(i)\text{score}^{(i)} \gets w^{(i)} \cdot \mathcal{C}^{(i)}
    • nkโ†NumToReveal(k,N,โˆฃMkโˆฃ)n_k \gets \text{NumToReveal}(k, N, |\mathcal{M}_k|)
    • Stโ†TopK(score,nk)\mathcal{S}_t \gets \text{TopK}(\text{score}, n_k) ย // select best tokens
    • for each index jโˆˆStj \in \mathcal{S}_t do xjโ†x^0jx^j \gets \hat{x}_0^j ย // reveal
  4. return xx

Hyperparameters

ParamMeaningRecommended
ฮป\lambda (--lambd)Positional bias strength: 0 = no bias, larger = stronger left-to-right0 (Sudoku), 0.25 (most tasks), 0.5 (Countdown)
ฮฑ\alpha (--alpha)Clipping threshold for the salience score10 (stable across tasks)
pDโ€ฒp_{\mathcal{D}'}Background frequency distribution from a reference corpussee data/baseline

๐Ÿ“Œ Citation

If you find UNCODE useful, please cite:

@inproceedings{huang-etal-2026-empirical,
    title     = "Empirical Analysis of Decoding Biases in Masked Diffusion Models",
    author    = "Huang, Pengcheng and Liu, Tianming and Liu, Zhenghao and Yan, Yukun and Wang, Shuo and Xiao, Tong and Chen, Zulong and Sun, Maosong",
    booktitle = "Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)",
    year      = "2026",
    publisher = "Association for Computational Linguistics",
    url       = "https://aclanthology.org/2026.acl-long.311/",
    pages     = "6853--6876",
}

๐Ÿ“ง Contact

Questions, suggestions, or bug reports are welcome โ€” please open an issue or email pengcheng.neu@outlook.com.