arXiv:2409.19128cs.CV2024-09被引 25

通过剪枝+重加权,用更少数据高效训练扩散模型。

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

  • 先用代理模型选关键数据,再按类别重加权提升性能。
  • 在CIFAR-10上速度提升2.34至8.32倍,生成质量接近全量数据。
  • 适用于像素级和潜在空间扩散模型,适合资源有限的训练场景。

尽管扩散模型(DMs)具备出色的生成能力,但其训练与推理仍极为耗时。现有研究多聚焦于加速采样过程,而对数据高效训练关注不足。本文从数据集剪枝角度探索高效训练方法。受生成对抗网络(GANs)数据高效训练启发,我们首先将GAN中使用的数据选择策略扩展至扩散模型训练:通过代理模型编码数据特征,并基于得分准则选取核心数据子集。为进一步提升生成性能,引入类别级重加权策略,利用分布鲁棒优化(DRO)在预训练参考扩散模型上推导类别权重。在CIFAR-10上的像素级扩散模型(DDPM)实验表明,该方法优于现有方法,在生成质量上接近全量数据训练模型的同时,实现2.34至8.32倍的速度提升。此外,该方法可推广至潜在扩散模型(如MDT、Stable Diffusion),在ImageNet上亦展现出竞争力。代码已公开。

原文摘要 · Abstract (English)

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).

扩散模型数据剪枝重加权高效训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。