提出雅可比强迫方法,实现高效且准确的并行生成。
Fast and Accurate Causal Parallel Decoding using Jacobi Forcing
- 通过自训练轨迹渐进式蒸馏,将自回归模型转为并行解码器。
- 在代码与数学任务上实现3.8倍时钟速度提升,性能损失极小。
- 适合追求低延迟推理的场景,尤其适用于大模型部署。
多标记生成已成为加速基于Transformer的大模型推理的有前景范式。近期工作主要探索扩散大语言模型(dLLMs)用于并行解码以降低推理延迟。为达到自回归(AR)模型的生成质量,许多技术将AR模型适配为dLLMs以支持并行解码,但受限于预训练与后训练分布不匹配,其加速效果有限。具体而言,后训练中的掩码数据分布与预训练时的真实世界分布存在显著差异,且dLLMs依赖双向注意力,与预训练中学习到的因果先验冲突,阻碍了精确键值缓存复用。为此,我们提出雅可比强迫(Jacobi Forcing),一种渐进式蒸馏范式,模型在自身生成的并行解码轨迹上进行训练,平稳地将AR模型转化为高效并行解码器,同时保留其预训练的因果推理特性。采用此范式的雅可比强迫模型在编码与数学基准上实现3.8倍的时钟速度提升,性能损失微乎其微。基于雅可比强迫模型的轨迹特征,我们引入多块解码与拒绝回收机制,使每轮迭代的可接受标记数提升至4.5倍,时钟速度提升接近4.0倍,有效以额外计算换取更低推理延迟。代码已开源:https://github.com/hao-ai-lab/JacobiForcing。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。