通过训练梯度动态提前终止扩散推理,显著减少步骤数且不降精度。
EDIT: Early Diffusion Inference Termination for dLLMs Based on Dynamics of Training Gradients
- 基于训练时的梯度信息构建推理路径映射,实时监测生成稳定性。
- 在多个推理任务中减少11.8%至68.3%的扩散步骤,精度基本不变或提升。
- 适合追求高效推理的模型部署者,尤其对资源受限场景友好。
基于扩散的大语言模型(dLLMs)通过迭代去噪逐步优化词元生成,但答案常在完成全部步骤前已稳定。本文提出编辑(EDIT),一种推理阶段的自适应终止机制,当检测到生成结果与训练时推理路径的匹配度足够高时即停止去噪。该方法监控词元激活与微调阶段通过AdamW聚合的LoRA更新所生成的推理图之间的对齐程度。训练过程中产生的优化动态蕴含丰富的参数重要性信息,以往通常被丢弃,本文将其压缩为学习路径的紧凑表示。推理时,对齐分数转化为当前未掩码词元上的分布,当连续两步间可见词元的KL散度低于阈值时判定收敛。在多个推理基准测试中,EDIT将扩散步骤减少11.8%至68.3%,多数情况下保持或提升准确率,存储开销仅约0.02%(8GB模型中32层块的QKV模块总计1.5-2MB)。利用训练梯度动态,本工作为降低dLLM推理时间和成本开辟了新方向。
原文摘要 · Abstract (English)
Diffusion-based large language models (dLLMs) refine token generations through iterative denoising, but answers often stabilize before all steps complete. We propose EDIT (Early Diffusion Inference Termination), an inference-time criterion that adaptively stops denoising once sufficient reasoning stability relative to training-time reasoning is detected. EDIT monitors the alignment between token activations and a reasoning map derived from AdamW-aggregated LoRA updates captured during supervised fine-tuning (SFT). During training, optimization dynamics generate rich metadata about parameter importance that in prior methods is typically discarded upon model release. We preserve this information as a compact representation of learned reasoning pathways. During inference, alignment scores are converted to a distribution over the tokens already unmasked at the current denoising step, and convergence is detected when KL divergence between consecutive steps falls below a threshold on the matched unmasked (visible) tokens. Across reasoning benchmarks, EDIT reduces diffusion steps by 11.8% to 68.3% while preserving or improving accuracy in most settings, with approximately 0.02% storage overhead (about 1.5-2 MB for all QKV modules across 32 blocks in an 8 GB model). By utilizing training-gradient dynamics, our work opens a new research direction for reducing dLLM inference time and cost.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。