PaperlayerView exact revision on arXiv

Source-linked AI summary

Fork-dLLM: Avoiding the Flexibility Trap in Diffusion Language Models

Stipe Frković, Metod Jazbec, Christian A. Naesseth

arXiv:2609.39859v1cs.LG

TL;DR

Confidence-based dLLM samplers can defer fork positions, reducing diversity and pass@k scaling, while AR sampling sacrifices parallelism. Fork-dLLM uses AR ordering only at uncertain fallback steps; ForkGRPO applies GRPO only there while preserving exact policy-likelihood ratios. Across mathematics and code, Fork-dLLM matches AR pass@k scaling with 2–3× fewer NFEs, while ForkGRPO matches or outperforms AR-based GRPO at lower training cost.

  • Problem

    Confidence-based sampling can defer uncertain, high-entropy fork positions, reducing generation diversity and pass@k scaling and limiting RL post-training gains.

    Source: §1 · p. 1

  • Method

    Fork-dLLM uses AR-style ordering at uncertain fallback steps while retaining parallel generation otherwise; ForkGRPO uses its rollouts and applies GRPO only at fallback steps.

    Source: §1 · p. 2

  • Results

    Across mathematics and code, Fork-dLLM matches AR sampling’s pass@k scaling with 2–3× fewer NFEs; ForkGRPO beats AR-based GRPO at equal compute and matches it at a fraction of the training cost.

    Source: §5 · p. 9

  • Takeaways & Limitations

    Avoiding the flexibility trap does not require abandoning confidence-based parallel decoding, and fallback-only GRPO preserves exact policy-likelihood ratios.

    Source: §1 · p. 2

    GRPO experiments use one LoRA configuration on LLaDA-8B-Instruct; evaluation on other models, with full fine-tuning, or with longer training remains future work.

    Source: §5 · p. 10

Abstract

from arXiv · show

Masked diffusion language models (dLLMs) have shown strong potential for faster inference through parallel token generation when combined with confidence-based samplers. However, recent work has shown that such methods can defer unmasking high-entropy fork positions at which multiple plausible continuations exist. This results in reduced generation diversity, as shown by worse pass@k scaling, and limits gains obtainable from RL post-training. To avoid this flexibility trap, prior work advocated for autoregressive (AR) sampling. Here, we show that discarding confidence-based sampling is unnecessary and, once inference cost is taken into account, wasteful. We first propose Fork-dLLM, a simple hybrid sampler that uses AR-style ordering only at uncertain fallback steps while retaining parallel generation otherwise. We then extend the same principle to post-training with ForkGRPO, which uses Fork-dLLM rollouts and applies the GRPO objective only at fallback steps, preserving exact policy-likelihood ratios while substantially reducing rollout and optimization cost. In our experiments, Fork-dLLM matches the strong pass@k scaling of AR sampling while being 2-3x more efficient, and ForkGRPO achieves downstream performance comparable to or better than AR-based GRPO baselines at a substantially lower training cost.

1 INTRODUCTION

Confidence-based sampling can reduce diversity at high-entropy fork positions, while autoregressive sampling sacrifices dLLM parallelism. The paper introduces fallback-only AR ordering for sampling and post-training, reporting strong scaling and lower costs.

1 INTRODUCTION

  • On GSM8K, ForkGRPO beats JustGRPO-Fast at equal training budget and reaches JustGRPO accuracy with 8× fewer H100-hours; overall, it outperforms AR-based GRPO at equal compute.The paper also reports that ForkGRPO matches AR-based GRPO at a fraction of the training cost.Source: Figure 1 · p. 1, §1 · p. 2
  • Confidence-based sampling can postpone uncertain, high-entropy fork positions, reducing generation diversity, pass@k scaling, and RL post-training gains; AR sampling sacrifices parallelism.Prior work terms this problem the flexibility trap and advocates AR sampling.Source: §1 · pp. 1–2
  • ForkGRPO uses Fork-dLLM rollouts and applies GRPO only at fallback steps, preserving exact policy-likelihood ratios while reducing rollout and optimization cost.Only one token is committed at each fallback step.Source: §1 · p. 2
  • Across mathematics and code benchmarks, Fork-dLLM matches AR sampling’s pass@k scaling while requiring 2–3× fewer NFEs.It applies AR-style ordering at uncertain fallback steps without training, extra model evaluations, or new hyperparameters.Source: §1 · p. 2
  • When tuned, other confidence-based samplers also retain strong pass@k scaling at a fraction of AR sampling’s cost.This supports avoiding the flexibility trap without abandoning confidence-based parallel decoding.Source: §1 · p. 2

2 BACKGROUND

The background describes masked dLLM training and iterative unmasking, then contrasts confidence-based parallel sampling with autoregressive rollouts used for exact-likelihood GRPO.

2 BACKGROUND

  • Masked dLLMs independently replace tokens with an absorbing [MASK] token and train a bidirectional transformer to reconstruct clean tokens at masked positions.The model predicts each clean token conditioned on the entire partially masked sequence.Source: §2 · p. 2
  • At each sampling step, the model proposes tokens for masked positions, and the chosen commit set determines generation order and parallelism.Fast-dLLM unmasks positions above a confidence threshold; if none qualify, it commits the highest-confidence position to ensure progress.Source: §2 · p. 3
  • Diffu-GRPO approximates sequence likelihood with a mean-field factorization across positions and estimates per-token likelihoods in one forward pass.GRPO is difficult to apply to dLLMs because their generation process lacks left-to-right policy-likelihood factorization.Source: §2 · p. 3
  • JustGRPO uses autoregressive rollouts and trajectory-level policy likelihoods for exact ratios, but incurs substantial costs in rollout generation and per-position likelihood evaluation.Its group-normalized advantage uses each rollout’s reward and the group’s reward mean and standard deviation.Source: §2 · pp. 3–4

3 METHODS

The methods alter only the Fast-dLLM fallback rule, treating rare high-uncertainty steps differently, and carry that strategy into GRPO by optimizing only at fallback steps.

3 METHODS

  • Fork-dLLM changes only Fast-dLLM’s fallback rule, and ForkGRPO applies the GRPO objective only at those fallback steps.The approach is presented as addressing the flexibility trap without AR sampling rollouts or ELBO-style likelihood approximations.Source: §3 · p. 4

3.1 FORK-DLLM

  • Approximately one eighth of committed tokens occur at fallback steps, but those tokens contribute nearly half of total token entropy.This motivates treating comparatively rare, high-uncertainty fallback steps differently while retaining confidence-based parallel decoding.Source: §3.1 · p. 4, Figure 2 · p. 4

3.2 FORKGRPO

  • ForkGRPO uses Fork-dLLM rollouts and evaluates the GRPO objective only at fallback steps, where single-token commits preserve exact policy-likelihood ratios.It avoids the likelihood approximations used by diffu-GRPO.Source: §3.2 · p. 5
  • Non-fallback tokens still affect rollout rewards and later fallback states, but they do not directly enter the policy objective, reducing training cost versus JustGRPO.ForkGRPO combines sparse policy-likelihood evaluation with parallel rollouts.Source: §3.2 · p. 5
  • Under a matched post-training budget, ForkGRPO substantially outperforms JustGRPO-Fast, which uses autoregressive rollouts and selects a fixed fraction of positions by entropy.ForkGRPO instead uses parallel rollouts and a variable, trajectory-dependent set of fallback positions.Source: §3.2 · p. 5

4 EXPERIMENTS

The experiments evaluate whether Fork-dLLM preserves pass@k scaling while reducing inference cost, then test ForkGRPO’s efficiency and downstream performance under fixed training budgets.

4 EXPERIMENTS

  • The sampler evaluation covered LLaDA-8B-Instruct and Dream-v0-Instruct-7B on mathematics and code benchmarks, using block length 32 and generation length 256.GRPO comparisons used fixed 4×H100, 12-hour budgets; ForkGRPO and FastGRPO used block length 32.Source: §4.1 · pp. 5–6, §4.2 · p. 8

4.1 PASS@k SCALING

  • All three confidence-based samplers remained close to AR pass@k scaling across datasets and substantially outperformed AR on pass@NFE.Fork-dLLM was generally the strongest parallel sampler, particularly at larger k.Source: §4.1 · p. 6, Figure 3 · p. 6
  • Fallback steps most often committed connectives such as “Since,” “First,” and “Thus,” while parallel steps mainly committed dataset-specific words following an established reasoning path.The authors interpret this contrast as consistent with fallback steps resolving forks directly.Source: §4.1 · p. 7, Figure 4
  • Making parallel generation greedy left scaling essentially unchanged, whereas replacing AR-style fallback ordering with lowest-confidence, highest-entropy, or random selection worsened scaling.The results indicate fallback ordering, rather than parallel-step stochasticity, drives most of Fork-dLLM’s gains.Source: §4.1 · p. 7, Figure 5 · p. 7

4.2 GRPO TRAINING

  • FastGRPO’s two runs initially matched ForkGRPO in reward but collapsed after roughly 350 and 220 steps, while ForkGRPO remained stable in an extended 24-hour run.The collapses coincided with sharp rises in fallback ratio and token entropy; the authors hypothesize that fallback selection independent of model confidence stabilizes training.Source: §4.2 · p. 8, Figure 6 · p. 8
  • At matched training steps, the methods reached similar accuracy, but JustGRPO-Fast required 1.6–1.9× more backward passes than ForkGRPO.The authors suggest Fork-dLLM’s fallback positions are more informative per update than entropy-percentile selection.Source: §4.2 · p. 8, Table 1 · p. 9
  • ForkGRPO completed more than twice as many training steps on every benchmark and improved over the base model by +6.7 vs. +2.1 points on GSM8K and +5.5 vs. +0.6 on HumanEval.On GSM8K, it completed 3.4× as many steps; its checkpoints also used fewer NFEs on GSM8K, MATH-500, and HumanEval at every threshold.Source: §4.2 · p. 8, Figure 7 · p. 9
  • With substantially less training compute, ForkGRPO reached comparable peak GSM8K accuracy to released JustGRPO and performed consistently better on MATH-500, at roughly 40% versus 37–38%.JustGRPO attained similar GSM8K accuracy at lower NFEs.Source: §4.2 · p. 9

5 CONCLUSION

Fork-dLLM preserves parallel decoding while matching AR pass@k scaling at lower inference cost, and ForkGRPO extends this approach to lower-cost RL post-training. The reported results and limitations are bounded by the sampler-tuning and GRPO settings described here.

5 CONCLUSION

  • Fork-dLLM matches AR sampling’s pass@k scaling across mathematics and code benchmarks with 2–3× fewer NFEs, while ForkGRPO matches or outperforms AR-based GRPO at lower training cost.The conclusion also reports that ForkGRPO substantially outperforms AR-based GRPO under the same compute budget.Source: §5 · p. 9
  • Confidence-based samplers were tuned on HumanEval with LLaDA-8B-Instruct and used unchanged elsewhere, although per-model and per-benchmark tuning may yield better operating points.This includes Fork-dLLM’s operating point.Source: §5 · pp. 9–10
  • GRPO experiments used one LoRA configuration on LLaDA-8B-Instruct, and testing other models, full fine-tuning, longer training, and other dLLM RL methods remains future work.Because of computational costs, the comparison used the released JustGRPO checkpoint rather than retraining JustGRPO under the same setup.Source: §5 · p. 10

A RELATED WORK

Related work addresses diversity in confidence-based sampling through several training-free methods, while dLLM RL work explores likelihood approximations, alternative objectives, and autoregressive rollouts. Fork-dLLM and ForkGRPO instead use confidence-based fallback steps to retain parallel decoding while targeting uncertain positions and exact likelihood ratios.

A RELATED WORK

  • Deterministic confidence ordering can reduce diversity across repeated generations (Ni et al., 2026; Olausson et al., 2026).Training-free alternatives include TCT’s stochastic ordering (Olausson et al., 2026), TAPS’s perturbed early states (Wu et al., 2026b), feature-space repulsion (Lamont et al., 2026), and Order-Token Search (Shen et al., 2026).Source: A RELATED WORK · p. 14
  • Fork-dLLM differs from these methods by changing Fast-dLLM’s fallback to commit the leftmost masked position when no position clears the confidence threshold.Low-confidence or high-entropy regions have been identified as important exploration points (Fu et al., 2025).Source: A RELATED WORK · p. 14
  • dLLM RL methods address the lack of tractable left-to-right sequence likelihood with approximate likelihoods, policy-gradient estimators, sequence-level ELBOs, likelihood-ratio estimators, or likelihood-free alternatives.Other approaches use denoising trajectories or optimize token and unmasking decisions; TraFL studies solution-coverage loss under reward-maximizing post-training.Source: A RELATED WORK · p. 14
  • JustGRPO uses AR rollouts for exact trajectory likelihoods but gives up parallel generation, whereas ForkGRPO applies GRPO only at sparse fallback steps with exact token-level likelihood ratios.ForkGRPO retains parallel Fork-dLLM rollouts.Source: A RELATED WORK · p. 14

B.1 STOCHASTIC SAMPLING

The appendix varies confidence threshold and sampling temperature to select settings for Fork-dLLM, Fast-dLLM, and TCT. Figures show pass@k and pass@NFE curves, and Table 2 reports the selected values.

B.1 STOCHASTIC SAMPLING

  • The appendix varies confidence threshold λ and sampling temperature τ for Fork-dLLM, Fast-dLLM, and TCT to find optimal hyperparameter values.Experiments used LLaDA-8B-Instruct on HumanEval; TCT’s additional position temperature was fixed at τpos = 0.1.Source: §B.1 · p. 15
  • Figures 8–10 show pass@k curves on top and pass@NFE curves below for Fork-dLLM, Fast-dLLM, and TCT, with AR and Fork-dLLM included as indicated baselines.Both rows plot the same data using different metrics.Source: §B.1 · p. 15, Figure 8 · p. 15, Figure 9 · p. 15
  • Table 2 lists the selected hyperparameter values for each sampling method.Source: §B.1 · p. 15, Table 2 · p. 15

B.2 GRPO TRAINING

The GRPO training appendix points to a full configuration table and clarifies two details of its settings: a JustGRPO-Fast-specific row and an omitted clipping range.

B.2 GRPO TRAINING

  • Table 3 lists the full GRPO configuration and its hyperparameters.Source: §B.2 · p. 16, Table 3 · p. 16
  • The starred row in Table 3 applies only to JustGRPO-Fast.Source: Table 3 · p. 16
  • The clipping range ϵ is omitted because a single policy update step per rollout means it never binds.Source: Table 3 · p. 16

C ADDITIONAL FIGURES

This appendix reports Dream-v0-Instruct-7B results and summarizes Fork-dLLM’s benchmark-specific scaling tradeoffs and comparisons with Fast-dLLM and TCT.

C ADDITIONAL FIGURES

  • On mathematics benchmarks, Fork-dLLM preserves AR scaling at a fraction of the cost; on coding tasks, it trades slightly worse pass@k scaling for better pass@NFE scaling.Figure 11 shows pass@k and pass@NFE curves on Dream-v0-Instruct-7B.Source: C ADDITIONAL FIGURES · p. 17, Figure 11 · p. 17
  • Fork-dLLM consistently matches or outperforms Fast-dLLM and TCT.Source: C ADDITIONAL FIGURES · p. 17
  • Tables 4 and 5 report Pass@1 and Pass@64 results for LLaDA-8B-Instruct and Dream-v0-Instruct-7B, respectively.Best results per column are bolded in each table caption.Source: Table 4 · p. 17, Table 5 · p. 17, C ADDITIONAL FIGURES · p. 17
  • Figure 12 plots Fork-dLLM and its ablations on HumanEval using LLaDA-8B-Instruct.Source: Figure 12 · p. 18
  • Figure 13 compares accuracy versus NFEs for final ForkGRPO and JustGRPO checkpoints on GSM8K and MATH-500, using Fast-dLLM and EB across thresholds.The JustGRPO checkpoint is from Ni et al. (2026).Source: Figure 13 · p. 18
Loading 2609.39859v1…