Schnelle und präzise kausale parallele Dekodierung mittels Jacobi-Forcing
Fast and Accurate Causal Parallel Decoding using Jacobi Forcing
December 16, 2025
papers.authors: Lanxiang Hu, Siqi Kou, Yichao Fu, Samyam Rajbhandari, Tajana Rosing, Yuxiong He, Zhijie Deng, Hao Zhang
cs.AI
papers.abstract
Die Multi-Token-Generierung hat sich als vielversprechendes Paradigma zur Beschleunigung der Inferenz transformerbasierter Großmodelle etabliert. Aktuelle Forschungsbemühungen untersuchen vorrangig diffuse Large Language Models (dLLMs) für paralleles Decodieren, um die Inferenzlatenz zu verringern. Um AR-ähnliche Generierungsqualität zu erreichen, passen viele Techniken AR-Modelle an dLLMs an, um paralleles Decodieren zu ermöglichen. Allerdings weisen sie im Vergleich zu AR-Modellen eine begrenzte Beschleunigung auf, bedingt durch eine Diskrepanz zwischen Pre-Training und Post-Training. Konkret weicht die maskierte Datenverteilung im Post-Training signifikant von der realen Datenverteilung während des Pre-Trainings ab, und dLLMs nutzen bidirektionale Aufmerksamkeit, was im Widerspruch zum kausalen Prior des Pre-Trainings steht und die Integration exakter KV-Cache-Wiederverwendung behindert. Um dies zu adressieren, führen wir Jacobi Forcing ein, ein Paradigma progressiver Distillation, bei dem Modelle anhand ihrer eigenen generierten parallelen Decodierungspfade trainiert werden. Dies wandelt AR-Modelle schrittweise in effiziente parallele Decodierer um und bewahrt dabei ihre vortrainierte kausale Inferenzeigenschaft. Die nach diesem Paradigma trainierten Modelle, Jacobi Forcing Models, erreichen eine 3,8-fache Echtzeit-Beschleunigung bei Coding- und Math-Benchmarks bei minimalem Leistungsverlust. Basierend auf den Trajektorie-Eigenschaften von Jacobi Forcing Models führen wir Multi-Block-Decodierung mit Rejection Recycling ein, die bis zu 4,5-mal mehr akzeptierte Tokens pro Iteration und nahezu eine 4,0-fache Echtzeit-Beschleunigung ermöglicht, wodurch effektiv zusätzliche Rechenleistung gegen geringere Inferenzlatenz getauscht wird. Unser Code ist verfügbar unter https://github.com/hao-ai-lab/JacobiForcing.
English
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.