Improved Immiscible Diffusion: Accelerate Diffusion Training by Reducing Its Miscibility

TL;DR

This paper introduces a generalized immiscible diffusion technique that reduces trajectory mixing in diffusion models, significantly accelerating training while preserving generative diversity across various tasks.

Contribution

It extends immiscible diffusion beyond linear assignment, proposing new implementations like KNN noise selection and image scaling to improve training efficiency.

Findings

Achieves over 4x faster training across multiple models and tasks.

Maintains generative diversity through bijective denoising process.

Provides analysis linking immiscibility with optimal transport to enhance understanding.

Abstract

The substantial training cost of diffusion models hinders their deployment. Immiscible Diffusion recently showed that reducing diffusion trajectory mixing in the noise space via linear assignment accelerates training by simplifying denoising. To extend immiscible diffusion beyond the inefficient linear assignment under high batch sizes and high dimensions, we refine this concept to a broader miscibility reduction at any layer and by any implementation. Specifically, we empirically demonstrate the bijective nature of the denoising process with respect to immiscible diffusion, ensuring its preservation of generative diversity. Moreover, we provide thorough analysis and show step-by-step how immiscibility eases denoising and improves efficiency. Extending beyond linear assignment, we propose a family of implementations including K-nearest neighbor (KNN) noise selection and image scaling to reduce miscibility, achieving up to >4x faster training across diverse models and tasks including unconditional/conditional generation, image editing, and robotics planning. Furthermore, our analysis of immiscibility offers a novel perspective on how optimal transport (OT) enhances diffusion training. By identifying trajectory miscibility as a fundamental bottleneck, we believe this work establishes a potentially new direction for future research into high-efficiency diffusion training. The code is available at https://github.com/yhli123/Immiscible-Diffusion.