[Paper Review] 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 Applications32 citations

TL;DR

Patch Diffusion introduces patch-wise conditional score matching with coordinate-conditioned inputs and random patch-size scheduling to accelerate diffusion-model training and improve data efficiency, achieving at least 2x faster training while maintaining or improving sample quality.

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.

Motivation & Objective

• Motivate democratizing diffusion-model training by reducing time and data requirements.
• Propose a patch-wise training framework that is agnostic to UNet backbones and sampling methods.
• Develop patch-level coordinate conditioning and patch-size scheduling to capture cross-region dependencies.
• Show that patch diffusion can halve training time while maintaining or improving generation quality, especially on small datasets.

Proposed method

• Define a patch-wise denoiser D_theta(x; sigma_t) trained via denoising score matching on random image patches.
• Condition the patch scores on patch location (i, j) and patch size s, sampled from a mixture of small and large patches.
• Introduce a pixel-level coordinate system by adding two coordinate channels to the input, encoding patch location.
• Use patch-size scheduling (stochastic or progressive) to mix patch sizes during training, including full-size images at a controllable ratio p.
• Maintain sampling as in standard diffusion models by providing full image coordinates at each reverse step.
• Demonstrate patch diffusion on top of EDM-DDPM++ and Latent Diffusion backbones, with CFG optional during evaluation.

Experimental results

Research questions

• RQ1 Does patch-wise score matching with location and size conditioning improve data efficiency and reduce training time for diffusion models?
• RQ2 How does patch-size scheduling (stochastic vs progressive) affect generation quality and training cost?
• RQ3 Can sampling be performed as in standard diffusion models without explicit patch-merging, while using patch-based training?
• RQ4 Is patch diffusion beneficial on small datasets and for fine-tuning pretrained diffusion models?

Key findings

• Patch Diffusion achieves at least 2x faster training while maintaining comparable or better generation quality on benchmark datasets.
• Using a mixture of small and large patches with coordinate conditioning improves cross-region consistency and FID scores.
• A sweet spot around p = 0.5 (ratio of full-size images during training) balances efficiency and quality.
• Patch Diffusion improves data efficiency, enabling competitive results on relatively small datasets (as few as ~5k images).
• Finetuning large pretrained models (e.g., ControlNet) with Patch Diffusion doubles training efficiency without sacrificing performance.
• On limited-size datasets, Patch Diffusion consistently outperforms baselines in FID with about 2x training-time reduction.