PaperlayerView exact revision on arXiv

Source-linked AI summary

A Few More Tokens Ahead: Reinforcing Parallelism in Diffusion LLMs

Jiahao Kuang, Yangfan Gao, Tao Gui, Yuanbin Wu, Qi Zhang, Xuanjing Huang, Tao Ji

arXiv:2610.18200v1cs.LGcs.AI

TL;DR

Fixed-quota evaluations show that reliable multi-token commitment remains a challenge for diffusion LLMs, despite their parallel prediction capability. RePara reinforces short consecutive blocks with group-relative policy optimization; on LLaDA and SDAR it improves mathematical accuracy over JustGRPO at higher decoding parallelism, with a measured training-step speedup.

  • Problem

    Flexible token prediction does not by itself establish reliable multi-token commitment, and performance drops sharply at higher tokens per diffusion step.

    Source: §1 · pp. 1–2

  • Method

    RePara uses group-relative policy optimization to reinforce short consecutive multi-token rollouts and their joint probabilities, reducing rollout and log-probability recomputation forward counts.

    Source: §1 · p. 2

  • Results

    At Diff#16, RePara improves LLaDA accuracy over JustGRPO by 40.34 percentage points on GSM8K and 14.00 pp. on MATH-500, with a 1.96× training-step speedup in the LLaDA/MATH setup.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Reinforcing short consecutive predictions can improve the accuracy–parallelism trade-off while accelerating training.

    Source: §1 · p. 2, §6 · p. 8

    Mathematical gains do not consistently extend to code generation, where direct application can substantially reduce accuracy.

    Source: Limitations · p. 8

Abstract

from arXiv · show

Parallel generation is widely regarded as an inherent advantage of diffusion large language models. Yet our fixed-quota evaluation of four representative models reveals severe degradation across mathematics and code, with all scores below 15% at 16 tokens per step. Token-level probes identify local predictive potential near the observed prefix, especially in mathematical text. We propose RePara, which reinforces short consecutive token blocks through group-relative policy optimization and reduces rollout and log-probability recomputation costs. On LLaDA and SDAR, \ours improves GSM8K and MATH-500 accuracy over JustGRPO under reduced denoising budgets, with a 1.96x training-step speedup in our LLaDA/MATH setup.

1 Introduction

Diffusion LLMs can predict tokens in parallel, but flexible generation order does not ensure reliable multi-token commitment. The paper identifies local predictive potential, especially in mathematical text, and introduces RePara to reinforce consecutive multi-token actions; experiments report improved high-parallelism mathematical accuracy over JustGRPO.

1 Introduction

  • At TPD 16, every evaluated model–benchmark combination scores below 15% under the fixed 512-token generation protocol.The paper calls the discrepancy between simultaneous token prediction and reliable multi-token commitment the Parallelism Gap.Source: §1 · p. 2
  • Ground-truth-prefix probes find that prediction quality persists farther into mathematical text than code, and two-token consecutive generation is competitive with confidence-based decoding on math.Code results are less consistent.Source: §1 · p. 2
  • RePara uses group-relative policy optimization on joint probabilities of consecutive-token blocks, reducing forward passes for rollout generation and log-probability recomputation.The paper reports improved mathematical reasoning on GSM8K and MATH-500 across both model families, with measured training-step acceleration.Source: §1 · p. 2
  • 40.34 percentage points on GSM8K and 14.00 pp. on MATH-500: RePara outperforms JustGRPO on LLaDA at Diff#16, with a 1.96× training-step speedup in the LLaDA/MATH setup.The reported speedup includes checkpoint and metric-preparation costs.Source: §1 · p. 2

2 A Few More Tokens Ahead

The section examines fixed-quota performance, token-level predictive behavior, and short consecutive decoding. It finds that high quotas sharply reduce accuracy, while local prediction extends farther in mathematical text and two-token consecutive generation is competitive on math.

2.1 The Parallelism Gap

  • Below 15%: every model–benchmark score falls under this level at Diff#16, after performance drops sharply from Diff#4 to Diff#16.The fixed-quota sweep uses four models and four benchmarks with a 512-token generation budget.Source: Figure 2 · p. 3, §2.1 · p. 3
  • At Diff#1–2, LLaDA GSM8K accuracy changes from 81.5% to 81.1%, while SDAR MATH-500 accuracy stays similar at 48.8% versus 49.4%.Code is more sensitive: LLaDA HumanEval Pass@1 falls from 45.1% to 33.5%, and SDAR from 77.4% to 46.3%.Source: §2.1 · p. 3

2.2 Local Predictive Potential

  • Math target-probability h80 boundaries are approximately 2.9–3.6 versus 1.9–2.1 in Code, and confidence boundaries are 3.8–5.7 versus 2.3–3.3.These are retention boundaries relative to first-position quality; probe measurements do not establish joint block correctness or a safe generation length.Source: §2.2 · p. 4

2.3 Short Consecutive Generation

  • Fix-2 matches Diff#2 on LLaDA GSM8K at 81.1% and closely tracks SDAR at 87.0% versus 87.3%, while code results can be substantially worse.On LLaDA2.1-mini, Fix-2 exceeds Diff#2 on GSM8K (82.9% vs. 79.5%) and MATH-500 (46.6% vs. 44.8%); on SDAR it trails by 15.8 pp. on HumanEval and 21.2 pp. on MBPP.Source: Figure 5 · p. 4, §2.3 · p. 4

3 RePara

RePara reinforces consecutive multi-token actions with GRPO while retaining left-to-right progression across blocks. Its rollout and learner replay require fewer forward passes as average block length increases.

3.1 Parallel Rollout Policy

  • RePara samples consecutive token blocks from a shared masked state in one forward pass, then advances left to right across blocks.Setting the maximum block length to one recovers JustGRPO’s single-token rollout structure.Source: Figure 6 · p. 5, §3.1 · p. 5
  • Adaptive selection commits the longest contiguous prefix meeting its score criterion, with a one-token fallback; the scores are not calibrated block-correctness probabilities.Fixed selection commits the full proposed block, while the adaptive variant uses normalized-entropy confidence.Source: §3.1 · p. 5

3.2 Group-Relative Optimization

  • GRPO applies each completion’s standardized correctness advantage to its blocks, reconstructing rollout states while keeping adaptive block boundaries fixed during updates.The learner computes block probabilities from the recorded rollout states.Source: §3.2 · p. 5

3.3 Training Efficiency

  • For an N-token completion with average block length M, rollout and learner replay each use N/M forwards rather than N for single-token generation.Synchronized rounds depend on the completion requiring the most blocks, and wall-clock speedup also depends on backward and optimizer costs, batching, and synchronization.Source: §3.3 · pp. 5–6
  • RePara leaves the backbone architecture unchanged, so inference can use parallel diffusion samplers independently of the training block-length rule.Source: §3.3 · p. 6

4 Experiments

Experiments compare RePara with JustGRPO on LLaDA and SDAR for mathematical reasoning at matched generation lengths and denoising budgets. RePara improves accuracy at higher parallelism, while code transfer is not consistent and training efficiency varies by policy.

4 Experiments

  • The study evaluates LLaDA and SDAR on GSM8K and MATH-500 using step-128 checkpoints, one response per problem, and generation budgets of 128, 256, or 512 tokens.Main comparisons use Fix2 with maximum block length B = 2; LLaDA evaluates fixed TPD through 16, while SDAR evaluates nominal TPD through 4.Source: §4 · p. 6

4.1.1 Accuracy Retention at Higher Parallelism

  • From Diff#8 to Diff#16, LLaDA GSM8K accuracy drops 12.73 pp. for RePara versus 46.86 pp. for JustGRPO; its advantage holds across three generation lengths.Source: §4.1.1 · p. 6, Figure 7 · p. 6
  • Mathematical gains do not consistently transfer to code generation.The paper directs readers to Appendix F for code-generation results and analysis.Source: §4.1.1 · p. 6
  • At Diff#16, LLaDA RePara reaches 68.54% GSM8K and 26.80% MATH-500, exceeding JustGRPO by 40.34 and 14.00 pp., and leads evaluated baselines at Diff#8 and Diff#16.Source: §4.1.1 · p. 6, Table 1 · p. 7

4.1.2 Robustness under Block-Parallel Decoding

  • At L = 512 and Diff#4, SDAR RePara reaches 81.80% GSM8K and 40.40% MATH-500, improving over JustGRPO by 35.17 and 5.00 pp.Source: Table 2 · p. 7, §4.1.2 · p. 7
  • On SDAR GSM8K, moving from Diff#1 to Diff#4 lowers accuracy by 3.57 pp. for RePara versus 35.48 pp. for JustGRPO; gains extend across evaluated lengths.Source: §4.1.2 · p. 7

4.2 Training Efficiency

  • In LLaDA/MATH, Fix2 cuts core training-step time by 49.5% for a 1.98× speedup, retaining a 1.96× speedup with checkpoint and metric-preparation costs.Adaptive2-0.95 reduces core time by 11.9%, giving a 1.14× speedup.Source: Table 3 · p. 7, §4.2 · p. 7

4.3.1 Effect of the Adaptive Training Threshold

  • At L = 512 on LLaDA/GSM8K, γ = 0.95 gives the highest Diff#1 accuracy at 88.48%, while γ = 0.80 is best at Diff#16 with 66.34%.At Diff#16, γ = 0.80 exceeds JustGRPO by 38.14 pp.; the preferred threshold depends on inference parallelism.Source: Table 4 · p. 7, §4.3.1 · p. 7

5 Related Work

Related work spans diffusion architectures, post-training methods, and parallel prediction or speculative decoding. RePara is distinguished by applying task-reward RL to the diffusion model’s own contiguous multi-token generation policy.

Diffusion language models

  • Masked diffusion recovers tokens from corrupted text, while block diffusion combines autoregressive conditioning across blocks with parallel denoising within blocks [Han et al., 2023].Source: §5 · p. 8

Post-training diffusion language models

  • Unlike JustGRPO’s left-to-right single-token rollouts, RePara samples contiguous multi-token blocks and evaluates their joint probabilities under recorded rollout states.Source: §5 · p. 8

Multi-token prediction and speculative decoding

  • Unlike multi-token prediction heads or speculative decoding with an autoregressive verifier, RePara trains the diffusion model’s own policy with task rewards [Chen et al., 2023].Source: §5 · p. 8

6 Conclusion

The paper finds that higher decoding parallelism can reduce accuracy, while token-level probes identify local predictive potential, broader in mathematical text. RePara reinforces consecutive multi-token rollouts and improves mathematical accuracy over JustGRPO at higher parallelism, but code transfer remains limited.

6 Conclusion

  • The evaluation finds substantial accuracy losses under reduced denoising budgets, while probes identify local multi-token predictive potential that is broader in mathematical text.Source: §6 · p. 8
  • On LLaDA and SDAR, RePara improves mathematical reasoning accuracy over JustGRPO at higher decoding parallelism while reducing rollout and log-probability recomputation forward counts.Source: §6 · p. 8
  • Code transfer remains limited: easier training data and JustGRPO initialization yield conditional gains without establishing consistent improvements.The paper relates this limitation to weaker parallel decoding and a shorter local predictive horizon on code.Source: §6 · p. 8

Limitations

The evaluation is limited to mathematical reasoning on GSM8K and MATH-500, with rollout and evaluation budgets that leave longer traces untested. Mathematical gains do not consistently extend to code, and key questions about coding transfer and multi-token training remain unresolved.

Limitations

  • Training uses 256-token rollouts and evaluations use at most 512 tokens, leaving longer reasoning traces untested.Source: Limitations · p. 8
  • Mathematical gains do not consistently extend to code, and the initialization study lacks a matched-budget JustGRPO continuation.These constraints limit conclusions about coding gains and the benefit of switching to multi-token training.Source: Limitations · p. 8
  • LLaDA accuracy still declines from eight to sixteen tokens per step, and the probes do not establish joint block correctness or explain coding regressions.Whether stronger initial multi-token capability enables reliable coding gains remains open.Source: Limitations · p. 8

Ethical Considerations

The appendix documents the ethical scope, forward-probe protocol, relative prediction boundaries, and decoding comparisons. Its probes find longer relative prediction retention in math than code, while emphasizing that these measures do not establish safe block lengths or explain downstream gains.

A Additional Analyses of Parallel Generation

  • The appendix supplies the forward-probe protocol, defines relative prediction boundaries, and documents decoding-comparison scope.Checkpoint and benchmark evaluation settings are in Appendix B.Source: A Additional Analyses of Parallel… · p. 12

A.1 Forward-Probe Protocol

  • Probes use 2,000 shared raw texts per domain and model, with 32-token masked blocks following ground-truth prefixes and no later text supplied.Source: §A.1 · p. 12
  • The AR reference uses teacher forcing and sees earlier ground-truth target tokens, unlike dLLMs; reported curves and heatmaps are unsmoothed observed means.Metrics are target-token probability, maximum-vocabulary confidence, and entropy in nats.Source: §A.1 · p. 12

A.2 Relative Prediction Boundaries

  • Across aggregate dLLM curves, entropy-based retention boundaries extend farther in Math than Code (h80 ≈1.7–2.2 vs. h80 ≈1.3–1.6).Longer prefixes shift the boundary but do not remove the decline farther into the masked block.Source: §A.2 · pp. 12–13
  • The boundaries describe relative mean marginal-token changes, not confidence intervals or calibrated safe block lengths; probes do not measure joint correctness or generated-prefix error accumulation.LLaDA and Dream probes use base checkpoints, unlike downstream instruction-tuned evaluations.Source: §A.2 · p. 13

A.3 Decoding Comparisons and Their Scope

  • Fixed-quota comparisons show code sensitivity on MBPP, while Dream’s small-quota trend is non-monotonic under a sampler differing from its native protocol.The authors characterize these results as specific to the fixed-quota setting, not native decoding performance.Source: §A.3 · p. 13
  • SDAR’s Fix-2 trails Diff#2 on HumanEval and MBPP, while LLaDA2.1-mini’s Fix-2 leads on HumanEval but trails on MBPP.The comparisons support short consecutive decoding as a mathematical rollout policy, with less consistent behavior on code.Source: §A.3 · p. 13
  • Matched token quotas do not isolate generation order or establish equal computation and latency, and the 32-token probe window is not an optimal rollout length.Source: §A.3 · p. 13

B.1 Common Evaluation Protocol

The evaluation protocol specifies datasets, prompts, sampling, scoring, and coding-study overrides. Model-specific decoding settings and baseline checkpoints differ across comparisons, while the training pipeline builds on JustGRPO.

B.1 Common Evaluation Protocol

  • Evaluation covers GSM8K, MATH-500, HumanEval, and MBPP, using backbone chat templates and benchmark-specific answer extraction or test execution.Default sampling uses temperature 0, one response per problem, seed 113, and batch size one per device; coding studies include specified overrides.Source: §B.1 · pp. 13–14

B.2 Model-Specific Decoding

  • Fixed-quota decoding commits the k highest-confidence masked positions, with window sizes and selection rules varying by model.SDAR’s main experiments instead use native four-token blocks and dynamic low-confidence selection.Source: §B.2 · p. 14

B.3 Baseline Checkpoints

  • Main mathematical comparisons use LLaDA-8B-Instruct and SDAR-8B-Chat, alongside released mathematical baselines re-evaluated under the shared protocol.Source: §B.3 · p. 14

C Training Details

  • The training pipeline builds on public data preparation and optimization settings following JustGRPO.Source: C Training Details · p. 14

C.1 Data and Training Configuration

Training uses separate GSM8K, MATH, and AceCode tasks with fixed-length rollouts and specified batch settings. RePara applies group-relative advantages to blocks, with a clipped objective and a single on-policy update per rollout batch.

C.1 Data and Training Configuration

  • GSM8K, MATH, and AceCode are trained as separate tasks, with MATH-trained models evaluated on MATH-500.For mathematical tasks, reference answers are used for reward computation but not included in prompts.Source: §C.1 · p. 14
  • Each global batch contains 64 prompts with 16 completions each, and training uses 256-token rollouts without EOS early stopping.Training uses all model parameters in bfloat16 with a constant learning rate.Source: §C.1 · pp. 14–15, Table 5 · p. 15
  • Main comparisons report individual step-128 checkpoint accuracies without estimating variability across independent training seeds.The coding interventions use single training seeds, and the initialization study lacks a matched-budget JustGRPO continuation.Source: §C.1 · p. 15, Table 6 · p. 15

C.2 Rewards and Optimization

  • Math completions receive +1 or −1 for answer correctness, while AceCode reward combines test-pass fraction and format score.Source: §C.2 · p. 15
  • Each sampled completion’s group-relative advantage is assigned to all its blocks, based on within-group reward standardization.The advantage uses the group reward mean and standard deviation, with ϵA = 10^-4 for numerical stability.Source: §C.2 · p. 15
  • The clipped surrogate normalizes by fixed generation length times the number of nonzero-advantage completions; batches with none contribute no policy gradient.Rollout block boundaries stay fixed during probability recomputation, including for Adaptive selection.Source: §C.2 · p. 15
  • Training uses one on-policy update per rollout batch without a KL penalty, rather than repeated updates against stored old-policy probabilities.Source: §C.2 · p. 16

D Detailed Results

The detailed results favor two-token over four-token training on MATH-500 under matched settings, and RePara reduces several training-phase costs. Adaptive block selection has a synchronization tail, while applying RePara directly to coding substantially degrades accuracy.

D Detailed Results

  • 2.8–5.2 pp. for LLaDA and 19.6–26.4 pp. for SDAR: Fix2 outperforms Fix4 on MATH-500 at Diff#1, Diff#2, and Diff#4 with L = 256.The comparison uses matched inference settings.Source: D Detailed Results · p. 16, Table 8 · p. 17
  • Continuing Fix2 for 128 more steps with B = 4 improves over direct Fix4 training but does not consistently surpass Fix2.The comparison favors shorter training blocks, with SDAR particularly sensitive to increasing B.Source: D Detailed Results · p. 16

E Training Efficiency Analysis

  • The detailed efficiency analysis uses LLaDA trained on MATH and evaluated on MATH-500; similar training speedups are observed in other settings.This is the representative case analyzed in detail.Source: E Training Efficiency Analysis · p. 16

E.1 Timing Protocol

  • 497 and 479 H200 GPU-hours are reported for 128-update, fixed-256-token-rollout MATH training on LLaDA and SDAR, respectively.Speedup is the JustGRPO mean divided by the corresponding method’s mean; Core and Step include different timing boundaries.Source: §E.1 · p. 16

E.2 Phase Costs

  • 49.6%, 51.2%, and 49.5%: Fix2 reduces rollout, log-probability recomputation, and backward time, respectively.Halving the block count saves rollout and replay forwards and the backward pass for each replayed block; similar checkpoint and metric costs keep Step speedup close to Core speedup.Source: §E.2 · pp. 16–17

E.3 Synchronization Tail

  • 153.23 logical blocks per completion still require 222.23 synchronized rounds per update; Adaptive’s round count falls 13.2% versus JustGRPO, alongside an 11.9% Core-time reduction.31.1% of Adaptive’s completion–round slots are assigned to finished rows; this measures inactive batch slots, not GPU idle time.Source: §E.3 · pp. 17–18

F.1 Applying RePara to Coding: Accuracy Losses

  • Directly applying RePara’s short-block rollouts and GRPO optimization to coding produces substantial accuracy degradation rather than reproducing the mathematical gains.The coding experiments use LLaDA-8B-Instruct and compare Fix2 and Adaptive2-0.95 with JustGRPO.Source: §F.1 · p. 19

F.2 Targeted Training Interventions

The coding interventions yield conditional rather than consistent gains: multi-token continuation improves higher-quota scores while reducing Diff#1 accuracy, and easier data does not produce consistent coding gains.

F.2 Targeted Training Interventions

  • Fix2 on AceCode-Medium reaches 30.80% full-pass success versus 7.80% on Hard500, yet remains near-zero on HumanEval/MBPP; Adaptive2-0.95 gains at Diff#4 but loses at Diff#1 and held-out Medium500.Easier data therefore does not yield consistent coding gains.Source: §F.2 · p. 20
  • The authors suggest weak initial multi-token rollouts may make task-reward learning difficult; stronger single-token JustGRPO initialization still leaves subsequent gains conditional.This interpretation is consistent with code's weaker parallel decoding and shorter local predictive horizon.Source: §F.2 · p. 20
  • At length 512, Adaptive2-0.95 raises HumanEval Diff#8 from 5.49% to 15.85%, while Diff#1 falls from 52.44% to 41.46% after multi-token continuation.The continuation starts from a single-token JustGRPO checkpoint.Source: §F.2 · p. 20, Figure 11 · p. 21

Related papers

Loading 2610.18200v1…