MosaicDiff: Training-free Structural Pruning for Diffusion Model Acceleration Reflecting Pretraining Dynamics

AI summary

Overview

Research area: Efficient generative modeling — specifically training-free structural pruning for diffusion model inference acceleration.

Technical level: Advanced. The paper assumes familiarity with diffusion forward/reverse processes, signal-to-noise ratio (SNR) schedules, Hessian-based (Optimal Brain Surgeon) pruning, and transformer/U-Net architectures.

Scope: The paper proposes MosaicDiff, a training-free pruning framework that divides the diffusion sampling trajectory into three stages using a noise-schedule-derived importance score and assigns each stage a different sparsity level, evaluated on DiT-XL/2 and SDXL-base-1.0.

What This Paper Is About

Diffusion models generate high-quality images but are expensive to sample from, and existing training-free acceleration methods (pruning, caching) apply the same compressed model at every denoising timestep. The authors observe that diffusion pretraining follows a slow-fast-slow learning trajectory, and they argue that this variation should be mirrored at inference time: timesteps that are learning fast deserve conservative pruning, while slow-learning timesteps can tolerate aggressive pruning. MosaicDiff is their attempt to derive these stage boundaries and sparsity levels purely from the noise schedule, without any pretraining or fine-tuning.

Key Contributions

1. A trajectory-aware pruning strategy that aligns training-free/post-training acceleration with the varying learning speeds of diffusion pretraining, while never actually performing pretraining.
2. A method for identifying stage-specific sparse subnetworks from a single dense parent diffusion network — aggressive pruning during slow-learning phases and conservative pruning during fast-learning phases.
3. An SNR-aware calibration dataset construction procedure and a second-order (Hessian-based) structural pruning pipeline that mirrors the inference setup, including null-label handling for classifier-free guidance.
4. Extensive experiments on DiT-XL/2 and SDXL claiming state-of-the-art generation and acceleration performance, particularly at high pruning ratios, outperforming Diff-Pruning, Learning-to-Cache, DiP-GO, EcoDiff, and DeepCache.

Main Findings

• Three learning stages exist in the reverse process. Using DiT-XL/2, the authors plot MSE between intermediate latents and the final latents: in the early stage T ∈ (600, 1000) MSE decreases slowly with images remaining largely noisy; in the middle stage T ∈ (200, 600) denoising accelerates and images converge rapidly; in the final stage T ∈ (0, 200) MSE reduction slows, indicating subtle perceptual refinements.

• Closed-form MSE and gradient expressions match empirical curves. Theorem 1 gives E[MSE(t)] = (1/d)[(1 − √ᾱ_t)²‖x̂₀‖²₂ + (1 − ᾱ_t)‖I‖²₂] and an analogous gradient expression involving δ_t := ᾱ_t − ᾱ_{t−1}. The theoretical and sampled curves align closely under the linear schedule.

• DiT-XL/2 results (50 sampling steps, DDIM): MosaicDiff Pruned-0.33 reaches FID 2.24, IS 267.8, Precision 82.01, Recall 57.31, 3.92T MACs, 2.90s latency, 1.32× speedup — versus Diff-Pruning-0.3 (FID 180.76, IS 4.68) and Learning-to-Cache (FID 2.27, IS 244.1). The vanilla 50-step baseline is FID 2.26.

• DiT-XL/2 at 20 and 10 sampling steps: MosaicDiff Pruned-0.30 achieves FID 3.20 at 20 steps (Diff-Pruning: 223.80; Learning-to-Cache: 3.46; vanilla: 3.48) and FID 12.28 at 10 steps (Diff-Pruning: 270.26; Learning-to-Cache: 12.79; vanilla: 12.38).

• Comparison with DiP-GO: At 100 steps, MosaicDiff Pruned-0.25 uses 8.52T MACs with FID 2.67 versus vanilla 11.86T MACs / FID 3.17 and DiP-GO Pruned-0.6 at FID 3.01. At 70 steps, MosaicDiff Pruned-0.25 uses 5.99T MACs with FID 3.01 versus DiP-GO Pruned-0.75 at 6.77T MACs / FID 3.14.

• SDXL results: At 10% sparsity MosaicDiff reaches FID 23.18, CLIP 0.32, SSIM 0.67 versus Diff-Pruning (FID 108.96, CLIP 0.22, SSIM 0.31) and EcoDiff (FID 33.75, CLIP 0.31, SSIM 0.53), with vanilla SDXL at FID 24.90, CLIP 0.32, SSIM 1. At 20% sparsity MosaicDiff gets FID 23.79 versus Diff-Pruning 404.87 and EcoDiff 34.41. At 30% sparsity MosaicDiff reaches FID 28.37, CLIP 0.31, SSIM 0.53.

• Comparison with DeepCache on SDXL: MosaicDiff at 50 steps uses 135.66T MACs with FID 23.73; at 25 steps 72.04T / FID 24.04; at 20 steps 57.63T / FID 24.17; at 10 steps 28.82T / FID 24.32. DeepCache-N=2/3/5 at 50 steps yield 93.23T/24.88, 59.66T/24.67, 37.26T/24.43 respectively.

• SNR-aware calibration helps substantially. Removing it drops IS from 266.1 to 227.3, raises FID from 3.22 to 3.96, and reduces Precision (81.48 → 76.42) and Recall (57.18 → 53.73). The text describes this as "boosting the IS by over 40 points and reducing the FID by 0.74."

• Stage division matters. With uniform sparsity 0.3 and 20 sampling steps on DiT-XL/2, no stage division gives FID 40.89, uniform stage division gives 33.21, and the proposed division gives the best results — with M=0.55 (dividers 550/100) yielding IS 266.7 / FID 3.20, better than M=0.45 (253.2 / 3.36) and M=0.70 (242.6 / 3.71).

• Sparsity placement matters. In the ablation with stage dividers at 450 and 900 and total sparsity 0.3 (20 sampling steps), increasing sparsity in Stage 2 (0.6/0.15/0.1, FID 10.20) degrades performance more than the final configuration (0.6/0.04/0.1, FID 3.22), while a "sparser Stage 1" setup (0.9/0.04/0.1) still reaches FID 3.53.

• Compression is fast. DiT-XL/2 (675M) is pruned in 0.5 GPU hours on an RTX 4090, versus 16.3 GPU hours for L2C on the same hardware. SD-1.5 (865M) takes 0.8 hours on an RTX 4090 versus 2.5 hours for DiP-GO on an AMD MI250. SDXL (2.6B, 1024×1024) takes 6.0 GPU hours.

• Compatibility with other accelerators. Combining MosaicDiff at sparsity 0.15 with DeepCache-N=3 reduces MACs from 59.66T to 50.71T and FID from 24.61 to 23.88; with DPM-solver++ it goes from 67.80T/24.96 to 57.63T/24.01; with SDXL-Turbo from 13.56T/30.93 to 11.53T/30.08.

• Stated limitation: MosaicDiff introduces slightly more memory overhead because multiple pruned sub-networks for different sampling stages must be stored.

Methodology in Plain English

The pipeline is called Divide, Prune and Conquer.

Divide. The authors compute a per-timestep importance score, score(t) = E[Grad(t)] + λ·ln SNR(t), where λ is a hyperparameter (set to 0.01) and SNR(t) = ᾱ_t / (1 − ᾱ_t). Because the SNR grows sharply in the final sampling steps, it is folded in to protect late timesteps where even small MSE changes affect image quality. A threshold, defined as M · max_t(score(t)) with M ∈ (0,1), splits the score curve into three segments. Each segment gets a sparsity sᵢ proportional to 1 − scorēᵢ, so segments with higher average scores are pruned less.

Prune. For each stage, the authors build an SNR-aware calibration set: images from a standard calibration set (ImageNet-1K for DiT, MS COCO 2017 training images with captions for SDXL) are resized and encoded into latents, then a random timestep within the target stage is chosen and noise is added per the forward equation so the latent lands in the desired SNR range. Under classifier-free guidance, latents and timesteps are duplicated and paired with a null label to mirror unconditional inference. These calibration inputs are fed through the pretrained model to compute per-layer Hessians H = XᵀX. Following an extension of Optimal Brain Surgeon, saliency scores are derived from the Hessian to select which columns to remove, and a weight update δ is computed to compensate for pruning.

Conquer. The per-stage pruned subnetworks are merged so that sampling uses the correct subnetwork at each timestep, keeping the accelerated trajectory consistent with the original model. Both transformer and U-Net models are pruned in two structures: attention heads in multi-head self-attention and the intermediate dimension of MLP modules. No fine-tuning or retraining is performed; the only cost is Hessian computation.

Why This Matters

Impact on research. The paper reframes diffusion acceleration as a problem of matching inference-time compression to training-time learning dynamics, rather than treating all timesteps as equally important. It also provides closed-form expressions for expected MSE and gradient as functions of the noise schedule, which the authors verify empirically, giving a theoretical handle for future stage-aware methods.

Real-world applications:

• Text-to-image services (e.g., SDXL-based products) that need lower per-image latency and GPU cost without retraining large checkpoints.
• On-device or edge image generation, where a smaller pruned subnetwork per stage reduces MACs and latency.
• Large-scale image generation pipelines (e.g., ImageNet-conditional generation with DiT) where 50,000-image evaluation budgets and inference throughput matter.
• Deployment stacks combining pruning with caching, fast solvers (DPM-Solver++), and step-distilled models (SDXL-Turbo).

Industry relevance. The method is training-free and reportedly compresses SDXL (2.6B parameters) in 6 GPU hours on a single RTX 4090, versus 16.3 hours of router training for Learning-to-Cache and 2.5 hours for DiP-GO's pruner on an AMD MI250. That cost profile makes it attractive for teams that cannot afford fine-tuning runs, and the reported compatibility with existing caches and samplers means it can slot into an existing inference stack rather than replacing it.

Future Directions

• Reducing the memory overhead of storing multiple sub-networks. The conclusion explicitly flags this as future work, since each sampling stage currently requires its own pruned subnetwork.
• Automating the choice of λ and M. The paper fixes λ = 0.01 and M = 0.55 for both models; whether these transfer to other schedules or architectures is not established by the reported experiments.
• Testing on more schedulers and architectures. The paper reports DiT under a linear schedule and SDXL under a scaled-linear schedule; generalization to other noise schedules and diffusion variants is left open.
• Combining stage-specific pruning with other acceleration families at higher sparsity. The compatibility experiments use sparsity 0.15 on SDXL, so how far the combination extends at more aggressive ratios is not reported.

Target Audience

Researchers and practitioners working on diffusion model efficiency — particularly those interested in structural pruning, training-free compression, or noise-schedule-aware inference. It is also relevant to engineers deploying DiT- or SDXL-class models under latency or compute constraints, and to readers interested in Optimal Brain Surgeon-style second-order pruning extended to iterative generative models. The density of formulas and the assumed background in diffusion math make it best suited to readers with prior exposure to the area.

Authors’ abstract

Diffusion models are renowned for their generative capabilities, yet their pretraining processes exhibit distinct phases of learning speed that have been entirely overlooked in prior post-training acceleration efforts in the community. In this study, we introduce a novel framework called MosaicDiff that aligns diffusion pretraining dynamics with post-training sampling acceleration via trajectory-aware structural pruning. Our approach leverages the observation that the middle, fast-learning stage of diffusion pretraining requires more conservative pruning to preserve critical model features, while the early and later, slow-learning stages benefit from a more aggressive pruning strategy. This adaptive pruning mechanism is the first to explicitly mirror the inherent learning speed variations of diffusion pretraining, thereby harmonizing the model's inner training dynamics with its accelerated sampling process. Extensive experiments on DiT and SDXL demonstrate that our method achieves significant speed-ups in sampling without compromising output quality, outperforming previous state-of-the-art methods by large margins, also providing a new viewpoint for more efficient and robust training-free diffusion acceleration.