Command Palette
Search for a command to run...
柔軟性の罠:拡散言語モデルにおける任意順序生成の価値の再考
柔軟性の罠:拡散言語モデルにおける任意順序生成の価値の再考
概要
拡散大規模言語モデル(dLLM)は、従来のLLMの厳格な左から右への生成制約を打破し、任意の順序でのトークン生成を可能にする。直感的には、この柔軟性は固定された自己回帰軌道を厳密に包含する解空間を示唆し、理論上は優れた推論能力を引き出す可能性がある。しかし本論文では、一般的な推論タスク(例:数学やコーディング)において、任意順序生成が実際にはdLLMの推論能力を制限する可能性があることを見出す。我々は、dLLMがこの順序柔軟性を悪用して、探索に重要であるが不確実性の高いトークンを回避する傾向があり、これが解の網羅性の早期崩壊につながる可能性があることを観察する。この観察は、dLLMのための強化学習アプローチの再考を促すものである。そこでは、組み合わせ軌道や扱いにくい尤度の処理など、かなりの複雑さがこの柔軟性を維持することに費やされることが多い。我々は、単に任意順序を放棄し、標準的なGroup Relative Policy Optimization(GRPO)を適用するだけで、効果的な推論を引き出せることを示す。我々のアプローチ、JustGRPOは、最小限でありながら驚くほど効果的であり(例:GSM8Kで89.1%の精度)、dLLMの並列デコーディング能力を完全に保持している。
One-sentence Summary
Researchers from Tsinghua University and Alibaba Group show that arbitrary-order generation in Diffusion LLMs (dLLMs) hinders reasoning by allowing models to bypass high-uncertainty tokens crucial for exploration, and propose JustGRPO, which forgoes this flexibility to apply standard Group Relative Policy Optimization, achieving 89.1% accuracy on GSM8K while retaining parallel decoding.
Key Contributions
- Identifies a limiting behavior of arbitrary-order generation in diffusion LLMs for general reasoning: the flexibility causes models to bypass high-uncertainty tokens essential for exploration, leading to premature collapse of solution coverage.
- Introduces JustGRPO, a minimalist alignment approach that forgoes arbitrary-order generation and applies standard Group Relative Policy Optimization (GRPO) directly, deliberately constraining exploration order to improve reasoning while preserving the parallel decoding capability of diffusion LLMs at inference.
- Achieves 89.1% accuracy on GSM8K, demonstrating that effective reasoning can be elicited without complex RL adaptations for order arbitrariness, and fully retains inference-time parallel decoding.
Introduction
Diffusion Large Language Models (dLLMs) are emerging as an alternative to autoregressive (AR) models by framing text generation as a discrete denoising process. They promise two theoretical advantages over fixed left-to-right decoding: efficient parallel decoding and the ability to generate tokens in arbitrary order. While the decoding speed benefits are well documented, the value of arbitrary-order generation for complex tasks such as mathematics and coding remains an open question. Prior work has largely assumed that this flexibility is essential, and consequently, reinforcement learning (RL) methods for dLLMs have grown complex, grappling with a combinatorial explosion of denoising trajectories, intractable marginal likelihoods, and reliance on inaccurate approximations. Some recent approaches introduce auxiliary policies or step-wise MDPs to handle arbitrary order, but all retain the mechanism without questioning whether it is actually helpful.
The authors challenge this assumption with a counter-intuitive finding: for general reasoning tasks, arbitrary-order generation may actually narrow, rather than expand, the reasoning potential accessible to RL. Using Pass@k as a proxy for the upper bound of solution space that RL can exploit, they show that restricting a dLLM to standard AR order yields higher Pass@k and a higher reasoning boundary than its flexible counterpart. The mechanism behind this is "entropy degradation": reasoning depends on sparse forking tokens such as "Therefore" or "Since," which steer the logical trajectory. AR decoding forces the model to confront the high uncertainty at these forks, preserving diversity in exploration. Arbitrary-order generation, by contrast, tends to fill in low-uncertainty tokens first, anchoring the reasoning path and prematurely resolving the ambiguity of the forking tokens before they are infilled.
Based on this insight, the authors propose JustGRPO, a return to simplicity. They show that eliciting reasoning ability does not require diffusion-specific RL adaptations; instead, the dLLM can be treated as an AR policy during RL training, allowing standard GRPO to be applied without modification. This AR formulation serves only as a training scaffold for better exploration; the model's native architecture, including bidirectional attention and discrete diffusion formulation, remains untouched. As a result, reasoning elicitation (which benefits from sequential exploration) and inference execution (which benefits from parallel decoding) are decoupled. JustGRPO achieves strong results, including 89.1% on GSM8K and 45.1% on MATH-500, competitive with methods that rely on intricate diffusion-specific RL adaptations, while preserving the parallel decoding efficiency of dLLMs at inference.
Method
The authors identify a "flexibility trap" in diffusion language models (dLLMs) where the theoretical flexibility of arbitrary-order generation inadvertently limits reasoning potential. This limitation stems from a phenomenon termed "entropy degradation," where the adaptive decoding process preferentially bypasses critical logical connectors, causing a sharp decrease in entropy at these branching points. As shown in the entropy analysis, while the global average token entropy remains comparable between decoding modes, the entropy at logical forks drops significantly under arbitrary order, effectively anchoring the trajectory to pre-determined conclusions early in the generation process.
To address this, the authors propose JustGRPO, a method that explicitly forgoes arbitrary-order generation during the reinforcement learning stage. By constraining the process to strict autoregressive order, the method transforms the dLLM from a chaotic sequence denoiser into a well-defined autoregressive policy. This design choice resolves the ambiguity in token-level credit assignment and the intractable sequence likelihood inherent in diffusion models.
The core of JustGRPO involves constructing a surrogate autoregressive policy, denoted as πθAR. For a given query q and history O<k, the method constructs a sequence x~k where the past tokens are observed and the future tokens are masked. The probability of the next token ok is defined as the softmax of the model logits at position k. This formulation allows the intractable marginalization over permutations to be converted into an exactly computable likelihood, expressed as the product of conditional probabilities.
With this surrogate policy established, the authors apply standard Group Relative Policy Optimization (GRPO). For each query, a group of outputs is sampled from the old policy. The advantage A^i,k is computed by standardizing the reward against the group mean and standard deviation. The optimization objective maximizes a clipped surrogate function with a KL regularization term, ensuring stable updates while leveraging the group-level reward normalization.
Despite its effectiveness, applying exact GRPO to dLLMs incurs a structural overhead because the model must evaluate each position independently rather than in a single causal forward pass. To mitigate this computational cost, the authors introduce JustGRPO-Fast. Motivated by the observation that reasoning is steered by a sparse set of high-entropy forking tokens, this variant computes the probability ratios ρi,k only at the top-25% highest-entropy positions. This optimization eliminates 75% of the forward evaluations required for ratio computation.
This targeted computation significantly improves the accuracy-to-wall-time trade-off, allowing JustGRPO-Fast to match the peak accuracy of approximation-based baselines in less time while continuing to improve performance. Crucially, the autoregressive constraint is applied exclusively during RL training, preserving the parallel sampling capability of the dLLM for efficient inference.
Experiment
The evaluation first contrasts arbitrary-order with autoregressive (AR) decoding in diffusion LLMs, showing that while arbitrary order matches AR at k=1, it scales worse with increasing k, covers a smaller effective solution space, and reduces Pass@k as the block size increases. The mechanism is entropy degradation: arbitrary order defers logical connectors, lowering their entropy upon decoding and constraining exploration. Applying JustGRPO, which enforces AR ordering during RL training, yields strong reasoning and coding performance competitive with prior diffusion-specific methods, preserves parallel decoding compatibility (with gains growing under higher parallelism), and offers a favorable accuracy/wall-clock tradeoff, further improved by JustGRPO-Fast which focuses updates on the highest-entropy positions.
JustGRPO, a simple application of standard GRPO without diffusion-specific adaptations, achieves competitive or higher accuracy on reasoning benchmarks like GSM8K and MATH-500 compared to prior diffusion-specific methods. It remains competitive with approaches that use auxiliary modules or larger private training data, despite heterogeneous baseline settings. The findings suggest that full combinatorial optimization over denoising trajectories may not be necessary for strong reasoning performance. JustGRPO matches or outperforms diffusion-specific RL methods on GSM8K and MATH-500 despite its simplicity. It remains competitive with baselines that rely on auxiliary modules or significantly larger private training data. The conclusion holds even when representative baselines are reproduced under a unified protocol aligned with JustGRPO.
Under a unified evaluation protocol with full fine-tuning and one token per decoding step, JustGRPO achieves the highest accuracy across all reported reasoning and code generation benchmarks, surpassing reproduced diffusion-specific baselines. The gains are consistent across mathematical reasoning and code tasks, supporting the effectiveness of treating diffusion LLMs as sequential generators during training. JustGRPO outperforms all reproduced baselines on GSM8K, MATH-500, HumanEval, and MBPP. The performance advantage holds across both mathematical reasoning and code generation tasks, with the largest margins on code benchmarks. The unified protocol confirms that JustGRPO remains ahead even when comparing against baselines that use auxiliary modules or larger private training data.
JustGRPO, a straightforward application of standard GRPO without diffusion-specific adaptations, outperforms or matches prior diffusion-specific reinforcement learning methods on reasoning benchmarks such as GSM8K and MATH-500. Under a unified evaluation protocol with full fine-tuning and one token per decoding step, it achieves the highest accuracy across all reported reasoning and code generation tasks, including HumanEval and MBPP, surpassing reproduced baselines. The findings indicate that full combinatorial optimization over denoising trajectories is not necessary for strong performance, and that treating diffusion LLMs as sequential generators during training is effective.