Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment (2406.12303v2)

Abstract: In this paper, we point out that suboptimal noise-data mapping leads to slow training of diffusion models. During diffusion training, current methods diffuse each image across the entire noise space, resulting in a mixture of all images at every point in the noise layer. We emphasize that this random mixture of noise-data mapping complicates the optimization of the denoising function in diffusion models. Drawing inspiration from the immiscibility phenomenon in physics, we propose Immiscible Diffusion, a simple and effective method to improve the random mixture of noise-data mapping. In physics, miscibility can vary according to various intermolecular forces. Thus, immiscibility means that the mixing of molecular sources is distinguishable. Inspired by this concept, we propose an assignment-then-diffusion training strategy to achieve Immiscible Diffusion. As one example, prior to diffusing the image data into noise, we assign diffusion target noise for the image data by minimizing the total image-noise pair distance in a mini-batch. The assignment functions analogously to external forces to expel the diffuse-able areas of images, thus mitigating the inherent difficulties in diffusion training. Our approach is remarkably simple, requiring only one line of code to restrict the diffuse-able area for each image while preserving the Gaussian distribution of noise. In this way, each image is preferably projected to nearby noise. Experiments demonstrate that our method can achieve up to 3x faster training for unconditional Consistency Models on the CIFAR dataset, as well as for DDIM and Stable Diffusion on CelebA and ImageNet dataset, and in class-conditional training and fine-tuning. In addition, we conducted a thorough analysis that sheds light on how it improves diffusion training speed while improving fidelity. The code is available at https://yhli123.github.io/immiscible-diffusion

Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment

The paper "Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment" by Yiheng Li et al. introduces an innovative method to accelerate the training of diffusion models, specifically targeting the inefficiencies associated with the noise-data mapping in current methodologies. Authors from UC Berkeley and Tsinghua University propose a concept called Immiscible Diffusion, which aims to mitigate these inefficiencies by reassigning noise to images in a manner inspired by the immiscibility phenomenon in physics.

Summary of Contributions

The primary contributions of this paper are encapsulated in the proposal and experimental validation of Immiscible Diffusion:

1. Identification of Suboptimal Noise-Data Mapping: The authors recognize that current diffusion models diffuse each image across the entire noise space, leading to a randomized mixture that complicates the optimization process.
2. Immiscible Diffusion Method: Drawing inspiration from immiscible fluid interactions, the authors propose an assignment-then-diffusion strategy. Prior to diffusing image data into noise, they assign diffusion target noise by minimizing the total image-noise distance within a mini-batch. This results in noise being assigned to nearby images, which simplifies the denoising function.
3. Quantized-assignment Strategy: To tackle the computational complexity of the assignment algorithm, the authors employ a quantized-assignment strategy. This significantly reduces overhead, making the approach computationally feasible even for large batch sizes and high-resolution images.

Experimental Validation

The authors validate their method on three diffusion model baselines—Consistency Models, DDIM, and Stable Diffusion—across multiple datasets including CIFAR-10, CelebA, and ImageNet. The results demonstrate substantial improvements in training efficiency and image quality:

• Consistency Model Improvements: Immiscible Diffusion achieved up to 3x faster training on the CIFAR-10 dataset and 1.3x on CelebA, with consistent reductions in Fréchet Inception Distance (FID) scores.
• DDIM Performance: For DDIM on CIFAR-10, the method not only improved training speed but also lowered the FID scores significantly, particularly when fewer inference steps were used.
• Stable Diffusion: While FID improvements for Stable Diffusion on ImageNet were not as pronounced, qualitative assessments indicated that images generated with Immiscible Diffusion were subjectively clearer and more detailed.

Analysis and Discussion

The theoretical foundation laid out in the paper includes a detailed analysis of the noise prediction task. The authors illustrate how Immiscible Diffusion makes the denoising task easier by ensuring that each noise point is mapped to nearby images in the noise space, as opposed to the traditional random assignment which leads to a miscible mixture.

Through mathematical illustration and thorough experimentation, the authors demonstrate that even a slight reduction in the image-noise pair distance (approximately 2%) can lead to significant gains in training efficiency. This finding is of particular relevance given the high dimensionality of the image and noise spaces involved.

Implications and Future Directions

The implications of this research are both practical and theoretical:

• Practical Implications: The significant reduction in training time and the enhancement in image quality make Immiscible Diffusion a highly attractive technique for accelerating the development of diffusion models. This has potential applications in any domain relying on quick iteration cycles, such as generative AI for image and video production.
• Theoretical Implications: The introduction of a noise assignment strategy opens new avenues for understanding and possibly improving other aspects of diffusion models. This includes extending these principles to different types of generative tasks and exploring alternative distance functions for noise assignment.

Conclusion and Future Work

This research presents Immiscible Diffusion as a simple yet effective method to accelerate diffusion model training, requiring minimal modifications to existing training pipelines. Future work could explore optimizing the assignment strategy further and scaling experiments to larger datasets like LAION. Additionally, the application of Immiscible Diffusion to text-to-image and text-to-video tasks represents an intriguing area for future exploration.

The findings here underscore the potential for even minor algorithmic innovations to yield significant gains in computational efficiency, a critical consideration as diffusion models continue to evolve and integrate into various AI systems.