Source-linked AI summary
Patch Diffusion: Faster and More Data-Efficient Training of Diffusion Models
Zhendong Wang, Yifan Jiang, Huangjie Zheng, Peihao Wang, Pengcheng He, Zhangyang Wang, Weizhu Chen, Mingyuan Zhou
TL;DR
Diffusion models are powerful but costly and data-hungry to train, limiting broader access. Patch Diffusion trains conditional score functions on coordinate-conditioned patches with diversified patch sizes, achieving generally faster training and improved small-data performance while retaining comparable or better generation quality.
Problem
Diffusion-model training requires prohibitive time, computation, and data, creating a bottleneck for democratizing access to competitive training.
Method
Patch Diffusion uses patch-level conditional score matching with patch-location coordinate channels and progressive or stochastic patch-size diversification.
Results
2× faster training is reported while maintaining comparable or better generation quality and improving performance in small-data regimes.
Takeaways & Limitations
Patch Diffusion provides a generally applicable, plug-and-play framework for reducing diffusion-model training costs and improving data efficiency.
Takeaways & Limitations
Patch-wise score learning may fail to capture global cross-region dependencies without strategies such as random patch sizes and large-patch guidance.
Abstract
from arXiv · showhide
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\times}$ 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$\times$64, 1.93 on AFHQv2-Wild-64$\times$64, and 2.72 on ImageNet-256$\times$256. We share our code and pre-trained models at https://github.com/Zhendong-Wang/Patch-Diffusion.
1 Introduction
Patch Diffusion addresses the high time and data costs of diffusion-model training with patch-wise conditional score learning. It uses coordinate conditioning and diversified patch sizes to improve efficiency while preserving global structure and generation quality.
- Diffusion models are expensive and data-hungry to train, with DDPM requiring approximately four days at 64×64 and over two weeks at 256×256 on eight V100 GPUs.
- The framework targets democratizing diffusion-model training by reducing prohibitive training time and data requirements for researchers with fewer resources.
- Patch Diffusion learns conditional score functions on image patches rather than full images, conditioning on patch location and size.Patch locations are encoded as additional coordinate channels, while patch sizes are diversified through progressive or stochastic scheduling.
- The method achieves comparable or better generation quality while reducing training costs and improving data efficiency.Sampling remains as easy as in the original diffusion model.
- Patch Diffusion generally halves training time while delivering competitive results and notable gains in small-data regimes.
2 Related Work
Related work has reduced diffusion-model costs through latent-space modeling and faster sampling, while prior patch-wise training approaches have focused mainly on other generative-model settings. Patch Diffusion is positioned as an orthogonal, plug-and-play training module.
- Earlier data-efficient generative-model work includes differentiable and adaptive augmentation for GANs trained with limited data.One cited GAN approach reports acceptable performance with only 10% data.
- Few-shot diffusion studies have mainly addressed adapting pretrained models using four images or even a single example.
- Coco-GAN introduced patch-wise training for GANs, but its discriminator still combines generated patches, limiting memory savings.
- Latent-space diffusion reduces training and inference cost, while improved samplers primarily accelerate inference rather than training.
- Patch Diffusion is described as orthogonal to latent-space and fast-sampling approaches and usable as a plug-and-play module.
- A concurrent patch-based diffusion model targets memory-efficient high-resolution synthesis and avoids boundary artifacts through feature collage and window sliding.
3 Patch Diffusion Training
Patch Diffusion trains conditional score functions on randomly cropped patches, using patch coordinates and varied sizes to reduce computation while preserving global coherence. The method retains simple global sampling and supports flexible UNet-based diffusion models.
- Sampling and Architecture Compatibility: Sampling remains global and as straightforward as in the original diffusion model, while convolutional UNets can handle patches of varying resolutions.The method concatenates full-image coordinates with sampled noise before reversing the diffusion chain.
- 3.1 Patch-wise Score Matching: Patch Diffusion replaces full-image score matching with conditional score matching on randomly sized patches.Each patch is conditioned on its location and size.
- 3.2 Progressive and Stochastic Patch Size Scheduling: Random patch sizes and occasional full-size images address missing cross-region dependencies by training the score function to combine local scores coherently.Large patches can be viewed as sequences of small patches, while full-size images help align the reverse process with the original data distribution.
- 3.2 Progressive and Stochastic Patch Size Scheduling: Patch-size scheduling can be stochastic or progressive, with progressive training moving from smaller patches to larger patches and then full-size images.The full-size phase occupies a ratio p of total training iterations.
- 3.2 Progressive and Stochastic Patch Size Scheduling: p = 0.5 with stochastic scheduling provides an empirical trade-off between training efficiency and generation quality.This setting is identified as a sweet point in the trade-off.
- 3.1 Patch-wise Score Matching: Patch coordinates are normalized to [−1, 1] and supplied as two additional input channels to the denoiser.The coordinate channels identify each patch’s position in the original image.
4 Experiments
Experiments evaluate Patch Diffusion across ablations, large and small datasets, high-resolution latent diffusion, and ControlNet finetuning. The method generally reduces training cost while preserving or improving generation quality, with stochastic patch scheduling and p = 0.5 providing a favorable trade-off.
- Experimental scope: Experiments cover ablations, standard benchmarks, latent diffusion, ControlNet finetuning, and datasets containing approximately 5k images.The evaluation uses CelebA, FFHQ, LSUN, ImageNet, AFHQv2-Cat, AFHQv2-Dog, and AFHQv2-Wild.
- Ablation study: p = 0.5 provides a favorable trade-off between generation quality and training efficiency.Larger p increases training cost and time, while p = 0.75 reaches the minimum FID level.
- Ablation study: 1.66 versus 2.05 FID on CelebA-64×64 and 3.11 versus 3.85 on FFHQ-64×64 favor stochastic over progressive patch-size scheduling.The comparison is reported as stochastic versus progressive scheduling.
- Large-scale datasets: Patch Diffusion generally maintains strong FID while significantly reducing training time on benchmark datasets.For high-resolution synthesis, Latent Patch Diffusion combines patch training with latent diffusion; the pretrained autoencoder may limit performance.
- Finetuning: Patch Diffusion can be applied to ControlNet finetuning without compromising performance after 20k training steps.The experiment finetunes ControlNet for HED-map-to-image generation from Stable Diffusion v1-5.
- Limited-size datasets: On three datasets with around 5k images each, Patch Diffusion consistently outperforms EDM-DDPM++ in FID while achieving ≥2× faster training.The models are trained from scratch for 75 million images.
5 Conclusion and Future Work
Patch Diffusion is a patch-level, coordinate-conditioned score-matching framework that diversifies patch sizes to capture dependencies across multiple scales. It targets faster training and improved data efficiency, while noting positional-embedding and convergence theory as future directions.
- Patch Diffusion trains diffusion models through coordinate conditioned score matching at the patch level.
- Randomized or progressive patch-size schedules capture cross-region dependencies at multiple scales.
- 2× faster training is reported while maintaining comparable or better generation quality.
- Patch Diffusion improves the data efficiency of models trained on relatively small datasets.
- Future work includes advanced positional embeddings and theoretical convergence proofs for general patch-wise score matching.
A Theoretical Interpretations
The theoretical interpretation views patch diffusion through Markov random fields and linear regression. Patch-wise score matching decomposes or projects the full-image problem, reducing computation but relying on image redundancy and symmetry for recovery.
- Markov Random Field: Markov random fields represent images as undirected pixel-grid graphs whose distributions factor through node and edge potentials.
- Markov Random Field: The MRF score decomposes into node and edge contributions, allowing patch training to learn local score components and combine them during inference.
- Markov Random Field: Coordinate conditioning models the location dependence of node and edge score components.
- Linear Regression: Patch distributions are marginals of the full-image distribution, represented in the Gaussian example by a selection matrix associated with patch location and size.
- Linear Regression: Patch-based and full-size score matching differ through the measurement matrix, trading computational cost against well-posedness.
- Linear Regression: Image redundancy and symmetry are proposed to make whole-distribution recovery from limited patch observations feasible, while general convergence remains future work.
B More Generated Images.
The appendix presents additional generated-image examples from Patch Diffusion and Latent Patch Diffusion across bedroom, church, and ImageNet settings, plus further ControlNet finetuning results.
- The appendix states that more generated images from Patch Diffusion are listed.
- Additional Patch Diffusion samples use an EDM-DDPM++ backbone trained on LSUN-Bedroom and LSUN-Church at 256×256 resolution.
- Additional Latent Patch Diffusion samples use an EDM-ADM backbone trained on ImageNet at 256×256 resolution.
- Further finetuning results for ControlNet with patch diffusion training are shown.
C Potential Social Implications
The authors identify potential social harms from computer-vision generation models, especially reduced trust in visual evidence and misuse for deceptive or abusive content.
- Generated or manipulated images may make visual evidence harder to assess as real, undermining credibility.
- Potential misuse includes false identities, deepfake pornography, propaganda, and disinformation campaigns.
- These risks may affect journalism, forensics, and other fields that rely on accurate visual representation.
D More Extrapolation Results
Patch Diffusion supports extrapolating beyond the training image boundary and produces ImageNet samples with or without classifier-free guidance. The extrapolation demonstration regenerates out-of-boundary regions at twice the training resolution.
- Patch Diffusion extrapolates beyond a 256 × 256 image boundary, generating a 512 × 512 sample while keeping the reference image fixed at the center.The out-of-boundary region is regenerated during the reverse process, although the model was trained only on 256 × 256 images.
- The extrapolation results demonstrate that the learned coordinate manifold can generate content beyond the original image boundary.
- Latent Patch Diffusion generates ImageNet-256×256 samples using an EDM-ADM backbone with either no classifier-free guidance or classifier-free guidance.The figure places samples without CFG on the left and samples with CFG on the right.