PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2601.15165v4cs.CLcs.AIcs.LG

TL;DR

The paper asks whether arbitrary-order generation actually expands the reasoning potential of dLLMs on general reasoning tasks. It compares this flexibility with AR-ordered training and introduces JustGRPO, which applies standard GRPO using AR order as a training scaffold while retaining the native diffusion model. The study finds that arbitrary order can constrain solution coverage, whereas JustGRPO achieves strong results, including 89.1% accuracy on GSM8K, without sacrificing parallel decoding at inference.

  • Problem

    For general reasoning tasks, the paper examines whether arbitrary-order generation expands dLLM reasoning potential or instead limits solution coverage.

  • Method

    JustGRPO treats the dLLM as an AR policy during RL training, enabling standard GRPO while preserving the native diffusion architecture.

  • Results

    Arbitrary-order generation may constrain reasoning potential, while JustGRPO achieves 89.1% accuracy on GSM8K and retains parallel decoding at inference.

  • Takeaways & Limitations

    Constraining exploration order during training can elicit dLLM reasoning capabilities without giving up parallel decoding during inference.

  • Takeaways & Limitations

    dLLM RL involves ambiguous token-level credit assignment because stochastic unmasking trajectories prevent a unique index-aligned conditional probability.

Abstract

from arXiv · show

Diffusion Large Language Models (dLLMs) break the rigid left-to-right constraint of traditional LLMs, enabling token generation in arbitrary orders. Intuitively, this flexibility implies a solution space that strictly supersets the fixed autoregressive trajectory, theoretically unlocking superior reasoning potential. However, in this paper, we find that for general reasoning tasks (e.g., mathematics and coding), arbitrary order generation may in fact limit the reasoning potential of dLLMs. We observe that dLLMs tend to exploit this order flexibility to bypass high-uncertainty tokens that are crucial for exploration, which can lead to a premature collapse of solution coverage. This observation motivates a rethink of RL approaches for dLLMs, where considerable complexities, such as handling combinatorial trajectories and intractable likelihoods, are often devoted to preserving this flexibility. We show that effective reasoning can be elicited by simply forgoing arbitrary order and applying standard Group Relative Policy Optimization (GRPO) instead. Our approach, JustGRPO, is minimalist yet surprisingly effective (e.g., 89.1% accuracy on GSM8K) while fully retaining the parallel decoding ability of dLLMs. Project page: https://nzl-thu.github.io/the-flexibility-trap

1. Introduction

Although arbitrary-order generation offers parallel decoding and flexible token ordering, the paper finds that it can narrow reasoning potential on general reasoning tasks. JustGRPO instead uses AR ordering as a training scaffold with standard GRPO, preserving parallel decoding at inference.

  • dLLMs offer efficient parallel decoding and arbitrary-order generation beyond the strict left-to-right constraint of AR models.
  • For mathematics and coding, arbitrary-order generation may narrow rather than expand the reasoning potential elicited by RL.The paper evaluates this potential using Pass@k as a measure of solution-space coverage.
  • AR decoding confronts high-uncertainty forking tokens, whereas arbitrary order can resolve easier tokens first and prune branching possibilities.The paper links this bypassing behavior to premature suppression of uncertainty at reasoning forks.
  • Preserving arbitrary-order flexibility complicates RL through combinatorial denoising trajectories and intractable marginal likelihoods.These difficulties can force reliance on inaccurate approximations, making diffusion-specific adaptations costly.
  • JustGRPO treats the dLLM as an AR policy during RL training and applies standard GRPO without modifying its native architecture.AR ordering serves as a training scaffold for exploration, while bidirectional attention and the discrete diffusion formulation remain intact.
  • 89.1% accuracy on GSM8K and 45.1% on MATH-500 demonstrate strong JustGRPO results while the trained model remains compatible with parallel decoding.

2. Preliminaries

dLLMs generate sequences by iteratively denoising masked states, while GRPO normalizes rewards within output groups for policy optimization. Pass@k measures whether correct solutions are sampled and therefore serves as a proxy for reasoning potential and its RL-accessible boundary.

  • dLLMs iteratively denoise a masked state x_t initialized from fully masked tokens, with continuous time t representing the masking ratio.The forward process independently masks each token with probability t.
  • A neural network estimates original token distributions at masked positions, and inference repeatedly unmasks selected tokens until completion.Selection can use confidence heuristics, and always unmasking the leftmost token yields autoregressive generation.
  • GRPO avoids value-function estimation by standardizing each output’s reward against the mean and standard deviation of a sampled group.Its objective uses a clipped surrogate with KL regularization.
  • Pass@k measures the probability that at least one correct solution appears among k independent samples.It is used to quantify the solution space accessible to reasoning optimization.
  • A high Pass@k indicates that correct reasoning trajectories lie within the model’s sampling distribution and can therefore provide positive exploration signals for RL.Consistently failing to sample a solution suggests an intrinsic reasoning boundary that standard RLVR cannot overcome.

3. The Flexibility Trap

Across general reasoning tasks, arbitrary-order decoding does not expand practically reachable solution coverage and instead shows weaker scaling than AR order. The paper attributes this flexibility trap to bypassing high-entropy logical forks, while reduced arbitrariness preserves exploration.

  • Empirical comparison: AR order uncovers more correct solutions as k increases, whereas arbitrary order often remains competitive only at k = 1.This flatter Pass@k scaling appears across three dLLMs and four reasoning benchmarks.
  • Empirical comparison: At k = 1024, arbitrary-order solutions largely overlap with AR solutions and form a smaller subset in practice.On HumanEval, 21.3% of problems are solved exclusively by AR, versus 0.6% exclusively by arbitrary order.
  • Empirical comparison: Pass@k decreases monotonically as semi-autoregressive block size B grows across k ∈ {8, 32, 128}.B = 1 recovers pure AR order, while larger B grants more freedom to choose the next position.
  • Mechanism: The Entropy Degradation: Arbitrary order preferentially resolves easy, high-confidence tokens while bypassing logical connectors and transition words that steer reasoning branches.AR order instead forces the model to confront uncertainty as it arises.
  • Mechanism: The Entropy Degradation: At logical forks, arbitrary order sharply reduces entropy despite comparable global average token entropy, converting open-ended branching into retrospective alignment.Future context resolves the bypassed connectors before they are filled, producing the paper’s entropy degradation phenomenon.
  • Mechanism: The Entropy Degradation: The resulting early anchoring constrains solution coverage, whereas AR ordering preserves high-entropy forks and supports diverse reasoning branches.The conclusion characterizes arbitrary order as favoring inference-time exploitation over exploration for general reasoning.

4. “Just GRPO” for dLLMs

JustGRPO removes arbitrary-order generation during RL training, replacing diffusion-specific optimization with a well-defined autoregressive scaffold and standard GRPO. This resolves key likelihood and credit-assignment difficulties while preserving the dLLM’s native parallel decoding capability.

  • 4.1. The Flexibility Tax in dLLMs’ RL: Arbitrary-order diffusion RL must handle combinatorial denoising trajectories, ambiguous token-level credit assignment, intractable sequence likelihoods, and sampler-learner mismatch.Exact likelihoods require marginalizing over O(N!) trajectories, while dLLMs lack a unique index-aligned conditional probability for standard importance ratios.
  • 4.2. JustGRPO: JustGRPO forgoes arbitrary-order generation during RL, turning the dLLM into a well-defined autoregressive policy for exploration and credit assignment.The autoregressive constraint is used only during training rather than imposed as a structural change to the diffusion model.
  • 4.2. JustGRPO: A surrogate autoregressive policy converts permutation marginalization into an exactly computable likelihood.The policy constructs inputs with observed history and masked future, then evaluates only the next-token prediction.
  • 4.2. JustGRPO: Standard GRPO applies directly once the surrogate policy supplies token-level probability ratios and group-standardized advantages.The objective uses the current-to-old policy ratio and the group-standardized advantage in the clipped GRPO formulation.
  • 4.2. JustGRPO: The autoregressive training scaffold refines the model distribution without imposing causal masking or compromising native parallel sampling.The resulting model remains compatible with parallel samplers at inference.

5. Experiments

Experiments evaluate JustGRPO on mathematical reasoning and coding benchmarks, testing whether autoregressive RL training elicits reasoning while preserving parallel inference. The method achieves strong benchmark performance, retains parallel decoding, and offers favorable training-efficiency trade-offs.

  • 5. Experiments: JustGRPO is evaluated on GSM8K, MATH-500, HumanEval, and MBPP using LLaDA-Instruct without additional task-specific SFT.The experiments use full-parameter fine-tuning and train on official mathematical splits and an AceCoder-87K coding subset.
  • 5.1. Main Results: 89.1% accuracy on GSM8K demonstrates strong performance from standard GRPO without diffusion-specific RL adaptations.JustGRPO also reaches 45.1% on MATH-500 and remains competitive with systems using auxiliary modules or larger private training data.
  • 5.1. Main Results: Table 1 comparisons are subject to heterogeneous baseline settings, while Table 2 reproduces representative baselines under a unified protocol.Differences include fine-tuning method, tokens unmasked per step, generation length, and whether results are quoted or reproduced.
  • 5.2. JustGRPO Preserves Parallel Decoding: +25.5% MBPP accuracy gap at approximately 5 tokens per step, versus +10.6% at 1 token per step, shows stronger gains under aggressive parallelism.The JustGRPO model remains more stable than the original LLaDA-Instruct model as parallelism increases.
  • 5.2. JustGRPO Preserves Parallel Decoding: JustGRPO preserves full parallel decoding compatibility while maintaining a favorable speed-accuracy trade-off.Inference uses the training-free Entropy Bounded sampler under varying numbers of tokens per step.
  • 5.3. Training Efficiency: JustGRPO achieves a better accuracy/wall-time trade-off than ESPO despite exact-likelihood overhead, and JustGRPO-Fast improves it further.JustGRPO-Fast computes probability ratios only at the top-25% highest-entropy positions, eliminating 75% of those forward evaluations.

6. Related Work

Related work has established diffusion language models, explored arbitrary-order generation, and adapted reinforcement learning to their denoising structure. These directions motivate the paper’s focus on the costs and assumptions surrounding order flexibility.

  • Diffusion language models: Diffusion language models evolved from embedding-space approaches toward masked diffusion models that operate directly in discrete text generation.The related work traces this development from continuous-domain diffusion to masked discrete denoising.
  • The value of order arbitrariness: Prior studies validated arbitrary-order generation in constrained tasks and extended it to general reasoning through sketch-first behaviors and optimized decoding schedules.These works underpin the perceived value of order arbitrariness for navigating complex reasoning paths.
  • Reinforcement learning for diffusion language models: Reinforcement learning for dLLMs faces optimization hurdles from the combinatorial explosion of denoising trajectories.Earlier token-level adaptations were limited by ill-defined state transitions and reliance on mean-field approximations.

7. Conclusion

The conclusion argues that arbitrary-order generation can constrain rather than expand reasoning potential on general reasoning tasks by narrowing solution coverage. JustGRPO addresses this with autoregressive alignment during training while preserving parallel decoding at inference.

  • 7. Conclusion: For mathematics and coding, arbitrary-order generation appears to prioritize refining a single trajectory over broader solution coverage.The conclusion characterizes this as a constraint on the reasoning potential accessible to RL.
  • 7. Conclusion: JustGRPO applies standard GRPO with left-to-right alignment during training and retains dLLM parallel decoding at inference.The conclusion presents this as a simpler alternative to complex adaptations designed to preserve arbitrary-order flexibility.

Impact Statement

The paper aims to advance diffusion language models and reinforcement learning for reasoning tasks. It reports that simpler training approaches may be effective and reduce computational costs.

  • The work targets diffusion language models and reinforcement learning for reasoning tasks.
  • The findings suggest that simpler training approaches can be effective and may reduce computational costs.
  • Its broader aim is to improve the understanding and training of language models.

A.2. Training Configuration

JustGRPO trains dLLMs directly with standard GRPO-compatible autoregressive rollouts, using task-specific rewards for mathematics and code generation. The setup uses uniform 125-step training for consistency across datasets.

  • Training setup: Reinforcement learning is applied directly to the dLLM without additional trainable modules, using exact autoregressive sampling for standard GRPO.This enables direct log-probability computation and supports fast, stable convergence.
  • Reward functions: Mathematical completions receive a binary reward of 1 only when the final answer is mathematically equivalent to ground truth.
  • Reward functions: Code reward combines correctness and format components, with correctness measured by unit-test pass rate when the output format is valid.
  • Training schedule: Training uses 125 steps uniformly across datasets, although GSM8K reaches approximately 89% accuracy after around 50 steps while coding tasks benefit from all 125.
  • Reward functions: Format reward scores valid Python Markdown blocks as 1.0, syntactically invalid blocks as 0.5, and invalid blocks as 0.0.
  • Additional analysis: Robustness analyses extend the reasoning-boundary investigation to HumanEval using LLaDA-Instruct, varying temperature and sampling strategies.

B.1. Temperature Analysis

Temperature and sampling analyses show that autoregressive decoding retains stronger reasoning potential than arbitrary-order decoding. Sampling methods with better scaling behavior more closely resemble autoregressive performance characteristics, but still do not catch up overall.

  • Temperature analysis: Autoregressive decoding peaks near T ≈0.6 and degrades above T > 1.0, whereas arbitrary-order decoding peaks at higher temperatures.The comparison extends Pass@k to k = 256 on HumanEval.
  • Temperature analysis: Even under optimized temperature settings, arbitrary-order decoding does not match autoregressive reasoning potential in these experiments.
  • Sampling algorithms: More sophisticated sampling algorithms improve Pass@k over default confidence-based sampling but still remain behind autoregressive order.
  • Sampling algorithms: Higher-Pass@k sampling methods show slightly worse Pass@1 than confidence-based sampling, bringing their overall curves closer to autoregressive Pass@k.
  • Sampling algorithms: Neg-Entropy achieves a 0.970 correlation with autoregressive performance characteristics, indicating that stronger scaling potential corresponds to stronger autoregressive similarity.

C. More Results on Entropy Degradation

Across a broad set of logical connectors used as reasoning forks, autoregressive decoding maintains higher entropy than arbitrary-order decoding. The results support consistent differences in exploration and branching behavior across these tokens.

  • Cross-token analysis: The entropy-degradation phenomenon persists across a comprehensive set of common logical connectors used as forking tokens.
  • Cross-token analysis: Autoregressive order maintains higher average entropy at reasoning forks, indicating active reasoning and decision-making.
  • Cross-token analysis: Arbitrary-order decoding consistently produces lower entropy at these forks, consistent with premature collapse of branching possibilities.
  • Evaluated connectors: The extended experiment evaluates connectors including Therefore, Thus, So, Since, However, First, Then, Finally, Because, Or, Consider, and Write.

D. Is Random Order an Alternative?

Fixed random decoding does not provide a useful alternative: it fails to improve solution coverage, produces structurally broken outputs, and remains weaker after RL. JustGRPO preserves general capabilities while retaining dLLM parallel decoding.

  • Random-order decoding: Random order does not improve solution coverage and can collapse single-shot accuracy relative to autoregressive decoding.Table 4 compares Pass@128 and Pass@1 for confidence-based arbitrary order, fully random order, and AR order on LLaDA-Instruct.
  • Random-order decoding: Random permutations often predict tokens before their surrounding context exists, producing structurally broken outputs.Such outputs rarely earn rewards, limiting the learning signal available to RL.
  • Random-order RL: 82.2% on GSM8K: JustGRPO-Random trails the 89.1% achieved by AR-order JustGRPO after RL.JustGRPO-Random keeps a single random permutation fixed throughout training while leaving the rest of the pipeline unchanged.
  • Capability preservation: JustGRPO preserves general capabilities, with only small task-dependent fluctuations across MMLU, MMLU-Pro, HellaSwag, and ARC-C.The evaluation compares the trained checkpoint with the original LLaDA-Instruct model.
  • Capability preservation: AR factorization is used only during training, so JustGRPO retains parallel decoding and leaves bidirectional refinement as a future direction.The paper connects this extension to corrective decoding and bidirectional-context settings.
Loading 2601.15165v4…