Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment

TL;DR

Immiscible Diffusion is a one-line training-time modification for diffusion models: before diffusing each image, assign it to a specific noise sample in the mini-batch by minimizing total image-noise pair distance (Hungarian matching, quantized to keep overhead negligible). The noise marginal stays Gaussian but each image is only ever paired with nearby noise, so the denoiser never has to disambiguate arbitrary image-noise pairings. Reports ~3× training speedup for Consistency Models and DDIM on CIFAR-10 and 1.3× on CelebA at matched FID, with better fidelity for Stable Diffusion on ImageNet. The technique is orthogonal to architectural / hyperparameter recipes and adds ~23ms per batch of 1024 on an A6000.

Key claims

Standard diffusion training implicitly requires the denoiser to map every noise point back to arbitrary images, and this random noise-data pairing is a major training-efficiency bottleneck — at large timesteps the optimal noise prediction collapses toward a constant (dataset mean), which provides no denoising signal [§3.3, Eq. 3-4, Fig. 3a].
Batch-wise Hungarian matching between images and noise samples, minimizing total L2 pair distance, preserves the Gaussian marginal while making each image’s diffuse-able region local and distinguishable — this is the entire method [§3.2, §3.3].
The assignment changes actual image-noise distances by only ~2% per batch, yet this small perturbation is sufficient to activate meaningful denoising signal at high noise levels [§4.3, Table 3].
Quantizing images and noise to int8 during the assignment step reduces Hungarian complexity to negligible cost — 22.8ms per batch of 1024 on A6000 [§3.3, Table 3].
Empirically: ~3× training speedup for Consistency Models on CIFAR-10, 1.3× on CelebA, 1.2× on tiny-ImageNet at matched FID; better fidelity for Stable Diffusion on ImageNet even at similar FID; L2 outperforms L1 for the distance metric [§4.2, Fig. 4-6, Table 2, Table 4].
The FID improvement is largest when inference steps are few (e.g., ΔFID = -2.57 at 20 steps vs -0.15 at 1000 steps for DDIM on CIFAR-10) — the method disproportionately helps few-step regimes, which is where the current field is heading [§4.3, Table 5].

Method

For each training batch of images x₁…x_B and freshly-sampled Gaussian noise ε₁…ε_B, run the Hungarian algorithm to find a permutation π that minimizes Σᵢ ‖xᵢ − ε_{π(i)}‖². Diffuse xᵢ toward ε_{π(i)} instead of the originally-sampled εᵢ. Because the noise samples themselves are unchanged (only their labeling to images shifts), the noise marginal remains exactly Gaussian and sampling at inference is unchanged. To avoid O(B³) Hungarian cost dominating throughput, images and noise are quantized to int8 for the assignment step only; the actual diffusion loss uses full-precision inputs. The trick is one line of code that replaces “sample noise per image” with “sample noise per batch and permute.”

Results

Consistency Model on CIFAR-10: ~3× faster to reach FID target thresholds (Table 2, Fig. 4).
DDIM on CIFAR-10: FID 5.28 vs baseline 6.30 at 70.2k steps, gap widens at fewer inference steps (Table 4-5).
CelebA Consistency Model: 1.3× training efficiency (Fig. 4).
Stable Diffusion on ImageNet: matched FID but visibly cleaner images; class-conditional variant improves FID 22.43 → 20.90 (§4.2, Fig. 7-8).
Assignment overhead: 22.8ms for batch 1024 on A6000, negligible fraction of total training step (Table 3).

Why it’s interesting

Immiscible Diffusion sits under Diffusion training efficiency as a training-time recipe that is orthogonal to the recipes currently dominating that concept page (external DINO supervision, latent design, perceptual losses, self-distillation). None of those touch the noise-data assignment; they all take random pairing for granted. It’s also historical context for the Terrance-flagged connection to Top-k Training of GANs: Improving GAN Performance by Throwing Away Bad Samples — Top-k GAN throws away the worst-scoring mini-batch samples during generator updates as a batch-level sample-selection efficiency trick, several years before diffusion models adopted a related batch-level idea. Whether the two mechanisms share a deeper “batch-level sample matters for gradient quality” thesis is an open question worth revisiting now that both live in one concept context.