Immiscible Diffusion: More aligned noise-image assignment

[Koratahiu]

This pull request implements the Immiscible Diffusion strategy, proposed by the paper:
"Improved Immiscible Diffusion: Accelerate Diffusion Training by Reducing Its Miscibility"

This method optimizes the pairing between training data and noise, effectively "straightening" the diffusion trajectories. By ensuring each image is paired with noise that is mathematically close to it, the model learns a more direct mapping, leading to faster convergence and better sample quality.

The new "Immiscible Noise" setting (Noise Oversampling) is available under the noise configuration settings.

• Set Noise Oversampling setting to 64, 128, 254, etc. (Higher values improve performance but may increase latency at larger batch sizes).

The Key Difference: Random vs. Optimized Noise Pairing

Standard Diffusion (Random Pairing):

• effective_noise = random_gaussian_noise()
• Each data point in a batch is assigned a random piece of noise. Because the noise is chosen arbitrarily, the "path" from the clean image to the noisy state is often tangled and complex.

Immiscible Diffusion (Optimal/Nearest Pairing):

• effective_noise = select_best_noise(source_tensor, candidates)
• Instead of random assignment, this method ensures that the noise is "assigned" to an image such that the distance between them is minimized.

Noise Oversampling Strategy

• The model generates k candidate noise tensors for every image in the batch.
• It calculates the Euclidean distance between the image and all k candidates.
• It selects the nearest candidate (the one with the minimum distance) and discards the rest.
• Recommended value: 64.

This implementation significantly reduces the complexity of the ODE/SDE trajectories that the model must learn. By making the diffusion paths "immiscible" (non-crossing), the model can achieve higher quality results in fewer training steps and produce sharper images during inference.

Implementation details

• Minimal changes in modules\modelSetup\mixin\ModelSetupNoiseMixin.py
• New file: modules\util\immiscible_diffusion.py with the paper logic.
• Universal support for all models, and it also works for flow-matching models.

Sources

v1 paper: Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment
v2 paper: Improved Immiscible Diffusion: Accelerate Diffusion Training by Reducing Its Miscibility
Official repo: https://github.com/yhli123/Immiscible-Diffusion