PaperlayerView exact revision on arXiv

Source-linked AI summary

Structural Pruning for Diffusion Models

Gongfan Fang, Xinyin Ma, Xinchao Wang

arXiv:2305.10924v3cs.LGcs.AIcs.CV

TL;DR

Diffusion models are powerful but computationally expensive, and existing work lacks a general way to compress and reuse pretrained models without heavy retraining. Diff-Pruning uses timestep-pruned Taylor expansions to identify important weights and remove non-critical parameters. It achieves substantial compression with 10% to 20% of pre-training cost while preserving generative behavior.

  • Problem

    Diffusion models incur substantial training and inference overhead, while existing methods lack a general compression approach for efficiently reusing pretrained models without heavy retraining.

  • Method

    Diff-Pruning uses a Taylor expansion over selected diffusion timesteps to aggregate informative gradients and identify parameters for structural removal.

  • Results

    50% FLOPs reduction is achieved on an LSUN Church model with only 10% of the original training cost, while preserving similar generation behavior.

  • Takeaways & Limitations

    Diff-Pruning provides lightweight diffusion models that retain behavior consistent with their pretrained counterparts and offers a baseline for future compression research.

Abstract

from arXiv · show

Generative modeling has recently undergone remarkable advancements, primarily propelled by the transformative implications of Diffusion Probabilistic Models (DPMs). The impressive capability of these models, however, often entails significant computational overhead during both training and inference. To tackle this challenge, we present Diff-Pruning, an efficient compression method tailored for learning lightweight diffusion models from pre-existing ones, without the need for extensive re-training. The essence of Diff-Pruning is encapsulated in a Taylor expansion over pruned timesteps, a process that disregards non-contributory diffusion steps and ensembles informative gradients to identify important weights. Our empirical assessment, undertaken across several datasets highlights two primary benefits of our proposed method: 1) Efficiency: it enables approximately a 50\% reduction in FLOPs at a mere 10\% to 20\% of the original training expenditure; 2) Consistency: the pruned diffusion models inherently preserve generative behavior congruent with their pre-trained models. Code is available at \url{https://github.com/VainF/Diff-Pruning}.

1 Introduction

Diff-Pruning is introduced as a structural-pruning method for compressing diffusion models, addressing their substantial computational overhead and the lack of a general compression method. It uses timestep-aware Taylor expansion and achieves substantial FLOPs reduction with limited training cost while preserving generative quality and consistency.

  • Diffusion Probabilistic Models offer broad applications but incur considerable computational overhead during training and inference.
  • Existing diffusion-model research has emphasized architectures, training methods, and sampling acceleration, while a general compression method remains absent.
  • Conventional pruning is challenged by diffusion models’ iterative generation, timestep sensitivity, and intricate diffusion-process interactions.
  • Diff-Pruning applies a Taylor expansion over pruned timesteps to account for image content, details, and the negative effects of noisy diffusion steps.
  • 50% FLOPs compression is achieved for an LSUN Church diffusion model with only 10% of the pre-training expenditure.
  • Diff-Pruning is positioned as an efficient compression baseline requiring 10% to 20% of pre-training costs and supporting future work on compressed-model quality and consistency.

2 Ralted Works

Prior diffusion efficiency research focuses on architectures, training procedures, and sampling acceleration. Diffusion pruning remains comparatively underexplored, motivating a dedicated structural-pruning method.

  • Efficient Diffusion Models: Existing diffusion-model efficiency methods primarily refine architectures, improve training procedures, or accelerate sampling.
  • Network Pruning: Network pruning research distinguishes structural pruning from unstructured pruning and has mainly focused on discriminative tasks.
  • Network Pruning: Only limited work has studied pruning for generative tasks, while structural pruning for diffusion models presents unique challenges.
  • Network Pruning: The paper introduces the first dedicated method explicitly designed for pruning diffusion models and positions it as a baseline for future work.

3 Diffusion Model Objectives

Diffusion models learn a generative distribution through forward noising and reverse denoising processes. Diff-Pruning seeks a lightweight noise predictor that produces similar outputs from the same initial noise.

  • Diffusion Model Objectives: Diffusion models aim to learn a generative distribution pθ(x) that approximates the data distribution q(x).
  • Diffusion Model Objectives: Latent variables x1 through xT form a joint distribution with learned Gaussian reverse transitions pθ(xt−1|xt).
  • Diffusion Model Objectives: The forward process adds noise according to a predefined variance schedule, while the reverse process denoises xt to obtain xt−1.
  • Diffusion Model Objectives: After training, synthetic images x0 are sampled iteratively from Gaussian noise xT ∼N(0, 1).
  • Diffusion Model Objectives: Diff-Pruning constructs a lightweight noise predictor by removing redundant parameters while aiming to produce similar x0 from the same xT.

4 Structrual Pruning for Diffusion Models

Diff-Pruning compresses pretrained diffusion models by estimating structural importance through Taylor-based loss disruption while accounting for the iterative, timestep-dependent generation process. It aggregates informative timestep gradients, excludes uninformative noisy steps, and uses the resulting scores to remove network sub-structures with limited retraining.

  • Pruning Objective: A lightweight parameter matrix θ′ is obtained by removing network sub-structures from a pretrained model while minimizing pruning-induced loss disruption.Structural pruning targets sparse parameter matrices whose retained structure preserves the original model’s performance.
  • Timestep Selection: Diff-Pruning weights timestep gradients with binary αt values, retaining informative partial steps to balance generated content, details, and noisy-step effects.Early steps emphasize local details, whereas later steps emphasize higher-level content; the method uses thresholding to trade off these factors.
  • Taylor-Based Importance: Taylor expansion estimates an individual weight’s importance from the first-order loss disruption produced by setting that weight to zero.For weight θik, the criterion is |θik · ∇θikLt(θ, x)|.
  • Taylor-Based Importance: Structural pruning scores an entire sub-structure by summing the Taylor influence of its scalar parameters and aggregating removal effects across the iterative diffusion process.This cumulative score accounts for the fact that small loss fluctuations can affect later generation steps.
  • Timestep Selection: Errors at larger timesteps can propagate through the reverse process and alter global image content, while smaller timesteps primarily refine generated details.The reverse-process chain can amplify prediction errors and create nonlinear disturbances in the final images.
  • Timestep Selection: Timesteps with relative loss below the threshold are treated as uninformative and excluded from importance estimation, avoiding converged gradients from noisy loss terms.Experiments report that loss terms rapidly approach 0 as t →T, and that full-step objectives can sometimes be worse than partial objectives.

5 Experiments

Experiments evaluate Diff-Pruning across datasets, model types, pruning criteria, resolutions, ratios, and timestep thresholds using efficiency, quality, and consistency metrics. Results show efficient convergence and generally preserved generation, while pruning sensitivity varies with dataset, timestep selection, and pruning ratio.

  • Datasets and Models: Experiments span six datasets and DDPMs or LDMs, using parameter count, MACs, FID, and SSIM to measure efficiency, quality, and consistency.SSIM compares outputs from pre-trained and pruned models under identical noise inputs.
  • Scratch Training v.s. Pruning: Pruning methods converge within approximately 100K steps and outperform scratch training on CIFAR-10 and CelebA-HQ in FID and SSIM.Scratch-trained compact models require relatively more steps to reach convergence.
  • Pruning Criteria: Diff-Pruning delivers stable improvements over Random, Magnitude, and Taylor baselines across CIFAR-10 and CelebA-HQ, surpassing the pre-trained model on CelebA-HQ after 100K optimizations.Magnitude has the best CIFAR-10 FID among the three baselines but performs poorly on CelebA datasets; CIFAR-10 degradation remains observable for Diff-Pruning.
  • Conditional LDMs on ImageNet: 30% channel sparsity reduces the ImageNet-1K conditional LDM’s U-Net to 189.43M parameters after importance estimation using 534 participating steps.The U-Net contains around 400M parameters, compared with 55M in the autoencoder.
  • Pruned Timesteps: Around 250 Taylor-expansion steps can suffice, but additional steps may slightly reduce quality; CelebA benefits from more steps for importance estimation.The reported explanation is that first-order Taylor estimates become inaccurate at converged points, whereas dataset behavior differs on CelebA.
  • Pruning Ratios: A 16% pruning ratio worsens FID from 4.19 to 4.62, showing that diffusion models are sensitive to model-size changes.The paper attributes this sensitivity to FID’s susceptibility to domain shift in generative models.
  • Visualization of Different Importance Criteria: Using timestep pruning with Diff-Pruning achieves an SSIM score of 0.905 after pruning, outperforming the no-timestep-pruning variant in the reported visual comparison.Figure 4 compares images from models pruned with different importance criteria and reports SSIM without post-training.

6 Conclusion

Diff-Pruning compresses diffusion models by using Taylor expansion over pruned timesteps to identify and remove non-critical parameters, producing lightweight models consistent with pre-trained ones at substantially lower cost.

  • Diff-Pruning uses Taylor expansion over pruned timesteps to identify and remove non-critical parameters.
  • The method crafts lightweight models from pre-trained diffusion models while preserving consistency with their original generative behavior.
  • 10% to 20% of pre-training cost is required to produce the compressed models.
Loading 2305.10924v3…