Decodificación causal paralela rápida y precisa mediante forzado de Jacobi

Resumen

La generación multi-token ha surgido como un paradigma prometedor para acelerar la inferencia de modelos grandes basados en transformadores. Los esfuerzos recientes exploran principalmente Modelos de Lenguaje Grandes de difusión (dLLMs) para la decodificación paralela con el fin de reducir la latencia de inferencia. Para lograr una calidad de generación comparable a los modelos autoregresivos (AR), muchas técnicas adaptan modelos AR en dLLMs para permitir la decodificación paralela. Sin embargo, adolecen de una aceleración limitada en comparación con los modelos AR debido a una discrepancia entre el preentrenamiento y el postentrenamiento. Específicamente, la distribución de datos enmascarados en el postentrenamiento se desvía significativamente de la distribución de datos del mundo real vista durante el preentrenamiento, y los dLLMs dependen de la atención bidireccional, lo que entra en conflicto con el conocimiento causal aprendido durante el preentrenamiento y dificulta la integración de la reutilización exacta de la caché KV. Para abordar esto, presentamos Jacobi Forcing, un paradigma de destilación progresiva donde los modelos se entrenan en sus propias trayectorias de decodificación paralela generadas, transformando suavemente los modelos AR en decodificadores paralelos eficientes mientras se preserva su propiedad de inferencia causal preentrenada. Los modelos entrenados bajo este paradigma, el Modelo de Forzamiento de Jacobi, logran una aceleración en tiempo real de 3.8x en benchmarks de código y matemáticas con una pérdida mínima de rendimiento. Basándonos en las características de la trayectoria de los Modelos de Forzamiento de Jacobi, introducimos la decodificación multi-bloque con reciclaje por rechazo, que permite hasta 4.5x más tokens aceptados por iteración y una aceleración en tiempo real de casi 4.0x, intercambiando efectivamente cálculo adicional por una menor latencia de inferencia. Nuestro código está disponible en 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.

Decodificación causal paralela rápida y precisa mediante forzado de Jacobi

Fast and Accurate Causal Parallel Decoding using Jacobi Forcing

Resumen

Support