EDIT: Early Diffusion Inference Termination for dLLMs Based on Dynamics of Training Gradients
By: He-Yen Hsieh, Hong Wang, H. T. Kung
Potential Business Impact:
Makes AI think faster and use less energy.
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.
Similar Papers
Reasoning in Diffusion Large Language Models is Concentrated in Dynamic Confusion Zones
Machine Learning (CS)
Teaches AI to learn better by focusing on tricky parts.
Beyond Fixed: Training-Free Variable-Length Denoising for Diffusion Large Language Models
Computation and Language
Makes AI write faster and smarter without wasting energy.
dInfer: An Efficient Inference Framework for Diffusion Language Models
Computation and Language
Makes AI write much faster and better.