Answer First, Reason Later: Commitment Order in Diffusion LLMs

Code and per-problem results for the paper "Answer First, Reason Later: Commitment Order in Diffusion LLMs."

Jewon Yeom¹, Jaewon Sok², Seonghyeon Park³, Jeongjae Park¹, Hwiyeong Lee¹, Taesup Kim¹* ¹Graduate School of Data Science, Seoul National University ²Department of Rural Systems Engineering, Seoul National University ³Department of Aerospace Engineering, Seoul National University *Corresponding author

• Paper: arXiv:2608.05687
• BibTeX:

  @article{yeom2026answer,
    title={Answer First, Reason Later: Commitment Order in Diffusion {LLM}s},
    author={Yeom, Jewon and Sok, Jaewon and Park, Seonghyeon and Park, Jeongjae and Lee, Hwiyeong and Kim, Taesup},
    journal={arXiv preprint arXiv:2608.05687},
    year={2026}
  }

Summary

Masked-diffusion LLMs (LLaDA-8B, Dream-7B) can commit the final answer before the reasoning is written, and on GSM8K they frequently collapse to an answer-only response with no chain of thought. We show (i) this collapse is a property of the decoding order, not of the prompt or the budget; (ii) chain of thought is causally attenuated — the effect of CoT depends on the decoder (a prompt×decoder interaction, ΔΔ ≈ +35 pp on LLaDA, +38 pp on Dream); and (iii) constraining commit order with a frontier window (w) recovers the lost accuracy, matching a semi-autoregressive decoder at any w ≤ 64. This repository reproduces every number, table, and figure in the paper from the included per-problem summaries (raw arrays via Release; see Data).

Install

python -m venv .venv && source .venv/bin/activate   # Python 3.10–3.12
pip install -r requirements.txt
# PyTorch: install the wheel matching your CUDA (run env used torch 2.11.0+cu128):
#   pip install torch==2.11.0 --index-url https://download.pytorch.org/whl/cu128

GPU: results were produced on a single NVIDIA A100-SXM4-40GB, batch size 1. The sampler loads an 8B model in bf16 (~17 GB VRAM). transformers==4.49.0 is required (LLaDA's trust_remote_code loader breaks on 5.x).

Minimal reproduction (~5 min after model download)

Compare pure diffusion vs. frontier-gated decoding on one GSM8K problem:

# pure diffusion (unconstrained commit order)  -> tends to answer-only collapse
python src/dllm_pc.py --family llada --dataset gsm8k --limit 1 \
    --gen-length 512 --steps 512 --block-length 512 \
    --out-root ./runs --run-id demo_pure

# frontier-gated (window w=16)                  -> writes reasoning, then the answer
python src/dllm_pc.py --family llada --dataset gsm8k --limit 1 \
    --gen-length 512 --steps 512 --block-length 512 --window 16 \
    --out-root ./runs --run-id demo_gated

Each writes runs/demo_*/problems/0000.npz + summary.csv (fields: correct, answer_span_start → reasoning length, answer_commit_step_norm → t*, …). The pure run typically commits the answer early with little reasoning; the gated run fills reasoning first.

Full reproduction

Every paper cell is one dllm_pc.py invocation; the exact flags for each are in configs/<cell>.yaml and the number→cell→command map is in docs/paper_numbers.md. Typical cell (250 problems, L=512, steps=512) is ~2 h on one A100. Runs are resumable: re-running the same --run-id --out-root skips finished problems (npz on disk); the local harvest also drives --skip-idx for VM-loss recovery.

• Phase 1 / 1.5 (collapse, 2×2, channels): analysis/analyze.py, analyze_2x2.py, channel_decomp.py, decompose.py, backfill.py, aimed_error.py.
• Phase 2 (frontier gating w-sweep, NFE): analysis/analyze_wsweep.py, analyze_2b.py.
• Dream replication / §4: src/dream_official_worker.py, analysis/analyze_dream.py.
• MATH-500: analysis/rescore_math_remote.py.
• Robustness R1/R2/R3: src/infra/runall_r1.ps1, runall_r2.ps1, src/r3_timing.py.
• Figures: analysis/fig_extract.py (reads frozen data, echo-verifies every number against results/mapping_internal.md) → analysis/make_figs.py.

Data root. Analysis scripts read logs/<cell>/problems/*.npz. The raw npz are not shipped (see Data); regenerate a cell with its configs/<cell>.yaml into ./logs/, or request the arrays from the authors. CSV-derivable numbers need no npz — run the analysis scripts against results/summary/.

Frontier gating (the intervention)

w is a manipulated variable, not a new sampler: at each step a masked position is eligible to be committed only if it lies within w of the current frontier (the left-most still-masked position). w ≥ gen_length is byte-identical to vanilla pure diffusion (regression gate: ndiff=0), so w cleanly isolates the causal effect of commit order. Implementation: src/dllm_pc.py → generate_logged, the window block (frontier eligibility + block-final flush).

Data

• In this repo (git): per-problem summaries results/summary/<cell>.csv (accuracy, reasoning length, commit step t*, PCI, gold difficulty …), all configs/<cell>.yaml, the figures, and results/manifest.json (SHA-256 of every frozen file).
• Raw arrays (npz, ~170 MB): not distributed (size). They hold the full per-position commit traces, snapshots, and EOS densities. They are (a) regenerable by re-running the sampler with the matching configs/<cell>.yaml (src/dllm_pc.py), or (b) available from the authors on request. A frozen file list with SHA-256 is in results/manifest.json.

Most headline numbers (accuracy, collapse rate, commit step t*, the prompt×decoder interaction, the w-sweep, the NFE table) are fully recomputable from the CSV summaries in this repo — no npz needed. Only two panels read the raw arrays: the Fig 1 commit raster and the Fig 2 EOS pressure↔commitment dissociation (per-step arrays not stored in the CSVs); regenerate those cells to reproduce them.

Limitations & notes

• Main tables use single-seed greedy decoding — which is precisely the official sampler at temperature 0 (add_gumbel_noise(logits, 0) is a no-op). Stochastic-sampling and prompt-paraphrase robustness (Appendix R1/R2) live in a separate namespace and do not feed the main tables.
• src/infra/ (the runall_*.ps1 / watch.ps1 orchestrators and the Colab CLI wrappers cx/bg/poll/harvest.ps1) is specific to our execution environment and is not needed to reproduce results — it is included only for transparency. These scripts contain no credentials (the Colab CLI reads its own config outside this repo).
• NFE (denoising steps) is the compute axis; wall-clock is exactly linear in NFE on A100 batch=1 (R²=1.00), so steps=128 is ~4× faster than steps=512.