1 Introduction↩︎

Enhancing the latent representation capability of Diffusion Transformers (DiTs) [1]–[4] during training has been demonstrated to accelerate convergence and improve generation quality [5]–[10]. Prior works, such as REPA [5], attempt to achieve this by aligning the internal features of DiTs with those of a frozen, pre-trained image encoder (e.g., DINOv2 [11]). However, this external alignment strategy often falls short in scenarios where a sufficiently powerful encoder is absent, or when scaling up the training data and model size for DiTs [10], [12]. To address this, recent research has pivoted toward representation alignment within the DiT itself [8], [10], [13]. Pioneer work like Self-Representation Alignment (SRA) [8], which aligns latent representations in earlier layers under higher noise conditions with those in deeper layers under lower noise levels of the same model to progressively reinforcing internal representation learning. Subsequently, Self-Flow [10] extends this self-representation alignment paradigm to multi-modal scenarios and larger scales (e.g, Text-to-Image, Text-to-Video, Text-to-Audio), demonstrating that it consistently outperforms the external alignment methods like REPA, and the self-representation alignment baseline SRA.

Notably, as illustrated in Figure 1, Self-Flow also adopts the Self-Representation Alignment method pioneered by SRA. The key distinction, however, lies in how the input samples are processed for the student. Specifically, Self-Flow introduces a dual-timestep scheduling, where a single input sample to the student model contains patches corrupted by two distinct noise-levels. Consequently, the performance gains achieved by Self-Flow over SRA are primarily attributed to this specific design. In Self-Flow paper, the explanation of the mechanism of this dual-timestep scheduling is: "by applying different noise levels to different tokens, the model is encouraged to use cleaner tokens to infer noisy tokens. This drives learning strong representations alongside generative capabilities." Nevertheless, we question that dose these improvements indeed stem from superior self-supervision achieved by interactions of different noise-level tokens?

In this work, we revisit the mechanism behind the gains of dual-timestep scheduling. Rather than attributing the improvement solely to better self-supervision by interactions, we argue that this design also functions as a form of data augmentation for diffusion training. Here, data augmentation does not directly alter the semantic content of the clean image[14], [15]; instead, it expands the effective training distribution along the noise dimension. By assigning different noise level to different token subsets, a single clean sample is presented to the model under more diverse noise states, allowing the model to observe more noise-conditioned variants of the same data within training, thus expands the effective training data for the model.

To verify this hypothesis, we introduce Attention Separation, as illustrated in Figure 2. The key idea is to preserve the same dual-timestep input as Self-Flow while removing the interaction between tokens at different noise levels. Specifically, tokens assigned to the same timestep can attend to each other, whereas tokens assigned to different timesteps are blocked from interacting. This creates a controlled setting: if the improvement of Self-Flow mainly comes from cleaner tokens guiding noisier tokens through attention, removing such interaction should degrade performance; if the gain remains, the dual-timestep scheduler is more likely acting as noise-state augmentation.

This observation further leads us to reinterpret Attention Separation itself as a form of data augmentation. When applied under single-timestep training, all tokens share the same noise level. Nevertheless, Attention Separation still improves training, we analyze that with such separation, each token group acts as a partial observation of the original image. These partial views are processed by the same shared-parameter model and optimized with the same denoising and self-alignment objectives in a single iteration as shown in Figure 2. Thus, each image yields multiple effective training samples with different content subsets to expand the effective training distribution.

We evaluate this interpretation through controlled ablations and system-level comparisons, our final training scheme improves over previous self-alignment baselines on most metrics and remains on par with, or better than, the external-encoder baseline on ImageNet.

In summary, our main contributions are as follows:

• We revisit the mechanism behind the improvement from SRA to Self-Flow and show that dual-timestep scheduling is better explained as data augmentation rather than self-supervision.

• We introduce Attention Separation, a controlled operation that blocks interactions between tokens at different noise levels, and further show that it can also serve as a data augmentation.

• We combine dual-timestep scheduling and Attention Separation within self-representation alignment, achieving stronger results than previous self-alignment baselines on most metrics and competitive performance with external-encoder alignment.

2 Related Work↩︎

2.1 Representation Alignment for Generation↩︎

Improved latent representations of diffusion models can accelerate convergence and enhance generation [5]–[10], [16]. One prominent avenue is leveraging discriminative priors from pretrained vision encoders for alignment [5]–[7], [17]. REPA [5] pioneered this paradigm by aligning intermediate diffusion features with representations from external visual encoders such as DINOv2 [11]. Building upon this paradigm of external representation alignment, subsequent studies have introduced refinements in alignment mechanisms [7], [17], training strategies [6], [18], etc.. Meanwhile, an alternative research avenue has emerged that dispenses with external encoders entirely, opting instead to perform representation alignment internally within the Diffusion Transformer [8]–[10], [13]. SRA [8] pioneered this paradigm by aligning latent representations in earlier layers of with those in deeper layers of the same model. Following work Self-Flow [10] extend this idea to larger-scale settings and more modalities (e.g., text-to-image, text-to-video, and text-to-audio). And shows that self alignment paradigm consistently outperforms the external alignment counterpart. In this work, we revisit this internal alignment paradigm and study whether the improvement of dual-timestep scheduling comes from stronger self-supervision by interactions or from the augmented noise states introduced during training.

2.2 Visual Self-Supervised Learning↩︎

Self-Supervised Learning (SSL) leverages pretext tasks to learn robust representations without manual labels [19]–[26]. For example, Masked Autoencoders (MAE) [21] reconstruct masked image patches to provide strong downstream initializations. MoCo [20] introduces a momentum encoder and a queue-based dictionary to learn instance-discriminative representations from augmented image views. DINO [19] instead adopts a self-distillation framework, where a student network learns to match a momentum teacher across different views without using negative samples. Follow-up works such as DINOv2 [11] further extend this idea to large-scale settings, producing strong general-purpose visual representations. In visual generation, Self-Supervised Learning is also applied in the training process like pre-training [8], [10], [27], [28] and post-training [29]. In this work, we further examine this self-supervised interpretation in diffusion training, showing that the benefit of dual-timestep scheduling does not primarily rely on self-supervision, but can be explained by its data augmentation effect.

2.3 Data Augmentation↩︎

Data augmentation enlarges the effective training data by constructing task-preserving variants of existing samples. In visual representation learning, classical strategies include Mixup [15], which interpolates images and labels, Manifold Mixup [30], which performs interpolation in hidden spaces, and CutMix [14], which replaces image regions across samples. In self-supervised learning, augmentations also define different views for representation learning, as in SimCLR [31]. For diffusion-based generation, data augmentation has recently been explored to improve both discriminative and generative training. Diffusion-generated synthetic images can improve ImageNet classification [32], [33], while Degeorge et al. [34] show that competitive text-to-image diffusion models can be trained from ImageNet alone with synthetic long captions and image augmentations such as CutMix and crop-based training. In this work, we focus on understanding the mechanism behind the gains of dual-timestep scheduling [10], and show through Attention Separation that it is better viewed as data augmentation.

3 Preliminary: SRA and Self-Flow↩︎

Since both SRA and Self-Flow are built upon the Flow Matching Models [2], [3], we begin by introducing the standard training objective of Flow Matching models. Subsequently, we elaborate on the formulations of SRA and Self-Flow, followed by an analysis of the key distinctions between these two methods.

3.0.0.1 Flow matching.

We consider a conditional flow-matching DiT parameterized by \(\theta\). Given a clean sample \(x_0 \sim p_{\mathrm{data}}\), condition \(c\), and Gaussian noise \(x_1 \sim \mathcal{N}(0,I)\), a noisy sample at timestep \(t \in [0,1]\) is obtained by the linear path [35], [36]: \[x_t = (1-t)x_0 + t x_1 ,\] where a larger \(t\) corresponds to a higher noise level. The model predicts the velocity field along this path and is trained with the standard generation objective: \[\mathcal{L}_{\mathrm{gen}} = \mathbb{E}_{x_0,x_1,t,c} \left[ \left\| v_\theta(x_t,t,c) - (x_1-x_0) \right\|_2^2 \right].\]

3.0.0.2 SRA.

SRA [8] introduces a self-alignment objective without relying on an external representation encoder. Let \(h_\theta^{m}(\cdot)\) denote the feature map from the \(m\)-th layer of the student DiT, and let \(h_{\bar{\theta}}^{n}(\cdot)\) denote the feature map from the \(n\)-th layer of its EMA teacher, where usually \(m \leq n\). For a high-noise timestep \(t\) and a lower-noise timestep \(s<t\), SRA feeds \(x_t\) into the student and \(x_s\) into the EMA teacher. The early-layer student representation is then projected by a lightweight head \(g_\psi\) and aligned to the stop-gradient teacher representation: \[\mathcal{L}_{\mathrm{SRA}} = \mathbb{E} \left[ \frac{1}{N} \sum_{i=1}^{N} d\left( g_\psi\!\left(h_\theta^{m}(x_t,t,c)\right)_i, \mathrm{sg}\!\left[h_{\bar{\theta}}^{n}(x_s,s,c)_i\right] \right) \right],\] where \(N\) is the number of tokens, \(i\) indexes a token, and \(d(\cdot,\cdot)\) is a feature distance such as cosine or \(\ell_2\) distance. The overall objective is: \[\mathcal{L}=\mathcal{L}_{\mathrm{gen}}+\lambda \mathcal{L}_{\mathrm{SRA}} .\] Thus, SRA constructs self-supervision from two asymmetries: the student observes a noisier input than the teacher, and an earlier student layer is encouraged to match a later teacher layer.

3.0.0.3 Self-Flow.

Self-Flow [10] follows the same EMA-based self-alignment principle introduced in SRA, but changes how the student input is constructed. Instead of assigning a single timestep to all tokens, it samples two timesteps \(t\) and \(s\), defines \(t_{\mathrm{hi}}=\max(t,s)\) and \(t_{\mathrm{lo}}=\min(t,s)\), and builds a token-wise timestep vector \(\boldsymbol{\tau}\in[0,1]^N\): \[\tau_i = \begin{cases} t_{\mathrm{hi}}, & i \in M,\\ t_{\mathrm{lo}}, & i \notin M, \end{cases}\] where \(M\) is a randomly sampled token mask. The student input is then mixed at the token level, \[x_{\boldsymbol{\tau},i} = (1-\tau_i)x_{0,i}+\tau_i x_{1,i},\] while the EMA teacher receives the cleaner input \(x_{t_{\mathrm{lo}}}\). Its representation objective can be written as: \[\mathcal{L}_{\mathrm{SF}} = \mathbb{E} \left[ \frac{1}{N} \sum_{i=1}^{N} d\left( g_\psi\!\left(h_\theta^{m}(x_{\boldsymbol{\tau}},\boldsymbol{\tau},c)\right)_i, \mathrm{sg}\!\left[h_{\bar{\theta}}^{n}(x_{t_{\mathrm{lo}}},t_{\mathrm{lo}},c)_i\right] \right) \right],\] with \(g_\psi\) set to the identity when no projection head is used.

This formulation highlights the essential difference between SRA and Self-Flow: SRA assigns a single noise level to the entire student input. In contrast, Self-Flow introduces dual-timestep scheduling, where tokens with different noise levels coexist within the same student input. Self-Flow attributes its improvement to interactions among tokens at different noise levels: cleaner tokens provide contextual cues for noisier tokens, thereby encouraging stronger self-supervised representation learning. However, this scheduler also changes the training data, the student observes multiple noise levels within one sample instead of a single global timestep, which exposes the model to more diverse noise-level instances and can be viewed as token-wise data augmentation for the denoising task. Therefore, we argue that the gain of Self-Flow over SRA may stem from two entangled factors: interactions for better self-supervision, and heterogeneous-noise training as data augmentation.

4 Isolate Effect by Attention Separation↩︎

To disentangle these two factors, we design an Attention Separation operation that preserves the dual-timestep noise while selectively removes token interaction from different noise-levels. Specifically, we keep the same dual-timestep scheduling as Self-Flow, which means that the model still observes heterogeneous noise levels within each training sample. The only modification is in the self-attention computation. Let \(r_i=\mathbb{1}[i\in M]\) be the group indicator of token \(i\). We construct a binary attention mask: \[A^{\mathrm{sep}}_{ij} = \mathbb{1}[r_i=r_j] = \mathbb{1}[\tau_i=\tau_j],\] which allows attention only between tokens of the same noise level. For a self-attention layer with queries, keys, and values \((Q,K,V)\), Attention Separation computes: \[\mathrm{Attn}^{\mathrm{sep}}(Q,K,V)_i = \sum_{j=1}^{N} \frac{ A^{\mathrm{sep}}_{ij}\exp(q_i^\top k_j/\sqrt{d}) }{ \sum_{l=1}^{N} A^{\mathrm{sep}}_{il}\exp(q_i^\top k_l/\sqrt{d}) } v_j .\] Equivalently, tokens from different timestep groups are assigned \(-\infty\) attention logits before the softmax. As shown in Figure 3, this turns the full attention matrix into a block-diagonal one: tokens assigned to the same noise level can attend to each other, whereas tokens from different noise levels are prevented from interacting through attention. This controlled setting preserves the heterogeneous-noise training but removes interactions. Therefore, comparing Self-Flow with its attention-separated counterpart allows us to isolate whether the observed gain mainly comes from self-supervision by interaction or from the dual-timestep noise as data augmentation.

5 Data Augmentation Matters More↩︎

To answer the question raised above, we conduct controlled ablations on ImageNet \(256\times256\) [37] using SiT-B [2] by default. We report FID-10K [38] and IS IS? at different training iterations. Unless otherwise specified, the training and inference hyperparameters settings follow the default choices used in SRA and Self-Flow [8], [10]. Our goal is to disentangle whether the gain mainly comes from token interaction or from heterogeneous-noise data augmentation.

5.0.0.1 Removing interaction does not weaken dual-timestep training.

As Attention Separation blocks attention between tokens assigned to different noise levels while preserving the same heterogeneous noise assignment. If the gain mainly came from cleaner tokens guiding noisier tokens through self-attention, this intervention should degrade performance. Table 1 shows the ablation results of whether isolates the role of token interaction of different noise-levels under dual-timestep scheduling. It can be seen that Attention Separation achieves comparable FID at 100K and improves both FID and IS at later stages, reducing FID from \(25.19\) to \(25.06\) and increasing IS from \(66.75\) to \(72.94\) at 800K. This result supports the interpretation that the dual-timestep benefit does not primarily rely on the interactions of tokens.

5pt

5.0.0.2 Dual-timestep scheduling primarily works as data augmentation.

Figure 4 compares single-timestep and dual-timestep training under matched attention settings. With full attention, dual-timestep training consistently improves FID from \(32.10/28.42/26.34\) to \(30.20/26.89/25.19\) at 400K/600K/800K iterations. More importantly, the same improvement remains under Attention Separation, where interactions between different noise-level tokens are explicitly blocked: dual-timestep training improves FID from \(32.45/28.30/25.81\) to \(29.89/26.97/25.06\) and also improves IS across all training stages. This indicates that the gain does not rely on stronger self-supervision induced by token interactions. Instead, dual-timestep scheduling changes the training data seen by the student: each image is decomposed into token subsets observed at different noise levels, exposing the model to more noise-state variants within the same training iteration, thus expand the effective training distribution. Therefore, its effect is better understood as data augmentation to expose the model to more data.

6 Attention Separation Is Also a Data Augmentation↩︎

Table 2 compares full attention and Attention Separation under the single-timestep setting. In this case, it do not has any interaction from tokens in different noise level, since all tokens share a timestep. It only partitions the image tokens into several non-interacting groups. However, We observe that Attention Separation still brings clear gains, especially in IS, even when all tokens share the same timestep. To investigate the source of the performance gains, we conducted the following equivalent substitution analysis.

As illustrated in Figure 2 and Figure 3, Attention Separation can be interpreted as converting one training image into multiple part-conditioned training views. Whether the two groups use different timesteps (\(t_1 \neq t_2\)) or the same timestep (\(t_1=t_2\)), the separation mask makes each token group behave like a partial observation of the original image. These partial views are processed by the same model with shared parameters and optimized with the same denoising and self-alignment objectives in one iteration. Equivalently, a single image provides multiple effective training samples, each containing a different subset of the full image. This expands the effective training distribution without introducing external data. Together with the results in Table 2, indicating that Attention Separation itself also acts as data augmentation along the sample view, while dual-timestep scheduling augments the sample along the noise-state dimension.

5pt

6.0.0.1 Effect of the mask ratio.

We further study how the mask ratio affects dual-timestep training with Attention Separation. Let \(\alpha=|M|/N\) denote the fraction of tokens assigned to one timestep group; thus, \(\alpha=0.25\) partitions an image into two groups with \(25\%\) and \(75\%\) tokens. Table 3 reports the ablation results.

5pt

When full attention is used, changing the mask ratio does not harm dual-timestep training, and the performance remains consistently better than the single-timestep baseline. This is expected under our augmentation interpretation: regardless of the exact partition ratio, the model is still exposed to more noise states within each image, while full-image attention allows every token to access the complete image context. Thus, changing the mask ratio mainly changes the relative amount of tokens in one group, but does not prevent the model from learning with global spatial context. However, the behavior changes once Attention Separation is applied. While \(\alpha=0.25\) achieves the best IS and comparable FID, larger ratios degrade FID substantially, especially at \(\alpha=0.50\). We hypothesize that this degradation comes from a stronger training–inference mismatch induced by overly balanced separation. During training, Attention Separation decomposes each image into two non-interacting token groups, so each attention component can only aggregate information from a partial view of the image. As the mask ratio approaches \(0.50\), both groups become incomplete views with similar size, and neither group consistently preserves most of the global image context. In contrast, inference uses standard full-image attention, where all tokens interact globally. This gap between part-level training and full-image inference becomes more severe at larger mask ratios, leading to the observed FID degradation.

To mitigate the training–inference mismatch at large mask ratios, we further mix single-timestep full-image samples into each training batch. Specifically, for a fraction \(\rho\) (in our experiments, we set \(\rho=0.25\)) of samples in a mini-batch, we disable the dual-timestep and Attention Separation and assign all tokens the same timestep, i.e., \(\tau_i=t\) for all \(i\), while using the standard full-attention. The remaining samples are trained with the original dual-timestep scheduling and Attention Separation. Therefore, each batch contains both separated dual-timestep samples, which preserve the heterogeneous-noise and part-level data augmentation effects, and full-image single-timestep samples, which expose the model to the same global attention pattern used at inference. This mixed setting preserves the augmentation effect from dual-timestep scheduling and Attention Separation for a subset of samples, while also exposing the model to the standard full-image, single-timestep attention pattern used at inference. As shown in Figure 5, this strategy substantially improves FID when the mask ratio is large. At \(\alpha=0.50\), where all-sample Attention Separation gives each attention component only half of the image context, adding full-image samples reduces the 800K FID from \(38.19\) to \(24.15\). The gain becomes smaller as the mask ratio decreases. This trend is consistent with our interpretation: when \(\alpha=0.25\), one token group already covers most of the image, so the separated training samples remain relatively close to the full-image inference condition. In this case, replacing part of the batch with vanilla single-timestep samples is less critical and may also weaken the augmentation effect, since those samples no longer receive either dual-timestep noise augmentation or Attention Separation.

7 Putting Things Together↩︎

The analyses above lead to a unified interpretation of the transition from SRA to Self-Flow. The key component that improves Self-Flow over SRA, dual-timestep scheduling, is not mainly explained by stronger self-supervision as suggested in the Self-Flow paper. By applying Attention Separation, we remove token interactions of different noise-level while preserving the same heterogeneous-noise input, yet the performance does not degrade and can even improve. This indicates that the benefit of dual-timestep scheduling mainly comes from data augmentation that expands the effective training data: the same image is observed under more diverse noise states. We further find that Attention Separation itself also acts as an augmentation mechanism: by splitting one training image into multiple independently optimized token groups under shared model parameters, it increases the number of effective training views derived from the same sample. In this sense, the answer to the question in our title is that the gain from SRA to Self-Flow is better understood as data augmentation, rather than as stronger self-supervision.

This interpretation naturally leads to our final training scheme. We retain the internal self-alignment objective of SRA, since it provides the representation-learning signal without relying on external encoders. On top of it, we use dual-timestep scheduling to augment each image along the noise-state dimension, and apply Attention Separation to further create part-conditioned training views. Both components are therefore used as augmentation mechanisms within the self-representation alignment framework for training.

8 System-Level Comparison↩︎

8.1 Setup↩︎

Unless specified otherwise, our training pipeline closely mirrors the configurations established in privious baselines [1], [2], [8], [10]. Specifically, we employ the AdamW optimizer [40] with a constant learning rate of 1e-4, zero weight decay, and a total batch size of 256, and uniform timestep sampling strategy. Latent representations are extracted utilizing the pre-trained Stable Diffusion VAE [41]. For the model backbone, we adopt the XL/2 SiT, all of which operate with a patch size of 2. For our method, we follow the setups of Self-Flow [10] and SRA [8], where the alignment layer for the student and teacher are 8 and 20, respectively. The teacher is obtained via the Exponential Moving Average (EMA) of the student with a decay of 0.9999, and the coefficient of the alignment loss is set to 0.5. The mask ratio is set to 0.25 as it yields the best performance (ablated in Table 3 and Figure 5). All experiments are conducted on 8 NIVIDA H20 GPUs.

To evaluate generation quality, we report Fréchet Inception Distance (FID [38]), sFID [42], Inception Score (IS IS?), along with precision and recall [43]. To ensure equitable comparisons with existing baselines, we compute these metrics using the official TensorFlow evaluation suite from ADM [44] with 50K generated samples and the standard reference statistics.

We benchmark our method against vanilla DiT and SiT [1], [2] as well as paradigms from both branches of representation alignment: namely, those with and without dependency on external models. Within each category, we benchmark against the representative method. Specifically, we select REPA [5] as the representative for external-model-assisted alignment, and SRA [8] and Self-Flow [10] for self-alignment approaches.

8.2 Results↩︎

Table 4 reports ImageNet \(256\times256\) results. Compared with the vanilla SiT-XL/2 trained for 7M steps, our method reaches a lower FID using 4M steps, improving FID from \(2.06\) to \(1.44\) and IS from \(270.3\) to \(315.3\). Among self-alignment methods, our method improves over SRA and Self-Flow in FID and IS, achieving the best IS and the second-best FID among all compared methods. Although REPA obtains a slightly lower FID with an external pretrained encoder, our method remains comparable while relying only on self-representation alignment inside the diffusion transformer.

Table 5 shows the results on ImageNet \(512\times512\). Our method matches the best FID of REPA at \(2.08\), outperforms both SRA and Self-Flow in FID and IS, and achieves the highest IS of \(282.7\). It also substantially improves over the vanilla SiT-XL/2 baseline trained for 3M steps, reducing FID from \(2.62\) to \(2.08\) with only 1M training steps. These results indicate that the augmentation interpretation developed in the controlled studies translates to stronger system-level performance, and that the resulting method remains effective when scaling to higher image resolution.

9 Conclusion↩︎

In this work, we revisit the transition from SRA to Self-Flow and study whether the improvement actually comes from. By introducing Attention Separation, we preserve the same heterogeneous-noise input while removing cross-noise token interaction. The resulting performance does not degrade and can even improve, indicating that the benefit of dual-timestep scheduling is better explained as noise-state data augmentation rather than cleaner-to-noisier token interaction alone. We further show that Attention Separation itself provides a part-level augmentation effect by splitting a single image into multiple effective training parts to expand the training data. Based on these findings, we combine dual-timestep scheduling and Attention Separation within the self-representation alignment framework. Experiments on ImageNet \(256\times256\) and \(512\times512\) show that this augmentation-based interpretation leads to a simple and effective training scheme, competitive with both external-encoder alignment and previous self-alignment methods.

References↩︎