Source-linked AI summary
Diffuse and Disperse: Image Generation with Representation Regularization
Runqian Wang, Kaiming He
TL;DR
Diffusion models commonly rely on regression objectives without explicit representation regularization. The paper introduces Dispersive Loss, which disperses internal representations without positive pairs or external resources, and reports consistent improvements across diffusion-model baselines and scales. The approach also extends to MeanFlow one-step generation, where the authors report state-of-the-art performance.
Problem
Diffusion generative models typically use regression-based objectives and lack explicit regularization of the representations learned for generation.
Method
Dispersive Loss regularizes internal representations by encouraging dispersion in hidden space alongside the standard regression loss.
Results
Dispersive Loss consistently improves DiT and SiT baselines across model scales and achieves state-of-the-art performance for one-step generation with MeanFlow.
Takeaways & Limitations
Representation regularization can enhance generative modeling without pre-training, extra parameters, or external data.
Takeaways & Limitations
The paper treats diffusion and flow-matching methods under one umbrella term, without distinguishing between them.
Abstract
from arXiv · showhide
The development of diffusion-based generative models over the past decade has largely proceeded independently of progress in representation learning. These diffusion models typically rely on regression-based objectives and generally lack explicit regularization. In this work, we propose \textit{Dispersive Loss}, a simple plug-and-play regularizer that effectively improves diffusion-based generative models. Our loss function encourages internal representations to disperse in the hidden space, analogous to contrastive self-supervised learning, with the key distinction that it requires no positive sample pairs and therefore does not interfere with the sampling process used for regression. Compared to the recent method of representation alignment (REPA), our approach is self-contained and minimalist, requiring no pre-training, no additional parameters, and no external data. We evaluate Dispersive Loss on the ImageNet dataset across a range of models and report consistent improvements over widely used and strong baselines. We hope our work will help bridge the gap between generative modeling and representation learning.
1 Introduction
Diffusion models have advanced largely independently of representation learning, relying mainly on regression objectives without explicit representation regularization. Dispersive Loss addresses this gap with a self-contained regularizer that disperses hidden representations and consistently improves diffusion-model baselines.
- Diffusion training typically emphasizes reconstruction or denoising through regression, without explicitly regularizing the representations learned for generation.
- Dispersive Loss adds an objective that encourages internal representations to spread through hidden space alongside the standard regression loss.The regression objective also serves as an alignment mechanism, so positive pairs need not be manually defined.
- Unlike contrastive learning, Dispersive Loss requires no positive pairs, two-view sampling, specialized augmentation, or additional encoder.
- Compared with REPA, the method requires no pre-training, additional model parameters, or external data.
- Dispersive Loss consistently improves DiT and SiT baselines across model scales and also benefits MeanFlow one-step generation.The authors report state-of-the-art performance for one-step diffusion-based generation with MeanFlow.
2 Related Work
Related work spans diffusion and flow-matching generative models, self-supervised representation learning, and auxiliary representation objectives. The paper positions Dispersive Loss as a positive-free, input-preserving alternative within this landscape.
- Diffusion Models: Diffusion and flow-matching models use progressive noising/denoising or velocity-based formulations, while regression-based training remains central to both families.
- Diffusion Models: This paper uses “diffusion models” as an umbrella term covering both diffusion methods and flow-matching methods.
- Self-Supervised Learning: Contrastive self-supervised learning encourages similarity between positive pairs and dissimilarity between negative pairs, whereas the proposed method is framed as positive-free contrastive learning.
- Self-Supervised Learning: Masked modeling applies more destructive input corruption than typical augmentations, while this method aims not to interfere with diffusion-training inputs.
- Representation Learning as Auxiliary Tasks: Auxiliary representation learning adds a representation objective alongside a main task, including supervised contrastive learning and self-supervised language-image pre-training.
- Representation Learning as Auxiliary Tasks: REPA aligns generative-model representations with a frozen, pre-trained encoder, and related methods extend this alignment approach structurally or adversarially.
3 Methodology
Dispersive Loss regularizes diffusion models by spreading intermediate representations while retaining the standard regression objective. It forms a general family of single-view, plug-and-play objectives derived by removing positive-pair terms from contrastive losses.
- 3.1 Dispersive Loss: Dispersive Loss adds a representation regularizer to the standard diffusion objective, encouraging intermediate representations to disperse in hidden space.The regularizer operates alongside the original diffusion loss rather than replacing it.
- 3.1 Dispersive Loss: The batch objective combines per-sample diffusion loss with batch-dependent Dispersive Loss weighted by λ, applied directly to intermediate representations without extra parameters.No projection head or additional layer is used.
- 3.1 Dispersive Loss: Dispersive Loss removes positive-pair terms from contrastive objectives, leaving repulsion or dispersion terms that can operate on independent single-view image batches.This avoids additional sampling views and data augmentation while preserving the regression targets supplied by diffusion training.
- 3.2 InfoNCE-Based Variant of Dispersive Loss: The InfoNCE-based variant computes dispersion over a batch of intermediate representations using a dissimilarity function such as cosine dissimilarity or squared ℓ2 distance.The batch-form loss has the same value for every sample and is computed once per batch.
- 3.3 Other Variants of Dispersive Loss: Hinge and covariance variants retain only negative-pair or off-diagonal covariance terms, extending Dispersive Loss beyond InfoNCE.All variants are simpler than their contrastive counterparts and function as plug-and-play regularizers without modifying the regression loss.
4 Experiments
Experiments on ImageNet show that Dispersive Loss consistently improves diffusion and flow-based baselines across loss variants, regularization locations, hyperparameters, model scales, and one-step generation settings.
- Loss Variants: All four Dispersive Loss variants outperform the SiT-B/2 baseline, demonstrating robustness across objective formulations.The variants are InfoNCE with ℓ2 or cosine dissimilarity, Hinge, and Covariance.
- Loss Variants: 4.14 FID points, or 11.35% relative improvement, is achieved by InfoNCE with ℓ2 distance over the baseline.The unnormalized ℓ2 formulation is used by default in subsequent experiments.
- Block Choice for Regularization: Applying Dispersive Loss to all blocks performs best, while applying it to any single block performs nearly as well.A single-block application also increases representation norms in later blocks where the loss is not directly applied.
- Hyperparameters: All studied combinations of regularization strength λ and temperature τ improve on the baseline FID of 36.49.The results are described as robust across a wide range of temperature values.
- Different Models: Dispersive Loss consistently improves DiT and SiT performance across four model sizes, with XL evaluated only for SiT.All models are trained on ImageNet for 80 epochs in this evaluation.
- Extended Settings: Dispersive Loss remains beneficial for SiT-XL/2 across longer training, classifier-free guidance, and both ODE- and SDE-based samplers, and improves one-step MeanFlow.The MeanFlow experiments report a new state of the art among compared one-step diffusion and flow-based models.
5 Conclusion
The paper proposes Dispersive Loss to regularize diffusion-model representations without added pre-training, parameters, or external data. It preserves the original diffusion training strategy while enhancing generative modeling.
- Dispersive Loss regularizes the internal representations of diffusion models without additional pre-training, extra parameters, or external data.
- The design minimizes interference with the original sampling process, allowing the original diffusion training strategy to be preserved.
- The authors suggest that the plug-and-play regularizer may also benefit applications beyond generative modeling, including image recognition.
A Implementation
The ImageNet SiT and DiT experiments use established codebases and fixed optimization and sampling configurations. MeanFlow experiments follow the authors’ codebase, with a compute-limited XL/2 training procedure.
- SiT and DiT Experiments: SiT and DiT experiments use AdamW with a constant learning rate of 1×10−4, no weight decay, and 32 TPU-v3/v4 cores.
- SiT and DiT Experiments: ODE sampling uses Heun with 250 steps, while SDE sampling uses Euler–Maruyama with SiT’s default drift and diffusion schedule.
- MeanFlow Experiments: MeanFlow evaluates B/4, B/2, and XL/2 models with and without Dispersive Loss, using the shared author codebase.
- MeanFlow Experiments: The MeanFlow-XL/2 model resumes from a checkpoint at 180 epochs and applies Dispersive Loss for the remaining 60 epochs.
- MeanFlow Experiments: MeanFlow regularization weights are λ=0.25 for B/4, 1.0 for B/2, and 1.5 for XL/2.
B.1 Inception Scores
The ImageNet evaluation includes Inception Scores for SiT models. Dispersive Loss improves these scores over baseline SiT models, where higher scores are better.
- Dispersive Loss improves SiT models’ Inception Scores over baseline SiT models on ImageNet.
- Higher Inception Scores indicate better performance in the reported SiT comparison.
B.2 CIFAR-10 Experiments
The CIFAR-10 experiments test Dispersive Loss with a Unet architecture using Flow Matching. Across training, the regularizer consistently improves FID over the reproduced baseline.
- The CIFAR-10 experiments use a Unet architecture and publicly available Flow Matching code with the original hyper-parameters.
- 2.13 FID is the reproduced Flow Matching baseline, compared with 2.07 FID reported by the original repository.
- 2.07 FID is achieved with Dispersive Loss versus 2.13 FID for the reproduced baseline.
- Dispersive Loss produces consistent gains throughout training when applied at Residual Block 15 of the Unet.