Быстрое и точное причинно-следственное параллельное декодирование с использованием метода Якоби

Fast and Accurate Causal Parallel Decoding using Jacobi Forcing

December 16, 2025

Авторы: Lanxiang Hu, Siqi Kou, Yichao Fu, Samyam Rajbhandari, Tajana Rosing, Yuxiong He, Zhijie Deng, Hao Zhang

cs.AI

Аннотация

Многотокенная генерация стала перспективной парадигмой для ускорения вывода больших моделей на основе трансформеров. Современные исследования в основном изучают диффузионные большие языковые модели (dLLM) для параллельного декодирования с целью снижения задержек вывода. Для достижения качества генерации на уровне авторегрессионных (AR) моделей многие методы адаптируют AR-модели в dLLM для обеспечения параллельного декодирования. Однако они страдают от ограниченного ускорения по сравнению с AR-моделями из-за несоответствия между предварительным и последующим обучением. В частности, маскированное распределение данных при последующем обучении существенно отклоняется от распределения реальных данных, наблюдаемых при предварительном обучении, а dLLM полагаются на двунаправленное внимание, что конфликтует с причинными априорными знаниями, усвоенными при предварительном обучении, и препятствует интеграции точного повторного использования KV-кэша. Для решения этой проблемы мы представляем Jacobi Forcing — прогрессивную парадигму дистилляции, в которой модели обучаются на собственных сгенерированных траекториях параллельного декодирования, плавно преобразуя AR-модели в эффективные параллельные декодеры с сохранением их причинных свойств вывода, усвоенных при предварительном обучении. Модели, обученные по этой парадигме (Jacobi Forcing Model), достигают 3.8-кратного ускорения в реальном времени на бенчмарках по программированию и математике при минимальной потере производительности. На основе характеристик траекторий Jacobi Forcing Models мы вводим многоблочное декодирование с рециклингом отклонений, которое позволяет увеличить количество принимаемых токенов за итерацию до 4.5 раз и достичь почти 4.0-кратного ускорения в реальном времени, эффективно обменивая дополнительные вычисления на снижение задержки вывода. Наш код доступен по адресу 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.