No Alignment Needed for Generation: Learning Linearly Separable Representations in Diffusion Models

## Abstract

Abstract Efficient training strategies for large-scale diffusion models have recently emphasized the importance of improving discriminative feature representations in these models.
A central line of work in this direction is representation alignment with features obtained from powerful external encoders, which improves the representation quality as assessed through linear probing . Alignment-based approaches show promise but depend on large pretrained encoders, which are computationally expensive to obtain. In this work, we propose an alternative regularization for training, based on promoting the L inear SEP arability (LSEP) of intermediate layer representations. LSEP eliminates the need for an auxiliary encoder and representation alignment, while incorporating linear probing directly into the network’s learning dynamics rather than treating it as a simple post-hoc evaluation tool. Our results demonstrate substantial improvements in both training efficiency and generation quality on flow-based transformer architectures such as SiTs, achieving an FID of 1.46 on 256 × 256 256\times 256 ImageNet dataset.

## 1 Introduction

Diffusion models have demonstrated remarkable generative quality in a wide range of visual tasks, motivating additional architectural innovations to further improve their performance (Dhariwal & Nichol, 2021; Karras et al., 2022; Peebles & Xie, 2023; Ma et al., 2024). These models inherently learn semantically meaningful representations through the noise-prediction objective used in denoising (Baranchuk et al., 2022; Xiang et al., 2023; Chen et al., 2025). However, their learned representations are not explicitly optimized for representation learning and are often less expressive than models trained specifically for this purpose (Xiang et al., 2023; Yu et al., 2025).

From a representational perspective, a growing body of research has examined how improving learned representations can benefit diffusion training (Zhu et al., 2024; Yu et al., 2025; Jiang et al., 2025; Leng et al., 2025; Yao et al., 2025). These studies have explored methods to improve the training efficiency and generative performance of transformer-based diffusion models through representation alignment. In particular, they have consistently shown that acquiring strong representations at specific network depths by aligning with high-quality internal (Zhu et al., 2024; Jiang et al., 2025) or external (Yu et al., 2025; Leng et al., 2025; Yao et al., 2025) representations significantly accelerates convergence and enhances the quality of generated outputs.

A recent powerful approach, representation alignment (REPA) (Yu et al., 2025), leverages high-quality representations extracted from large-scale pre-trained transformer models, such as DINOv2 (Oquab et al., 2024) and CLIP (Radford et al., 2021). REPA explicitly aligns early-stage diffusion transformer features with clean image features from external encoders, thereby encouraging stronger representational capacity. This alignment enables deeper layers to concentrate on high-frequency content, ultimately improving generative performance. However, training such pre-trained visual encoders demands access to large-scale datasets and involves substantial computational costs.

An alternative approach leverages self-representation alignment (Zhu et al., 2024; Jiang et al., 2025) without external encoders. In this framework, self-representation alignment (SRA) (Jiang et al., 2025) adapts teacher–student discriminative pair structures within diffusion models to enable self-supervised knowledge distillation (Zhu et al., 2024). Using the observation that deeper layers produce richer representations (Xiang et al., 2023; Yu et al., 2025), the teacher is assigned to deeper layers with lower noise, while the student is associated with the earlier layers. This strategy has been shown to enhance training effectiveness without relying on pre-trained encoders. Nevertheless, the approach remains fundamentally constrained by the representational capacity of the diffusion model itself.

To assess the effectiveness of learned representations, prior work commonly employs linear probing (Alain & Bengio, 2016). In this evaluation protocol, the trained diffusion model is frozen, and features from specific layers are extracted to train a lightweight linear probe (classifier). The classifier’s accuracy serves as a measure of the *degree of linear separability* of the feature space. Linear probing accuracy has been shown to correlate strongly with both the training efficiency and the generative quality of diffusion transformers (Leng et al., 2025; Yu et al., 2025). These findings suggest that a key factor underlying the success of prior approaches is the refinement of feature representations, thereby promoting *linear separability*.

#### Our Approach:

This work is motivated by a simple yet fundamental question:

*“Can diffusion models learn highly linearly separable representations that improve training efficiency while producing higher-quality outputs without representation alignment or external encoders?”*

To this end,
we introduce Linear SEParability (LSEP), a training regularization strategy that incorporates a linear probe into diffusion models to simultaneously improve the separability of early-stage hidden representations and optimize the denoising objective.

In particular, we insert a trainable linear probe into an intermediate layer of the diffusion model, as shown in  Fig. [1](https://arxiv.org/html/2509.21565v1#S0.F1), similar in spirit to linear probing evaluations, but *without freezing* the model parameters. Unlike prior alignment-based approaches, our method does *not require an external visual encoder or explicit representation alignment*, while still promoting stronger linear separability at the intended depth of the diffusion model. Experiments show that LSEP substantially boosts both the training efficiency and the generative output quality on flow-based transformer SiTs (Ma et al., 2024).

Our main contributions are as follows:

- •
We introduce *Linear SEParability (LSEP)*, a framework that integrates a linear probe (classifier) into generative model training to simultaneously enhance linear separability and the standard denoising objective.
- •
To enable the two objectives to mutually reinforce each other, we propose novel training techniques for the classification term while keeping the denoising training unchanged: (1) classification-specific conditioning for the linear probe branch, (2) random cropping to enhance patch-level linear separability, and (3) time-dependent weighting of the classification loss.
- •
We demonstrate that LSEP substantially enhances both training efficiency and generative quality in flow-based transformer architectures, without relying on *an external visual encoder or explicit representation alignment*.
- •
SiT-XL with LSEP converges to lower FID significantly faster than the baseline SiT-XL, reaching performance comparable to the alignment-based model REPA, as shown in Fig. [1](https://arxiv.org/html/2509.21565v1#S0.F1) (b). Moreover, it achieves an FID of 1.46 with classifier-free guidance using the guidance interval, establishing the best performance among models without relying on external encoder architectures.
- •
Finally, we show that LSEP synergistically combines with alignment-based methods to enhance linear separability of representations through distinct mechanisms, and to further improve both training efficiency and generative performance.

## 2 Preliminaries

### 2.1 Training Flow-based Diffusion Transformer

Flow-based approaches (Lipman et al., 2023) learn a velocity field $\mathbf{v}_{\theta}(\mathbf{x}_{t},t)$ that defines a probability flow ordinary differential equation (PF-ODE), characterizing the deterministic evolution of a data point $\mathbf{x}_{t}$. SiT (Ma et al., 2024) adopts this framework to model a continuous-time forward process:

$$ $\mathbf{x}_{t}=\alpha_{t}\mathbf{x}_{0}+\sigma_{t}\bm{\epsilon},\quad\bm{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I}),$ (1) $$

where $\alpha_{t}$ and $\sigma_{t}$ are time-dependent functions such that $\alpha_{t}$ decreases and $\sigma_{t}$ increases with $t\in[0,T]$, satisfying $\alpha_{0}=\sigma_{T}=1$ and $\alpha_{T}=\sigma_{0}=0$. The PF-ODE is given by
$\frac{d\mathbf{x}_{t}}{dt}=\mathbf{v}_{\theta}(\mathbf{x}_{t},t),$
where the distribution of the ODE solution at time $t$ matches the marginal distribution of the forward process.
To train the velocity model $\mathbf{v}_{\theta}(\mathbf{x}_{t},t)$, the following velocity matching loss is minimized:

$$ $\mathcal{L}_{\text{velocity}}=\mathbb{E}_{\mathbf{x}_{0},\bm{\epsilon},t}\left[\left\|\mathbf{v}_{\theta}(\mathbf{x}_{t},t)-\dot{\alpha}_{t}\mathbf{x}_{0}-\dot{\sigma}_{t}\bm{\epsilon}\right\|^{2}\right],$ (2) $$

where $\dot{\alpha}_{t}=\frac{d\alpha_{t}}{dt}$ and $\dot{\sigma}_{t}=\frac{d\sigma_{t}}{dt}$ denote the time derivatives of $\alpha_{t}$ and $\sigma_{t}$, respectively.

In practice, to incorporate class guidance via classifier-free guidance (CFG) (Ho & Salimans, 2021), the velocity model $\mathbf{v}_{\theta}$ is trained conditionally on a label $\mathbf{c}\in\mathcal{C}$, resulting in a model of the form $\mathbf{v}_{\theta}(\mathbf{x}_{t},t\mid\mathbf{c})$ (Peebles & Xie, 2023; Ma et al., 2024; Yu et al., 2025). Here, $\mathcal{C}$ includes both the dataset class labels $\{\mathbf{c}_{\text{class}}\}$
and an unconditional label $\mathbf{c}_{\varnothing}$, *i.e*. $\mathcal{C}=\{\mathbf{c}_{\text{class}}\}\cup\mathbf{c}_{\varnothing}$. Both types of conditions are employed during training, where the unconditional label is used with a small probability, $\rho_{D}\ll 1.$

### 2.2 Representation Alignment for Generative Models

To improve training efficiency and generation quality of diffusion transformers, REPA (Yu et al., 2025) aligns a model’s hidden-layer representations with large-scale pretrained self-supervised visual encoders such as DINOv2 (Oquab et al., 2024) and CLIP (Radford et al., 2021). Specifically, REPA introduces a regularization term that maximizes patch-wise similarity between features $\mathbf{y}_{\mathrm{clean}}$, extracted from the pretrained external encoder, and $k^{\text{th}}$ hidden layer features $\mathbf{h}_{t}^{k}$ of the diffusion transformer encoder at timestep $t$. The representation alignment loss is defined as:

$$ $\mathcal{L}_{\textrm{repa}}=-\mathbb{E}_{\mathbf{x}_{\textrm{clean}},\epsilon,t}\left[\frac{1}{N}\sum_{n=1}^{N}\textrm{sim}_{\textrm{cos}}\!\left(\mathbf{y}_{\textrm{clean}}^{n},\,\textrm{MLP}(\mathbf{h}_{t}^{k,n})\right)\right],$ (3) $$

where $N$ is the total number of patches, $\epsilon$ is the noise level at timestep $t$, and MLP is a trainable multilayer perceptron that projects $\mathbf{h}_{t}^{k,n}$ to adaptively align it with the representation space of $\mathbf{y}_{\textrm{clean}}^{n}$, while $\textrm{sim}_{\textrm{cos}}(\cdot,\cdot)$ denotes the cosine similarity function. This regularization term is incorporated into the original diffusion-based objective in Eq. [2](https://arxiv.org/html/2509.21565v1#S2.E2) during training, yielding the combined loss:

$$ $\mathcal{L}_{\text{REPA}}=\mathcal{L}_{\text{velocity}}+\lambda\cdot\mathcal{L}_{\text{repa}}.$ (4) $$

Large-scale pretrained encoders (*e.g*., DINOv2) are optimized not only with image-level but also with patch-level objectives, enabling them to learn more robust representations (Zhou et al., 2022; Oquab et al., 2024). From the perspective of the patch-level manifold space (Hao et al., 2022), REPA leverages patch-wise representation alignment to capture these finer-grained structures, thereby maximizing the effectiveness of alignment with such pretrained encoders.

### 2.3 Linear Probing

Linear probing evaluations (Alain & Bengio, 2016) were originally proposed to analyze deep neural networks by measuring the degree of linear separability of features across layers. In this approach, the model under investigation is frozen. Given an input $\mathbf{x}$, the neural network encoder extracts features $\mathbf{h}^{k}$ at the $k^{\text{th}}$ layer, analogous to Section [2.2](https://arxiv.org/html/2509.21565v1#S2.SS2), which are subsequently reduced via global pooling to yield the representation $\mathbf{g}^{k}$. A linear classifier $f^{k}$ is then trained on these pooled features to solve the classification task by minimizing the cross-entropy loss:

$$ $\mathcal{L}_{\text{class}}=-\mathbb{E}_{(\mathbf{x},\mathbf{c})}\big[\,\mathbf{c}_{\textrm{gt}}^{\top}\log\big(f^{k}(\mathbf{g}^{k})\big)\big],$ (5) $$

where $\mathbf{c}_{\textrm{gt}}$ is ground-truth label and $f^{k}(\mathbf{g}^{k})=\text{softmax}(\mathbf{W}\mathbf{g}^{k}+\mathbf{b})$ with $(\mathbf{W},\mathbf{b})$ as probe parameters.

Recent studies have demonstrated that both transformer-based and CNN-based diffusion models learn semantically meaningful representations through the denoising task in training (Baranchuk et al., 2022; Chen et al., 2025; Xiang et al., 2023), as evidenced by evaluations of their linear separability via linear probing. Furthermore, alignment-based methods have employed linear probing to quantify the improvements in the linear separability of the target layers’ representations, demonstrating how such enhancements correlate with training efficiency (Jiang et al., 2025; Yu et al., 2025).

## 3 LSEP: Promoting Linear Separability of Denoising Networks

We introduce a simple yet novel trainable linear probe into an intermediate layer of the diffusion model, in the spirit of linear probing evaluations, but train it *without freezing* its model parameters. We conceptualize our proposed method as a harmonious integration of diffusion and linear classification tasks, designed to enhance both simultaneously, as described in Fig. [1](https://arxiv.org/html/2509.21565v1#S0.F1). All intermediate layers, along with a linear probe inserted at a pre-specified depth, are simultaneously trained to enhance linear separability of the intermediate representations.
The entire diffusion model leverages these well-separated features to perform effective denoising.

Consequently, our model should be carefully designed to accommodate the coexistence of these two distinct tasks. To this end, we introduce several novel training strategies: (1) design of the linear probe branch, (2) task-specific conditioning of the shared intermediate layers, (3) random cropping for improving patch-level linear separability, and (4) time-dependent weighting scheduling of the linear probing loss. Each subsequent subsection provides details on these components, and their respective results are presented in Section [4.3](https://arxiv.org/html/2509.21565v1#S4.SS3).

### 3.1 Linear Classifier Built on Diffusion Model

Figure: Figure 2: Architectures of the linear classifier.
Refer to caption: Figures/linear_probe.jpg

As illustrated in Fig. [2](https://arxiv.org/html/2509.21565v1#S3.F2), our linear classifier consists of a normalization module followed by a linear layer, similar in spirit to prior works (Alain & Bengio, 2016; He et al., 2022; Chen et al., 2025; Yu et al., 2025). Unlike these works, which may or may not employ batch normalization, we adopt layer normalization to ensure stable training in the presence of the denoising objective. Thus, intermediate features extracted from the SiT model are globally aggregated via average pooling for dimensionality reduction, followed by layer normalization and a linear classification head. The diffusion transformer is trained up to a specified target depth jointly with the classifier. This joint optimization encourages early-stage representations to become more linearly separable, guided by supervision from the classifier. The classification loss is given in Eq. [5](https://arxiv.org/html/2509.21565v1#S2.E5), with the only modification being the addition of layer normalization.

Our proposed total loss, which incorporates the scaling factor $\omega_{\text{class}}$ for the classification term to control the trade-off between classification and denoising tasks, together with Eq. [2](https://arxiv.org/html/2509.21565v1#S2.E2), is defined as:

$$ $\mathcal{L}_{\text{LSEP}}=\mathcal{L}_{\text{velocity}}+\omega_{\text{class}}\cdot\mathcal{L}_{\text{class}}.$ (6) $$

### 3.2 Conditioning to Linear Probe Branch

In practice, conditioning of the diffusion transformer consists of time and class embeddings, which are summed and applied through projection operators. For time embedding, we use the same timestep conditioning for both tasks, which helps the classifier to align with the temporal dynamics of the denoising process.
However, we apply a different class conditioning to the linear probe branch, as indicated by the blue path in Fig. [1](https://arxiv.org/html/2509.21565v1#S0.F1) to enhance linear separability of intermediate features. This is because the class embedding is a function of the class label, $\mathbf{c}_{\textrm{class}}$. Thus, this information may cause shortcut learning, where the linear classifier relies directly on class conditioning to perform classification rather than learning linearly separable representations.

To mitigate this, we assign the unconditional class label ($\mathbf{c}_{\varnothing}$) as $\mathbf{c}_{\textrm{class}}$ with probability $\rho_{L}$, chosen to be close to 1. This withholds the class information from the classifier and prevents it from relying solely on class conditioning. With probability $1-\rho_{L}$, we assign another class label ${\bf c}\in{\cal C}-\{\mathbf{c}_{\varnothing}\}$ as $\mathbf{c}_{\textrm{class}}$, which exposes the classifier to non–affine-transformed features. Note this approach is different from the class conditioning used for the denoising task of the diffusion model, where non-null class information is included in the class conditioning with higher probability, $1-\rho_{D}$. Our combined strategy encourages feature representations in the early stages to capture meaningful semantics.

### 3.3 Random Cropping for Improving Patch-level Linear Separability

Figure: Figure 3: t-SNE 3D projection visualization of the patch-level manifold space of SiT-XL using baseline, REPA, and our method with and without random cropping (400K iters, $t$=0.7). Each patch within the $8^{th}$ layer intermediate features is represented by a square, with distinct colors for 7 classes.
Refer to caption: Figures/3d_patch_level.jpg

Our linear probe uses a summary statistic of the intermediate features, obtained via global pooling, to perform its classification task. While this ensures linear separability of the pooled feature vector, the patch-level features may not be as optimally separated in the patch-level manifold space (Hao et al., 2022). To further investigate this, we performed a t-SNE (Maaten & Hinton, 2008) 3D visualization in the patch-level manifold space, as illustrated in Fig. [3](https://arxiv.org/html/2509.21565v1#S3.F3). In Fig. [3](https://arxiv.org/html/2509.21565v1#S3.F3) (b), the alignment-based method (REPA) forms well-separated clusters by leveraging patch-wise alignment with an extensively pre-trained external encoder. Our method (Fig. [3](https://arxiv.org/html/2509.21565v1#S3.F3) (c)), optimized with globally pooled features, achieves improved clustering compared to the baseline (a). However, certain patches (in light orange and light blue) exhibit a dispersed distribution.

To address this, we randomly crop the intermediate feature map into $n\times n$ patches and compute the mean over this subset before feeding it to the linear classifier as illustrated in Fig. [2](https://arxiv.org/html/2509.21565v1#S3.F2). Concretely, we first reshape the intermediate features from $\mathbf{h}^{k}\in\mathbb{R}^{T\times D}$ to $\mathbb{R}^{t\times t\times D}$, where $T$ is the length of features, $D$ is the channel dimension, and $t^{2}=T$. We then randomly crop the features to $\mathbb{R}^{n\times n\times D}$, with $n\leq t$. This strategy enhances the separability of both the mean vectors and subsets of patch-level features, while also introducing diversity similar to data augmentation. As illustrated in Fig. [3](https://arxiv.org/html/2509.21565v1#S3.F3) (d), this promotes clearer cluster formation and further contributes to improved denoising training.

### 3.4 Time-Dependent Weighting Scheduling of ℒ c ​ l ​ a ​ s ​ s \mathcal{L}_{class}

One of the major differences between standard classification and classification within our designed diffusion model is that the inputs are combined with different noise levels, resulting in a multitude of input distributions. Although the time embedding effectively tracks these variations and guides the corresponding denoising tasks, a single linear probe has limitations in classifying across such diverse distributions. While incorporating a time-dependent embedding into the classifier or using multiple classifiers for different time steps may be viable, our focus is not on enhancing the classifier head. Instead, we maintain a single classifier and introduce diversity through weight scheduling to improve the linear separability of intermediate transformer blocks.

To this end, we apply a time-dependent piecewise constant weighting scheme to the linear probe loss, allowing the linear probe to assign different weights according to the noise level and thus learn more effectively. This approach divides the different noise level distributions into $k$ groups, allowing the linear probe to be optimized for each group. It also assigns larger weights to higher noise levels, which helps the classifier perform more effectively on noisier inputs. The time-dependent weight, $\omega_{\text{class}}(t)$ with $k$ bins is defined as:

$$ $\omega_{\text{class}}(t,k)=\omega_{\text{start}}+\lfloor t\cdot k\rfloor\cdot\Delta\omega,\quad t\in[0,1]$ (7) $$

where $\Delta\omega=(\omega_{\text{max}}-\omega_{\text{min}})/k$, and $\omega_{\text{min}}$ and $\omega_{\text{max}}$ are the minimum and maximum weighting values, respectively.
We denote this as
$\omega_{\text{class}}(t,k)=[\omega_{\text{start}},\omega_{\text{end}}]_{k\text{ bins}}$
for future reference.

**Table 1: Performance comparison of the proposed key training strategies on the SiT-L/2 model. Each strategy is distinguished by a different color, and dark colors indicate the best option within each component. The symbols $\uparrow$ and $\downarrow$ indicate that higher and lower values are better, respectively.**
| \rowcolorblack!10<br>Iter. | Uncond. Prob. | Target | Rand. Crop ($n\!\times\!n$) | $\omega_{\text{class}}$ | FID$\downarrow$ | sFID$\downarrow$ | IS$\uparrow$ | Pre$\uparrow$ | Rec$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| \rowcoloryellow!10<br>400K | Baseline SiT-L/2 | 18.5 | 5.20 | 72.2 | 0.64 | 0.63 |  |  |  |
| \arrayrulecolorgray<br>400K | \cellcolorblue!10 0.1 | 8 | $n=16$ | 0.03 | 20.6 | 5.47 | 70.2 | 0.62 | 0.65 |
| \cellcolorblue!10 0.8 | 8 | 13.1 | 5.33 | 95.2 | 0.67 | 0.63 |  |  |  |
| \cellcolorblue!20 0.9 | 8 | 12.9 | 5.37 | 96.2 | 0.67 | 0.64 |  |  |  |
| \cellcolorblue!10 1.0 | 8 | 14.1 | 5.34 | 91.0 | 0.66 | 0.64 |  |  |  |
| \arrayrulecolorgray<br>400K | 0.9 | \cellcolororange!10 6 | $n=16$ | 0.03 | 12.8 | 5.41 | 96.2 | 0.67 | 0.64 |
| \cellcolororange!30 7 | 12.7 | 5.34 | 98.1 | 0.67 | 0.64 |  |  |  |  |
| \cellcolororange!10 8 | 12.9 | 5.37 | 96.2 | 0.67 | 0.64 |  |  |  |  |
| \cellcolororange!10 9 | 13.9 | 5.37 | 91.5 | 0.66 | 0.65 |  |  |  |  |
| \cellcolororange!10 10 | 13.4 | 5.43 | 93.0 | 0.66 | 0.64 |  |  |  |  |
| \arrayrulecolorgray<br>400K | 0.9 | 7 | $n=16$ | 0.03 | 12.7 | 5.34 | 98.1 | 0.67 | 0.64 |
| 7 | \cellcolorgreen!10 $n\in[11,16]\cap\mathbb{Z}$ | 0.03 | 13.0 | 5.33 | 95.8 | 0.67 | 0.64 |  |  |
| 7 | \cellcolorgreen!30 $\mathbf{n}\in[\mathbf{12},\mathbf{16}]\cap\mathbb{Z}$ | 0.03 | 12.5 | 5.34 | 98.8 | 0.67 | 0.64 |  |  |
| 7 | \cellcolorgreen!10 $n\in[13,16]\cap\mathbb{Z}$ | 0.03 | 13.2 | 5.33 | 95.2 | 0.67 | 0.64 |  |  |
| 7 | \cellcolorgreen!10 $n\in[14,16]\cap\mathbb{Z}$ | 0.03 | 12.5 | 5.37 | 97.6 | 0.68 | 0.64 |  |  |
| \arrayrulecolorgray<br>400K | 0.9 | 7 | $n=16$ | 0.03 | 12.7 | 5.34 | 98.1 | 0.67 | 0.64 |
| 7 | $n=16$ | \cellcolorcyan!10 $[0.02,0.03]_{\text{10 Bins}}$ | 12.5 | 5.33 | 97.4 | 0.67 | 0.64 |  |  |
| 7 | $n=16$ | \cellcolorcyan!10 $[0.02,0.03]_{\text{5 Bins}}$ | 12.5 | 5.36 | 98.6 | 0.67 | 0.64 |  |  |
| 7 | $n\in[12,16]\cap\mathbb{Z}$ | \cellcolorcyan!10 $[0.02,0.025]_{\text{10 Bins}}$ | 12.5 | 5.32 | 97.9 | 0.67 | 0.63 |  |  |
| 7 | $n\in[12,16]\cap\mathbb{Z}$ | \cellcolorcyan!10 $[0.0275,0.0325]_{\text{5 Bins}}$ | 12.6 | 5.41 | 97.7 | 0.67 | 0.64 |  |  |
| 7 | $n\in[12,16]\cap\mathbb{Z}$ | \cellcolorcyan!30 $[\mathbf{0.0275},\mathbf{0.0325}]_{\text{10 Bins}}$ | 12.3 | 5.40 | 98.3 | 0.68 | 0.64 |  |  |
| \arrayrulecolorblack |  |  |  |  |  |  |  |  |  |

## 4 Experimental Setup

### 4.1 Implementation details

We closely follow the experimental setup of SiT (Ma et al., 2024) and REPA (Yu et al., 2025), unless stated otherwise. All models are trained and evaluated on ImageNet (Deng et al., 2009), where images are preprocessed to a resolution of 256$\times$256 following the data preprocessing protocol of ADM (Dhariwal & Nichol, 2021). We employ the Base, Large, and X-Large model variants introduced in SiT (Ma et al., 2024), all configured with a patch size of 2. We insert a linear classifier after the 4^th, 7^th, and 8^th layers in the SiT-B/2, SiT-L/2, and SiT-XL/2 models, respectively.

To ensure fair comparison with prior work (Ma et al., 2024; Yu et al., 2025), we train all models using a consistent global batch size of 256. Optimization is performed using AdamW (Kinga et al., 2015; Loshchilov & Hutter, 2019) with a constant learning rate of $1\times 10^{-4}$ for training the diffusion model. We provide the detailed hyperparameters in Appendix [A](https://arxiv.org/html/2509.21565v1#A1).

### 4.2 Evaluation Protocol

For evaluating image generation quality, we strictly follow the ADM evaluation protocol (Dhariwal & Nichol, 2021). We report several standard metrics, including Fréchet Inception Distance (FID) (Heusel et al., 2017), Structural FID (sFID) (Nash et al., 2021), Inception Score (IS) (Salimans et al., 2016), and Precision and Recall (Kynkäänniemi et al., 2019), all computed using 50K generated samples. Following the sampling procedures from SiT (Ma et al., 2024) and REPA (Yu et al., 2025), we adopt the SDE-based Euler–Maruyama sampler with 250 steps. All evaluations are performed on 50K validation images from the ImageNet dataset (Deng et al., 2009), resized to 256$\times$256 resolution.

### 4.3 Analysis of Key Strategies

We present detailed experimental results for the SiT-L/2 model in Tab. [3.4](https://arxiv.org/html/2509.21565v1#S3.SS4), effectively serving as an ablation study for the training innovations outlined in Section [3](https://arxiv.org/html/2509.21565v1#S3).

Class Conditioning for Linear Probe
The unconditioning probability $\rho_{L}$ indicates the proportion of the null class $\mathbf{c}_{\varnothing}$, as in Section [3.2](https://arxiv.org/html/2509.21565v1#S3.SS2). When the unconditioning probability is set to 0.1, matching that of the denoising model, the linear probe relies on the conditioning label for shortcut learning. This not only fails to improve the linear separability of intermediate features but also degrades generative performance. Conversely, setting it to $\rho_{L}=1.0$ (*i.e*., using only $\mathbf{c}_{\varnothing}$) prominently improves generative performance. However, this also causes the class conditioning for the classifier and the denoiser to be learned independently, leading to inconsistencies in class representation. Therefore, we find that using $\rho_{L}=0.9$ provides the optimal conditioning ratio. Further analysis is provided in Appendix [B.1](https://arxiv.org/html/2509.21565v1#A2.SS1).

Target Depth Numerous works have demonstrated that efficient training of diffusion models depends on representation quality in early layers (Yu et al., 2025; Jiang et al., 2025). Our results align with these insights, showing that the optimal depth for incorporating the linear classifier corresponds to the shallow layers of the model. When the depth is shifted toward the middle layers (*e.g*., 9 or 10 for SiT-L/2), the generative performance degrades, as the reduced capacity for denoising limits overall effectiveness. Conversely, connecting to very early layers also hampers denoising training. We empirically find that layer 7 yields the best results for SiT-L, and adopt this configuration in our experiments. For SiT-B and SiT-XL, the optimal depths are found to be layers 4 and 8, respectively.

Random Cropping
To enhance linear separability in the patch-level manifold space as shown in Fig. [3](https://arxiv.org/html/2509.21565v1#S3.F3), we investigated the effects of random cropping with varying box sizes. We found that selecting a box size randomly from the range 12 to 16 yielded the best results. This introduces variability, providing more diverse feature samples for the linear classifier and improving separability not only of pooled features but also within the patch-level manifold space.

Time-dependent $\omega_{\text{class}}$
The first three rows compare a time-dependent $\omega_{\text{class}}$ to a constant $\omega_{\text{class}}$ in the absence of random cropping. This demonstrates that assigning different weights to $\mathcal{L}_{\text{class}}$ over time, using a piecewise constant schedule with multiple bins for training a linear classifier improves the generative results. The last three rows further examine different weight intervals and numbers of bins when time-dependent $\omega_{\text{class}}$ is combined with random cropping approach. The results show $[0.0275,0.0325]_{\text{10 Bins}}$ achieves the best performance, and is used for the remainder of the study.

## 5 Results

### 5.1 Linear Separability

Figure: Figure 4: (a) Linear probing evaluation results on baseline SiT-XL and SiT-XL+LSEP with $t=0.1$. (b) t-SNE 3D projections for baseline SiT-XL (left), alignment-based REPA (middle), and proposed method (right). Intermediate features are extracted from the $8^{\text{th}}$ layer of each pre-trained SiT-XL model with $t=0.1$. Each feature is globally pooled and sampled from 20 randomly selected classes, with 100 samples per class. The representations are then projected into 3D space, with each class visualized using a distinct color. For all methods, the 400K-iteration checkpoint was used.
Refer to caption: Figures/linear_sep.jpg

We evaluate the effect of promoting linear separability using linear probing, 3D projections, and PCA visualization for SiT-XL with our proposed method. Fig. [4](https://arxiv.org/html/2509.21565v1#S5.F4) (a) shows that LSEP achieves higher overall linear probing accuracy across layers, particularly at the early stages. Unsurprisingly, linear probing performance is substantially improved using the LSEP training strategy that jointly targets a linear classification task. In Fig. [4](https://arxiv.org/html/2509.21565v1#S5.F4) (b), our proposed method produces well-defined clusters, comparable to those observed in REPA, which indicates that the classes are clearly linearly separable, even though no external encoder information is used. We also provide PCA visualizations in Appendix [B.3](https://arxiv.org/html/2509.21565v1#A2.SS3), which demonstrate that LSEP preserves clearly separable components across varying noise levels.

### 5.2 Quantitative Results

#### Analysis Across Model Sizes.

Tab. LABEL:tab:table_results2 summarizes generative performance of models of different sizes with different training paradigms. Note the number of parameters for REPA *excludes* the additional parameters related to the external encoder. Our method consistently improves the generation FID compared to the baseline SiT baseline trained for the same 400K iterations. Notably, FID of 12.3 (Tab. LABEL:tab:table_results2) and IS of 98.3 (Tab. [3.4](https://arxiv.org/html/2509.21565v1#S3.SS4)) achieved by our LSEP (SiT-L) outperform the 12.5 and 90.7, respectively, reported for REPA (SiT-L) (Yu et al., 2025) with pretrained external MAE-L encoder (He et al., 2022), which contains an additional 304M parameters in addition to the transformer diffusion parameters.

Figure: Table 2: FID comparisons on $256\times 256$ ImageNet, without employing CFG. $\downarrow$ indicates lower values are better.

While LSEP shows substantial gains in early stages of trainings, we emphasize that LSEP continues to steadily improve FID throughout training, with consistent gains: by 4M iterations, its performance approaches that of REPA, while remaining well beyond the reach of baseline training at 7M iterations, as illustrated in Fig. [1](https://arxiv.org/html/2509.21565v1#S0.F1) (b). Furthermore, our approach achieves these promising results *without relying on any pre-trained external models or explicit representation alignment*. Selected qualitative results for SiT-XL/2 using our method are presented in Fig. [5](https://arxiv.org/html/2509.21565v1#S5.F5). More qualitative results are provided in Appendix [D](https://arxiv.org/html/2509.21565v1#A4).

**Table 3: System-level comparison on ImageNet $256\times 256$ with CFG. Within each architecture, the best on each metric are bolded. Arrows indicate whether lower $(\downarrow)$ or higher $(\uparrow)$ values are better.**
| \rowcolorblack!10 Model | Epochs | Tokenizer | FID$\downarrow$ | sFID$\downarrow$ | IS$\uparrow$ | Pre.$\uparrow$ | Rec.$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Pixel diffusion |  |  |  |  |  |  |  |
| ADM-U (Dhariwal & Nichol, 2021) | 400 | - | 3.94 | 6.14 | 186.7 | 0.82 | 0.52 |
| VDM++ (Kingma & Gao, 2023) | 560 | - | 2.40 | - | 225.3 | - | - |
| Simple diffusion (Hoogeboom et al., 2023) | 800 | - | 2.77 | - | 211.8 | - | - |
| CDM (Ho et al., 2022) | 2160 | - | 4.88 | - | 158.7 | - | - |
| \arrayrulecolorgray       Latent diffusion, U-Net |  |  |  |  |  |  |  |
| LDM-4 (Rombach et al., 2022) | 200 | LDM-VAE | 3.60 | - | 247.7 | 0.87 | 0.48 |
| \arrayrulecolorgray       Latent diffusion, Transformer *with* pre-trained external encoder |  |  |  |  |  |  |  |
| SiT + REPA (Yu et al., 2025) | 800 | SD-VAE | 1.42 | 4.70 | 305.7 | 0.80 | 0.65 |
| LightningDiT (Yao et al., 2025) | 800 | VA-VAE | 1.35 | 4.15 | 295.3 | 0.79 | 0.65 |
| REPA-E (Leng et al., 2025) | 800 | VA-VAE | 1.26 | 4.11 | 314.9 | 0.79 | 0.66 |
| \arrayrulecolorgray       Latent diffusion, Transformer *without* pre-trained external encoder |  |  |  |  |  |  |  |
| DiT-XL/2 (Peebles & Xie, 2023) | 1400 | SD-VAE | 2.27 | 4.60 | 278.2 | 0.83 | 0.57 |
| SiT-XL/2 (Ma et al., 2024) | 1400 | SD-VAE | 2.06 | 4.50 | 270.3 | 0.82 | 0.59 |
| SD-DiT (Zhu et al., 2024) | 480 | SD-VAE | 3.23 | - | - | - | - |
| MaskDiT (Zheng et al., 2024) | 1600 | SD-VAE | 2.28 | 5.67 | 276.6 | 0.80 | 0.61 |
| DiT + TREAD (Krause et al., 2025) | 740 | SD-VAE | 1.69 | 4.73 | 292.7 | 0.81 | 0.63 |
| SiT + SRA (Jiang et al., 2025) | 800 | SD-VAE | 1.58 | 4.65 | 311.4 | 0.80 | 0.63 |
| \rowcolorbrown!30 SiT + LSEP (ours) | 800 | SD-VAE | 1.46 | 4.94 | 296.8 | 0.80 | 0.64 |

#### Combining Alignment-Based Methods with LSEP.

REPA improves the representations by performing token-wise alignment, while our LSEP enhances linear separability by mainly utilizing mean representation vectors. Both methods individually improve linear separability; however, when used together, they are expected to provide more powerful representation geometry and improved training efficiency. As shown in Tab. [2](https://arxiv.org/html/2509.21565v1#S5.T2), combining REPA with LSEP consistently drives FID even lower, proving that this synergy not only further accelerates training but also enhances generative quality. Additional analysis using patch-level 3D projection visualizations is presented in Appendix [B.2](https://arxiv.org/html/2509.21565v1#A2.SS2).

#### System-Level Comparison.

In Tab. [5.2](https://arxiv.org/html/2509.21565v1#S5.SS2.SSS0.Px1), we report evaluation results using CFG (Ho & Salimans, 2021) with the guidance interval (Kynkäänniemi et al., 2024) (please see Appendix [C](https://arxiv.org/html/2509.21565v1#A3) for details). Our method achieves a best FID of 1.46 among models that do not use an external visual encoder and is comparable to alignment-based models, such as REPA, which rely on pretrained external encoders.

## 6 Discussion

Limitations. The scope of this paper does not cover advanced generative settings such as text-to-image generation, video generation, and 2K higher-resolution images. Further investigation is warranted to evaluate the applicability and scalability of our framework in these scenarios.

Future Directions.
Similar to how REPA variants (Yao et al., 2025; Leng et al., 2025) improve generative performance by tuning the VAE through representation alignment, LSEP inherently offers an alternative means to achieve this. Furthermore, since most vision foundation models are transformer-based, prior studies on representation alignment have focused mainly on transformer diffusion models. In contrast, LSEP can also be applied to CNN-based models, broadening its applicability. These ideas will be investigated in future studies.

## 7 Conclusion

In this paper, we proposed LSEP, a regularization approach that independently enhances the linear separability of generative models without relying on large-scale pre-trained encoders or representation alignment. We demonstrated that linear separability is a core principle of diffusion training and repurposed linear probing, which is typically used for post-hoc evaluation, as an effective training tool. Our results show that LSEP improves both training efficiency and generative performance, achieving a state-of-the-art FID score of 1.46 with a single model, without the need for any alignment.

## References

- Alain & Bengio (2016)
Guillaume Alain and Yoshua Bengio.
Understanding intermediate layers using linear classifier probes.
In *Proc. Int. Conf. Learn. Represent.*, 2016.
- Baranchuk et al. (2022)
Dmitry Baranchuk, Andrey Voynov, Ivan Rubachev, Valentin Khrulkov, and Artem Babenko.
Label-efficient semantic segmentation with diffusion models.
In *Proc. Int. Conf. Learn. Represent.*, 2022.
- Chen et al. (2025)
Xinlei Chen, Zhuang Liu, Saining Xie, and Kaiming He.
Deconstructing denoising diffusion models for self-supervised learning.
In *Proc. Int. Conf. Learn. Represent.*, 2025.
- 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 *Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recog.*, pp. 248–255, 2009.
- Dhariwal & Nichol (2021)
Prafulla Dhariwal and Alexander Nichol.
Diffusion models beat GANs on image synthesis.
In *Proc. Adv. Neural Inf. Process. Syst.*, pp. 8780–8794, 2021.
- Hao et al. (2022)
Zhiwei Hao, Jianyuan Guo, Ding Jia, Kai Han, Yehui Tang, Chao Zhang, Han Hu, and Yunhe Wang.
Learning efficient vision transformers via fine-grained manifold distillation.
In *Proc. Adv. Neural Inf. Process. Syst.*, 2022.
- 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 *Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recog.*, pp. 16000–16009, 2022.
- Heusel et al. (2017)
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.
GANs trained by a two time-scale update rule converge to a local nash equilibrium.
In *Proc. Adv. Neural Inf. Process. Syst.*, 2017.
- Ho & Salimans (2021)
Jonathan Ho and Tim Salimans.
Classifier-free diffusion guidance.
In *Proc. NeurIPS Workshop DGMs Appl.*, 2021.
- Ho et al. (2022)
Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans.
Cascaded diffusion models for high fidelity image generation.
*J. Mach. Learn. Res.*, 23(47):1–33, 2022.
- Hoogeboom et al. (2023)
Emiel Hoogeboom, Jonathan Heek, and Tim Salimans.
Simple diffusion: End-to-end diffusion for high resolution images.
In *Proc. Int. Conf. Mach. Learn.*, pp. 13213–13232, 2023.
- Jiang et al. (2025)
Dengyang Jiang, Mengmeng Wang, Liuzhuozheng Li, Lei Zhang, Haoyu Wang, Wei Wei, Guang Dai, Yanning Zhang, and Jingdong Wang.
No other representation component is needed: Diffusion transformers can provide representation guidance by themselves, 2025.
arXiv:2505.02831.
- Karras et al. (2022)
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine.
Elucidating the design space of diffusion-based generative models.
In *Proc. Adv. Neural Inf. Process. Syst.*, pp. 26565–26577, 2022.
- Kinga et al. (2015)
Diederik Kinga, Jimmy Ba Adam, et al.
A method for stochastic optimization.
In *Proc. Int. Conf. Learn. Represent.*, 2015.
- Kingma & Gao (2023)
Diederik Kingma and Ruiqi Gao.
Understanding diffusion objectives as the elbo with simple data augmentation.
In *Proc. Adv. Neural Inf. Process. Syst.*, pp. 65484–65516, 2023.
- Krause et al. (2025)
Felix Krause, Timy Phan, Ming Gui, Stefan Andreas Baumann, Vincent Tao Hu, and Björn Ommer.
TREAD: Token routing for efficient architecture-agnostic diffusion training, 2025.
arXiv:2501.04765.
- Kynkäänniemi et al. (2019)
Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila.
Improved precision and recall metric for assessing generative models.
In *Proc. Adv. Neural Inf. Process. Syst.*, 2019.
- Kynkäänniemi et al. (2024)
Tuomas Kynkäänniemi, Miika Aittala, Tero Karras, Samuli Laine, Timo Aila, and Jaakko Lehtinen.
Applying guidance in a limited interval improves sample and distribution quality in diffusion models.
In *Proc. Adv. Neural Inf. Process. Syst.*, 2024.
- Leng et al. (2025)
Xingjian Leng, Jaskirat Singh, Yunzhong Hou, Zhenchang Xing, Saining Xie, and Liang Zheng.
REPA-E: Unlocking VAE for end-to-end tuning with latent diffusion transformers, 2025.
arXiv:2504.10483.
- Lipman et al. (2023)
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le.
Flow matching for generative modeling.
In *Proc. Int. Conf. Learn. Represent.*, 2023.
- Loshchilov & Hutter (2019)
Ilya Loshchilov and Frank Hutter.
Decoupled weight decay regularization.
In *Proc. Int. Conf. Learn. Represent.*, 2019.
- 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 *Proc. Eur. Conf. Comput. Vis.*, pp. 23–40, 2024.
- Maaten & Hinton (2008)
Laurens van der Maaten and Geoffrey Hinton.
Visualizing data using t-SNE.
*J. Mach. Learn. Res.*, 9(Nov.):2579–2605, 2008.
- Nash et al. (2021)
Charlie Nash, Jacob Menick, Sander Dieleman, and Peter W Battaglia.
Generating images with sparse representations.
*Proc. Int. Conf. Mach. Learn.*, 2021.
- Oquab et al. (2024)
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy V. Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel HAZIZA, Francisco Massa, Alaaeldin El-Nouby, Mido Assran, Nicolas Ballas, Wojciech Galuba, Russell Howes, Po-Yao Huang, Shang-Wen Li, Ishan Misra, Michael Rabbat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Herve Jegou, Julien Mairal, Patrick Labatut, Armand Joulin, and Piotr Bojanowski.
DINOv2: Learning robust visual features without supervision.
*Trans. Mach. Learn. Res.*, 2024.
ISSN 2835-8856.
- Peebles & Xie (2023)
William Peebles and Saining Xie.
Scalable diffusion models with transformers.
In *Proc. IEEE/CVF Int. Conf. Comput. Vis.*, pp. 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 *Proc. Int. Conf. Mach. Learn.*, pp. 8748–8763, 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 *Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recog.*, pp. 10684–10695, 2022.
- Salimans et al. (2016)
Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen.
Improved techniques for training GANs.
In *Proc. Adv. Neural Inf. Process. Syst.*, 2016.
- Xiang et al. (2023)
Weilai Xiang, Hongyu Yang, Di Huang, and Yunhong Wang.
Denoising diffusion autoencoders are unified self-supervised learners.
In *Proc. IEEE/CVF Int. Conf. Comput. Vis.*, pp. 15802–15812, 2023.
- Yao et al. (2025)
Jingfeng Yao, Bin Yang, and Xinggang Wang.
Reconstruction vs. generation: Taming optimization dilemma in latent diffusion models.
In *Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recog.*, pp. 15703–15712, 2025.
- Yu et al. (2025)
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.
In *Proc. Int. Conf. Learn. Represent.*, 2025.
- Zheng et al. (2024)
Hongkai Zheng, Weili Nie, Arash Vahdat, and Anima Anandkumar.
Fast training of diffusion models with masked transformers.
In *Trans. Mach. Learn. Res.*, 2024.
- Zhou et al. (2022)
Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong.
Image BERT pre-training with online tokenizer.
In *Proc. Int. Conf. Learn. Represent.*, 2022.
- Zhu et al. (2024)
Rui Zhu, Yingwei Pan, Yehao Li, Ting Yao, Zhenglong Sun, Tao Mei, and Chang Wen Chen.
SD-DiT: Unleashing the power of self-supervised discrimination in diffusion transformer.
In *Proc. IEEE/CVF Conf. Comput. Vis. Pattern Recog.*, pp. 8435–8445, 2024.

## Appendix A Implementation details

**Table 4: Hyperparameter setup for different SiT models and LSEP.**
|  | SiT-B/2 | SiT-L/2 | SiT-XL/2 |
| --- | --- | --- | --- |
| Architecture |  |  |  |
| Input dim. | $32\times 32\times 4$ | $32\times 32\times 4$ | $32\times 32\times 4$ |
| Num. layers | 12 | 24 | 28 |
| Hidden dim. | 768 | 1,024 | 1,152 |
| Num. heads | 12 | 16 | 16 |
| LSEP |  |  |  |
| Uncon. Prob. | 90% | 90% | 90% |
| Target depth | 4 | 7 | 8 |
| Rand. Crop | [14, 16] | [12, 16] | [12, 16] |
| $\omega(t)$ | $[\text{0.005},\text{0.01}]_{10\ \mathrm{bins}}$ | $[\text{0.0275},\text{0.0325}]_{10\ \mathrm{bins}}$ | $[\text{0.0225},\text{0.03}]_{10\ \mathrm{bins}}$ |
| lr for linear probes. | 0.03 | 0.0001 | 0.0001 |
| Optimization |  |  |  |
| Training iteration | 400K | 400K | 400K / 4M (Fig. 1) |
| Batch size | 256 | 256 | 256 |
| Optimizer | AdamW | AdamW | AdamW |
| lr | 0.0001 | 0.0001 | 0.0001 |
| $(\beta_{1},\beta_{2})$ | (0.9,0.999) | (0.9,0.999) | (0.9,0.999) |
| Interpolants |  |  |  |
| $\alpha_{t}$ | $1-t$ | $1-t$ | $1-t$ |
| $\sigma_{t}$ | $t$ | $t$ | $t$ |
| $w_{t}$ | $\sigma_{t}$ | $\sigma_{t}$ | $\sigma_{t}$ |
| Training objective | v-prediction | v-prediction | v-prediction |
| Sampler | Euler-Maruyama | Euler-Maruyama | Euler-Maruyama |
| Sampling steps | 250 | 250 | 250 |
| Guidance | - | - | $\omega_{cfg}=1.8,\text{Interval : }[0,0.65]$ (Tab.3) |

Strictly following the experimental setups of SiT (Ma et al., 2024) and REPA (Yu et al., 2025) for the denoising loss, we integrate the linear probe into our model and employ the hyperparameters listed in Tab. LABEL:tab:table_hyperparameters. We use an increased learning rate for the linear classifier in the SiT-B/2 model, which significantly improves both training efficiency and generative performance. On the other hand, for the Large and X-Large models, this adjustment does not lead to improvements in generative results. All experiments were conducted using 4 NVIDIA A100 GPUs.

## Appendix B Additional Analysis

### B.1 Conditioning for linear probe branch

Figure: Figure 6: Comparison of $\mathcal{L}_{\text{class}}$ across training epochs with varying conditioning ratios in the linear probe branch.
Refer to caption: Figures/Appdx_iter_vs_loss_class.jpg

As discussed in Section [4.3](https://arxiv.org/html/2509.21565v1#S4.SS3), conditioning for linear probing plays an important role in improving both linear separability and generation quality. When the unconditioning probability $\rho_{L}$ in Section [3.2](https://arxiv.org/html/2509.21565v1#S3.SS2) is set to 0.1, matching that of the denoising model, the linear probe tend to rely on shortcut learning, as illustrated by the yellow curve in Fig. [6](https://arxiv.org/html/2509.21565v1#A2.F6). Moreover, as the unconditioning ratio increases, the classifier’s reliance on class information gradually diminishes, leading to a more progressive learning process. However, as mentioned earlier, when the conditioning is fully assigned to the unconditional class (*i.e*., $\rho_{L}$ = 1.0), a mismatch arises with the conditioning of the denoising model, which prevents achieving optimal performance. Therefore, incorporating a small proportion of the conditional class $\mathbf{c}_{\text{class}}$ is necessary to obtain the best results.

### B.2 Further Analysis of Combining LSEP and REPA

Figure: Figure 7: t-SNE 3D projection visualization of baseline SiT-XL, SiT-XL with our method, SiT-XL with REPA, and SiT-XL with REPA combined with our method in the patch-level manifold space (400K iterations, $t=0.7$), using the same settings as in Fig. [3](https://arxiv.org/html/2509.21565v1#S3.F3).
Refer to caption: Figures/Appdx_REPA_Ours.jpg

In Section [5.2](https://arxiv.org/html/2509.21565v1#S5.SS2), we showed that combining LSEP with the alignment-based method (REPA) further improves training efficiency and generative performance of the latter. The patch-level linear separability provided by REPA, together with the mean-vector linear separability from LSEP, enhances the learned representations and thereby strengthens training. As shown in Fig. [7](https://arxiv.org/html/2509.21565v1#A2.F7), in t-SNE 3D projection, the combination of REPA and LSEP results in the strongest class-wise separability.

### B.3 PCA visualization

Figure: Figure 8: PCA visualizations at various noise levels and layers within the SiT-XL models. Checkpoints at 400K iterations of baseline SiT-XL and SiT-XL + LSEP are used for the visualizations.
Refer to caption: Figures/pca_vis.jpg

We plot the PCA visualization of intermediate features. These demonstrate that LSEP preserves distinctly separable components across different noise levels., with particularly clear separation in the early-stage layers and under higher noise conditions (*e.g*., layers up to 14, $t=0.6$ and $0.8$).

**Table 5: Evaluation results of SiT-XL + LSEP at 4M iteration with different CFG weight and guidance interval.**
| Model | #Params | Iter. | $w_{cfg}$ | Interval | FID$\downarrow$ | sFID$\downarrow$ | IS$\uparrow$ | Pre.$\uparrow$ | Rec.$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| \rowcoloryellow!10 SiT-XL/2 (Ma et al., 2024) | 675M | 7M | [0, 1] | 1.50 | 2.06 | 4.50 | 270.3 | 0.82 | 0.59 |
| + LSEP (ours) | 675M | 4M | 1.95 | [0, 0.75] | 2.18 | 4.55 | 344.3 | 0.83 | 0.60 |
| + LSEP (ours) | 675M | 4M | 1.95 | [0, 0.70] | 1.81 | 4.68 | 328.6 | 0.82 | 0.62 |
| + LSEP (ours) | 675M | 4M | 1.95 | [0, 0.65] | 1.59 | 4.87 | 312.2 | 0.81 | 0.63 |
| + LSEP (ours) | 675M | 4M | 1.95 | [0, 0.60] | 1.49 | 5.15 | 296.7 | 0.79 | 0.64 |
| + LSEP (ours) | 675M | 4M | 1.90 | [0, 0.60] | 1.50 | 5.15 | 291.8 | 0.79 | 0.64 |
| + LSEP (ours) | 675M | 4M | 1.85 | [0, 0.60] | 1.56 | 5.10 | 282.6 | 0.79 | 0.65 |
| + LSEP (ours) | 675M | 4M | 1.80 | [0, 0.65] | 1.46 | 4.94 | 296.8 | 0.80 | 0.64 |
| + LSEP (ours) | 675M | 4M | 1.80 | [0, 0.60] | 1.55 | 5.17 | 280.6 | 0.79 | 0.65 |

## Appendix C Detailed evaluation results with the different CFG weights scheduling

We provide detailed evaluation results of SiT-XL + LSEP under different classifier-free guidance schedules in Tab. [5](https://arxiv.org/html/2509.21565v1#A2.T5). With $\omega_{cfg}$ in the interval [0, 0.65], our LSEP achieves the best performance.

## Appendix D Additional Qualitative Results

Figure: Figure 9: Uncurated samples from SiT-XL/2 + LSEP ($\omega_{\text{CFG}}=4.0$, class = loggerhead turtle (33))
Refer to caption: Figures/Appdx_collage_33_loggerhead_turtle.png

Figure: Figure 10: Uncurated samples from SiT-XL/2 + LSEP ($\omega_{\text{CFG}}=4.0$, class = scorpion (71))
Refer to caption: Figures/Appdx_collage_71_scorpion.png

Figure: Figure 11: Uncurated samples from SiT-XL/2 + LSEP ($\omega_{\text{CFG}}=4.0$, class = golden retriever (207))
Refer to caption: Figures/Appdx_collage_207_golden_retriever.png

Figure: Figure 12: Uncurated samples from of SiT-XL/2 + LSEP ($\omega_{\text{CFG}}=4.0$, class = border collie(232))
Refer to caption: Figures/Appdx_collage_232_border_collie.png

Figure: Figure 13: Uncurated samples from SiT-XL/2 + LSEP ($\omega_{\text{CFG}}=4.0$, class = arctic fox (279))
Refer to caption: Figures/Appdx_collage_279_Arctic_fox.png

Figure: Figure 14: Uncurated samples from SiT-XL/2 + LSEP ($\omega_{\text{CFG}}=4.0$, class = wood rabbit (330))
Refer to caption: Figures/Appdx_collage_330_wood_rabbit.png

Figure: Figure 15: Uncurated samples from SiT-XL/2 + LSEP ($\omega_{\text{CFG}}=4.0$, class = panda (388))
Refer to caption: Figures/Appdx_collage_388_panda.png

Figure: Figure 16: Uncurated samples from SiT-XL/2 + LSEP ($\omega_{\text{CFG}}=4.0$, class = castle (483))
Refer to caption: Figures/Appdx_collage_483_castle.png

Figure: Figure 17: Uncurated samples from SiT-XL/2 + LSEP ($\omega_{\text{CFG}}=4.0$, class = valley (797))
Refer to caption: Figures/Appdx_collage_979_valley.png

Figure: Figure 18: Uncurated samples from SiT-XL/2 + LSEP ($\omega_{\text{CFG}}=4.0$, class = space shuttle (812))
Refer to caption: Figures/Appdx_collage_812_space_shuttle.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: See weaknesses above
soundness: 3
strengths: - The paper tackles a concrete and important problem of improving the quality of features learnt by generative models. 
- The intervention is architecturally minimal and easy to adopt, since it adds a single linear head and combines it with the standard denoising loss via a simple weight.
- The experiments show consistent FID gains over baseline SiT and demonstrate complementarity with REPA, with clear tables and figures that place the numbers in context and avoid relying solely on a single configuration. 
- The writing is generally clear about the two-branch setup and provides a useful hyperparameter grid for reproduction.
summary: - The paper proposes LSEP, a simple regularization that inserts a trainable linear probe at an intermediate layer of a diffusion transformer and co-optimizes it with the standard denoising loss, with the goal of directly promoting linear separability of intermediate features during training rather than measuring it post-hoc. 
- On ImageNet 256×256 with SiT backbones, LSEP reduces FID versus the SiT baseline, improves linear-probe accuracy at early layers, and shows complementarity with alignment methods like REPA
- The paper positions LSEP as a no-alignment alternative that can also be combined with alignment (e.g., REPA) to further speed training and improve generation quality, and it discusses limitations such as the lack of experiments on text-to-image, video, or higher resolutions.
weaknesses: - The goal of pretraining is to learn general-purpose, task-agnostic features that transfer across datasets, label spaces, and objectives (e.g., detection, segmentation, retrieval, open-set recognition), but this method explicitly supervises a linear probe on ImageNet class labels during pretraining and pressures intermediate features to be separable with respect to that specific taxonomy. As a result, the learned geometry is optimized for class-conditional ImageNet generation and linearly decoding those same labels, which does not guarantee that the features capture fine-grained structure needed by unrelated downstream tasks and may bias representations toward dataset-specific shortcuts.
- The probe and denoiser use different class-conditioning schedules, which creates gradient-level tension; the paper notes that $ρ_L=1$ causes a mismatch but does not analyze interference between the two schedules or test mitigations such as stop-gradient into shared embeddings or separate class embeddings.
- The paper trains its linear probe on pooled features (and sometimes pooled crops), which can make the average representation linearly separable while leaving individual patch tokens messy. Without token-level evaluations, the claim that LSEP improves patch-space geometry remains unclear and could reflect pooling shortcuts 

- The paper sets the probe’s unconditioning probability very high but below one, yet it does not test whether the method remains stable under label noise or whether those occasional non-null conditioning steps make the probe exploit conditioning artifacts rather than visual content. It will be interesting to see if the reported gains survive mislabeled classes and still hold when the class-conditioning pathway is perturbed or ablated.

---

confidence: 4
contribution: 1
initial_rating: 4
presentation: 3
questions: 1. Please answer the Weaknesses.

2. Table 2 shows SiT + REPA + LSEP. However, in Table 3, the authors report only SiT + LSEP. Can SiT + REPA + LSEP achieve SOTA performance?
soundness: 2
strengths: 1. The idea is simple yet effective: It can even be used with REPA to further accelerate the training of the diffusion transformer.

2. Extensive analysis of the design choices facilitates the practitioners to use this approach.
summary: This paper aims to accelerate the training of diffusion transformers without additional computational costs by using a linear classifier. Specifically, instead of aligning the intermediate features of the diffusion transformers with a pretrained vision encoder, it uses intermediate features to be linearly separable to classify the category of the given image. The authors demonstrate that the proposed method (LSEP) shows competitive results with REPA.
weaknesses: 1. LEAP is limited only to the dataset that has human annotations (i.e., labels). How can we use this method to accelerate training diffusion models that do not have ground-truth labels, e.g., T2I generation (MS-COCO)?

2. Even though REPA uses an external encoder as conditioning, the reported REPA results do not use any external labels (i.e., unsupervised), but LSEP needs labels (i.e., supervised). In some sense, LSEP uses more information to generate images, but shows worse performance and slower training than REPA (in terms of training time).

---

confidence: 4
contribution: 3
initial_rating: 4
presentation: 3
questions: Please refer to the weaknesses regarding the correlation between linear separability and performance, and the generalizability to other conditioning signals.

The current manuscript also reads like a strong empirical study with good performance but leaves many questions unanswered. For example, will other well-known techniques in supervised learning (e.g., label smoothing) be helpful? Will other discriminative tasks (other than a simple classification) also be beneficial? Why does a classifier trained on diffusion backbone lag behind supervised ViTs?
soundness: 2
strengths: - The paper's primary strength lies in its novel and straightforward idea: a linear classifier can directly be an effective training regularizer. This provides a new, simple alternative to the dominant paradigm of external representation alignment.
- Although the idea of supervised training is simple, several tricks like patch-level random cropping are clever and effective.
- The empirical results are convincing. The method demonstrates consistent improvements over baselines. The finding that LSEP and REPA are complementary is also valuable.
summary: This paper proposes a regularization method, Linear SEParability (LSEP), to improve the training efficiency and generation quality of diffusion models, without relying on external pre-trained encoders.

The idea is to incorporate a linear classifier, inspired by linear probing in post-hoc evaluation, directly into diffusion training as an auxiliary objective, encouraging the model to learn more linearly separable intermediate representations. The approach is straightforward but also includes several specific training strategies and hyper-parameter choices.

Experiments on the ImageNet 256x256 dataset show that LSEP accelerates convergence and improves final performance across various model sizes of SiT using the original SD-VAE. It also demonstrates a synergistic effect when combined with existing alignment-based methods like REPA.
weaknesses: - Training cost: The paper claims improvements in training efficiency, demonstrated by training iterations vs. FID scores.
  - However, the proposed method appears to require two separate forward passes through the first few layers, as the two branches require different class-conditioning (Figure 1, Section 3.2).
  - This could introduce significant computational overhead, potentially doubling the cost for a portion of the model, and this overhead is not discussed. A wall-clock time (or compute) vs. FID plot would be helpful.
- On linear separability: The paper is motivated by a claim that enhancing linear separability is the mechanism driving the improvements in generation. However, evidence for this link is sparse.
  - The main results in Table 1 do not report the linear probing accuracy for each configuration, making it impossible to directly correlate the degree of separability with the final FID. If we plot the (acc, FID) points for all 21 configurations in Table 1, will it demonstrate a clear, strong correlation?
  - In Table 2, it is also unclear how REPA and LSEP independently (and jointly) improve linear separability.
  - While Figure 4(a) shows an increase in accuracy, the absolute value remains relatively low (below 60%, much weaker than supervised learning), which warrants further discussion.
- Generalizability to other settings: The method, in its current form, explicitly relies on class labels for the auxiliary loss. This dependency limits its applicability to other tasks such as text-to-image or unconditional generation.
  - It would strengthen the paper to demonstrate how LSEP could be adapted for such scenarios, for instance, by training on class or attribute labels as a regularizer while still performing text-conditional or unconditional generation for the primary task.
  - Can the input conditions and labels for auxiliary classification be decoupled?
- Loss weighting design: Eq. (7) appears somewhat hand-crafted and might be sensitive to changes in the training-time timestep sampling (e.g., LogNorm distribution in SD3 and VA-VAE).
  - The ablation study in Table 1 shows that this parameter is carefully tailored to a very narrow range of values ([0.0275, 0.0325]).
  - The paper also does not show how the general magnitude of this value (around 0.03) is determined, as the cross entropy loss itself is already low enough (<0.1, according to Figure 6).

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 3
questions: - The linear probe module is already conditioned on timestep embeddings. Could the authors elaborate on why additional time-dependent weight is still necessary? While Table 1 shows empirical improvements, a deeper explanation of its role would be helpful.  

- The authors mention that LSEP continues to improve and approaches REPA. Does this trend persist with longer training (e.g., beyond 4M steps)?

-  All experiments are conducted only on SiTs. Is this also applicable to DiTs?
soundness: 3
strengths: - The proposed training framework is simple yet efficient, achieving both faster convergence and improved generation quality without relying on external modules.

- The paper is well-written and easy to follow. The motivation and ideas are clearly articulated, figures and visualizations make the concepts easy to understand.

- The ablation studies on each design components are well-designed and detailed which effectively validates each design choice.
summary: This paper presents Linear Separability (LSEP), a regularization strategy for diffusion transformers which explicitly encourages linearly separable representations without external alignment. The high-level idea is to jointly train a linear probe with the denoising objective thereby enhancing feature separability in the early layer. The authors introduce three design components (1) classification specific conditioning, (2) random cropping for patch-level separability, and (3) time-dependent loss weighting. Experiments on SiT models show that LSEP improves both training speed and generation quality achieving comparable results to alignment-based methods.
weaknesses: - The study mainly focus on linear probing as the main indicator of representational quality. It is unclear whether this approach generalizes beyond class-labeled settings (e.g., text-to-image or video generation). It would be helpful to discuss how linear separability could be defined or leveraged in such contexts.

- Although linear separability is a useful cue, good linear separation does not necessarily imply better representation. The setup therefore feels somewhat constrained and overly tied to the classification context. When combined with REPA (Table 2 and Figure 7), it achieves much higher performance and clearer separability, showing that linear probing alone does not fully explain the improvements in generation quality.

- The rationale for selecting the weighting parameter $\omega_{\text{class}}$ is not well explained. Its value (e.g., 0.03) appears empirically chosen; an ablation study or sensitivity analysis would strengthen this aspect.  

- Higher resolution (e.g., 512x512) experiments would help demonstrate scalability.


Decision: Reject