Jacobi Forcing Advances Causal Parallel Decoding, Delivering 4.5x Faster Inference
Summarize this article with:
Generating text with large language models typically happens one token at a time, creating a bottleneck for speed, but recent work explores generating multiple tokens simultaneously to accelerate the process. Lanxiang Hu, Siqi Kou from Shanghai Jiao Tong University, and Yichao Fu, along with colleagues, tackle a key challenge in this area, namely the loss of quality when adapting conventional language models for parallel decoding.
The team introduces Jacobi Forcing, a novel training method that allows models to learn from their own generated outputs, effectively bridging the gap between traditional, sequential generation and efficient, parallel decoding while maintaining high quality. This approach achieves substantial speedups, up to 3. 8x on coding and mathematical tasks, and, when combined with a new decoding strategy, delivers nearly 4. 0x faster inference, representing a significant step towards real-time language generation.
Jacobi Parallel Decoding for LLM Inference This research details advancements in accelerating Large Language Model (LLM) inference, focusing on maximizing the number of tokens processed per second while respecting hardware limitations.
The team investigated Jacobi Parallel Decoding, a technique that processes sequences in blocks, generating a draft and then verifying and correcting it in parallel, allowing for increased concurrency. Optimizing parameters like block size and verification size is crucial for achieving peak performance on different GPUs. The research demonstrates that simply increasing parallelism does not always improve performance, as it must be balanced against computational and memory constraints. The ideal block size and verification size are hardware-dependent, varying with the specific GPU used due to differing performance limits. Understanding these hardware limits through roofline analysis is essential for identifying bottlenecks and achieving optimal configuration.
The team discovered that a smooth surface, created through interpolation of discrete measurements, effectively explores the configuration space and identifies optimal parameters. Experiments revealed that smaller block sizes and moderate verification sizes perform best on A100 GPUs, while H200 and B200 GPUs can handle larger block sizes without significant performance degradation.
Jacobi Forcing Accelerates Transformer Inference and Quality Researchers have developed Jacobi Forcing, a novel progressive distillation technique that accelerates transformer-based inference while maintaining high generation quality. This approach addresses limitations in existing methods by smoothly transitioning autoregressive models into efficient parallel decoders, preserving crucial causal inference properties. The core innovation lies in training models on their own generated parallel decoding trajectories, overcoming discrepancies between pre-training and post-training data distributions. The training process was further refined with a progressive consistency loss, optimizing token prediction within each block and aggregating losses to reduce variance. This loss is combined with a conventional autoregressive loss, safeguarding output quality and ensuring high-fidelity generation. To enhance training efficiency, the researchers introduced a noise-aware causal attention mechanism and a sequence packing technique, reducing the number of forward and backward passes required for loss computation. These combined innovations achieve a 3. 8x wall-clock speedup on coding and math benchmarks with minimal performance loss, and further optimization with multi-block decoding and rejection recycling enables up to a 4. 5x higher token acceptance count per iteration and nearly a 4. 0x wall-clock speedup, effectively trading compute for lower inference latency.
Jacobi Forcing Accelerates Transformer Inference Dramatically A significant breakthrough in accelerating transformer-based inference has been achieved through Jacobi Forcing, a technique focused on efficiently training models to generate multiple tokens simultaneously without sacrificing accuracy. Experiments demonstrate a 3. 8x wall-clock speedup on coding and math benchmarks, representing a considerable improvement over existing approaches. The core of this achievement lies in a progressive distillation paradigm where the model learns from its own generated outputs, smoothly transitioning from standard sequential generation to efficient parallel decoding while preserving its ability to understand causal relationships in the data. Further refinement with multi-block decoding and rejection recycling pushes performance even higher, enabling up to 4. 5times more tokens to be accepted per iteration and achieving nearly a 4. 0x wall-clock speedup. Detailed performance metrics on benchmarks like GSM8K and MATH reveal substantial gains, with Jacobi Forcing consistently outperforming the standard sequential approach, delivering a 3. 7x speedup while maintaining competitive accuracy. On the MATH benchmark, the method achieves a solve rate of 77. 4%, a slight improvement over the baseline model, while simultaneously increasing tokens processed per second to 150. 7. On the B200 GPU, the team achieved a 3. 95x speedup, demonstrating the scalability of the technique.
Progressive Distillation Speeds Autoregressive Decoding This work presents Jacobi Forcing, a new progressive distillation technique for training autoregressive models as faster and more accurate parallel decoders. Unlike existing methods that directly train models to predict large blocks of tokens, Jacobi Forcing employs a progressively more difficult learning objective, achieved through a carefully designed noise schedule and sequence packing strategy. This approach allows for parallel token prediction while maintaining the benefits of causal inference learned during pre-training. The resulting model demonstrates a significant speedup of 3. 8times on coding and mathematical benchmarks, with minimal loss in performance. Further improvements, including multi-block decoding and rejection recycling, increase the token acceptance rate by 4. 5times and achieve nearly a 4-fold speedup on the HumanEval benchmark using both A100 and B200 GPUs. Analysis of the model’s generated trajectories reveals its ability to produce high-quality draft tokens, particularly towards the end of sequences. 👉 More information 🗞 Fast and Accurate Causal Parallel Decoding using Jacobi Forcing 🧠 ArXiv: https://arxiv.org/abs/2512.14681 Tags: