Technical Summary: Improved Immiscible Diffusion

1. Problem Statement

The training of diffusion models is hindered by substantial computational costs and time requirements, which scale significantly with model size and dataset complexity. A fundamental bottleneck identified in this work is trajectory miscibility. In vanilla diffusion training, random Gaussian noise is added to images, causing the diffusion paths (flows) of different images to mix at the same points in the noise space, particularly as the noise level increases ($t \to T$).

When multiple distinct image trajectories converge to the same noisy point $x_t$, the denoising network $\epsilon_\theta(x_t, t)$ faces an ambiguous learning objective. Instead of predicting a single specific noise vector to recover a unique image, the model is forced to predict an average of conflicting noise vectors. As the noise level approaches the maximum, the denoising target converges to a constant (the average of the dataset), rendering the learning process ineffective for recovering specific images. This "mixing" creates irreversible and confusing trajectories that degrade training efficiency.

Previous attempts to address this, such as applying batch-wise Optimal Transport (OT) or linear assignment to pair images with closer noise points, have shown limited success. These methods reduce the image-noise distance by only $\approx 2\%$ and the standard deviation of the denoising function by $\approx 4\%$. Furthermore, the linear assignment approach has a computational complexity of $O(n^3)$, making it infeasible for large batch sizes, and raises concerns about whether restricting image-noise pairings compromises the diversity of generated images.

2. Methodology

The authors propose Improved Immiscible Diffusion, a refined framework that redefines the core concept of immiscibility not merely as a specific image-noise pairing strategy, but as a general reduction in the mixture of diffusion trajectories from different images.

2.1 Theoretical Foundation: Stable Image-Noise Correlation

The work first establishes that generated images are stably correlated with their noise origins. Through perturbation experiments on vanilla DDIM, the authors demonstrate that adding small amounts of noise (up to 20% perturbation) to a specific noise point results in generated images that remain semantically identical. Only with significant perturbation (30%) do object-level changes occur. This finding validates that restricting diffusion to specific noise regions (immiscibility) does not inherently destroy generation diversity, as the mapping from noise to image is robust.

2.2 Mechanism of Action

The core mechanism involves clarifying the denoising goal for the model. By reducing the overlap of diffusion trajectories:

1. Noise Clusters Separate: Noise points assigned to different images form distinct, non-overlapping clusters in the noise space.
2. Denoising Goals Clarify: The denoising function no longer needs to average conflicting targets. Even at high noise levels ($\tau = S$), the model can effectively predict the specific noise added to a specific image.
3. Feature Space Improvement: t-SNE analysis shows that under immiscible diffusion, predicted images across denoising steps form distinct, non-intersecting paths, whereas vanilla diffusion paths are chaotically tangled.

2.3 Implementations

The paper introduces a family of implementations to achieve immiscibility, moving beyond the computationally expensive linear assignment:

• Batch-wise Linear Assignment (Baseline): Assigns each image to the nearest noise point via $L_2$ minimization. While effective at separating trajectories, it suffers from $O(n^3)$ complexity.
• K-Nearest Neighbor (KNN) Noise Selection: For each image, $k$ random noise points are sampled, and the one with the minimum $L_2$ distance is selected. This method is computationally efficient ($O(n)$) and scales well to large batch sizes. Analysis shows that the resulting noise distribution remains nearly Gaussian (KL divergence is negligible compared to strict Gaussian).
• Image Scaling: A novel approach that does not rely on image-noise pairing. By multiplying image pixel values by a constant factor $>1$ (e.g., 2 or 4) before diffusion, the $L_2$ distance between the centers of different images' diffusion areas increases. This naturally reduces the intersection of diffusion trajectories without altering the noise space or requiring explicit pairing.

3. Key Contributions

1. Conceptual Refinement: The authors extend immiscible diffusion from a specific linear assignment technique to a broader concept defined by the reduction of diffusion trajectory mixture. They demonstrate that this reduction does not harm generation diversity due to the intrinsic stability of image-noise correlations.
2. Efficient Implementations: The paper proposes KNN Noise Selection and Image Scaling as highly efficient alternatives to linear assignment. These methods achieve immiscibility with significantly lower computational overhead, enabling their use in large-scale training.
3. Comprehensive Validation: The work provides extensive empirical evidence across diverse models (Consistency Models, Flow Matching, DDIM, Stable Diffusion), datasets (CIFAR-10, ImageNet, MSCOCO), and tasks (unconditional/conditional generation, fine-tuning, image editing, and robotics planning).
4. Mechanistic Analysis: Through step-by-step feature analysis, the paper elucidates how immiscibility accelerates training: by activating the denoising capabilities of the noisiest layers, which are typically the most difficult to train in vanilla diffusion.

4. Experimental Results

The experiments demonstrate consistent and significant improvements in training efficiency and final performance:

• Training Speed: The immiscible diffusion family achieves up to 4$\times$ faster training to reach the same performance level (FID) compared to vanilla baselines.
• Performance Gains:
  • Unconditional Generation: On CIFAR-10, KNN-implemented immiscible DDIM improved FID from 6.63 (Vanilla) to 4.45.
  • Conditional Generation: In Stable Diffusion training on ImageNet and MSCOCO, immiscible diffusion consistently outperformed vanilla baselines in both training-from-scratch and fine-tuning scenarios.
  • Image Editing: In-painting and out-painting tasks showed lower FID scores (e.g., In-painting FID reduced from 18.35 to 17.32), indicating better preservation of source information.
  • Robotics Planning: Applied to the PushT task, immiscible diffusion improved the average area coverage of the robot's success from 79.56% to 82.83%.
• Efficiency: The KNN method reduced execution time for noise selection from milliseconds (Linear Assignment) to sub-millisecond levels (e.g., 0.2ms for batch size 256), making it scalable.
• Diversity: CLIPScore evaluations confirmed that immiscible diffusion maintains prompt-image correspondence and generation diversity comparable to vanilla models.

5. Significance and Claims

The paper posits that trajectory miscibility is a fundamental bottleneck in diffusion training that has been overlooked in favor of other efficiency techniques like flow matching or rectified flow. While flow matching linearizes trajectories, it does not prevent them from mixing (being miscible).

The authors claim that by addressing the root cause of denoising ambiguity—trajectory mixing—they establish a new direction for high-efficiency diffusion training. The work suggests that the "miscibility problem" is a general issue affecting various diffusion architectures and that reducing it offers a robust, model-agnostic path to faster convergence and better performance. The paper concludes that this approach provides a potentially new research direction for future studies in efficient diffusion training, moving beyond simple architectural changes to fundamentally altering the diffusion process dynamics.