Mixture of Efficient Diffusion Experts Through Automatic Interval and Sub-Network Selection

TL;DR

This paper introduces DiffPruning, a method to create a mixture of specialized diffusion model experts for different denoising intervals, reducing sampling costs while maintaining high quality.

Contribution

It proposes a novel approach to prune and fine-tune diffusion models into interval-specific experts with learned routing, optimizing computational efficiency.

Findings

Effective across multiple datasets including LSUN, FFHQ, and ImageNet.

Reduces sampling time while preserving sample quality.

Demonstrates the benefits of interval-specific expert specialization.

Abstract

Diffusion probabilistic models can generate high-quality samples. Yet, their sampling process requires numerous denoising steps, making it slow and computationally intensive. We propose to reduce the sampling cost by pruning a pretrained diffusion model into a mixture of efficient experts. First, we study the similarities between pairs of denoising timesteps, observing a natural clustering, even across different datasets. This suggests that rather than having a single model for all time steps, separate models can serve as ``experts'' for their respective time intervals. As such, we separately fine-tune the pretrained model on each interval, with elastic dimensions in depth and width, to obtain experts specialized in their corresponding denoising interval. To optimize the resource usage between experts, we introduce our Expert Routing Agent, which learns to select a set of proper network configurations. By doing so, our method can allocate the computing budget between the experts in an end-to-end manner without requiring manual heuristics. Finally, with a selected configuration, we fine-tune our pruned experts to obtain our mixture of efficient experts. We demonstrate the effectiveness of our method, DiffPruning, across several datasets, LSUN-Church, LSUN-Beds, FFHQ, and ImageNet, on the Latent Diffusion Model architecture.