Sparse-to-Sparse Training of Diffusion Models

Technical Summary: Sparse-to-Sparse Training of Diffusion Models

Problem Statement
Diffusion Models (DMs) have achieved state-of-the-art performance in generative tasks, including image synthesis, natural language processing, and temporal data modeling. However, their widespread adoption is hindered by significant computational costs during both training and inference. While prior research has largely focused on accelerating the inference stage (e.g., through distillation or faster sampling), the training process remains extensive and resource-intensive. Furthermore, existing network compression techniques, such as post-training quantization or pruning pre-trained dense models, still require the initial expensive training of full-scale dense networks. There is a need for methods that reduce computational costs and memory footprints from the outset of the training process.

Methodology
This paper introduces the paradigm of sparse-to-sparse training to Diffusion Models, aiming to train sparse networks from scratch rather than compressing pre-trained dense ones. The study focuses on unstructured sparsity and evaluates three specific training strategies applied to two state-of-the-art DM architectures: Latent Diffusion (for continuous image data) and ChiroDiff (for discrete spatiotemporal sketch data).

The three proposed methods are:

1. Static-DM (Static Sparse Training): The network connectivity is fixed at initialization based on a sparsity ratio. The authors employ Erdős–Rényi (ER) and ERK strategies to allocate non-zero weights, ensuring larger layers have higher sparsity. The connectivity pattern remains unchanged throughout training.
2. RigL-DM (Dynamic Sparse Training): This method utilizes a dynamic prune-and-grow scheme. Weights are pruned based on their magnitude, and new connections are regrown based on the magnitude of the gradients (gradient growth).
3. MagRan-DM (Dynamic Sparse Training): Similar to RigL-DM regarding pruning (magnitude-based), but regrowth is determined randomly rather than by gradients.

The experiments cover six datasets (LSUN-Bedrooms, CelebA-HQ, Imagenette, QuickDraw, KanjiVG, and VMNIST) across various sparsity levels ($S \in \{0.1, 0.25, 0.5, 0.75, 0.9\}$). The authors also investigate the impact of the prune and regrowth rate ($p$) and the exploration frequency ($\Delta T_e$) for dynamic methods.

Key Contributions

• First Application of Sparse-to-Sparse Training to DMs: The paper establishes the feasibility of training unconditional Diffusion Models from scratch using sparse connectivity, covering both static and dynamic strategies.
• Comprehensive Empirical Evaluation: The study benchmarks three sparse training methods against dense baselines across two distinct model architectures and six datasets, providing a broad view of sparsity effects on DM performance.
• Identification of Optimal Hyperparameters: The research identifies that while dynamic sparse training (DST) generally performs well, the choice of the prune and regrowth rate is critical. A conservative rate ($p=0.05$) yields better results in high-sparsity regimes compared to the more aggressive rate ($p=0.5$) often used in other contexts.
• Performance vs. Efficiency Trade-offs: The work demonstrates that sparse DMs can match or exceed the performance of dense counterparts while significantly reducing the number of trainable parameters and Floating Point Operations (FLOPs).

Experimental Results

• Performance: In most cases, at least one sparse configuration outperformed its dense counterpart. For instance, on the CelebA-HQ dataset, RigL-DM with 25% sparsity achieved a lower FID score (32.12) than the dense model (32.74) while reducing parameters by 25% and FLOPs by 10%. On QuickDraw, MagRan-DM with 90% sparsity maintained comparable performance to the dense model while reducing FLOPs by 88%.
• Sparsity Levels: Moderate sparsity levels (25%–50%) generally yielded the best balance between performance and efficiency. Extremely high sparsity (e.g., 90%) often led to performance degradation, though this varied by dataset and model architecture.
• Dynamic vs. Static: While Dynamic Sparse Training (DST) methods (RigL-DM, MagRan-DM) often outperformed Static-DM in moderate sparsity regimes, Static-DM showed competitive or superior results in very high sparsity settings ($S > 0.5$) when using the aggressive $p=0.5$ rate. However, when using the conservative $p=0.05$ rate, DST methods regained superiority even in high sparsity regimes.
• Training Dynamics: Sparse models exhibited stable training dynamics similar to their dense counterparts, suggesting that the sparse-to-sparse paradigm does not destabilize the diffusion training process.
• Conditional Generation: Preliminary experiments on class-conditional models (ImageNet and QuickDraw) indicated that sparse-to-sparse training is also effective in conditional settings, though dense models generally held a slight edge in ImageNet.

Significance and Claims
The paper claims that sparse-to-sparse training offers a promising pathway to improve the efficiency of Diffusion Models without sacrificing generation quality. By training sparse networks from scratch, the approach addresses both training and inference costs, potentially reducing the carbon footprint and memory requirements associated with DMs. The authors emphasize that selecting the appropriate method (Static vs. Dynamic), sparsity level, and hyperparameters (specifically the prune/regrowth rate) is crucial for success. The work suggests that with the right configuration, sparse DMs can be a viable alternative to dense models, enabling broader implementation and faster iteration in generative AI research. The authors note that while current hardware often limits the practical speedup of unstructured sparsity, the reduction in parameter count and FLOPs remains a significant theoretical and practical advantage, especially as hardware support for sparse operations evolves.