Diffuse and Disperse: Image Generation with Representation Regularization

TL;DR

Dispersive Loss is a plug-and-play regularizer for diffusion transformers that encourages the model’s intermediate representations to spread out in hidden space — a “contrastive loss without positive pairs.” Unlike REPA, it needs no pretrained encoder, no extra parameters, no external data, no second view, and no augmentation; the existing denoising regression serves as the implicit alignment term. Across DiT and SiT at S/B/L/XL sizes on ImageNet-256, it consistently improves FID, and applied to MeanFlow it sets a new one-step SOTA (FID 3.21 at XL/2). The work is from Kaiming He’s group and directly engages the REPA-vs-internal-supervision fork that the wiki already tracks.

Key claims

Dispersive Loss adds a single regularizer term L_disp to the standard diffusion loss with weight λ; when λ = 0 it reduces exactly to the baseline — no other implementation changes [§3.1, Eq. 1].
The loss is “contrastive without positive pairs”: derived by dropping the positive (alignment) term from an InfoNCE/Hinge/Covariance contrastive loss and keeping only the repulsion term, so it operates on a single-view batch with no two-view sampling or augmentation [§3.1, §3.2, Table 1].
Three variants are derived (InfoNCE, Hinge, Covariance) and all improve over the baseline; InfoNCE with ℓ² distance (no normalization, no projection head) is best, improving SiT-B/2 FID 36.49 → 32.35 at 80 epochs (–11.35%) [§3.3, Table 2].
Contrastive loss (with positive pairs) using independent noise per view degrades FID by 14-20% across all variants; only when the two views are restricted to near-identical noise levels does contrastive training help — and even then it underperforms the positive-free Dispersive variant [Table 2].
Applying Dispersive Loss to any single Transformer block helps; applying it at all blocks is best (FID 32.05 vs 36.49 baseline), and the representation-norm effect at one regularized block propagates to all other blocks even when they aren’t directly regularized [§4.2, Table 3, Fig. 3].
The two hyperparameters (loss weight λ, temperature τ) are robust: every (λ, τ) combination in a 4×3 sweep improves over baseline; results are surprisingly insensitive to τ in contrast to standard contrastive learning [§4.2, Table 4].
Improvements are consistent across DiT and SiT at S, B, L, XL sizes; relative gain grows with the strength of the baseline (SiT > DiT, L > B > S), interpreted as larger models benefiting more from regularization [§4.2, Fig. 4].
On SiT-XL/2 at 800 epochs with CFG + SDE sampler, Dispersive Loss improves FID 2.27 → 2.09; at 1200 epochs FID 2.06 → 1.97 — gains shrink with longer training but persist [Table 5].
Head-to-head with REPA on SiT-XL/2: REPA achieves FID 1.80 with 1500-ImageNet-epoch DINOv2-distilled-from-1.1B-params teacher trained on 142M curated images; Dispersive Loss achieves FID 1.97 with zero pretraining, zero extra params, zero external data [Table 6].
Applied to one-step MeanFlow: MF-XL/2 + Dispersive improves FID 3.43 → 3.21, a new one-step SOTA versus iCT (34.24), Shortcut (10.60), IMM (7.77) [§4.3, Table 7].

Method

Two losses combined: the standard diffusion regression loss L_diff (DDPM denoising or flow-matching velocity) on each sample, plus a per-batch dispersive term L_disp on intermediate representations of the same batch — no additional forward pass, no augmented view, no projection head. For the recommended InfoNCE-ℓ² variant, L_disp = log(E_{i,j}[exp(-‖z_i - z_j‖² / τ)]) where z are unnormalized intermediate features and the expectation is over all pairs in the batch. This is mathematically a “uniformity loss” in the Wang & Isola (2020) sense, but applied to intermediate (not output) representations and unpaired with any alignment term. Algorithm 1 in the paper shows the implementation is a few lines of PyTorch. By default the loss is applied to the block at the first quarter of total depth (e.g. block 6 of 24 for XL/2) with λ = 0.5 and τ = 0.5. Applied to MeanFlow the recipe is unchanged.

Results

SiT-B/2 ImageNet-256, 80 epochs, no CFG: FID 36.49 (baseline) → 32.35 (Dispersive InfoNCE-ℓ²), –11.35% [Table 2].
SiT-XL/2 800 epochs SDE w/ CFG: FID 2.27 → 2.09 [Table 5].
SiT-XL/2 1200 epochs SDE w/ CFG: FID 2.06 → 1.97 [Table 5].
SiT-XL/2 vs REPA: 1.97 vs 1.80 — REPA wins by 0.17 FID but uses a teacher distilled from a 1.1B backbone on 142M curated images for ~1500 epochs equivalent; Dispersive uses nothing external [Table 6].
MeanFlow-XL/2 one-step: FID 3.43 → 3.21 (–6.4%); new one-step SOTA on ImageNet-256 [Table 7].
Contrastive baseline (independent noise): InfoNCE-ℓ² FID 43.66, +19.65% over baseline — a strict regression [Table 2]. This is the paper’s evidence that the positive-free design is the key, not contrastive learning per se.

Why it’s interesting

This is the canonical Kaiming-He-group entry in the Diffusion training efficiency internal-vs-external-supervision fork the wiki is already tracking, and it lands almost exactly between the two camps the page currently splits:

It’s a sibling to No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves (SRA) (SRA, May 2025) — both argue REPA’s external-DINO substrate is dispensable and both are “self-contained, no pretraining” recipes. But SRA still constructs an internal teacher view (later-layer-cleaner) and student view (earlier-layer-noisier); Dispersive Loss is even simpler — there is no teacher at all, just a repulsion term on a single batch’s representations. The two papers reach the same conclusion (REPA’s external supervision isn’t necessary) via opposite mechanisms: SRA replaces external alignment with internal alignment, Dispersive Loss replaces external alignment with no alignment, only repulsion.
It directly attacks the REPA cost/benefit accounting that REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training (HASTE) was patching from the other direction. HASTE keeps REPA but turns it off when gradients flip; Dispersive Loss never turns it on. The paper’s Table 6 — making the pretraining/data/param overhead of REPA fully explicit (1.1B teacher, 142M images, 1500 ImageNet-epoch-equivalents) — is exactly the bookkeeping that lets Luma compare “0.17 FID vs all that cost” honestly.
The MeanFlow result is the practical hook: a +0.22 FID improvement on a one-step model with literally a few lines of code is the sort of regularizer that ships. Combined with the existing Diffusion Distillation thread on the wiki (SoFlow, Causal Forcing, AnyFlow), this is a candidate “free” auxiliary loss to layer on top of distilled / few-step setups where REPA’s two-view requirement is awkward.

The open question this paper sharpens: if the regression loss is itself a sufficient “alignment” signal (the paper’s central claim), then what exactly does internal self-distillation (SRA, Self-Flow) buy beyond pure repulsion? A clean ablation — Dispersive Loss vs SRA vs Self-Flow at compute parity, same backbone, same epoch count — would resolve the wiki’s current “structural difference” open question for diffusion-training-efficiency.