Décodage Causal Parallèle Rapide et Précis par Forçage de Jacobi

papers.abstract

La génération multi-jetons est apparue comme un paradigme prometteur pour accélérer l'inférence des grands modèles basés sur les transformers. Les travaux récents explorent principalement les modèles de langage de grande taille à diffusion (dLLM) pour le décodage parallèle afin de réduire la latence d'inférence. Pour atteindre une qualité de génération équivalente aux modèles autorégressifs (AR), de nombreuses techniques adaptent les modèles AR en dLLM pour permettre le décodage parallèle. Cependant, ils souffrent d'une accélération limitée par rapport aux modèles AR en raison d'une inadéquation entre le pré-entraînement et le post-entraînement. Concrètement, la distribution de données masquées lors du post-entraînement s'écarte significativement de la distribution de données réelles observée pendant le pré-entraînement, et les dLLM s'appuient sur une attention bidirectionnelle, ce qui entre en conflit avec l'a priori causal appris lors du pré-entraînement et entrave l'intégration de la réutilisation exacte du cache KV. Pour résoudre ce problème, nous introduisons Jacobi Forcing, un paradigme de distillation progressive où les modèles sont entraînés sur leurs propres trajectoires de décodage parallèle générées, transformant progressivement les modèles AR en décodeurs parallèles efficaces tout en préservant leur propriété d'inférence causale pré-entraînée. Les modèles entraînés selon ce paradigme, Jacobi Forcing Model, atteignent une accélération en temps réel de 3,8x sur des benchmarks de codage et de mathématiques avec une perte de performance minimale. Sur la base des caractéristiques de trajectoire des modèles Jacobi Forcing, nous introduisons le décodage multi-bloc avec recyclage par rejet, qui permet jusqu'à 4,5x plus de jetons acceptés par itération et une accélération en temps réel de près de 4,0x, échangeant efficacement un calcul supplémentaire contre une latence d'inférence réduite. Notre code est disponible à l'adresse 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.

Décodage Causal Parallèle Rapide et Précis par Forçage de Jacobi

Fast and Accurate Causal Parallel Decoding using Jacobi Forcing

papers.abstract

Support