[论文解读] Patch Diffusion: Faster and More Data-Efficient Training of Diffusion Models

Zhendong Wang, Yifan Jiang+6 more|arXiv (Cornell University)|Apr 25, 2023

Advanced Neuroimaging Techniques and Applications被引用 32

一句话总结

Patch Diffusion 引入基于补丁的条件分数匹配，结合坐标条件输入与随机补丁大小调度，以加速扩散模型训练并提高数据效率，在保持或提升样本质量的同时实现至少2×更快的训练速度。

ABSTRACT

Diffusion models are powerful, but they require a lot of time and data to train. We propose Patch Diffusion, a generic patch-wise training framework, to significantly reduce the training time costs while improving data efficiency, which thus helps democratize diffusion model training to broader users. At the core of our innovations is a new conditional score function at the patch level, where the patch location in the original image is included as additional coordinate channels, while the patch size is randomized and diversified throughout training to encode the cross-region dependency at multiple scales. Sampling with our method is as easy as in the original diffusion model. Through Patch Diffusion, we could achieve $\mathbf{\ge 2 imes}$ faster training, while maintaining comparable or better generation quality. Patch Diffusion meanwhile improves the performance of diffusion models trained on relatively small datasets, $e.g.$, as few as 5,000 images to train from scratch. We achieve outstanding FID scores in line with state-of-the-art benchmarks: 1.77 on CelebA-64$ imes$64, 1.93 on AFHQv2-Wild-64$ imes$64, and 2.72 on ImageNet-256$ imes$256. We share our code and pre-trained models at https://github.com/Zhendong-Wang/Patch-Diffusion.

研究动机与目标

• 通过降低时间和数据需求来推动扩散模型训练的民主化。
• 提出一个对 UNet 骨架和采样方法不依赖的基于补丁的训练框架。
• 开发补丁级坐标条件和补丁大小调度，以捕捉跨区域依赖。
• 证明补丁扩散可以使训练时间减半，同时保持或提高生成质量，特别是在小数据集上。

提出的方法

• 通过对随机图像补丁进行去噪分数匹配来训练的补丁级去噪器 D_theta(x; sigma_t)。
• 将补丁分数条件化为补丁位置 (i, j) 和补丁大小 s，来自小补丁和大补丁的混合采样。
• 通过在输入中添加两个坐标通道来引入像素级坐标系，编码补丁位置。
• 使用补丁大小调度（随机或渐进）在训练期间混合补丁大小，包括以可控比例 p 使用全尺寸图像。
• 通过在每个反向步骤提供全图坐标，保持与标准扩散模型的采样方式。
• 在 EDM-DDPM++ 和 Latent Diffusion 骨架上演示补丁扩散，在评估阶段可选 CFG。

实验结果

研究问题

• RQ1 基于位置和大小条件的基于补丁的分数匹配是否能提高数据效率并减少扩散模型的训练时间？
• RQ2 补丁大小调度（随机 vs 渐进）如何影响生成质量和训练成本？
• RQ3 在不进行显式补丁合并的情况下，是否可以像标准扩散模型那样进行采样，同时使用基于补丁的训练？
• RQ4 在小数据集上，补丁扩散是否有利于对预训练扩散模型进行微调？

主要发现

• 补丁扩散在基准数据集上实现了至少2×更快的训练，同时保持或达到相当或更好的生成质量。
• 使用小补丁与大补丁混合并结合坐标条件，提升跨区域一致性与FID分数。
• 当训练时全尺寸图像的比例 p 约为 0.5 时，达到了效率与质量的平衡点。
• 补丁扩散提高数据效率，即使在相对较小的数据集上也能获得有竞争力的结果（约 5k 图像左右）。
• 微调大型预训练模型（如 ControlNet）使用补丁扩散在不牺牲性能的情况下提高训练效率两倍。
• 在数据集规模有限时，补丁扩散在 FID 上持续优于基线，训练时间约降低2×。