Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment

We thank you for some of the constructive feedback on our work. We would like to present following updates:

For W1, we have revised our mathematical proof part in rebuttal to reviewer-neka - W1, taking into consideration your constructive suggestions.

For W1, we compute the average of the images on the ImageNet dataset (1M images), the result is shown in Figure R1, demonstrating that the average of images does not contain much meaningful information, as we stated in the revised mathematical proof.

For W2, we have included an additional discussion on our difference to OT-family methods in the global response for your reference.

For L1, we have included a discussion of the influence of the batch size in rebuttal to review-neka-Q1 for your reference.

For L3, we have provided the performance of conditional generation of immiscible diffusion in the global response for your reference.

We sincerely thank you for your acknowledgement of our work. Besides, we also hope to emphasize that our method is one simple implementation to address the miscibility problem demonstrated to be important, into which we hope to inspire more work to address. We hope this work can benefit the diffusion model society beyond the implementation itself. We hope to address your concerns below from our understanding:

W1 - Diversity

We sincerely thank you for the comment. Balancing the model fitting for FID and diversity is truly hard. When there are no image-noise relationships, as shown in Fig. 3(a), the denoising of noisy layers can not perform meaningful denoising functions due to miscibility so the model is hard to be optimized to fit the data; however when the relation comes too tight like “fixing and memorizing”, there can logically be diversity problems as you raised.

Our method balances between these two extremes: we perform image-noise assignment in a mini-batch. In this way, we avoid significant miscibility and make noisy denoising layers effective (as shown in Fig. 3(b)) by encouraging the diffusion of each image to the surrounding area. From the other hand, considering the low batch size (~256 for DDIM in CIFAR) and the very high image dimension (32 * 32 * 3=3072 for CIFAR), such encouragement is very weak. Tab. 3 shows that the distance between image and noise only decreases ~2% after assignment (the stddev of it is >10%). Therefore, the diffuseable area for each image is still very flexible and the diversity can be kept.

Our image generation experiments further confirm these claims. Firstly, as suggested by your 3rd point, we show in Fig. R7 that immiscible diffusion does not significantly affect the diversity. Furthermore, we employ immiscible diffusion to the conditional generation problem with Stable Diffusion on ImageNet. Results in Fig. R2 shows that image quality as well as the FID are improved, which further supports that our immiscible diffusion balances the miscibility and the diversity well.

As a conclusion, we believe that immiscible diffusion takes the balance between miscibility and diversity, thus providing faster training without significantly sacrificing the diversity.

W2 - Theory for Performance Enhancement

We really appreciate your question. Immiscible diffusion works by making images less miscible in the noise space, which is proved in Fig. 3: where we see that the denoiser learns which image to reconstruct after immiscible training. Our further theoretical experiment in Fig. R5 shows that such immiscibility is mainly achieved by pushing away images far from each other in the noise space, as the fitted line only has a significant positive absolute value in the far right of the image. In this way, we achieve more immiscible diffusion while avoiding major disturbance to the diversity. The image-noise distance change of ~2% is to prove how minor our disturbance is, which supports that immiscible diffusion would not significantly affect the diversity as discussed in the last question.

W3 - 2 Proposed Experiments

We are so grateful for your offering of experiment designs, which inspire us a lot. We perform both sampling experiments with the setting of Unconditional DDIM + CIFAR10 + Batch Size 256 + 187k Training Steps + 50 Sampling Steps:

i) We generate images with a fixed noise, altered on 10/96/320 dimensions (10 pixel = 30 dimensions due to 3 channels), which are shown in Fig. R7. We found that altering 10 dimensions would not significantly alter the image for both the immiscible and the vanilla DDIM. With 96 dimensions altered, some minor changes are observed for both models. And when we altered 320 dimensions (~10%), both models show class alternations in generated images. In conclusion, we don’t observe a significant diversity difference between the vanilla and the immiscible DDIMs.

ii) The results can be found in Fig. R8 We observe that with the same noise, the global seed does not seem to have an impact on the generated image. We think that is caused by little randomness of the DDIM sampling process. For DDPM, diverse classes of images are generated, as randomness in DDPM can also come from the sampling path, in addition to the initial noise.

W4 - Scaling Law Thanks for the comment. We studied the scaling law from theoretical analysis. For theoretical analysis, as illustrated above, our method almost does not change the assumption that $P(X_T)$ should be distributed in Gaussian while we push the image-noise distribution away from each other. This means our method still preserves the manifold of diffusion training while making the optimization on the high-noise level easier (see Fig. 3 of our main paper). In that case, our method should not hurt the scaling-law facts for Diffusion Models. We also conducted experiments on the whole ImageNet dataset and found that our immiscible-Stable Diffusion can improve the performance of the vanilla Stable Diffusion model. This further demonstrates that our method can work on large-scale datasets (see general response).

Thanks for your acknowledgement on the potential of our method. Besides, we also hope to emphasize that our method is one implementation to address the miscibility problem, into which we hope to inspire more work to work on. Below we will address your concerns:

W1 - Mathematical Proof

We re-write the part 3.3 below:

In DDPM, we know for any image data-point $x_0$, when it comes to the last diffusion step $T$, i.e. $t = T$, the image is sufficiently wiped out and nearly only Gaussian noise remains. Therefore,

Utilizing Bayes' Rules and Equation 1, we can find that for a specific $x_T$:

which indicates that the distributions of the corresponding images for any noise data-point are the same as the distribution of all images.

The simplified training objective in DDPM is the added noise $\epsilon(x_t, t)$. However, we find that for a specific point $x_T$ in the noise space at diffusion step $T$,

where $a = -\frac{\sqrt{\overline{\alpha_t}}}{\sqrt{1-\overline{\alpha_t}}}$ and $b = \frac{1}{\sqrt{1-\overline{\alpha_t}}}$ are constants, and $\bar{x_0}$ is an average of images in the dataset. When the number of images is large enough, $\bar{x_0}$ contains little meaningful information.

However, in our Immiscible Diffusion, while we still have

for each specific data point $x_T$ or $x_0$, the conditional noise distribution does not follow the Gaussian distribution because of the batch-wise noise assignment

Instead of the Gaussian distribution, we assume that the predicted noise with noise assignment has a distribution:

where $f$ is a function denoting the influence of assignment on the conditional distribution of $x_T$. According to the definition of the linear assignment problem, $f$ decreases when $x_T - x_0$ increases its norm, say L2 norm as our default setting.

Therefore, from Eqn. 2 and 6, we have

which means that for a specific noise data-point, the possibility of denoising it to the nearby image data-point would be higher than to a far-away image.

For the noise prediction task, we see that

where $\overline{x_0 f(x_T - x_0, ...)}$ is the weighted average of $x_0$ with more weights on image data-points closer to the noisy data-point $x_t$. Therefore, the predicted noise leads to the average of nearby image data-points, which makes more sense than pointing to a constant...

W2 - Evidence on Miscibility

The evidence is in Fig. 3(a), which shows the predicted noise in each step, and the image generated solely with the predicted noise for this step, both from a trained DDIM. In this 20 steps sampling, T=20 is the step denoising from pure noise while T=1 is the step outputting the image. We see that the predicted noise at T=20 looks messy, and the image generated with this noise has little meaningful info. These support that the model can not determine the image to be denoised to, which supports the miscibility. Mathematically, in vanilla diffusion models, the noise added to each image during training is N(0, I). Therefore, each image would be projected to the noise space with the same possibility distribution N(0, I), which is miscible.

W3&4 - Typo on Eqn. 2 and Eqn. 7

We appreciate your comment. We change the notation t to T, which means the last diffusion step, so that Eqn. 2 would approx. hold considering the new Eqn. 1.

Q1 - Discussion on Batch Size (BS) To see its impact, we vary the BS on DDIM with CIFAR with 50 sampling steps. Results in Fig. R9 show that larger BS enjoys better FID, but the influence of immiscibility significantly outweighs that of BS, and such influence does not vary significantly across selected BS.

Q2 - Applying Assignment to Which Steps

Yes! We further apply immiscible diffusion on large diffusion steps only, on conditional generation with Stable Diffusion (SD) and ImageNet, BS = 2048 and training steps = 20k. We compare: 1) Vanilla diffusion; 2) Immiscible diffusion on all steps; 3) Immiscible diffusion on t > 50% total diffusion steps;4) Immiscible diffusion on t > 75% total diffusion steps. The FID results are 1) 22.44; 2) 20.90; 3) 21.23; 4) 21.51. We find that denoising at large steps indeed improves the performance comparing to vanilla diffusion. This matches findings in Fig.3 that our method improves large diffusion steps. Assigning the image-noise across more steps can further improve the performance, which saves us the efforts to finetune the hyperparameters of steps.

Q3 - Finetuning

Yes. We finetune conditional Stable Diffusion with pre-trained stable-diffusion-v1.4 model on ImageNet dataset with a batch size of 512. We see that after 2.5k steps, the FID for immiscible and vanilla stable diffusion is 11.10 and 12.19 respectively, supporting that our method works effectively for fine-tuning.

Q4 - Runtime for Assignment

We refer to Tab. 3 for the detailed time necessary for the assignment with different batch sizes. In a typical situation (DDIM, batch size=256), each training step is ~460ms, where the 6.7ms assignment time is truly negligible (~1.5% overhead).

Reviewer nEkA:

It is only one day away, please provide your feedback to the authors' rebuttal.

AC

I appreciate the reviewers' clarification and have decided to increase my score.

We are really excited to hear that you agree with our strengths listed above, and we sincerely hope that the proposed perspective of image-noise matching can be realized and further discovered by our diffusion society. We hope to address your concerns below:

W1-Practicality

We thank you for the concern on whether immiscible diffusion would work on datasets without simple classification distributions, and we highly appreciate your suggestion on training on LAION. We train conditional Stable Diffusion model on a subset of LAION (laion-10k) with a batch size of 512. We see that after 25k steps, the FID for immiscible and vanilla diffusion is 124.55 and 146.69 respectively, showing that immiscible diffusion can work on the training efficiency of datasets with complex distributions. Note that the training steps are limited by the time and computational resources.

W1-Generalization

We can’t agree more to explore on this concern, which we have been seriously considering after submission. A typical example and an important application here is conditional image generation, where noise is randomly sampled and paired with the text prompt. We performed class-conditional image generation on Stable Diffusion with ImageNet dataset. Our results in Fig. R2 showed that immiscible diffusion outperformed the vanilla one, specifically the FID for immiscible and vanilla diffusion is 20.90 and 22.44 respectively, and the training steps necessary to go below the FID <23 is 12.5k for vanilla and 7.5k for immiscible diffusion. This effectively suggests that the generalization problem is not significant for current immiscible diffusion methods, which is not surprising as the distance between image and noise shrinks only ~2% after the assignment. We hope this experiment can help to improve our discussion on generalization.

W2 Image Correlation according to Classes

We appreciate your interesting questions, but the answer is no. The images are more correlated according to its own structure rather than the category. In our response to Reviewer y1v3 - W3 - (1), we perform an experiment showing that when enough perturbation is added onto a noise, the image would change to another image in another class, which proves our claim, as shown in Fig. R7. Furthermore, we had an experiment in W1-Practicality on LAION-10k, which nearly has no classifications and immiscible diffusion preliminarily works. This also proves that classification is not a prerequisite for immiscible diffusion.

Thanks for your rebuttal, here is my reply to your rebuttal:

W1-Practicality

Your reply partially addresses my concerns about the effectiveness of this method on complex datasets (rather than simple classification distributions). However, the training data for this experiment is only 10k, and I think the data size should be at least about 10% ImageNet (the dataset used for conditional generation in the paper) to prove its effectiveness.

In addition, could you provide the results of how much the improvement in FID on ImageNet is compared to the improvement on FID on LAION under the same training setting (10k data + 25k training steps)? It would be interesting to explore how much performance differences your method has on training sets with different data distributions.

W1-Generalization

The authors may have misunderstood my concerns. FID is good, but it is not an accurate measure of the quality and generalization of generated images. Lower FID does not mean your generative model is better in performance and generalization

• For performance, you can also report CLIP-Score [1] and CMMD [2] for conditional generation (class name as the text prompt).
• Regarding generalization, I am worried that after training with your method, the model will sample pictures with the same pattern for different random initialization noises, thereby damaging the diversity of generated images. To prove this, you can refer to this representation learning method [3] to see whether the features corresponding to the generated images of the same category are evenly distributed on the feature plane, and calculate their quantized values ​​to determine whether there is enough diversity between the generated images.

[1] Learning Transferable Visual Models From Natural Language Supervision, ICML 2021

[2] Rethinking FID: Towards a Better Evaluation Metric for Image Generation, CVPR 2024

[3] Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere, ICML 2020

W2-Image Correlation according to Classes

I am satisfied with the answers and have no further questions or doubts.

About Rating

I tend to keep my score as is so far and will adjust the final score based on further explanations from the authors.

Thanks to the author for his efforts and reply. After reading the author's rebuttal carefully, I decided to increase my score to Weak Accept, but I hope the author can explain in detail the impact of the change on generalization in the subsequent version.

Specifically, you can generate 100 images for each class, then test the image features corresponding to these 100 images, and measure the distance between these image features to determine whether these images are diverse. If the generated images are diverse, it means that there are differences between these image features, and the distance between them is not very close, and there will be no mode collapse phenomenon similar to that in GAN (the features of all images are very close and collapse to a point in the feature space), and I believe t-SNE or the mentioned paper can better help you verify this problem.

We are sincerely grateful for the reviewer's time and efforts on reviewing our work! We hope to provide more information to clarify our findings on the reviewer's concerns:

W1 - Practicality

1. Larger Complex Image Dataset

We thank the reviewer for raising the concern. We are actively seeking to address this concern via training on large-scale LAION dataset. However, we found that the full LAION dataset has been removed due to the legal policy so we can not download it yet. We can only take the LAION-10K dataset which is legally available on Huggingface "wtcherr/LAION10K". To try our best to do reviewer h6mk a favor, we conduct experiments on the LAION-10K dataset, the results on LAION 10K show that we improve the baseline by 22.14 FID, demonstrating the effectiveness of our method. Despite the relatively small scale, we believe it still effectively shows the sign of life for our method on complex conditional generation.

2. Comparison Experiments on LAION-10k and on 10k Images in ImageNet

We thank the reviewer for proposing the experiment. We perform an additional experiment on 10k ImageNet data + 25k training steps with vanilla and immiscible conditional SD. The results are nearly the same as the experiments on LAION-10k dataset: the FIDs for vanilla SD on ImageNet (10k images) and LAION are 148.90 and 146.69 respectively, and those for immiscible SD are 128.66 and 124.55 respectively. The FID improvements are 20.24 and 22.14 respectively. The improvement is even a little bit larger on LAION-10k than on ImageNet (10k images), supporting that immiscible diffusion does not lose practicality in complex datasets.

W1 - Generalization

We understand the concern on the FID, and hope to present some additional evaluations to address your concerns:

For Performance:

We appreciate the suggestions of providing CLIP-Score and CMMD as other evaluation metrics for evaluating the image-prompt correlation and the image quality, respectively. We evaluate these metrics on 50k images generated from conditional Stable Diffusion models trained for 20k steps, which corresponds to the settings reported in Fig. R2.

The CLIP-Score for baseline and the immiscible model is both 28.55, with standard deviation of 0.02 and 0.01 respectively. We measured it for 3 times as the scores are so close, and our results further validate that there are no differences concerning the CLIP-Score, indicating that the image-prompt correlation is not damaged by immiscible diffusion.

For CMMD, the value for baseline and immiscible model is 1.436 and 1.385 respectively (the lower the better). This also demonstrates that our immiscible diffusion outperforms the vanilla class-conditional SD.

For generalization:

We sincerely thank the reviewer for proposing the additional experiment to prove the diversity of images under the same class-condition. We carefully read the paper provided, and we actively try to design experiments to show the image diversity with representation learning methods. However, the previous representation learning is studied on top of the feedforward models, and we find it unclear how to investigate the visual representation for a multi-step diffusion models. Nevertheless, we have submitted new proofs in Fig. R10, which contains images generated with the same prompt on our immiscible class-conditional SD without cherry picking, to the ACs with anonymous links. It clearly shows that the images are diverse in each class. We hope the release of those additional images can address the reviewer’s concerns.

W2 - Image Correlation according to Classes

We thank the reviewer for acknowledging our response, and are sincerely grateful for the reviewer’s efforts in reviewing our work.

Thank you for your rebuttal. Here is my answer to your global rebuttal:

Conditional Generation

The authors experimented on conditional generation by training a Stable Diffusion model to generate ImageNet images conditionally on their class. The authors show the FID score decreases faster with their method, which is encouraging.

However, I have concerns related to this experiment:

• The authors only show 5 images (potentially cherry-picked) as qualitative examples. This concern was already raised in the initial reviews (Related to reviewer Da24 W4.1 Cherry-picking)
• It is possible that the model ignores the class name. The results shown in Figure R2 do not show that the generated images have the correct class. Furthermore, consider this thought experiment where you would have only two classes, “bright image” and “dark image”. With the proposed data-noise assignment pair, bright images will only be diffused to the brighter noise samples and dark images will only be diffused to the darker noise samples. Then the model will likely lose its ability to transform a “dark initial noise” (randomly sampled at inference time) into a “bright image” when conditioned on the class “dark image”. The results shown in the rebuttal do not address this concern. This concern was already raised in the initial reviews (reviewer h6mk W1-Generalization)
• The main argument mentioned by the authors to support their method is that the first denoising direction predicted by the model (at timestep T) points to the average of the data distribution, which the authors say “does not provide any meaningful information” (L161). By “immiscibility” approach, the first denoising direction now points to the “average of nearby data-points” (L181). While this argument may hold for unconditional generation, it does not make sense anymore as (without immiscibly) it is a good thing to first predict a direction point to the average of images of the class. This average of images of the class is informative (see https://web.mit.edu/torralba/www/ne3302.pdf Figure 1, see also first column of Figure 5 of https://arxiv.org/pdf/2305.08891 ). Predicting this average in the first denoising step allows to guide the generation towards images of the desired class.

Comparison to OT-CFM and related works

The argument given by authors in the rebuttal are not convincing:

• To differentiate from OT-CFM, the authors now propose to redefine immiscibility as “immiscibility of diffusion destinations without necessarily minimizing image-noise distance in a mini-batch”, in other words: coupling (CFM) but no optimal transport (OT-CFM). However, this new definition still makes the proposed method very similar to CFM. The explanation in Theoretical Difference only explains the difference between OT-CFM and normal CFM. This concern was already mentioned in the initial reviews (Da24, W2.1 Coupling).
• Additional Experiments Splitting Two Proposals/Domination of Immiscible Effect to OT in Noise Assignment. The paper [1] (referenced in the rebuttal) already showed that CFM (new definition of “immiscibility”) can have better FID that standard FM.
• Immiscible Diffusion with Image Projection: Multiplying all pixels by a factor of 2 or 4 in the training preprocessing is equivalent to reducing all noise levels by this factor. Some studies (https://arxiv.org/abs/2305.08891, https://www.crosslabs.org/blog/diffusion-with-offset-noise, https://arxiv.org/abs/2309.15842 ) showed however that the noise level used in Stable Diffusion (v1 and v2) are already too low, preventing the model to generate non-greyish images. The paper https://arxiv.org/pdf/2305.08891 showed (table 3) that increasing the noise levels fixes this issue and improve the FID score. I am very surprised that the authors here (Fig R6) find that decreasing all noise levels by a factor of 2 or 4 also improves the FID. Some discussion would be needed here, at least by showing qualitative examples.