Bridging Diffusion Pruning and Step Distillation with Teacher-Aligned Repair

2026-07-07

AI Summary

Diffusion models suffer from high inference costs due to their large network size and multi-step denoising process, and existing compression methods struggle to simultaneously achieve structural simplicity and performance retention. This work proposes a synergistic compression framework that integrates structured pruning with lightweight teacher-aligned restoration and single-step distillation (SiDA), enabling seamless integration of step distillation without requiring retraining after pruning for the first time. Built upon the EDM2-XS architecture, the pruned model achieves an FID of 3.12 on ImageNet-512 with only 98.8M parameters and a single forward pass at 20% sparsity, and an FID of 4.26 at 30% sparsity—significantly outperforming current baselines while substantially reducing inference overhead without compromising generation quality.

Abstract

Diffusion models generate high-quality images, but their inference cost comes from two sources: large denoising networks and repeated denoising steps. Existing compression pipelines usually attack these costs separately. Pruning reduces the network, but most pruning methods still rely on a long post-pruning retraining stage to recover a many-step sampler. Step distillation reduces the number of denoising steps, but it usually assumes a student that can already follow the teacher well enough to receive useful distillation gradients. This paper asks whether post-pruning retraining can be replaced by step distillation. We find that the direct replacement fails: after pruning an EDM2-XS teacher, starting SiDA from the pruned checkpoint produces unusable samples. We introduce a short teacher-alignment repair stage as a bridge between pruning and step distillation. The bridge matches the pruned generator to the teacher on noisy real-image latents, then hands the repaired checkpoint to one-step distillation. On ImageNet-512, the original EDM2-XS baseline uses 124.713M parameters and 63 network evaluations, reaching an FID of 3.53. With a suitable distillation objective, our 20% pruned one-step generator uses 98.826M parameters and one network evaluation, reaching an FID of 3.12. With 30% pruning, the model uses 88.029M parameters and one network evaluation, with an FID of 4.26.

Problem

• Diffusion Pruning
• Step Distillation
• Model Compression
• Inference Cost
• Teacher-Student Alignment

Innovation

• diffusion pruning
• step distillation
• teacher-aligned repair
• one-step generation
• model compression