Pruning then Reweighting: Towards Data-Efficient Training of Diffusion Models

TL;DR

This paper proposes a data-efficient training method for diffusion models by combining dataset pruning with class-wise reweighting, resulting in faster training with comparable image synthesis quality.

Contribution

It introduces a novel data selection and reweighting approach for diffusion models, improving training efficiency without sacrificing generation quality.

Findings

Achieves 2.34x to 8.32x speed-up on CIFAR-10 with comparable quality

Extends the method to latent diffusion models like SD and MDT

Demonstrates effectiveness on ImageNet dataset

Abstract

Despite the remarkable generation capabilities of Diffusion Models (DMs), conducting training and inference remains computationally expensive. Previous works have been devoted to accelerating diffusion sampling, but achieving data-efficient diffusion training has often been overlooked. In this work, we investigate efficient diffusion training from the perspective of dataset pruning. Inspired by the principles of data-efficient training for generative models such as generative adversarial networks (GANs), we first extend the data selection scheme used in GANs to DM training, where data features are encoded by a surrogate model, and a score criterion is then applied to select the coreset. To further improve the generation performance, we employ a class-wise reweighting approach, which derives class weights through distributionally robust optimization (DRO) over a pre-trained reference DM. For a pixel-wise DM (DDPM) on CIFAR-10, experiments demonstrate the superiority of our methodology over existing approaches and its effectiveness in image synthesis comparable to that of the original full-data model while achieving the speed-up between 2.34 times and 8.32 times. Additionally, our method could be generalized to latent DMs (LDMs), e.g., Masked Diffusion Transformer (MDT) and Stable Diffusion (SD), and achieves competitive generation capability on ImageNet. Code is available here (https://github.com/Yeez-lee/Data-Selection-and-Reweighting-for-Diffusion-Models).