SAI
← All ICML 2026 orals

The Flexibility Trap: Rethinking the Value of Arbitrary Order in Diffusion Language Models

Zanlin Ni, Shenzhi Wang, Yang Yue, Tianyu Yu, Weilin Zhao, Yeguo Hua, Tianyi Chen, Jun Song, Cheng Yu, Bo Zheng, Gao Huang

OralReplication score 20%Paper PDFCode repoOpenReview

The Flexibility Trap: Rethinking the Value of Arbitrary Order in Diffusion Language Models

SAI replication review · Referee report

Summary

The paper questions a widely-held premise about diffusion language models (dLLMs) — that generating tokens in arbitrary order is a source of reasoning power — and argues, against the intuition that a strictly larger trajectory space should help, that arbitrary-order generation narrows the reasoning solution space accessible to RL on general reasoning (math and code). The empirical case rests on Pass@k as a proxy for the post-RL reasoning boundary: across three dLLMs and four benchmarks, restricting decoding to strict left-to-right yields flatter-shape but higher-scaling Pass@k curves than confidence-based arbitrary-order decoding, a per-problem set-difference showing AR-solved problems dominate AO-solved ones, and a monotone block-size sweep in which more arbitrariness means less coverage. The proposed mechanism is entropy degradation: confidence-based samplers systematically defer logical connectives ("Therefore", "Thus", "Since") whose entropy is what carries branching, so by the time those tokens are unmasked the surrounding future has resolved their ambiguity. From this diagnosis the paper draws a refreshingly minimal methodological move — treat the dLLM as an AR policy only during RL, apply vanilla GRPO on the exactly-computable AR factorisation, keep bidirectional attention and parallel sampling at inference. JustGRPO reports 89.1% on GSM8K and 45.1% on MATH-500 and, in the paper's own unified-protocol reproductions of d1, ESPO and SPG, retains clear margins on math; a JustGRPO-Fast variant restricts the per-position ratio computation to top-entropy positions and cuts wall-clock further. Running the released checkpoints through the shipped eval.py reproduces the JustGRPO benchmark numbers closely, but every figure and table carrying the paper's scientific argument (Pass@k curves, coverage sets, entropy-at-forks, block-size sweep, random-order controls, general-capability retention) sits outside the shipped code. The main conceptual limitation is that the argument leans heavily on a Pass@k-of-base-model sharpening view of RLVR that is itself contested, and the "entropy degradation" diagnosis relies on a measurement guaranteed by information theory (conditional entropy dropping when future context is fixed) that does not separate harmful commitment from beneficial informed decision.

Strengths

  • Conceptual contribution. A clean, counter-intuitive claim — arbitrary order can limit reasoning potential on general reasoning tasks — supported by three distinct empirical angles: Pass@k curve shape (Figure 3), per-problem set-difference at Pass@1024 (Figure 4), and a monotone block-size sweep (Figure 5). The block-size continuum is especially compelling as a "dose-response" argument.
  • Mechanistic hypothesis. Tying the coverage gap to a specific behaviour — confidence-based samplers deferring high-entropy connectives — is a testable, falsifiable story rather than a vague appeal to inductive bias, and Figure 7's fork-vs-global entropy contrast is a nice targeted measurement.
  • Methodological restraint. JustGRPO is a genuinely minimal recipe: no auxiliary policy, no ELBO surrogate, no importance-weighted trajectory reweighting. The framing as "AR is a training-time scaffold, not a structural constraint" cleanly separates optimisation from inference and is well-defended by the parallel-decoding evidence in Figure 8.
  • Head-to-head reproductions. Table 2 reruns d1, ESPO, and SPG under a unified protocol rather than only quoting original numbers, which is the right move given how heterogeneous the reported settings across dLLM-RL papers are. Even if the baseline code is not shipped, the effort raises the bar for the comparison.
  • Efficient variant. JustGRPO-Fast is a natural application of the paper's own mechanistic story to its own algorithm, and the fact that pruning ratio computations to the top-entropy quartile improves accuracy/wall-clock is a nice validation of the fork-token diagnosis.
  • Analytical breadth. Appendices B–D substantially firm up the story: temperature sweeps, alternative samplers (Neg-Entropy, Top-k Margin), a broader entropy-vs-fork-tokens plot (Figure 12), a random-order control, and a general-capability check on MMLU/HellaSwag/ARC-C.

Weaknesses

  • Contested proxy for the central claim. The Section 4 argument that arbitrary order limits the achievable-after-RL boundary depends on base-model Pass@k being a valid proxy for post-RL potential. The sharpening view of RLVR is contested in exactly the papers arbitrary-order RL methods would cite (including ProRL in the paper's own references). A post-RL Pass@k for both AR- and arbitrary-order-trained models on at least one benchmark would substantially firm up the argument; without it, the reader is asked to grant the very premise the arbitrary-order camp would reject.
  • Confound in the AR-vs-arbitrary Pass@k comparison. Semi-AR block-32 arbitrary-order decoding commits multiple tokens per step, while AR order (256 steps for 256 tokens) makes one stochastic decision per token. Part of the AR Pass@k advantage could therefore be attributed to more Monte-Carlo variance rather than to "confronting forks". A one-token-per-step arbitrary-order run would isolate the order effect from the stochastic-granularity effect.
  • Entropy-degradation diagnosis is partly tautological. Conditioning on more of the surrounding future can only decrease conditional entropy on average, so Figure 7's entropy drop at forks is guaranteed by information theory — not diagnostic of pathology on its own. The Pass@k evidence carries the coverage argument; the entropy measurement is interpretive scaffolding, and the paper leans on it more than the measurement can bear. A control against equally-frequent non-fork tokens at matched positions would help.
  • JustGRPO-Random does not test the entropy-degradation mechanism. The paper's mechanism is that confidence-based decoding systematically defers connectives. A single fixed random permutation does not systematically defer connectives — it places them uniformly. So the JustGRPO-Random vs JustGRPO gap tests "AR order specifically helps" but not "AR helps because it avoids fork-deferral". An ablation on an anti-confidence order (defer low-entropy first) would sit closer to the mechanism.
  • "Random order does not help coverage" is contradicted by Table 4. On MBPP (71.8 random vs 67.1 confidence-based) and GSM8K (97.5 vs 97.0), random beats confidence-based at Pass@128. The correct statement is "random still falls short of AR", not "random does not help coverage", and the current phrasing weakens rather than strengthens the mechanistic story: if random can beat confidence-based, then confidence-based's specific deferral is the culprit worth naming.
  • "Fails" is not a fair label for a +3.6-pp RL gain. JustGRPO-Random raises GSM8K from 78.6 to 82.2, a real if smaller improvement. Framing this as "random order also fails in RL post-training" is a false dichotomy; the honest reading is "random order underperforms AR order".
  • No decoupling ablation of AR-during-training vs AR-at-inference. The decoupling story is central, but the paper does not report the 2x2 grid (train AR/train AO × infer AR/infer AO). Without this, the training-time-scaffold interpretation of Figure 8 is inferential rather than tested.
  • No SFT-only baseline. Given that the paper's main takeaway is "simple is enough", a supervised-fine-tune baseline on the same data would sharpen the attribution of gains to GRPO vs. to any AR-order training signal.
  • Widening-gap parallel-decoding claim is not an absolute-robustness claim. Figure 8 shows the gap between JustGRPO and LLaDA-Instruct widening with parallelism, but the caption already notes the baseline degrades sharply. "More resilient to the approximations inherent in parallel sampling" reads as an absolute-robustness claim; the evidence is a widening-gap claim. Absolute accuracy curves under parallel decoding — and comparison against at least one arbitrary-order-trained RL baseline decoded the same way — are needed.
  • No seed/variance reporting. All main-table numbers, figures, and ablations are single-run point estimates. The +2.2/+3.3 pp margins over SPG* on math are the paper's most direct head-to-head evidence; without seed variance a reader cannot judge how much of that margin is signal.
  • Block-size sweep is single-benchmark. Figure 5's monotone result is on HumanEval / LLaDA-Instruct only, yet the prose reads as a general claim. Replicating on a math benchmark and at least one other dLLM would keep the reach honest.
  • Overstated framing of the RL objective. "Transforming an otherwise intractable optimization into a well-defined task" and "defining an AR policy for dLLMs" oversell what is really a substitution: the marginal-likelihood objective remains intractable, and JustGRPO optimises a different, well-defined AR surrogate. The paper acknowledges this in a couple of places but the abstract-level framing goes further.
  • Underspecified JustGRPO-Fast. How the other 75% of positions are treated in the loss, whether top-25% is per-sequence or per-batch, and whether selection uses rollout-time or current-policy entropy are all left implicit. This matters both for reproducibility and for whether the wall-clock claim maps to the objective a reader will assume.
  • Bypassing analysis lacks an operational definition and a position-matched control. "Bypassed" is never quantitatively defined, and connectives sit at particular positions in reasoning chains (typically after a full clause), so the Figure 6 word-cloud can partly reflect position-in-sequence rather than a specific deferral of forks. A deferral distribution for equally-frequent non-connective tokens matched by position would sharpen the diagnosis.
  • Transfer of "forking token" concept from AR to dLLM is asserted, not verified. The "connectives are high-entropy forks" line comes from AR-model analyses; whether the same tokens carry branch selection in bidirectional dLLMs under confidence-based decoding is not directly demonstrated.
  • "General capabilities preserved" oversells small drops without variance. Table 6 shows ARC-C dropping 88.5 → 87.5 and MMLU-Pro 37.0 → 36.7. Without any variance signal, "preserved" is stronger than the numbers warrant.
  • Sequence-length non-monotonicity on code passes uncommented. MBPP drops 52.4 → 49.0 and HumanEval 49.4 → 48.7 when moving from gen len 256 to 512 for JustGRPO, yet the paper does not comment.
  • Table 1 layout obscures attribution. Eight methods × three sequence lengths worth of numbers are flattened into a single column per benchmark; a careful reader cannot easily map values to methods/lengths, and the paper's cross-references (e.g. SPG on GSM8K vs. the reproduced SPG* = 86.9) are hard to verify.
  • Minor stumbles. A^i,k\hat A_{i,k} carries a token index despite being sequence-level; "smaller subset" (Section 4.1) implies strict set inclusion contradicted by the AO-only slices in Figure 4; and the RL rollout temperature (T=1.0T = 1.0, Table 3) differs from the Section 4 analysis temperature (T=0.6T = 0.6) without justification.

Reproducibility & code

  • Headline JustGRPO checkpoints reproduce closely at the settings we could actually run. Executing the shipped eval.py against the released JustGRPO-GSM8K, JustGRPO-Math500, and JustGRPO-Code checkpoints on 200-example subsets at gen_length = 128, block_length = 32, steps = 128 yields GSM8K = 87.5, MATH-500 = 41.5, HumanEval = 39.6, MBPP = 50.8 — all within 1–3 pp of the paper's own gen_length = 128 column (83.8 / 39.0 / 37.8 / 50.6) and directionally consistent with the headline gen_length = 256 values (89.1 / 45.1 / 49.4 / 52.4). The gen_length = 256 configuration itself, LoRA gen_length = 256 (Section A.2 Table), and end-to-end GRPO training were not exercised for compute-budget reasons; those cells are inferred from the shorter-budget proxy runs rather than directly verified.
  • The analytical infrastructure behind the paper's scientific claims is not shipped. The release covers training and greedy pass@1/accuracy on GSM8K/MATH-500/HumanEval/MBPP; it does not ship Pass@k or Pass@1024 sampling drivers (Figures 2–5, 10, 11; Tables 4, 5), a per-position entropy-logging hook or forking-token aggregation script (Figures 6, 7, 12), the EB parallel sampler (Figure 8), the top-25%-entropy variant of grpo.py (Figure 9, JustGRPO-Fast), a fixed-permutation variant (Table 5), or the Neg-Entropy and Top-k Margin remaskers used in Figure 11. "Bypassed" is not operationally defined in either the paper or the code. This is the material that turns "we observed" into a checkable claim, and it is essentially absent from the release.
  • Unified-protocol baseline reproductions in Table 2 are not backed by shipped baseline code. The +2.2/+3.3 pp margins over SPG* on GSM8K/MATH-500 are load-bearing for the head-to-head comparison, but the repo ships no training code, configuration, or checkpoints for d1, ESPO, or SPG under the paper's unified protocol. A reproducer would have to fork three external repositories and port each from a one-sentence caption description.
  • Code-benchmark training depends on an externally hosted 21K subset. The AceCoder-Hard filtering ("select 21K challenging samples from AceCoder-87K") is not implemented in the repo; users are asked to download a pre-filtered JSONL. The 49.4 HumanEval / 52.4 MBPP numbers can be evaluated from the released checkpoints but cannot be re-trained without trusting an externally hosted file and an undocumented selection pipeline.
  • Documentation gaps that only surface at code level. data/math.py applies a len(problem) < 1500 filter during training but eval.py does not; the paper does not say whether the reported 45.1% is on the full MATH-500 or a filtered subset. The code-side format-reward regex in data/code.py is strict enough to shift HumanEval/MBPP by several points and is not documented in the paper.
  • General-capability harness is not shipped. eval.py supports only the four reasoning benchmarks; MMLU, MMLU-Pro, HellaSwag, and ARC-C (Table 6) cannot be re-evaluated from the release.

Recommended Changes

Essential

  • Report post-RL Pass@k under both decoding modes. Directly address the base-model-Pass@k-as-proxy concern by comparing post-RL Pass@k for an AR-trained and an arbitrary-order-trained model on at least GSM8K. Speaks to "Contested proxy for the central claim" and "Widening-gap parallel-decoding claim".
  • Add a stochastic-granularity control to the AR-vs-arbitrary Pass@k comparison. Rerun arbitrary-order Pass@k at one-token-per-step (block_length = gen_length) so the order effect is not confounded with the number of independent sampling opportunities. Speaks to "Confound in the AR-vs-arbitrary Pass@k comparison".
  • Release the analytical infrastructure or scope the reproducibility statement. Ship (or clearly designate as out-of-scope) the Pass@k / Pass@1024 driver, per-position entropy-logging hook, EB parallel sampler, top-25%-entropy grpo.py variant, fixed-permutation grpo.py variant, and the MMLU/HellaSwag/ARC-C eval harness. Speaks to "The analytical infrastructure behind the paper's scientific claims is not shipped".
  • Ship baseline configs (or checkpoints) for the unified-protocol reproductions. Release per-baseline hyperparameter YAML/JSON and, ideally, the reproduced d1/ESPO/SPG checkpoints so Table 2's +2.2/+3.3 pp margins can be independently verified. Speaks to "Unified-protocol baseline reproductions in Table 2 are not backed by shipped baseline code".
  • Report seed variance on Table 2 (GSM8K + MATH-500 at minimum). A 2-seed run would put the +2.2/+3.3 pp margins on statistically firmer footing. Speaks to "No seed/variance reporting".
  • Rewrite "random order does not help coverage" and "random order also fails in RL post-training". State that random beats confidence-based on MBPP and GSM8K at Pass@128 but still falls short of AR, and that JustGRPO-Random still yields +3.6 pp over the base model but is dominated by JustGRPO. Speaks to "'Random order does not help coverage' is contradicted by Table 4" and "'Fails' is not a fair label for a +3.6-pp RL gain".
  • State the MATH-500 eval-time filtering protocol. Explicitly say whether the 45.1% is on the full MATH-500 or the len(problem) < 1500 subset, and apply the same filter both places or justify the asymmetry. Speaks to "Documentation gaps that only surface at code level".

Suggested

  • Add the 2x2 train-order × infer-order ablation. Even on GSM8K alone, the (AR-train, AR-infer) / (AR-train, AO-infer) / (AO-train, AR-infer) / (AO-train, AO-infer) grid would isolate the training-scaffold effect from the inference-mode effect. Speaks to "No decoupling ablation of AR-during-training vs AR-at-inference".
  • Add a plain-SFT baseline on the same training data. Quantify how much of the JustGRPO gain is attributable to GRPO vs. to AR-order training exposure. Speaks to "No SFT-only baseline".
  • Add an anti-confidence ablation. Train a variant that decodes high-entropy positions first (or fixes a schedule placing connectives early) to test the entropy-degradation causal story rather than the coarser "AR helps" claim. Speaks to "JustGRPO-Random does not test the entropy-degradation mechanism".
  • Add a matched-position non-fork control to Figure 7. Show that entropy drops more at logical forks than at equally-frequent, similarly-positioned non-forks — to distinguish pathology from the generic H(XY)H(X)H(X \mid Y) \le H(X) effect. Speaks to "Entropy-degradation diagnosis is partly tautological".
  • Fully specify JustGRPO-Fast. Add a loss equation; state whether skipped positions are dropped or held at ρ=1\rho = 1, whether top-25% is per-sequence or per-batch, whether the driving entropy comes from πθold\pi_{\theta_\mathrm{old}} or πθ\pi_\theta, and whether the effective batch of updated tokens is compensated by a step-count or learning-rate change. Speaks to "Underspecified JustGRPO-Fast".
  • Add an explicit forward-pass sketch for the AR policy on a bidirectional dLLM. State how many forward passes per sequence per policy-update step JustGRPO performs, and point at the corresponding grpo.py lines. Speaks to "Overstated framing of the RL objective" and "Underspecified JustGRPO-Fast".
  • Restructure Table 1 with per-length sub-columns. Break the flattened 8-method × 3-length column into rows or sub-columns so values can be attributed to methods and sequence lengths. Speaks to "Table 1 layout obscures attribution".
  • Soften the "general capabilities preserved" framing. Rephrase to "broadly preserved, with minor drops on ARC-C and MMLU-Pro", or add seed variance for the base model on these benchmarks. Speaks to "'General capabilities preserved' oversells small drops without variance".
  • Comment on MBPP/HumanEval non-monotonicity at gen_len 512. One sentence on likely cause (saturation, off-distribution boilerplate). Speaks to "Sequence-length non-monotonicity on code passes uncommented".
  • Explicitly scope "AR-is-better" to math and code. One sentence in the abstract/introduction and conclusion noting the claim is on general reasoning tasks and that arbitrary order remains preferable on constrained tasks such as sudoku/zebra puzzles. Speaks to "Transfer of forking token concept from AR to dLLM is asserted, not verified".
  • Add a bypass-vs-position control and an operational definition of "bypassed". Define bypassing quantitatively (e.g., mean deferral steps) and show that connectives are deferred more than same-position non-connectives. Speaks to "Bypassing analysis lacks an operational definition and a position-matched control".
  • Tighten the JustGRPO-vs-ESPO wall-clock phrasing. Choose one framing — "reaches ESPO's plateau in comparable time and continues improving" — with explicit numbers, rather than "matches ESPO within comparable time" alongside "ESPO saturates early". Speaks to "Minor stumbles" (Figure 9 discussion).
  • Justify the T=1.0T = 1.0 rollout temperature. One sentence relating the RL rollout temperature (Table 3) to the analysis temperature (Section 4, T=0.6T = 0.6). Speaks to "Minor stumbles".
  • Document the format-reward regex and prompt templates for code benchmarks in the paper. These are strict enough to shift HumanEval/MBPP scores by a few points and currently live only in data/code.py. Speaks to "Documentation gaps that only surface at code level" and "Code-benchmark training depends on an externally hosted 21K subset".
  • Add an independently-runnable gen_length = 256 checkpoint-eval script. Even a small wrapper that runs the released JustGRPO checkpoints at the exact headline configuration would let a reader verify the 89.1/45.1/49.4/52.4 numbers directly rather than through proxy runs. Speaks to "Headline JustGRPO checkpoints reproduce closely at the settings we could actually run".