PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast and Accurate Causal Parallel Decoding using Jacobi Forcing

Lanxiang Hu, Siqi Kou, Yichao Fu, Samyam Rajbhandari, Tajana Rosing, Yuxiong He, Zhijie Deng, Hao Zhang

arXiv:2512.14681v1cs.CL

TL;DR

Adapting AR models for parallel decoding faces a pretrain-to-posttrain mismatch involving masked data and bidirectional attention. Jacobi Forcing progressively distills AR models on their own generated trajectories while preserving causal attention, achieving 3.8× speedup on coding and math benchmarks with largely preserved accuracy.

  • Problem

    AR-to-dLLM adaptation has limited speedup under quality constraints amid a mismatch between pretrained causal attention and post-training data and attention patterns.

    Source: §1 · p. 1

  • Method

    Jacobi Forcing progressively distills AR models on their own generated parallel-decoding trajectories without modifying causal attention.

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

  • Results

    3.8× speedup while largely preserving accuracy, with Jacobi Forcing Models serving as efficient parallel decoders across coding and math benchmarks.

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

  • Takeaways & Limitations

    Rejection recycling and multi-block decoding use future-token n-grams to further boost speedup to around 4×.

    Source: §1 · p. 2

Abstract

from arXiv · show

Multi-token generation has emerged as a promising paradigm for accelerating transformer-based large model inference. Recent efforts primarily explore diffusion Large Language Models (dLLMs) for parallel decoding to reduce inference latency. To achieve AR-level generation quality, many techniques adapt AR models into dLLMs to enable parallel decoding. However, they suffer from limited speedup compared to AR models due to a pretrain-to-posttrain mismatch. Specifically, the masked data distribution in post-training deviates significantly from the real-world data distribution seen during pretraining, and dLLMs rely on bidirectional attention, which conflicts with the causal prior learned during pretraining and hinders the integration of exact KV cache reuse. To address this, we introduce Jacobi Forcing, a progressive distillation paradigm where models are trained on their own generated parallel decoding trajectories, smoothly shifting AR models into efficient parallel decoders while preserving their pretrained causal inference property. The models trained under this paradigm, Jacobi Forcing Model, achieves 3.8x wall-clock speedup on coding and math benchmarks with minimal loss in performance. Based on Jacobi Forcing Models' trajectory characteristics, we introduce multi-block decoding with rejection recycling, which enables up to 4.5x higher token acceptance count per iteration and nearly 4.0x wall-clock speedup, effectively trading additional compute for lower inference latency. Our code is available at https://github.com/hao-ai-lab/JacobiForcing.

1 INTRODUCTION

The introduction identifies pretrain-to-posttrain mismatch as a barrier to adapting AR models for parallel decoding, then presents Jacobi Forcing and trajectory-based decoding optimizations. The authors report speedups across coding and math benchmarks.

1 INTRODUCTION

  • Open pretrained dLLMs underperform AR models, while adapting AR models to dLLMs yields limited speedup under quality constraints.The introduction attributes the dLLM quality gap mainly to the NELBO objective and describes mismatch in both masked data distributions and bidirectional attention.Source: §1 · pp. 1–2
  • Jacobi Forcing progressively distills AR models on their own generated trajectories while preserving causal attention.Noise-aware causal attention trains predictions of each block’s converged point conditioned on previous unconverged blocks; training repeats with larger blocks and noisier data.Source: §1 · p. 2
  • Jacobi Forcing models generate more correct future tokens from noisy context, and their trajectories provide higher-quality future-token n-grams.The introduction contrasts this with prior trajectory-based consistency methods, where correctly decoded tokens per iteration stay essentially constant as block size increases.Source: §1 · p. 2
  • Rejection recycling reuses high-quality consecutive tokens discarded in earlier Jacobi iterations, while multi-block decoding refines later blocks before preceding blocks converge.Together, these methods verify multiple candidate branches in one iteration and exploit the model’s ability to predict from noisy context.Source: §1 · p. 2
  • 3.8× generation speedup across coding and math benchmarks, with rejection recycling and multi-block decoding boosting speedup to nearly 4.0×.The approach also produces higher-quality future-token n-grams for use during decoding.Source: §1 · p. 2

2 PRELIMINARY

The preliminary section introduces Jacobi decoding and consistency distillation as approaches to accelerate AR-model decoding. Jacobi decoding iteratively updates blocks to an AR-equivalent fixed point, while consistency distillation trains predictions from trajectory states.

2.1 JACOBI DECODING

  • Jacobi decoding addresses sequential AR generation by expressing token generation as a system of nonlinear equations solved with Jacobi fixed-point iteration.Standard AR decoding requires n forward passes to generate n tokens, whereas Jacobi updates the sequence iteratively.Source: §2.1 · p. 3
  • Jacobi decoding updates an initialized n-token sequence in parallel using a causal attention mask, requiring one model forward pass per iteration.Iteration stops at a fixed point that is identical to greedy AR decoding; a long response is generated by concatenating fixed points from consecutive blocks.Source: §2.1 · p. 3

2.2 CONSISTENCY DISTILLATION

  • Consistency distillation trains a model to predict multiple correct tokens from arbitrary states along a Jacobi trajectory.The objective samples a block and a trajectory state, with KL divergence aggregated across the block’s n tokens.Source: §2.2 · p. 3
  • CLLMs combine consistency loss with standard AR loss after collecting Jacobi trajectories, achieving up to 2× speedup while maintaining quality.Related training objectives have also been used for action prediction in VLA models.Source: §2.2 · p. 4

3 METHODOLOGY

Jacobi Forcing progressively distills autoregressive models on their own Jacobi decoding trajectories while preserving causal attention. The method combines progressive noise scheduling and efficient training with trajectory-based inference optimizations.

3.1 JACOBI FORCING

  • Jacobi Forcing splits large blocks into smaller blocks with cyclically increasing noise ratios, preserving partially clean context and shortening noisy-token dependencies.The progressive schedule reduces the longest consecutive noisy-input span from O(nN) under an all-noisy random schedule to O(⌈tn⌉).Source: §3.1 · p. 4
  • Progressive consistency loss aggregates correct-token prediction losses across blocks, while a conventional autoregressive loss is included to safeguard output quality.The authors report that consistency-only training produces the same effect motivating the autoregressive-loss addition.Source: §3.1 · p. 4
  • Sequence packing and block-wise sparse attention reduce the forward and backward passes needed to compute the training objective from O(N) to O(1).The packed sequence supports computation of progressive consistency and autoregressive losses in one forward pass.Source: §3.1 · p. 5, Figure 2 · p. 5
  • An additional distillation round on trajectories from the Jacobi Forcing Model yields 20% further speedup with minor performance degradation.This round uses progressively larger block sizes after speedup from training on original-model trajectories saturates.Source: §3.1 · p. 5

3.2 INFERENCE OPTIMIZATION

  • Jacobi Forcing Model trajectories develop extending fixed-point segments among unconverged noisy tokens, motivating rejection recycling for inference optimization.The figure distinguishes accepted tokens matching the fixed point at their positions from noisy tokens and highlights long runs matching regardless of position.Source: §3.2 · p. 5, Figure 4 · p. 6
  • Rejection recycling extends candidate n-grams from noisy Jacobi states, verifies branches in parallel, and selects the candidate with the most newly accepted tokens.Multi-block decoding maintains and refines multiple blocks.Source: §3.2 · pp. 5–6

4 EXPERIMENTS

Experiments evaluate coding and math performance and speed against diffusion and autoregressive decoding baselines, then ablate training and inference choices.

4.1 EVALUATION SETTINGS

  • Evaluation uses coding and math models trained on curated OpenCodeInstruct and Openthought2 prompts, and compares against diffusion and AR-based parallel-decoding baselines.Training includes two successive rounds with block sizes 16 and 32; speedup comparisons include dLLMs, Jacobi decoding, and CLLM.Source: §4.1 · pp. 6–7

4.2 RESULTS

  • 3.6× speedup over the AR baseline on coding and 3.70× on GSM8K and MATH, with competitive accuracy on both benchmark groups.On MATH, solve rate rises from 77.0% to 77.4%.Source: §4.2 · pp. 7–8, Table 1 · p. 7, Table 2 · p. 8
  • On B200, multi-block decoding with rejection recycling achieves the highest fast-forwarded tokens per iteration, especially at larger block sizes.The comparison uses block size 128 and verification size 4.Source: §4.2 · p. 8, Table 3 · p. 8

4.3 ABLATION STUDY

  • With window size 8, linear progressive noise outperforms random and reverse progressive schedules; reverse progressive is significantly worse in the block-size-256 ablation.The schedule ablation uses checkpoints trained on 10k randomly sampled instances, not the full Table 1 data scale, and tests only one window size.Source: §4.3 · p. 8, Table 4 · p. 9
  • Noise-conditioned masks are more effective than masks conditioned on a clean intrawindow context for speedup while maintaining generation quality.The comparison is evaluated on HumanEval with A100 GPUs.Source: §4.3 · pp. 8–9, Table 5 · p. 9
  • On H200, Jacobi decoding with block sizes up to 64 incurs no latency penalty, while size 128 causes only minor degradation, supporting parallel decoding of up to 126 tokens without significant overhead.This pattern holds across accepted-token counts fixed at 2, 3, 4, and 5.Source: §4.3 · p. 9
  • For block size 64, initialization threshold r = 0.85 is consistently optimal across verification sizes 2–8; additional blocks’ gains saturate at block size 2.The selected maximum-utilization setting is block size 64 and verification size 4, with stable wall-clock speedup until parallel decoding exceeds 256 tokens.Source: §4.3 · pp. 9–10

5 RELATED WORK

Related work covers discrete text diffusion, Jacobi decoding, and speculative decoding as approaches to parallel generation or faster autoregressive inference.

Discrete Text Diffusion

  • Discrete text diffusion replaces left-to-right causality with iterative denoising for parallel multi-token generation, while open-source systems remain less efficient than autoregressive decoding.The cited work also reports throughput gains and competitive quality for closed-source systems.Source: §5 · p. 10

Jacobi Decoding

  • Jacobi decoding treats AR generation as parallel fixed-point updates whose convergence is linked to greedy AR, with follow-up work refining decoding or training parallel decoders.Examples include CLLMs’ consistency distillation and methods targeting robotics, reasoning, sampling, or retrieval-augmented settings.Source: §5 · p. 10

Speculative Decoding

  • Speculative decoding has a lightweight drafter propose future tokens for target-model verification, while other methods create candidates from context or future tokens without a separate drafter.The passage identifies self-speculative decoding as closely connected to the Jacobi decoding used in this work.Source: §5 · p. 10

6 CONCLUSION

The conclusion presents Jacobi Forcing as progressive distillation for AR models, with trajectory-based training and inference optimizations.

6 CONCLUSION

  • Jacobi Forcing uses progressive noise, sequence packing, and a noise-aware causal mask to train AR models as parallel decoders.Iterative training regenerates trajectories with progressively larger block sizes.Source: §6 · p. 10
  • Jacobi Forcing Model achieves 3.8× speedup while largely preserving accuracy.The conclusion reports high-quality draft tokens toward the tails of generated trajectories.Source: §6 · p. 10
  • Rejection recycling and multi-block decoding raise tokens accepted per iteration to 4.5× as high and bring nearly 4× speedup on HumanEval across A100 and B200 GPUs.Source: §6 · p. 10

ETHICS STATEMENT

The statement reports that the work used no human subjects or sensitive personal data and that the authors take responsibility for its conclusions.

ETHICS STATEMENT

  • The work involved no human subjects, sensitive personal data, or experiments with potential to cause harm.Source: ETHICS STATEMENT · p. 11
  • The authors report using no confidential or proprietary data.Source: ETHICS STATEMENT · p. 11
  • The paper discusses potential societal impacts, including limitations and biases of large language models, and assigns responsibility for its conclusions to the authors.Source: ETHICS STATEMENT · p. 11

REPRODUCIBILITY STATEMENT

The authors report providing public datasets, documented preprocessing, model and hyperparameter details, and ablations to support independent reproduction and robustness checks.

REPRODUCIBILITY STATEMENT

USE OF LLM

The authors report using an LLM for grammar and clarity edits, while retaining responsibility for the final text. Their decoding algorithm combines rejection recycling and multi-block decoding, which together improve throughput over vanilla Jacobi decoding on B200.

USE OF LLM

  • An LLM was used to refine grammar and clarity, and the authors reviewed and revised its outputs and take responsibility for the final content.Source: USE OF LLM · p. 11

A DETAILED DECODING ALGORITHM

  • The two decoding techniques together improve throughput by over 30 TPS versus vanilla Jacobi decoding on a B200 GPU.Rejection recycling reuses high-quality consecutive tokens discarded in earlier Jacobi iterations, while multi-block decoding refines several blocks jointly.Source: A DETAILED DECODING ALGORITHM · p. 16
  • Rejection recycling and multi-block decoding are orthogonal and can be combined; multi-block decoding can recover later-block tokens while earlier blocks remain unconverged.Source: A DETAILED DECODING ALGORITHM · p. 16

B FURTHER BASELINE COMPARISONS

The appendix compares Jacobi Forcing Model with distilled diffusion and speculative-decoding baselines, and details how progressive noise scheduling constructs paired noisy and clean training blocks.

B FURTHER BASELINE COMPARISONS

  • The appendix extends comparisons to distilled discrete diffusion models and state-of-the-art speculative-decoding baselines, including dParallel, EAGLE-3, and HASS.Table 6 also includes Fast-dLLM v2 and SDAR alongside other dLLM variants.Source: B FURTHER BASELINE COMPARISONS · pp. 16–17
  • On HumanEval, Jacobi Forcing Model (MR) has 29% higher accuracy and more than 80% higher TPF and TPS than dParallel.On GSM8K, it improves accuracy by 8 absolute points with about 20% higher TPF and TPS; the authors describe its speed–quality profile as more favorable at comparable scale.Source: B FURTHER BASELINE COMPARISONS · pp. 16–17, Table 6 · p. 17

C MAPPING NOISE SCHEDULE TO TRAINING SEQUENCE FOR PROGRESSIVE CONSISTENCY DISTILLATION

  • For each partitioned target block, the training procedure selects a Jacobi-trajectory state whose unconverged-token fraction is closest to its scheduled noise ratio.The sequence is partitioned into N = ⌈L/n⌉ blocks, and the schedule is applied over a window of w blocks.Source: C MAPPING NOISE SCHEDULE TO TRAINING… · p. 17
  • Interleaving each noisy block with its clean counterpart lets one forward pass compute AR loss on clean blocks and consistency loss on noisy blocks.The progressive schedule limits the longest consecutive noisy span within a block to O(⌈tn⌉), rather than the naive O(nN) worst case.Source: C MAPPING NOISE SCHEDULE TO TRAINING… · p. 17, Figure 6 · p. 18

D UNDERSTANDING TPF AND FLOPS TRADE-OFF

Hardware latency profiles show a platform-specific knee that bounds useful parallelism, while a grid search identifies a moderate block and verification size as a high-throughput configuration.

D UNDERSTANDING TPF AND FLOPS TRADE-OFF

E INFERENCE CONFIGURATION SEARCH

  • The inference search sweeps block sizes from 8 to 256 and n-gram verification sizes from 1 to 12, then fits and interpolates a smoothed surface over measured TPS.The configuration must respect the target GPU’s FLOPs budget rather than simply maximize block size or verification depth.Source: E INFERENCE CONFIGURATION SEARCH · p. 19
  • On B200, tokens-per-second peaks near block size 64 and n-gram verification size 4, supporting that final configuration under the GPU’s FLOPs budget.The analysis finds that very small settings underuse FLOPs, while larger settings can degrade wall-clock latency.Source: E INFERENCE CONFIGURATION SEARCH · p. 19, Figure 8 · p. 19
Loading 2512.14681v1…