Fast and Accurate Causal Parallel Decoding using Jacobi Forcing (2512.14681v1)

Abstract: Multi-token generation has emerged as a promising paradigm for accelerating transformer-based large model inference. Recent efforts primarily explore diffusion LLMs (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.

• The paper introduces Jacobi Forcing, a progressive distillation framework that enables fast causal parallel decoding while preserving the autoregressive model structure.
• It employs a progressive noise schedule, noise-aware causal masks, and multi-block decoding to overcome AR and diffusion LLM mismatches and enhance inference efficiency.
• Empirical results demonstrate up to 3.8x wall-clock speedup with only marginal accuracy degradation, validating the method's effectiveness on modern accelerator architectures.

Fast and Accurate Causal Parallel Decoding using Jacobi Forcing

Introduction

The paper "Fast and Accurate Causal Parallel Decoding using Jacobi Forcing" (2512.14681) presents a comprehensive study on accelerating transformer-based large model inference with multi-token generation. The authors identify critical limitations in the adaptation of autoregressive (AR) models to diffusion-based LLMs (dLLMs) for parallel decoding, specifically emphasizing the pretrain-to-posttrain mismatch arising from masked data distributions and conflicting attention mechanisms. To address these issues, the paper introduces Jacobi Forcing (JF): a progressive distillation framework that enables parallel decoding for AR models while preserving the underlying causal structure, yielding substantial speedups with minimal generation quality degradation.

Background: Parallel Decoding and Its Bottlenecks

Jacobi Decoding and Consistency Distillation

Jacobi Decoding reconceptualizes AR generation as a parallel fixed-point computation via Jacobi-style nonlinear equation solving, where $n$ tokens can be produced in parallel over each block using causal masks. In practice, standard Jacobi Decoding rarely yields more than one correct token per iteration, resulting in only marginal speed improvements over standard AR generation. Consistency distillation, as established in prior work, introduces KL-divergence-based losses along Jacobi trajectories to teach models to converge more rapidly—yielding speedups around $2\times$.

Limitations of Diffusion LLMs and Blockwise Parallelization

Simple adaptation of AR models to dLLMs via block-wise bidirectional masking and NELBO training introduces a significant distributional mismatch, undermining both training and inference efficiency. Furthermore, bidirectional masking prohibits the efficient reuse of key/value (KV) caches, which is essential for scalable execution on accelerators.

Jacobi Forcing: Progressive Distillation for Causal Parallel Decoding

Jacobi Forcing (JF) is designed to directly address the aforementioned bottlenecks by leveraging a progressive noise schedule and a corresponding causal attention mask. This paradigm preserves the causal reasoning structure learned during AR pretraining.

Progressive Noise Schedule and Sequence Packing

During training, each sequence is partitioned into blocks; for each block, a trajectory is sampled where noisy (unconverged) tokens are injected according to a scheduled noise ratio that grows within a windowed schedule. This reduces the dependency length on noisy tokens, facilitating more effective parallel prediction (Figure 1).

Figure 1: Illustration of the progressive noise schedule and training sequence packing; noisy and clean blocks are interleaved for joint optimization of AR and consistency losses.

A noise-aware causal mask ensures that predictions within each block are properly conditioned, with preceding blocks providing clean context and intra-block non-converged positions delivering noisy conditioning. This enables a single forward pass to compute both autoregressive (AR) and progressive consistency losses (Figure 2).

Figure 2: Sequence packing with clean and noise-aware attention masks to enable efficient computation of progressive consistency and AR losses.

Training Objective and Iterative Distillation

The loss function is a weighted sum of the progressive consistency loss (encouraging fast fixed-point convergence) and the standard AR loss (maintaining generation quality). An iterative distillation regime is employed: after the initial training using AR trajectories, JF is further trained on its own trajectories with increasing block sizes, leading to further speedups with negligible additional quality degradation.

Inference Optimizations: Rejection Recycling and Multi-block Decoding

JF's learned trajectories are characterized by the frequent emergence of long correct (fixed-point) n-grams, even within unconverged blocks (Figure 3).

Figure 3: Visualization of JF’s Jacobi decoding trajectory; blue = accepted/fixed-point tokens, red = >3-token correct n-grams within noisy regions.

This observation motivates two inference-time optimizations:

• Rejection Recycling: High-quality n-grams from prior Jacobi iterations are stored and verified in parallel in subsequent steps, enabling rapid commitment of multiple tokens per iteration.
• Multi-block Decoding: Multiple blocks are maintained in parallel. A real-active block appends confirmed tokens to the committed sequence, while others (pseudo-active) are prepared for rapid verification when promoted. This allows the system to exploit high-quality drafts beyond the effective boundary of the current KV cache.

Collectively, these techniques significantly amplify speedup and token acceptance rates, especially on contemporary accelerator architectures (Figure 4).

Figure 4: Effect of block size choices on fast-forward count and wall-clock speedup—wall-clock speedup tracks fast-forward count until hardware parallelism saturates.

Empirical Results

JF achieves up to $3.8\times$ wall-clock speedup compared to strong AR baselines on coding (HumanEval, MBPP) and mathematical (GSM8K, MATH) benchmarks, with only marginal drops (typically 4–5 points) in strict generation accuracy. When combined with multi-block/rejection-recycling decoding, up to $4.5\times$ more tokens are accepted per iteration, and nearly $4.0\times$ wall-clock speedup is observed, consistently outperforming both Jacobi-based and dLLM-based alternatives. Notably, state-of-the-art dLLMs exhibit $2\times$–$10\times$ lower generation quality at similar throughput, even when equipped with blockwise KV caching, indicating that the Jacobi Forcing paradigm dramatically improves the speed–quality tradeoff for parallel decoding.

Ablation and Practical Considerations

• Progressive Schedules: Linear progressive noise schedules outperform both random and reverse progressive schedules for training, yielding a better balance of speed and accuracy.
• Mask Types: Noise-aware causal masks are substantially more effective than alternatives using intra-window clean contexts in preserving both speedup and robustness.
• FLOPs Utilization: On modern accelerators (H200, B200), optimal speedup is attained with block sizes around 64–128, saturating hardware efficiency without incurring wall-time penalties due to excessive parallel token decoding per iteration.

Implications and Future Directions

The findings demonstrate that it is possible to realize efficient, high-quality parallel token decoding without sacrificing the causal reasoning structure or the advantages of KV cache reuse—limitations that have hindered dLLMs and blockwise speculative architectures. The framework is theoretically robust (recovering AR decoding in the degenerate case) and is empirically validated at scale. The multi-block and recycling inference techniques introduced here are orthogonal and may be incorporated into other parallel decoding regimes.

This suggests that, as accelerators continue to scale in FLOPs and memory, techniques such as Jacobi Forcing will be increasingly valuable for lowering generation latency, especially for agentic, multi-step, or interactive LLM workloads. Additionally, progressive consistency distillation regimes may inspire new curricula for blockwise or diffusion-based training, potentially yielding universal decoders that interpolate flexibly between sequential and parallel extremes.

Conclusion

Jacobi Forcing offers a principled and practical solution to causal parallel decoding in transformers, facilitating multi-token generation with minimal accuracy loss and strong empirical speedup. The approach closes the principal gap between AR and dLLM-based paradigms for parallel decoding, and its dual focus on training curriculum and inference optimization marks a robust advance in the scalable, efficient deployment of LLMs.