Diffuse and Disperse: Image Generation with Representation Regularization

## Abstract

Abstract 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 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. Our code is available at https://github.com/raywang4/DispLoss .

## 1 Introduction

Diffusion generative models [35, 38, 18] have demonstrated remarkable performance in modeling complex data distributions, but their success remains largely disconnected from advances in representation learning.
The training objectives of diffusion models typically consist of a regression term focused on reconstruction (*e.g.*, denoising), yet lack an explicit regularization term on the representations learned for generation. This paradigm for image generation stands in contrast to its counterpart in image recognition, where representation learning has been a core topic and a driving force over the past decade [3].

In the field of representation learning, self-supervised learning has made significant progress in learning general-purpose representations applicable to a wide range of downstream tasks (*e.g.*, [15, 5, 16, 30]). Among these approaches, contrastive learning [7, 14, 29, 15, 5] offers a conceptually simple yet effective framework for learning representations from sample pairs. Intuitively, these methods encourage attraction between sample pairs considered similar (“positive pairs”) and repulsion between those considered dissimilar (“negative pairs”). Representation learning via contrastive learning has been shown useful across a variety of recognition tasks, including classification, detection, and segmentation. However, the effectiveness of these learning paradigms for generative modeling remains an underexplored problem.

Figure: Figure 1: Dispersive Loss for Generative Modeling. Left: A standard diffusion-based model (*e.g.*, SiT [26] or DiT [31]) with a regression-driven diffusion loss, augmented by a Dispersive Loss applied to an intermediate block. Right: A zoom-in on the first few blocks with Dispersive Loss. Our loss encourages intermediate representations to disperse in the hidden space. It operates on the same input batch of noised images, shares the existing network blocks and computations, introducing negligible overhead, no additional parameters, and no reliance on external data.
Refer to caption: x1.png

In light of the potential of representation learning for generative modeling, Representation Alignment (REPA) [41] has been proposed to leverage the capabilities of pre-trained, off-the-shelf representation models. This method trains a generative model while encouraging alignment between its internal representations and external pre-trained representations. As a pioneering effort, REPA has revealed the importance of representation learning in the context of generative modeling; however, its instantiation depends on additional pre-training, extra model parameters, and access to external data.
Developing a self-contained and minimalist approach to representation-based generative modeling is still an essential topic in this line of research.

In this paper, we propose Dispersive Loss, a flexible and general plug-and-play regularizer that integrates self-supervised learning into diffusion-based generative models.
Our core idea is simple: alongside the standard regression loss on the model output, we introduce an objective that regularizes the model’s internal representations ([Fig. 1](https://arxiv.org/html/2506.09027v2#S1.F1)).
Intuitively, Dispersive Loss encourages internal representations to spread out in the hidden space, analogous to the repulsive effect in contrastive learning. Meanwhile, the original regression loss (*e.g.*, denoising) naturally serves as an alignment mechanism, eliminating the need to manually define positive pairs as in contrastive learning.

In a nutshell, Dispersive Loss behaves like a “contrastive loss without positive pairs”—and as such, unlike contrastive learning, it requires neither two-view sampling, specialized data augmentation, nor an additional encoder. The training pipeline can follow standard practice used in diffusion-based models (and their flow-based counterparts)(^1^11In the context of this work, we do not distinguish between diffusion methods and flow matching methods [23, 25, 1], and refer to both under the umbrella term of “diffusion models”.), with the only distinction being an additional regularization loss that has negligible overhead. In comparison to the REPA mechanism, our method requires no pre-training, no additional model parameters, and no external data. With a self-contained and minimalist design, our method clearly demonstrates that representation learning can benefit generative modeling without relying on external sources of information.

We evaluate the effectiveness and generality of Dispersive Loss through extensive experiments. Our results show that Dispersive Loss consistently improves upon strong and widely used baselines in diffusion models, namely, DiT [31] and SiT [26] ([Fig. 2](https://arxiv.org/html/2506.09027v2#S1.F2)), across a wide range of model scales. We also find that various formulations of Dispersive Loss are consistently beneficial, demonstrating the robustness and generality of our approach; by comparison, incorporating a contrastive loss can degrade performance.
Beyond multi-step diffusion models, we further apply Dispersive Loss to the recent MeanFlow framework [12], achieving state-of-the-art performance in one-step diffusion-based generation. We hope that the simplicity, effectiveness, and generality of Dispersive Loss will help bridge the gap between generative modeling and representation learning.

Figure: Figure 2: Effectiveness of Dispersive Loss. We show FID-50k on ImageNet 256$\times$256 at different training epochs, comparing the standard SiT-XL/2 baseline [26] with its counterpart trained with Dispersive Loss. Classifier-free guidance is not used in this plot. More results and details are in [Section 4](https://arxiv.org/html/2506.09027v2#S4).
Refer to caption: x2.png

## 2 Related Work

#### Diffusion Models.

Diffusion probabilistic models [35] formulate generative modeling as a progressive noising/denoising process, trained via a variational lower bound. This process is substantially simplified and improved by reformulating it as a denoising-based reconstruction objective [18]. Subsequent work [36, 28, 10, 19] further advances this direction by improving noise schedules and sampling strategies.
Flow matching models [23, 25, 1] extend the diffusion paradigm by focusing on velocity-based formulations and modeling deterministic trajectories. Despite these advancements, the regression-based training objective remains central to diffusion-based models and their variants.

#### Self-Supervised Learning.

In parallel with advances in generative modeling, self-supervised representation learning has made steady progress in visual recognition.
Contrastive learning [7] has become a dominant self-supervised paradigm which encourages similarity between positive pairs and dissimilarity between negative pairs. Modern contrastive learning methods [15, 5], driven by advances such as information-based contrastive losses [29] and large-scale negative sampling, have demonstrated encouraging results that can rival those of supervised representation learning. Interestingly, there is a specific line of research on “negative-free" contrastive learning [13, 6], which lead to representation learning results that are competitively strong. Within this context, our method can be interpreted as “positive-free" contrastive learning, which is an underexplored direction in standard self-supervised learning.

Beyond contrastive learning, masked modeling [9] has been established as another effective self-supervised learning paradigm in vision [16, 2, 43]. Compared to contrastive learning, masked modeling methods apply more destructive corruption to the input data than typical data augmentations. In contrast, our method aims to introduce no interference to the input samples for diffusion training.

#### Representation Learning as Auxiliary Tasks.

Beyond the common pre-training/fine-tuning paradigm, representation learning can be performed as an auxiliary task jointly with the main objective. Supervised contrastive learning [20] extends classical supervised models (*e.g.*, classification) with extra contrastive objectives.
Self-supervised language-image pre-training (SLIP) [27] augments its contrastive counterpart (CLIP [32]) with a self-supervised objective trained in parallel.

In the context of image generation, REPA [41] explores a line of research on enhancing generative modeling with auxiliary representation learning. REPA aims to align the intermediate representations of a generative model to those from a frozen, high-capacity, pre-trained encoder, which may itself be trained using external data and diverse objectives.
SARA [4] further advances this approach by introducing structural and adversarial representation alignment. In the multimodal setting, SoftREPA [22] extends REPA by aligning noisy image representations with soft text embeddings.
While REPA and its extensions yield substantial gains in practice, they rely on additional pre-training overhead and, more notably, on external sources of information.
It is difficult to disentangle whether the improvements from representation alignment arise from a self-supervised objective or primarily from increased compute and access to external data.

## 3 Methodology

### 3.1 Dispersive Loss

At the core of our method is the idea of regularizing a generative model’s internal representations by encouraging their dispersion in the hidden space.
We refer to the original regression loss in diffusion-based models as the diffusion loss, and to our introduced regularization term as “Dispersive Loss”.
Formally, denoting $X=\{x_{i}\}$ as a batch of noisy images $x_{i}$, the objective function of this batch is given by:

$$ $\mathcal{L}(X)=\mathbb{E}_{x_{i}\in X}[\mathcal{L}_{\text{Diff}}\bigl{(}x_{i} \bigr{)}]\leavevmode\nobreak\ +\leavevmode\nobreak\ \lambda\mathcal{L}_{\text{ Disp}}\bigl{(}X\bigr{)}.$ (1) $$

Here, $\mathcal{L}_{\text{Diff}}(x_{i})$ is the standard diffusion loss of one sample, and $\mathcal{L}_{\text{Disp}}(X)$ is the dispersive loss term that depends on the entire batch, with a weighting scale $\lambda$.
In our practice, we do not apply any additional layer (*e.g.*, a projection head [5]), and Dispersive Loss is applied directly on the intermediate representations, requiring no extra learnable parameters.

Our method is self-contained and minimalist. In particular, it does not alter the implementation of the original $\mathcal{L}_{\text{Diff}}$ term: it introduces no additional sampling views, no extra data augmentation, and when $\lambda$ is zero it reduces exactly to the baseline diffusion model. This design is possible because the introduced Dispersive Loss $\mathcal{L}_{\text{Disp}}(X)$ only depends on the already-computed intermediate representations of the same input batch. This is unlike standard contrastive learning (e.g., [5]), where the additional augmentation and view may interfere with the per-sample regression objective.

Intuitively, Dispersive Loss behaves like a “contrastive loss without positive pairs”. In the context of generative modeling, this formulation is plausible because the regression terms provide predefined targets for training, making the use of “positive pairs” unnecessary. This is consistent with prior research on self-supervised learning [39], where positive terms are interpreted as alignment objectives, and negative terms as forms of regularization. By eliminating the need for positive pairs, the loss terms can be defined on any standard batch of (independent) images.

Conceptually, Dispersive Loss can be derived from any existing contrastive loss by appropriately removing the positive terms. In this regard, the term “Dispersive Loss” does not refer to a specific implementation, but rather to a general class of objectives that encourage dispersion. We introduce several variants of Dispersive Loss functions in the following.

### 3.2 InfoNCE-Based Variant of Dispersive Loss

InfoNCE [29] is a widely used and effective variant of contrastive loss in self-supervised learning (*e.g.*, [15, 5]). As a case study, we introduce the dispersive counterpart of the InfoNCE loss.
Formally, let $z_{i}=f(x_{i})$ denote the intermediate representations of the generative model for an input sample $x_{i}$, where $f$ represents the subset of layers used to compute the intermediate representations. The original InfoNCE loss [29] can be interpreted as a categorical cross-entropy objective that encourages high similarity between positive pairs and low similarity between negative pairs:

$$ $\mathcal{L}_{\text{Contrast}}=-\log\!\frac{\exp\!\bigl{(}-\mathcal{D}\!\bigl{( }z_{i},z_{i}^{+}\bigr{)}/\tau\bigr{)}}{\sum_{j}\exp\!\bigl{(}-\mathcal{D}\! \bigl{(}z_{i},z_{j}\bigr{)}/\tau\bigr{)}}.$ (2) $$

Here, $(z_{i},z_{i}^{+})$ denotes a pair of positive samples (*e.g.*, obtained by data augmentation of the same image), and $(z_{i},z_{j})$ denotes any pair of samples which include the positive pair and all negative pairs (*i.e.*, $i\neq j$). $\mathcal{D}$ denotes a dissimilarity function (*e.g.*, distance), and $\tau$ is a hyper-parameter known as the temperature. A commonly used form of $\mathcal{D}$ is the negative cosine similarity [40]: $\mathcal{D}(z_{i},z_{j})=-{z_{i}^{\top}z_{j}}/{(\|z_{i}\|\|z_{j}\|)}$.

**Table 1: Variants of Dispersive Loss functions, compared against their contrastive counterparts. Each variant of Dispersive Loss can be viewed as the contrastive loss without positive pairs.**
| variant | contrastive | dispersive |
| --- | --- | --- |
| InfoNCE | $D(z_{i},z_{i}^{+})/\tau+\log\sum_{j}\exp(-D(z_{i},z_{j})/\tau)$ | $\log\mathbb{E}_{i,j}[\exp(-D(z_{i},z_{j})/\tau)]$ |
| Hinge | $D(z_{i},z_{i}^{+})^{2}+\mathbb{E}_{j}[\max(0,\epsilon-D(z_{i},z_{j}))^{2}]$ | $\mathbb{E}_{i,j}[\max(0,\epsilon-D(z_{i},z_{j}))^{2}]$ |
| Covariance | $(1-\operatorname{Cov}_{mm})^{2}+w\,\sum_{n\neq m}\operatorname{Cov}_{mn}^{2}$ | $\sum_{m,n}\operatorname{Cov}_{mn}^{2}$ |

Inside the logarithm in [Eq. 2](https://arxiv.org/html/2506.09027v2#S3.E2), the numerator involves only the positive pair $(z_{i},z_{i}^{+})$, whereas the denominator includes all pairs in the batch.
Following [13, 39], we can rewrite [Eq. 2](https://arxiv.org/html/2506.09027v2#S3.E2) equivalently as:

$$ $\mathcal{L}_{\text{Contrast}}=\,\mathcal{D}\!\left(z_{i},z_{i}^{+}\right)/\tau \leavevmode\nobreak\ +\leavevmode\nobreak\ \log\!\sum_{j}\exp\!\bigl{(}- \mathcal{D}\!\bigl{(}z_{i},z_{j}\bigr{)}/\tau\bigr{)}.$ (3) $$

Here, the first term is similar to a regression objective, which minimizes the distance between $z_{i}$ and its target $z_{i}^{+}$.
On the other hand, the second term encourages any pair of $(z_{i},z_{j})$ to be as distant as possible.

To construct the Dispersive Loss counterpart, we keep only the second term:

$$ $\mathcal{L}_{\text{Disp}}\leavevmode\nobreak\ =\leavevmode\nobreak\ \log\!\sum _{j}\exp\!\bigl{(}-\mathcal{D}\!\bigl{(}z_{i},z_{j}\bigr{)}/\tau\bigr{)}.$ (4) $$

This formulation can also be viewed as a contrastive loss ([Eq. 3](https://arxiv.org/html/2506.09027v2#S3.E3)), where each positive pair consists of two identical views $z_{i}^{+}=z_{i}$, making $\mathcal{D}\!\left(z_{i},z_{i}^{+}\right)$ a constant.
[Eq. 4](https://arxiv.org/html/2506.09027v2#S3.E4) is equivalent to:

$$ $\mathcal{L}_{\text{Disp}}\leavevmode\nobreak\ =\leavevmode\nobreak\ \log \mathbb{E}_{j}\!\bigl{[}\exp\!\bigl{(}-\mathcal{D}\!\bigl{(}z_{i},z_{j}\bigr{) }/\tau\bigr{)}\bigr{]},$ (5) $$

up to a constant $\log$(batch size) that does not impact optimization.
Conceptually, this loss definition is based on a reference sample $z_{i}$. To have a form defined on a batch of samples $Z=\{z_{i}\}$, we follow [39] and redefine it as:

$$ $\mathcal{L}_{\text{Disp}}=\log\mathbb{E}_{i,j}\!\bigl{[}\exp\!\bigl{(}- \mathcal{D}\!\bigl{(}z_{i},z_{j}\bigr{)}/\tau\bigr{)}\bigr{]}.$ (6) $$

This loss function has the same value for all samples within a batch and is computed only once per batch. In our experiments, in addition to the cosine dissimilarity, we also study the squared $\ell_{2}$ distance: $\mathcal{D}\!\bigl{(}z_{i},z_{j}\bigr{)}=\|z_{i}-z_{j}\|_{2}^{2}$. When using this $\ell_{2}$ form, the Dispersive Loss can be easily computed by a few lines of code, as shown in [Algorithm 1](https://arxiv.org/html/2506.09027v2#LSTx1).

The InfoNCE-based Dispersive Loss as defined in [Eq. 6](https://arxiv.org/html/2506.09027v2#S3.E6) is similar to the uniformity loss in [39] (though we do not $\ell_{2}$-normalize the representations). In the context of contrastive representation learning considered in [39], the uniformity loss is applied to the output representation and must be paired with an alignment loss (*i.e.*, the positive terms). Our formulation goes a step further by removing the alignment term on the intermediate representations, and thereby focusing solely on the regularization perspective.

We note that we do not need to explicitly exclude the term $\mathcal{D}(z_{i},z_{j})$ when $j=i$. Since we do not use multiple views of the same image in one batch, this term always corresponds to a constant and minimal dissimilarity, *e.g.*, 0 in the $\ell_{2}$ case and –1 in the cosine case. As such, this term acts as a constant bias inside the logarithm, and its contribution becomes small when the batch size is sufficiently large. In practice, it is not necessary to exclude this term, which also simplifies the implementation.

### 3.3 Other Variants of Dispersive Loss

The concept of Dispersive Loss naturally extends to a broad class of contrastive loss functions beyond InfoNCE.
Any objective that encourages the repulsion of negative samples may be considered a dispersive objective and instantiated as a variant of Dispersive Loss. We introduce two additional variants based on other types of contrastive loss functions.
[Table 1](https://arxiv.org/html/2506.09027v2#S3.T1) summarizes all three variants and compares the contrastive and dispersive counterparts, as described below.

#### Hinge Loss

. In the classical formulation of contrastive learning [7], the loss function is defined as a sum of independent loss terms, each corresponding to a positive or negative pair.
The loss term for a positive pair is $\mathcal{D}(z_{i},z_{i}^{+})$;
the loss term for a negative pair is formulated as a squared hinge loss, *i.e.*, $\max(0,\epsilon-\mathcal{D}(z_{i},z_{j}))^{2}$, where $\epsilon>0$ is the margin. To construct the Dispersive Loss counterpart, we simply discard the loss terms for positive pairs and compute only the terms for negative pairs. See [Table 1](https://arxiv.org/html/2506.09027v2#S3.T1) (row 2).

#### Covariance Loss

. Another class of (generalized) contrastive loss functions operates on the cross-covariance matrix of the representations. This class of loss functions encourages the cross-covariance matrix to be close to the identity matrix. As an example, we consider the loss defined in “Barlow Twins” [42], which calculates a cross-covariance matrix between the normalized representations of two augmented views of a batch. Denote the $D{\times}D$ cross-covariance as “$\operatorname{Cov}$" with elements indexed by $(m,n)$.
The loss encourages the diagonal elements $\operatorname{Cov}_{mm}$ to be one, using a loss term $(1-\operatorname{Cov}_{mm})^{2}$, and off-diagonal elements $\operatorname{Cov}_{mn}$ ($\forall m\neq n$) to be zero, using a loss term $w\sum_{m\neq n}\operatorname{Cov}_{mn}^{2}$ for a weight $w$.

In our dispersive counterpart, we consider only the off-diagonal elements $\operatorname{Cov}_{mn}$. Since we do not use augmented views, the cross-covariance reduces to the covariance matrix computed over a single-view batch. In this case, the diagonal elements $\operatorname{Cov}_{mm}$ automatically equal one when the representations are $\ell_{2}$-normalized, and thus do not need to be explicitly addressed in the loss function.
The resulting Dispersive Loss is simply $\sum_{m,n}\operatorname{Cov}_{mn}^{2}$. See [Table 1](https://arxiv.org/html/2506.09027v2#S3.T1) (row 3).

Table: Algorithm 1 Dispersive Loss (InfoNCE, $\ell_{2}$ dist.) It takes as input the flattened intermediate representations $Z$ of shape $N{\times}D$ (with $D{=}H{\times}W{\times}C$).

### 3.4 Diffusion Models with Dispersive Loss

As summarized in [Table 1](https://arxiv.org/html/2506.09027v2#S3.T1), all variants of Dispersive Loss exhibit simpler formulations than their contrastive counterparts. More importantly, all Dispersive Loss functions are applicable to a single-view batch, eliminating the need for multi-view augmentations. As a result, they can serve as plug-and-play regularizers within existing generative models, without modifying the implementation of the regression loss.

In practice, incorporating Dispersive Loss requires only minimal adjustments: (i) specifying the intermediate layer on which to apply the regularizer, and (ii) computing the Dispersive Loss at this layer and adding it to the original diffusion loss. [Algorithm 2](https://arxiv.org/html/2506.09027v2#LSTx2) presents the training pseudo-code, with a specific form of Dispersive Loss as defined in [Algorithm 1](https://arxiv.org/html/2506.09027v2#LSTx1). We believe that such simplicity greatly facilitates the practical adoption of our methodology, enabling its application across various generative models.

## 4 Experiments

### 4.1 Experiment Setup

The majority of our experiments are conducted on standard models: DiT [31] and SiT [26], which respectively serve as the de facto diffusion-based and flow-based baselines. We faithfully follow the original implementations in [31, 26] and conduct experiments on the ImageNet dataset [8] at 256$\times$256 resolution. The generative models are trained on a 32$\times$32$\times$4 latent space produced by a VAE tokenizer [33]. Sampling is performed using the ODE-based Heun sampler with 250 steps, following [31, 26].
In our ablation experiments, unless specified, we train the models for 80 epochs [26] and without classifier-free guidance (CFG) [17].
By default, the weight scale $\lambda$ in [Eq. 1](https://arxiv.org/html/2506.09027v2#S3.E1) is 0.5, and the temperature $\tau$ in [Eq. 6](https://arxiv.org/html/2506.09027v2#S3.E6) is 0.5.
The full implementation details are in [Appendix A](https://arxiv.org/html/2506.09027v2#A1).

### 4.2 Main Observations on Diffusion Models

#### Dispersive vs. Contrastive.

In [Table 2](https://arxiv.org/html/2506.09027v2#S4.T2), we compare different variants of Dispersive Loss with their contrastive counterparts. To apply a contrastive loss, two views are sampled for each training example to form a positive pair. We study two strategies for adding noise to both views: (i) sampling noise independently for each view, following the generative model’s noising policy; or (ii) sampling noise for the first view according to the same policy, and then restricting the second view’s noise level to differ by at most 0.005. Moreover, to avoid doubling training epochs due to two-view sampling, we only apply the denoising loss to the first view, for fair comparisons with both the baseline and our method.

[Table 2](https://arxiv.org/html/2506.09027v2#S4.T2) shows that when using independent noise, the contrastive loss fails to improve generation quality in all cases studied. We hypothesize that aligning two views with substantially different noise levels impairs learning. As evidence, [Table 2](https://arxiv.org/html/2506.09027v2#S4.T2) shows that contrastive loss with restricted noise leads to improvements over the baseline in three of the four evaluated variants. These experiments suggest that contrastive learning is sensitive to the choice of data augmentation (consistent with observations in self-supervised learning [5]). Noising, which serves as a built-in form of augmentation in diffusion models, further complicates the problem.. While contrastive learning can be mildly beneficial (as shown in [Table 2](https://arxiv.org/html/2506.09027v2#S4.T2)), the design of the additional view and the coupling of two views may limit its application.

By comparison, [Table 2](https://arxiv.org/html/2506.09027v2#S4.T2) shows that Dispersive Loss yields consistent improvements over the baseline, while avoiding the complications introduced by two-view sampling. Despite its simplicity, Dispersive Loss achieves better results than its contrastive counterpart, even when the latter uses carefully tuned noise. We note that Dispersive Loss is applied to a single-view batch, and its influence on training arises solely from its regularizing effect on the intermediate representations.

#### Variants of Dispersive Loss.

The experiments in [Table 2](https://arxiv.org/html/2506.09027v2#S4.T2) involve four variants of Dispersive Loss (see [Table 1](https://arxiv.org/html/2506.09027v2#S3.T1)): InfoNCE ($\ell_{2}$ or cosine dissimilarity), Hinge, and Covariance. As discussed, all variants of Dispersive Loss outperform the baseline, highlighting the generality and robustness of the approach.

Among these variants, InfoNCE with the $\ell_{2}$ distance performs best: it improves the FID by a substantial margin of 4.14, or relative 11.35%. This is in contrast to common practice in self-supervised learning, where cosine dissimilarity is typically preferred. We note that we do not apply normalization to the representations before computing this InfoNCE loss, and as a result, the distance between two samples can be arbitrarily large.
We hypothesize that this design can encourage the representations to be more dispersed, leading to stronger regularization. We use the $\ell_{2}$-based InfoNCE in other experiments by default.

**Table 2: Variants of Dispersive Loss, compared against their contrastive counterparts. We report FID-50k of SiT-B/2 on ImageNet, trained for 80 epochs, without CFG. The “baseline” entry refers to original SiT-B/2 without any contrastive or dispersive loss. The contrastive loss involves two views per training sample, with “independent” or “restricted” noise (the latter case restricts the noise level to differ by up to $0.005$). Values in brackets denote the relative improvement or degradation with respect to the baseline.**
| variant | baseline | contrastive<br>(independent noise) | contrastive<br>(restricted noise) | dispersive |
| --- | --- | --- | --- | --- |
| none | 36.49 | – | – | – |
| InfoNCE, $\ell_{2}$ | – | 43.66 (+19.65%) | 36.57 (+0.22%) | 32.35 (–11.35%) |
| InfoNCE, cosine | – | 41.62 (+14.06%) | 34.83 (–4.55%) | 34.33 (  –5.92%) |
| Hinge | – | 43.02 (+17.89%) | 35.14 (–3.70%) | 33.93 (  –7.02%) |
| Covariance | – | 37.85 (  +3.73%) | 35.87 (–1.70%) | 35.82 (  –1.84%) |

**Table 3: Block Choice for Regularization. We apply Dispersive Loss at different blocks of the SiT-B/2 model (12-block). FID is evaluated on ImageNet. The regularizer is beneficial in all cases.**
| baseline | 36.49 |
| --- | --- |
| block 1 | 33.64 (  –7.81%) |
| block 2 | 32.98 (  –9.62%) |
| block 3 | 32.35 (–11.35%) |
| block 4 | 32.65 (–10.52%) |
| block 8 | 32.75 (–10.24%) |
| block 12 | 33.06 (  –9.93%) |
| all blocks | 32.05 (–12.17%) |

Figure: Figure 3: Evolvement of Representation Norm. Dispersive Loss is applied only at Block 3. Dispersive Loss significantly increases the representation norm compared to the baseline—even in layers where it is not directly applied. The model is SiT-B/2 on ImageNet.
Refer to caption: x3.png

#### Block Choice for Regularization.

In [Table 4](https://arxiv.org/html/2506.09027v2#S4.T4), we investigate the effect of Dispersive Loss at different layers (*i.e.*, Transformer blocks). Overall, the regularizer improves over the baseline by a substantial margin in all cases studied, showing the generality of our approach.
Applying Dispersive Loss to all blocks yields the best result, while applying it to any single block performs nearly as well.

To take a closer look, [Fig. 3](https://arxiv.org/html/2506.09027v2#S4.F3) shows the $\ell_{2}$ norm of the representations in a model where Dispersive Loss is applied only at Block 3. Notably, our regularizer yields a larger representation norm at Block 3 and propagates this effect to all other blocks, even though it is not directly applied to them. This helps explain the consistent gains observed in [Table 4](https://arxiv.org/html/2506.09027v2#S4.T4) wherever Dispersive Loss is applied.
In our other experiments, Dispersive Loss is applied to the single block at the first quarter of the total blocks.

#### Loss weight λ 𝜆 \lambda italic_λ and temperature τ 𝜏 \tau italic_τ

. Our method only involves two hyper-parameters that need to be specified: $\lambda$, which controls the regularization strength, and $\tau$, the temperature in InfoNCE. In [Table 4](https://arxiv.org/html/2506.09027v2#S4.T4), we study their effect on the results. Overall, all configurations studied improve upon the baseline (FID 36.49), further suggesting that the regularizer is broadly effective. Notably, we compare a wide range of temperature values $\tau$ and find that the results are surprisingly robust: this is in contrary to contrastive self-supervised learning, where performance degrades significantly when deviating from the optimal temperature [5].

Figure: Figure 4: Dispersive Loss for Different Models. We evaluate DiT and SiT baselines across four model sizes, both with and without Dispersive Loss. Each subplot corresponds to a specific model size (XL is only for SiT due to limited computation). All models are trained on ImageNet for 80 epochs.
Refer to caption: x4.png

**Table 5: More Experiments on SiT-XL/2. We train the largest SiT model (XL) for more epochs, with and without using CFG, to match the setup in the SiT paper [26]. Results are reported as FID-50k on ImageNet 256$\times$256. The sampler used is the ODE-based Heun method, except for the last two rows, which use the SDE-based Euler method (following [26]). In the last row, the baseline SiT results are taken from the original paper [26], which reports 1400-epoch training; our results with Dispersive Loss in the last row are based on 1200-epoch training. In all cases, Dispersive Loss yields substantial improvements.**
|  | w/o CFG | w/ CFG |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| epochs | baseline | dispersive | $\Delta$ | baseline | dispersive | $\Delta$ |
| 80 | 18.46 | 15.95 | –2.51 (–13.6%) | 6.02 | 5.09 | –0.93 (–15.4%) |
| 140 | 14.06 | 12.08 | –1.98 (–14.0%) | 3.95 | 3.42 | –0.53 (–13.4%) |
| 200 | 12.18 | 10.64 | –1.54 (–12.6%) | 3.30 | 2.90 | –0.40 (–12.1%) |
| 400 | 10.11 | 8.81 | –1.30 (–12.8%) | 2.69 | 2.39 | –0.30 (–11.1%) |
| 800 | 8.99 | 8.08 | –0.91 (–10.1%) | 2.46 | 2.12 | –0.34 (–13.8%) |
| 800 (w/ SDE) | 8.58 | 7.71 | –0.87 (–10.1%) | 2.27 | 2.09 | –0.18 (–7.93%) |
| $\geq$1200 (w/ SDE) | 8.26 | 7.43 | –0.83 (–10.0%) | 2.06 | 1.97 | –0.09 (–4.36%) |

**Table 6: System-level comparison with REPA. While Dispersive Loss yields a smaller gain compared to its REPA counterpart, our method is self-contained and does not rely on any external models. In contrast, REPA uses a pre-trained DINOv2 model [30], whose overhead is summarized in this table.**
| method | pre-training | additional params | external data | FID (SiT-XL/2) |
| --- | --- | --- | --- | --- |
| REPA [41] | ✓ (1500 ImgNet epochs) | ✓ (1.1B params) | ✓ (142M images) | 1.80 |
| Dispersive Loss | ✗ | ✗ | ✗ | 1.97 |

#### Dispersive Loss for Different Models.

Thus far, our ablations have focused on the SiT-B/2 model. In [Fig. 4](https://arxiv.org/html/2506.09027v2#S4.F4), we extend our evaluation to both DiT [31] and SiT [26] across four commonly used model sizes (S, B, L, XL). Once again, Dispersive Loss consistently improves performance over the baseline in all scenarios.

Interestingly, we observe that both the relative and even absolute improvements tend to be larger when the baseline is stronger. For each specific model size, both the relative and absolute improvements upon the SiT baseline are larger than those upon the DiT baseline—even though the SiT baselines are stronger than their DiT counterparts. Similarly, L-size models exhibit greater relative improvements compared to B- or S-size models (this trend plateaus at the XL size, likely because there is less room for further improvement). Overall, this trend provides strong evidence that the primary effect of Dispersive Loss lies in regularization. As larger and more capable models are more prone to overfitting, they tend to benefit more from effective regularization.

#### More Experiments on SiT-XL/2.

In [Table 5](https://arxiv.org/html/2506.09027v2#S4.T5), we train SiT-XL/2 while faithfully following all practices described in the original SiT paper [26]. Specifically, we train the model for more epochs, both with and without classifier-free guidance (CFG) [17]. Both ODE-based and SDE-based samplers are investigated. Overall, our method proves beneficial across all settings, even as the baselines become stronger.
Some example images generated by this model are in [Fig. 5](https://arxiv.org/html/2506.09027v2#S4.F5).

#### System-level Comparison with REPA

. Our method is related to REPA [41]. While our regularizer operates directly on the model’s internal representations, REPA aligns them with those from an external model. As such, for a fair comparison, both the additional computational overhead and the external sources of information should be taken into account, summarized in [Table 6](https://arxiv.org/html/2506.09027v2#S4.T6). REPA relies on a pre-trained DINOv2 model [30], which itself is distilled from a 1.1B-parameter backbone trained on 142 million curated images for the equivalent of 1500 ImageNet epochs. In comparison, our method is entirely self-contained, requiring no pre-training, no external data, and no additional model parameters.
Our method is readily applicable when scaling up training to larger models and datasets, and we expect that the regularization effect will be desirable in such regimes.

Figure: Figure 5: Qualitative results. We present curated samples generated from SiT-XL/2 with Dispersive Loss.
Refer to caption: extracted/6650736/tile_26_08.png

**Table 7: Dispersive Loss for One-Step Generation. All results are FID-50k on ImageNet 256$\times$256. Left: Dispersive Loss consistently improves the MeanFlow [12] baseline. Right: Comparison of state-of-the-art one-step diffusion/flow-based models.**
| model | epochs | baseline | dispersive | $\Delta$ |
| --- | --- | --- | --- | --- |
| MF-B/4 | 80 | 18.78 | 17.61 | –6.23% |
| MF-B/2 | 80 | 9.77 | 8.97 | –8.18% |
| MF-B/2 | 240 | 6.17 | 5.69 | –7.77% |
| MF-XL/2 | 240 | 3.43 | 3.21 | –6.41% |

### 4.3 One-step Generation Models

Our method can be directly generalized to one-step diffusion-based generative models. In [Table 7](https://arxiv.org/html/2506.09027v2#S4.T7) (left), we apply Dispersive Loss to the recent MeanFlow model [12] (implementation details in [Appendix A](https://arxiv.org/html/2506.09027v2#A1)), and observe consistent improvements. [Table 7](https://arxiv.org/html/2506.09027v2#S4.T7) (right) compares these results with the latest one-step diffusion/flow-based models, showing that our method enhances MeanFlow and achieves a new state of the art.

## 5 Conclusion

In this work, we have proposed Dispersive Loss, an objective that regularizes the internal representations of diffusion models. Without relying on additional pre-training, extra parameters, or external data, our approach demonstrates that representation regularization can effectively enhance generative modeling.

A key principle guiding our design is to introduce minimal or no interference with the sampling process of the original training objective. This allows us to fully preserve the original diffusion training strategy, which can be critical for maintaining generative performance. Nevertheless, a standalone, plug-and-play regularizer is not only desirable for generative modeling; its favorable properties may also benefit other applications, including image recognition.
We hope our regularizer will generalize to broader scenarios, a direction we leave for future exploration.

## Acknowledgement

We greatly thank Google TPU Research Cloud (TRC) for granting us access to TPUs.
Runqian Wang is supported by the MIT Undergraduate Research Opportunities Program (UROP).
We thank Qiao Sun, Zhicheng Jiang, Hanhong Zhao, Yiyang Lu, and Xianbang Wang for their help on the JAX and TPU implementation. We thank Zhengyang Geng for sharing the MeanFlow codebase.

## References

- Albergo and Vanden-Eijnden [2023]
Michael Samuel Albergo and Eric Vanden-Eijnden.
Building normalizing flows with stochastic interpolants.
In *International Conference on Learning Representations (ICLR)*, 2023.
- [2]
Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei.
Beit: Bert pre-training of image transformers.
In *International Conference on Learning Representations*.
- Bengio et al. [2013]
Yoshua Bengio, Aaron Courville, and Pascal Vincent.
Representation learning: A review and new perspectives.
*IEEE transactions on pattern analysis and machine intelligence*, 35(8):1798–1828, 2013.
- Chen et al. [2025]
Hesen Chen, Junyan Wang, Zhiyu Tan, and Hao Li.
SARA: Structural and adversarial representation alignment for training-efficient diffusion models.
*arXiv preprint arXiv:2503.08253*, 2025.
- Chen et al. [2020]
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.
A simple framework for contrastive learning of visual representations.
In *International conference on machine learning*, pages 1597–1607. PmLR, 2020.
- Chen and He [2021]
Xinlei Chen and Kaiming He.
Exploring simple siamese representation learning.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pages 15750–15758, 2021.
- Chopra et al. [2005]
Sumit Chopra, Raia Hadsell, and Yann LeCun.
Learning a similarity metric discriminatively, with application to face verification.
In *2005 IEEE computer society conference on computer vision and pattern recognition (CVPR’05)*, pages 539–546. IEEE, 2005.
- Deng et al. [2009]
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.
ImageNet: A large-scale hierarchical image database.
In *2009 IEEE conference on computer vision and pattern recognition*, pages 248–255. Ieee, 2009.
- Devlin et al. [2019]
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.
Bert: Pre-training of deep bidirectional transformers for language understanding.
In *Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume 1 (long and short papers)*, pages 4171–4186, 2019.
- Dhariwal and Nichol [2021]
Prafulla Dhariwal and Alexander Nichol.
Diffusion models beat gans on image synthesis.
*Advances in neural information processing systems*, 34:8780–8794, 2021.
- [11]
Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel.
One step diffusion via shortcut models.
In *The Thirteenth International Conference on Learning Representations*.
- Geng et al. [2025]
Zhengyang Geng, Mingyang Deng, Xingjian Bai, J Zico Kolter, and Kaiming He.
Mean flows for one-step generative modeling.
*arXiv preprint arXiv:2505.13447*, 2025.
- Grill et al. [2020]
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al.
Bootstrap your own latent-a new approach to self-supervised learning.
*Advances in neural information processing systems*, 33:21271–21284, 2020.
- Hadsell et al. [2006]
Raia Hadsell, Sumit Chopra, and Yann LeCun.
Dimensionality reduction by learning an invariant mapping.
In *IEEE Conference on Computer Vision and Pattern Recognition (CVPR)*, pages 1735–1742. IEEE, 2006.
- He et al. [2020]
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick.
Momentum contrast for unsupervised visual representation learning.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pages 9729–9738, 2020.
- He et al. [2022]
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick.
Masked autoencoders are scalable vision learners.
In *IEEE Conference on Computer Vision and Pattern Recognition (CVPR)*, 2022.
- Ho and Salimans [2022]
Jonathan Ho and Tim Salimans.
Classifier-free diffusion guidance.
*arXiv preprint arXiv:2207.12598*, 2022.
- Ho et al. [2020]
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
Denoising diffusion probabilistic models.
*Advances in neural information processing systems*, 33:6840–6851, 2020.
- Karras et al. [2022]
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine.
Elucidating the design space of diffusion-based generative models.
*Advances in neural information processing systems*, 35:26565–26577, 2022.
- Khosla et al. [2020]
Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan.
Supervised contrastive learning.
*Advances in neural information processing systems*, 33:18661–18673, 2020.
- Krizhevsky [2009]
Alex Krizhevsky.
Learning multiple layers of features from tiny images.
2009.
- Lee et al. [2025]
Jaa-Yeon Lee, Byunghee Cha, Jeongsol Kim, and Jong Chul Ye.
Aligning text to image in diffusion models is easier than you think.
*arXiv preprint arXiv:2503.08250*, 2025.
- [23]
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le.
Flow matching for generative modeling.
In *The Eleventh International Conference on Learning Representations*.
- Lipman et al. [2024]
Yaron Lipman, Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky TQ Chen, David Lopez-Paz, Heli Ben-Hamu, and Itai Gat.
Flow matching guide and code.
*arXiv preprint arXiv:2412.06264*, 2024.
- [25]
Xingchao Liu, Chengyue Gong, et al.
Flow straight and fast: Learning to generate and transfer data with rectified flow.
In *The Eleventh International Conference on Learning Representations*.
- Ma et al. [2024]
Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden, and Saining Xie.
SiT: Exploring flow and diffusion-based generative models with scalable interpolant transformers.
In *European Conference on Computer Vision*, pages 23–40. Springer, 2024.
- Mu et al. [2022]
Norman Mu, Alexander Kirillov, David Wagner, and Saining Xie.
Slip: Self-supervision meets language-image pre-training.
In *European conference on computer vision*, pages 529–544. Springer, 2022.
- Nichol and Dhariwal [2021]
Alexander Quinn Nichol and Prafulla Dhariwal.
Improved denoising diffusion probabilistic models.
In *International Conference on Machine Learning (ICML)*. PMLR, 2021.
- Oord et al. [2018]
Aaron van den Oord, Yazhe Li, and Oriol Vinyals.
Representation learning with contrastive predictive coding.
*arXiv preprint arXiv:1807.03748*, 2018.
- Oquab et al. [2023]
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al.
Dinov2: Learning robust visual features without supervision.
*arXiv preprint arXiv:2304.07193*, 2023.
- Peebles and Xie [2023]
William Peebles and Saining Xie.
Scalable diffusion models with transformers.
In *Proceedings of the IEEE/CVF international conference on computer vision*, pages 4195–4205, 2023.
- Radford et al. [2021]
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.
Learning transferable visual models from natural language supervision.
In *International Conference on Machine Learning (ICML)*, pages 8748–8763. PmLR, 2021.
- Rombach et al. [2022]
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.
High-resolution image synthesis with latent diffusion models.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pages 10684–10695, 2022.
- Ronneberger et al. [2015]
Olaf Ronneberger, Philipp Fischer, and Thomas Brox.
U-net: Convolutional networks for biomedical image segmentation.
In *Medical image computing and computer-assisted intervention (MICCAI)*, 2015.
- Sohl-Dickstein et al. [2015]
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.
Deep unsupervised learning using nonequilibrium thermodynamics.
In *International conference on machine learning*, pages 2256–2265. pmlr, 2015.
- [36]
Jiaming Song, Chenlin Meng, and Stefano Ermon.
Denoising diffusion implicit models.
In *International Conference on Learning Representations*.
- [37]
Yang Song and Prafulla Dhariwal.
Improved techniques for training consistency models.
In *The Twelfth International Conference on Learning Representations*.
- Song and Ermon [2019]
Yang Song and Stefano Ermon.
Generative modeling by estimating gradients of the data distribution.
*Neural Information Processing Systems (NeurIPS)*, 2019.
- Wang and Isola [2020]
Tongzhou Wang and Phillip Isola.
Understanding contrastive representation learning through alignment and uniformity on the hypersphere.
In *International conference on machine learning*, pages 9929–9939. PMLR, 2020.
- Wu et al. [2018]
Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin.
Unsupervised feature learning via non-parametric instance discrimination.
In *Proceedings of the IEEE conference on computer vision and pattern recognition*, pages 3733–3742, 2018.
- Yu et al. [2024]
Sihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong, Jonathan Huang, Jinwoo Shin, and Saining Xie.
Representation alignment for generation: Training diffusion transformers is easier than you think.
*arXiv preprint arXiv:2410.06940*, 2024.
- Zbontar et al. [2021]
Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny.
Barlow twins: Self-supervised learning via redundancy reduction.
In *International conference on machine learning*, pages 12310–12320. PMLR, 2021.
- Zhang et al. [2022]
Qi Zhang, Yifei Wang, and Yisen Wang.
How mask matters: Towards theoretical understandings of masked autoencoders.
*Advances in Neural Information Processing Systems*, 35:27127–27139, 2022.
- Zhou et al. [2025]
Linqi Zhou, Stefano Ermon, and Jiaming Song.
Inductive moment matching.
*arXiv preprint arXiv:2503.07565*, 2025.

## Appendix A Implementation

#### SiT and DiT Experiments.

We faithfully follow the SiT/DiT codebase for ImageNet experiments. We employ the AdamW optimizer with a constant learning rate of $1\times 10^{-4}$, $(\beta_{1},\beta_{2})=(0.9,0.95)$, and no weight decay. For ODE sampling, we use the Heun solver with 250 steps, following the SiT paper [26]; for SDE sampling, we apply an Euler–Maruyama solver with the default drift and diffusion schedule from SiT [26].
We re-implement both SiT and DiT in JAX. All experiments are run on 32 TPU-v3/v4 cores. The configuration details are in [Table 8](https://arxiv.org/html/2506.09027v2#A0.T8).

#### MeanFlow Experiments.

Our MeanFlow experiments use the exact same codebase shared by the authors of [12]. In [Table 7](https://arxiv.org/html/2506.09027v2#S4.T7), we evaluate three models: B/4, B/2, and XL/2, with and without Dispersive Loss. The B/4 and B/2 models are all trained from scratch. For the XL/2 model, due to limited computation resources, we take the MeanFlow-XL/2 checkpoint at 180 epochs and then apply Dispersive Loss in the remaining 60 epochs; we expect training this model from scratch will lead to better results.

For the MeanFlow baselines, all configurations of MF-B/2 and MF-XL/2 follow those in [12]. For MF-B/4, we adopt the setting of $\omega^{\prime}=2.0$ and $\kappa=0.5$ using the notation of [12]. Dispersive Loss is applied on the same single block as in our SiT experiments. We set the regularization weight $\lambda$ as 0.25 for B/4, 1.0 for B/2, and 1.5 for XL/2.

## Appendix B Additional Experiments

### B.1 Inception Scores

**Table 9: Inception Scores on ImageNet.**
|  |  | w/o CFG | w/ CFG |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| model | epochs | baseline | dispersive | $\Delta$ | baseline | dispersive | $\Delta$ |
| B/2 | 80 | 42.79 | 48.79 | +14.02% | 93.26 | 108.23 | +16.05% |
| XL/2 | 800 | 141.28 | 147.35 | +4.11% | 270.44 | 281.26 | +4.00% |

In [Table 9](https://arxiv.org/html/2506.09027v2#A2.T9) we report the Inception Scores of SiT models on ImageNet. Similar to observations on FID ([Table 5](https://arxiv.org/html/2506.09027v2#S4.T5)), Dispersive Loss improves Inception Scores (higher being better) over the baseline SiT models.

### B.2 CIFAR-10 Experiments

Figure: Figure 6: CIFAR-10 Result.
Refer to caption: x5.png

We also evaluate our approach on non-transformer architectures in the CIFAR-10 dataset [21], where the commonly used network architecture is Unet [34].
The experiments are based on the publicly available code of Flow Matching [24].(^2^22[https://github.com/facebookresearch/flow_matching](https://github.com/facebookresearch/flow_matching))
We use the same hyper-parameters as the original codebase, and our rerunning of the Flow Matching baseline has 2.13 FID. As a reference, the original repo reported 2.07 FID.

[Figure 6](https://arxiv.org/html/2506.09027v2#A2.F6) represents the evolvement of the FID during training, comparing the baseline against the one with Dispersive Loss, using the same random seed. The Dispersive Loss is applied at Residual Block 15 of the Unet. It can be observed that our regularizer yields consistent gains over the baseline throughout training. Our final result is 2.07 FID, which is better than the 2.13 FID of our reproduced baseline. While the absolute numbers may depend on subtle variations in the runtime environment, [Figure 6](https://arxiv.org/html/2506.09027v2#A2.F6) shows that the gains are persistent when evaluated under a controlled environment.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: - Although CIFAR-10 is included, other generative domains (text-to-image or high-res synthesis) could further establish generality. It would be interesting to see whether Dispersive Loss also benefits conditional or multimodal diffusion models.
- While Figure 3 shows increased representation norms, further qualitative or visualization-based analyses (e.g., embedding t-SNEs) would make the mechanism more intuitive.
- Some missing works on representation learning using/within diffusion models should be included [5, 6, 7, 8, 9] in related works.

Overall, I think this is a good paper and I would be happy to raise the score if my comments are addressed properly.

[5] Wang, Yingheng, et al. "Infodiffusion: Representation learning using information maximizing diffusion models." International conference on machine learning. PMLR, 2023.

[6] Mittal, Sarthak, et al. "Diffusion based representation learning." International conference on machine learning. PMLR, 2023.

[7] Hudson, Drew A., et al. "Soda: Bottleneck diffusion models for representation learning." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.

[8] Zhang, Zijian, Zhou Zhao, and Zhijie Lin. "Unsupervised representation learning from pre-trained diffusion probabilistic models." Advances in neural information processing systems 35 (2022): 22117-22130.

[9] Yang, Xingyi, and Xinchao Wang. "Diffusion model as representation learner." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.
soundness: 4
strengths: - The idea of removing positive pairs while retaining the repulsive regularization aspect is conceptually appealing and practically justified by diffusion models’ intrinsic alignment objective.
- Comprehensive experiments across multiple architectures (DiT, SiT, MeanFlow) and scales (S/B/L/XL) show consistent improvements in FID and Inception Scores.
- The improvement trend scales with model size, indicating the loss acts as an effective regularizer for large-capacity models prone to overfitting.
- Dispersive Loss outperforms all contrastive baselines even with careful tuning of noise schedules.
- The plug-and-play simplicity (no multi-view augmentation or external encoders) is convincingly demonstrated.
- The authors provide ablations for hyperparameters ($\lambda$, $\tau$), layer placement, and different loss variants (Table 2–4), showing robustness across configurations.
- Implementation details are transparent (Algorithm 1–2, Table 8).
- The inclusion of MeanFlow and CIFAR-10 experiments supports generality across diffusion and flow-matching paradigms.
- Figures (e.g., Fig. 1–4) clearly illustrate how Dispersive Loss integrates into existing architectures with negligible computational overhead.
- Comparisons with REPA (Table 6) highlight the system-level efficiency advantage (no 1.1B-parameter pre-trained model, no 142M external images).
- The paper adheres to reproducibility standards (code, README included) and presents results with careful quantitative analysis.
- The method provides a bridge between generative and representation learning, a frontier that is conceptually and practically valuable for the field.
summary: This paper introduces Dispersive Loss, a simple yet effective plug-and-play regularizer designed to improve diffusion-based generative models by encouraging dispersion of internal representations in hidden space. Unlike prior work such as REPA [1], which requires pre-trained external encoders and additional parameters, the proposed method is self-contained, requiring no pre-training, no extra parameters, and no external data.

Conceptually, Dispersive Loss can be interpreted as a “contrastive loss without positive pairs.” It regularizes the hidden representations by penalizing excessive clustering and promoting diversity, inspired by the repulsive component of contrastive learning. The authors provide theoretical motivation, multiple instantiations (InfoNCE-, Hinge-, and Covariance-based), and efficient implementations requiring only a few lines of code.

The paper provides extensive empirical evaluation across several diffusion backbones, including SiT [2], DiT [3], and MeanFlow [4]. Experiments on ImageNet 256×256 consistently show significant FID improvements (up to ~11–13%) over strong baselines and even outperform contrastive variants that require two-view sampling. The method is also shown to generalize well to one-step generation and smaller datasets like CIFAR-10, demonstrating broad applicability.

[1] Yu, Sihyun, et al. "Representation alignment for generation: Training diffusion transformers is easier than you think." arXiv preprint arXiv:2410.06940 (2024).

[2] Ma, Nanye, et al. "Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers." European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2024.

[3] Peebles, William, and Saining Xie. "Scalable diffusion models with transformers." Proceedings of the IEEE/CVF international conference on computer vision. 2023.

[4] Geng, Zhengyang, et al. "Mean flows for one-step generative modeling." arXiv preprint arXiv:2505.13447 (2025).
weaknesses: - The method is motivated intuitively but lacks a formal analysis of why dispersion improves generation quality. A deeper information-theoretic or geometric argument (e.g., on latent coverage or mutual information bounds) would strengthen the theoretical grounding.
- While FID and Inception Scores are strong indicators, evaluation on semantic diversity, perceptual similarity, or representation quality (e.g., CLIP-based metrics) could better reveal what aspects of representation regularization improve.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 2
questions: See "Weaknesses"
soundness: 3
strengths: 1. The proposed regularizer can be directly integrated into diffusion models with intermediate representations and requires little additional computational effort. 
2. It elegantly incorporates concepts from self-supervised learning into diffusion model training in a straightforward and theoretically sound manner. 
3. Experiments on a real-world image dataset demonstrate that adding the regularizer significantly improves performance, and the experimental results are comprehensive.
summary: The paper introduces a plug-in regularizer designed to enhance diffusion-based models by promoting the dispersion of internal representations in the hidden space. To demonstrate its effectiveness, the authors apply the proposed regularizer to two representative diffusion-based models, DiT (Peebles & Xie, 2023) and SiT (Ma et al., 2024). Experiments conducted on the ImageNet dataset show that the regularizer effectively improves model performance.
weaknesses: 1. It would be helpful to clarify the scope of applicability. Can the proposed regularizer be applied to all diffusion models with intermediate representations? 
2. Qualitative comparisons between images generated with and without the proposed regularizer would make the improvements more intuitive and visually convincing. 
3. Although experiments explore different blocks, loss weights, and temperatures, it would be beneficial to provide systematic guidance or heuristics for selecting these hyperparameters.

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 3
questions: Please see the weakness part above.
soundness: 3
strengths: 1. The authors aim to address a fundamental problem, namely that the generation task should not be left to stand alone with representation learning, and offer insightful perspectives.

2. The paper features a clear structure and coherent logic.

3. The proposed method does not rely on a pretrained encoder.
summary: This authors try to address a key limitation of generative diffusion models with representation learning: diffusion models trained with regression-based denoising objectives often lack explicit regularization for the feature space, and therefore are limited in performance. To tackle this, the authors introduce a Dispersive Loss, which eliminates the need for explicit positive sample pairs, compared to the contrastive learning counterpart.
weaknesses: 1. The authors used limited evaluation metrics. As they don't rely on a pretrained encoder and claim the importance of representation learning in the generative task, the authors should evaluate how the method performs with metrics like linear probing.

2. It's not clear why the major improvements were made in the case without CFG, while the performance with CFG only achieves very limited improvements. The authors should provide deeper analyses to explain this discrepancy and not rely only on FID (which is not in favor of the proposed method when evaluated with CFG).

3. It's not clear how the proposed method scales with the diversity and the size of datasets. Some empirical verification or theoretical motivation would be useful.

4. There is a slight redundancy in the mathematical derivation process of Section 3.2.

---

confidence: 5
contribution: 1
initial_rating: 4
presentation: 2
questions: 1. Could the authors provide theoretical or empirical evidence to clarify what specific kinds of high-level semantic features are required by image generation models, particularly diffusion models?

2. how does the proposed dispersive loss specifically facilitate the learning of such required semantic features?

3. Given that contrastive learning frameworks typically rely on heavy image augmentations to learn high-level semantics, how is the Dispersive Loss able to learn similar semantic information without relying on these image augmentations?

4. To support the main claim, shouldn't experiments be conducted to explicitly demonstrate the model's ability to learn rich semantic features, perhaps including evaluations on downstream tasks to provide valuable support?

5. Does a high representation norm truly indicate the capture of semantic information, and could the authors validate whether the model has learned semantic features, similar to the REPA method, by using linear probing?

6. Could the authors include feature similarity comparisons, such as Centered Kernel Alignment (CKA), with other feature extractors to better demonstrate the improved semantics?

7. To further validate the effectiveness of dispersive loss, would it be beneficial to test it on more complex generative tasks, such as text-to-image generation or text-to-video generation?

8. Since the method claims removing alignment term in line 231, shouldn't there be some investigation into why the alignment term deteriorates performance?

9. What happens when augmented images are used for positive pairs (i.e., should an ablation study compare the use of the same image versus an augmented image)?

10. How does down-weighting the alignment term in the contrastive loss affect the overall performance (i.e., how does performance change when moving from a standard contrastive loss toward the dispersive loss)?

11. Does this additional regularization lead to features that are usable for other discriminative tasks (even though this is not the main focus, it could demonstrate a strength of the learned features)?
soundness: 2
strengths: 1. Simplicity and Plug-and-Play Design: The proposed Dispersive Loss is highly appealing due to its self-contained and minimalist nature. It is a simple regularization term that can be added to the standard diffusion loss without requiring pre-training, external data, or additional visual encoders. This makes it easy to integrate into existing diffusion pipelines.

2. Demonstrated Training Acceleration: The paper shows an intriguing result of training acceleration without the need for external data or modules (unlike methods like REPA). While the gain is marginal, the fact that a simple internal regularization can accelerate convergence is a valuable finding.
summary: This paper introduces the Dispersive Loss as a plug-and-play regularization term for diffusion-based generative models. The Dispersive Loss encourages the hidden feature representations to disperse (or spread out) in the latent space, analogous to the non-alignment term of contrastive learning, but uniquely requires no positive sample pairs. The authors claim this enhances the quality of generated images and improves training efficiency. They report consistent improvements over strong baselines across various models on the ImageNet dataset.
weaknesses: ## 1. Fundamental Gaps in Theoretical Justification and Representation Analysis

The paper's core weakness lies in the disconnect between its SSL-inspired motivation and its analysis.

- Lack of Deep Theoretical Insight: The paper does not provide sufficient theoretical insights into why the Dispersive Loss leads to improvement, nor does it explore what kind of high-level representations generative models fundamentally require. The mechanism of dispersion is asserted to be beneficial without a rigorous explanation of its effect on the learned noise-prediction function or the resulting generative manifold.

- Missing Representation Utility Analysis: Since the method is explicitly motivated by SSL, it should verify the "rich semantic features" from SSL methods. Analysis techniques like linear probing on the average-pooled features (at the intermediate or final layer) for tasks like classification or retrieval are essential to validate the core hypothesis that the regularization improves semantic content.

- Unclear Design Rationale: It is asserted that the alignment (positive pair) term of contrastive loss is detrimental to the generation task, leading to its exclusion, as in Table 2. This is a strong claim that is not clearly supported. Intuitively, the alignment term would, at worst, be redundant; the paper must theoretically or empirically demonstrate why it would be actively harmful.

## 2. Lack of Experimental Rigor and Insufficient Evidence

The experimental section suffers from issues of completeness, methodology, and evidence that undermine the paper's claims.

- Missing Empirical Evidence. The proposed method has not been thoroughly investigated. While the paper evaluates various settings of the regularizers, e.g., different models and depths, it does not sufficiently analyze why this method is effective. 

- Missing Scalability and High-Resolution Tests: The paper lacks experiments in high-resolution settings, such as ImageNet-512x512 or text-to-image generation. This omission makes it impossible to assess the method's scalability and effectiveness on datasets where modern generative models typically face their greatest architectural challenges.

## 3. Marginal Performance and Missing Training Dynamics Analysis

The reported performance gains are modest, and the paper fails to investigate the interaction between its components.

- Marginal Performance Gains: The overall performance gains achieved by Dispersive Loss are marginal. The resulting models neither consistently match SOTA generative performance (e.g., REPA) nor demonstrate effectiveness in any downstream context.

- Missing Integration with Faster Frameworks: While the paper notes its potential for acceleration, it is unclear whether the proposed framework can be combined with existing faster training frameworks like REPA. An experiment or intuition on this synergy is needed.

- Unanalyzed Training Dynamics: The paper completely lacks an in-depth analysis of the training dynamics involving the new loss term. The authors should show the evolution of the Dispersive Loss over the course of training to see if it plateaus or exhibits different behavior from the standard denoising loss. An analysis of the conflict or synergy between the denoising loss and the Dispersive Loss (e.g., by analyzing the direction of their respective gradients) is necessary to clarify the training process. Furthermore, an investigation into how the dispersive loss varies across diffusion timesteps is crucial, as its impact is intuitively expected to be different for high-noise versus low-noise regimes.

- Unfair Comparisons: The paper lacks comparison with fairer versions of SOTA models. For example, when comparing against REPA (800 epochs), the comparison should be normalized for total training time or resources consumed to provide a true measure of efficiency (Disperse loss uses >=1200 epochs in Table 6).

- Incomplete Results. The paper only compares with SiT and SiT-REPA (except for one-step generation) and uses FID as their single evaluation metric. Authors should conduct comprehensive comparisons with current state-of-the-art methods that share the same idea, such as Lighting-DiT [`1], DC-AE1.5 [2], DDT [3], etc. Besides, more evaluation metrics, such as Inception-Score, should be included in the paper.




**I recommend that the authors conduct a deeper investigation into the learned representations and strictly refine their experimental validation and theoretical justification during rebuttal period.**


[1] Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models

[2] DC-AE 1.5: Accelerating Diffusion Model Convergence with Structured Latent Space

[3] DDT: Decoupled Diffusion Transformer


Decision: Reject