Patch Diffusion: Faster and More Data-Efficient Training of Diffusion Models

Summary

This paper presents PatchDiffusion, a novel framework designed to address the scalability challenges faced by most diffusion models in terms of training and sampling. The proposed framework adopts a patch-wise training approach, where a denoising network is trained on image patches rather than the entire high-resolution images. To generate patches at the target resolution, PatchDiffusion leverages progressive or stochastic scheduling techniques that utilize different patch sizes throughout the training process. Experimental results on a small-scale dataset demonstrate that PatchDiffusion not only enhances the quality of generated samples but also reduces the overall training time. Furthermore, from the experiment results for medium-scale datasets, PatchDiffusion could be a resource-efficient solution for diffusion training methods.

Strengths

This paper tackles a significant challenge in training diffusion models, which often involve substantial computational costs. To address this issue, the authors propose PatchDiffusion as an optional solution that can be seamlessly integrated into any diffusion model pipeline, regardless of the chosen backbone, sampler, or other modules within the pipeline. I think that PatchDiffusion operates as a form of data augmentation, thereby enhancing the generation quality of the model. This approach provides additional benefits beyond reduced computational costs.

Weaknesses

[Limited experiments - important baseline missing and low performance] The primary motivation behind patch-wise training is to reduce computation costs during training. Given this motivation, it might be worth considering a more efficient backbone instead of U-Net. Recently, successful approaches have replaced U-Net with ViT, as demonstrated in the papers "Scalable Diffusion Models with Transformers" (arXiv'22) and "All are Worth Words: A ViT Backbone for Diffusion Models" (CVPR'23). In the case of DiT, reducing the latent resolution by patchifying with a 2x2 patch size has shown improvements in training scalability and performance. It is worth exploring whether applying patch-wise training to DiT and U-ViT backbones could yield better results. However, the potential gains from such an approach are uncertain. Additionally, it is important to note that in class-conditional image generation on ImageNet-1K, the FID score of PatchDiffusion appears to be significantly worse than that of DiT and other related works. While the current state-of-the-art methods achieve an FID score of less than 4, this work reports a score of around 7.65. [Writing needs to be improved. ] The writing quality of the paper could be improved further, especially in terms of highlighting the comparison between this work and the baselines. It would be beneficial to include a figure illustrating the trade-off between FID (or any other measure of generation quality) and FLOPs. This would help in understanding how PatchDiffusion and other efficient diffusion backbones compare to each other. The authors can refer to similar trade-off figures presented in the DiT and U-ViT papers. Evaluating the quality of writing can be subjective, and I am open to hearing other reviewers' comments on this matter.

Questions

#1. The primary objective of this work is to minimize the computational cost of training/inference through patch-wise training. However, it could be worthwhile to consider an alternative solution by employing a more efficient backbone inherited from ViT, which incorporates a patch-fying (i.e., tokenizing) module. Including a comparison of this work with DiT or U-ViT in the paper would further highlight the benefits of this approach. Moreover, presenting training compute vs. FID plots for the comparison would greatly assist in determining the most effective approach. #2. In my understanding, the true advantages of this work are likely to be demonstrated through experiments on fine-tuning. This approach has the potential to reduce the fine-tuning cost for any pre-trained diffusion backbone. In this context, incorporating LoRA-type methods could complement this approach. Including empirical analysis of this nature in the paper would enhance the understanding of the benefits of this work. #3. The GAN community has explored patch-wise training in various ways. Recently, Any-resolution GAN (ECCV’22) has been introduced as a promising solution for training generative models with variable-size images. Although this paper primarily aims to reduce training costs rather than utilizing multiple size images in the training procedure, it would be interesting to explore whether the proposed framework can be extended to generate variable-size images. Doing so would further highlight the benefits of this framework. #4. A minor comment regarding the inclusion of "Appendix A." To provide a concise overview of the theoretical interpretation of the patch-wise training scheme, it would be beneficial to incorporate a brief summary of these observations within the main body of the paper.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

The limitations of this work were not explicitly outlined, and I didn’t observe any discussion regarding potential negative societal impacts. However, there is no clear negative societal impact, since all experiments are conducted in controlled benchmark datasets.

Summary

This paper presents a new training technique that improves the training speed of diffusion models. Instead of training the diffusion model on the entire image, the authors propose training on sampled patches of the image. This approach maintains the theoretical foundation of the diffusion model by keeping the training objective function mostly unchanged. By combining this approach with a fully-convolutional U-Net architecture, the computational complexity is reduced, resulting in faster training. To minimize the quality difference between the partial image approach and the conventional whole image approach, a stochastic/progressive patch size scheduling was proposed, and the corresponding ablation study was conducted. This study investigates the optimal probability of using the whole image, considering the trade-off between training time and generation quality. Summarizing, the proposed method enables faster learning while preserving the theoretical foundation and generation quality of traditional diffusion models.

Strengths

This paper presents two significant advantages of training diffusion models using partial images instead of the entire images. Firstly, it reduces model complexity, leading to faster training, which is especially beneficial for state-of-the-art baseline diffusion models that require extensive GPU hours for training. This approach offers the potential for energy-efficient training by effectively reducing the overall training time. Secondly, training with partial images proves effective in scenarios with limited datasets, outperforming traditional methods. In cases where the dataset is insufficient, diffusion models struggle to accurately predict the true data distribution due to overfitting (limited data is essentially a sparse sampling of true data distribution). By partitioning the image into patch units, the training process simulates a larger dataset, providing the model with more training samples. This enables the diffusion model to estimate a more accurate data distribution, resulting in improved generation quality. The paper conducted a series of experiments encompassing large-scale datasets, limited-size datasets, and fine-tuning scenarios, to demonstrate the effectiveness of training images in patch units. This approach maintains performance while significantly improving training efficiency.

Weaknesses

The key idea of this paper is to modify the input data format while maintaining the training process of the diffusion model. This approach aligns with similar strategies proposed in previous works like COCO-GAN. Considering the large overlap in ideas, the authors should present various case studies, such as showcasing the application of this technique to diffusion models, in order to compensate for the limited novelty of the paper. Firstly, there is a lack of analysis concerning spatial conditions. The authors convert the traditional three-channel format (R, G, B) to a five-channel format (R, G, B, i, j) incorporating location information. Meanwhile, other methods such as positional encoding in modern Transformer structures or Fourier feature methods in Alias-Free GAN have been proposed for spatial conditioning. Multiple experiments in various literature have demonstrated that positional encoding is more effective than raw methods like (i, j). Conducting an ablation study on different spatial conditioning methods and providing an analysis of the most suitable conditioning approach in a diffusion setting would strengthen the paper's credibility. Secondly, there is insufficient analysis regarding the ability of patch diffusion to generate structural diversity. The experimental results in Table 1 and Table 2 show a notable improvement in the quality of patch diffusion for datasets with weak common structures (e.g., Bedroom, Church) compared to those with strong common structures (e.g., FFHQ, CelebA). Including an analysis of the factors contributing to this quality improvement, such as visualizing attention layers for patch images, would help readers' understanding of the method.

Questions

1. Patch-wise generation: Is it feasible to generate the entire image by combining generated parts instead of generating it all at once? 2. Patch diffusion for image extrapolation: What would happen if channels for (i, j) were provided with a range of (-1.2, 1.2)? Can patch diffusion extrapolate the image using this approach? 3. Figure quality: To enhance the readability and visual appeal of Figure 1, I recommend refining its quality. Currently, the overall figure appears hastily created, resembling a PowerPoint slide. I kindly request aligning each object for the camera-ready submission to improve its structure. The objects currently occupy space without significant value. Also, the range of i, j is -1~1, but the crop has a different scale, such as 16x16. Please provide an example of the actual value when a 16x16 crop is performed. Furthermore, the font size in the figure is very small, making it difficult to read. Considering the available white space, increasing the font size throughout would greatly improve legibility. 4. Typo: Line 86: "e.g.,." should be corrected to "e.g."

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The authors adequately addressed the limitations and potential negative societal impact.

Summary

The paper introduces a path-wise diffusion algorithm for faster training. The authors propose patch coordinate conditioned diffusion models and present a patch-size conditioning scheduling technique for efficient training. The method has similar motivation with patch based GAN such as COCOGAN, but it is applied to diffusion model.

Strengths

The method presented in the paper is simple yet effective, successfully reducing the training time by half. The motivation behind the approach aligns with that of COCOGAN, and it can be seen as a reasonable extension for the diffusion model to reduce computational requirements.

Weaknesses

To learn global structure, the algorithm still need a portion of full-resolution diffusion which introduces a bottleneck in terms of time and memory costs. Additionally, the patch-size scheduling approach appears to be manually designed, potentially limiting its adaptability and automation in optimizing the training process. More controlled experiments on various positional encoding or different hyperparameters would provide better intuition about the work.

Questions

1. Is it possible to extrapolate an images as shown in COCOGAN? 2. What is FID scores of baseline model with same train-time without patch training? (e.g. 24h FIDs of EMM-DDPM++ for CelebA)

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

The author provides some limitations on conclusion section.

Summary

The authors propose a new formulation of training diffusion models by sampling different-sized patches from the training data. The models trained with this formulation have comparable FID scores to models trained on full images on many datasets.

Strengths

The authors proposed a way to train faster diffusion models, which can cut down training time in half with almost similar performance. Well-written and to the point paper.

Weaknesses

Training with patches of image with same guidance might not work when the training data is more heterogeneous like LAION. In datasets where the subject might not be centered in image, or occupy small portion of the image or existence of multiple object in the scene, I wonder if model can still perform well.

Questions

Do you have any results on models trained on LAION or MS COCO?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

Summary

The paper proposes a novel training framework for diffusion models, that significantly reduces the training time, while improving data efficiency. For the first time, the proposed method suggests patch-wise diffusion training, which can be deployed to any UNet-based diffusion models. Experimental results show that the patch-wise diffusion training can halve the training time while maintaining comparable or better image quality than the baseline models. On small scale datasets, the proposed method outperformed other baselines, validating the data efficiency of path-wise training.

Strengths

- The paper is well organized and easy to follow. The motivation of the paper is very clear, which is to shorten the training time of diffusion models while maintaining image generation quality. - Proper ablation studies are delivered to fully validate roles of different components of the model, and affect of different parameters. - 2x speedup of training time is non-negligible, considering long training time of conventional diffusion models. - The data efficiency followed by patch-wise training framework is considered a significant discovery. - The proposed method can be applied to any UNet-based diffusion models in a plug and play method.

Weaknesses

- Albeit the empirical evidences, the theoretical proof of convergence of patch-wise score matching is missing. - Experiments on high resolution image synthesis, beyond 256x256 resolution is missing. According to Table1 and 2, compared to the baseline method, the proposed patch diffusion showed better performance boost on the larger scale datasets (LSUN-Bedroom&Church), than the smaller scale datasets (CelebaA and FFHQ). Thus, it might imply that the patch diffusion can be more beneficial in high resolution image synthesis scenarios. Therefore, validation on high resolution image datasets, beyond 256x256 resolution could be interesting.

Questions

- While training, when entering the denoiser (UNet), has the small patches resized into the original image size? If not, will there be a difference between using resizing and not?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

Limitations are addressed by the authors.

Thank you

Thank you for the response. I stick to my current score.

Increasing my score from BR to BA.

Thank you for your detailed response. It has clarified most of the concerns I initially had about the paper. I had misunderstood certain aspects, particularly regarding the FiD score of DiT. Beyond the clarifications, I find the extrapolation (out-painting) results particularly interesting, since it's very practical that this diffusion model, trained using patch-wise methods, inherently offers the out-painting capability. As a result, I have revised my evaluation to BA (borderline accept). Regarding the manuscript, I agree with Reviewer Vqdw's observations about the subpar quality of the main figure. I recommend updating both the main figure and the manuscript itself.