PaperlayerView exact revision on arXiv

Source-linked AI summary

From SRA to Self-Flow: Data Augmentation or Self-Supervision?

Dengyang Jiang, Mengmeng Wang, Harry Yang, Jingdong Wang

arXiv:2607.02508v1cs.CV

TL;DR

The mechanism behind Self-Flow’s improvement over SRA remains unclear: it may reflect cross-noise token interaction or noise-dimension data augmentation. The paper introduces Attention Separation to preserve heterogeneous-noise inputs while blocking those interactions, finding that augmentation better explains the gain and supports stronger ImageNet results.

  • Problem

    The source of Self-Flow’s improvement over SRA—cross-noise token interaction or another effect of dual-timestep scheduling—remains under-examined.

  • Method

    The paper introduces Attention Separation, which blocks attention between differently noised tokens while combining dual-timestep scheduling with self-representation alignment.

  • Results

    FID improves from 2.06 to 1.44 and IS from 270.3 to 315.3 versus vanilla SiT-XL/2 on ImageNet 256 × 256.

  • Takeaways & Limitations

    The findings support interpreting dual-timestep scheduling and Attention Separation primarily as noise-state and part-level data augmentation rather than cross-noise self-supervision.

Abstract

from arXiv · show

Representation alignment has become an effective way to accelerate diffusion transformer training and improve generation quality. Recent self-alignment methods, such as SRA and Self-Flow, further remove the dependency on external pretrained encoders by constructing alignment within the diffusion model itself. However, the mechanism behind the improvement from SRA to Self-Flow, dual-time scheduling, remains under-examined: Self-Flow attributes its gain to interactions between tokens at different noise levels, where cleaner tokens help infer noisier ones. In this work, we revisit this explanation and ask whether the gain instead comes from data augmentation along the noise dimension. To disentangle these factors, we introduce Attention Separation, which preserves the same dual-timestep input as Self-Flow while blocking attention between tokens assigned to different noise levels. Surprisingly, removing such interaction does not degrade performance and can even improve it, suggesting that the improvement from SRA to Self-Flow mainly comes from data augmentation. Furthermore,We show that Attention Separation itself provides an augmentation effect by splitting a single image into multiple effective training parts to expand the training data. Based on these observations, we combine self-representation alignment with dual-timestep and attention-separation augmentation, and demonstrate the effectiveness of this design on ImageNet.

1. Introduction

The paper revisits why Self-Flow improves over SRA, arguing that dual-timestep scheduling primarily augments noise states rather than strengthening cross-noise self-supervision. Attention Separation tests this claim by blocking cross-timestep attention while retaining the heterogeneous input.

  • SRA versus Self-Flow: SRA assigns one noise level to every student token, whereas Self-Flow places tokens at two noise levels in the same image.Both methods retain SRA’s self-representation alignment; their key difference is student input processing.
  • SRA versus Self-Flow: Self-Flow attributes its gains to cleaner tokens providing contextual cues that help infer noisier tokens through attention.The paper questions whether these interactions, rather than the altered noise distribution, explain the improvement.
  • Data augmentation interpretation: Dual-timestep scheduling can augment training by presenting one clean image through more diverse noise states without changing its semantic content.Different token subsets receive different noise levels, expanding the effective training distribution along the noise dimension.
  • Attention Separation: Attention Separation blocks attention between different noise levels while preserving dual-timestep inputs, creating a controlled test of the interaction hypothesis.Tokens at the same timestep can still attend to one another.
  • Attention Separation: Attention Separation also creates multiple part-conditioned views, allowing one image to yield several effective training parts under shared objectives.This provides an additional augmentation effect by expanding the effective training distribution.
  • Results: The final combination of self-representation alignment, dual-timestep scheduling, and Attention Separation outperforms prior self-alignment baselines on most metrics.It remains competitive with external-encoder alignment while avoiding dependence on a pretrained encoder.

2. Related Work

Prior work improves diffusion training through external representation alignment, self-supervised representation learning, and data augmentation. This paper positions its contribution at their intersection by reconsidering dual-timestep scheduling as augmentation rather than interaction-driven supervision.

  • Representation alignment: External alignment methods match diffusion features to representations from frozen pretrained vision encoders such as DINOv2.REPA is presented as a prominent example of this paradigm.
  • Representation alignment: Self-alignment methods instead construct representation targets within the diffusion model, avoiding reliance on external encoders.The paper studies whether Self-Flow’s dual-timestep gain comes from self-supervision or augmented noise states.
  • Self-supervised learning: Self-supervised learning uses pretext tasks or multiple views to learn representations without manual labels.Examples include masked reconstruction, momentum-encoder contrastive learning, and self-distillation.
  • Data augmentation: Data augmentation enlarges effective training data through task-preserving variants, including Mixup, Manifold Mixup, CutMix, and view-based augmentation.The paper extends this perspective to diffusion training along the noise dimension.

3. Preliminary: SRA and Self-Flow

SRA and Self-Flow share EMA-based self-representation alignment but differ in how student inputs are noised. SRA uses a single timestep, while Self-Flow mixes tokens assigned to distinct timesteps.

  • Flow matching: Flow Matching constructs a noisy sample from a clean sample, Gaussian noise, and a timestep along a linear path.Larger t corresponds to a higher noise level, and the DiT predicts the velocity field.
  • SRA: SRA aligns an earlier student-layer representation from a noisier input with a later EMA-teacher representation from a cleaner input.The student and teacher therefore differ in both noise condition and network depth.
  • SRA: The SRA representation loss compares token features using a distance such as cosine or ℓ2 distance.N denotes the number of tokens and i indexes a token.
  • Self-Flow: Self-Flow retains SRA’s EMA self-alignment principle but assigns multiple timesteps across the student’s tokens.Its dual-timestep input allows different noise levels to coexist within one student sample.
  • Self-Flow: Self-Flow’s proposed mechanism is that cleaner tokens provide contextual information for noisier tokens through cross-noise interactions.The paper later tests whether this interaction is necessary for the observed gains.

4. Isolate Effect by Attention Separation

Attention Separation isolates whether Self-Flow gains require cross-noise token interaction by preserving dual-timestep inputs while blocking attention between timestep groups.

  • Controlled intervention: Attention Separation preserves Self-Flow’s heterogeneous dual-timestep input while removing attention between tokens assigned to different noise levels.Only tokens sharing a timestep can interact, creating a controlled test of the proposed cleaner-to-noisier guidance mechanism.
  • Controlled intervention: Comparing standard Self-Flow with the separated variant isolates token interaction from the dual-timestep noise assignment.Both settings retain heterogeneous noise; their difference is whether cross-group attention is allowed.
  • Mask design: The binary mask permits attention only among tokens with the same noise level.Different-timestep pairs receive effectively blocked attention, producing a block-diagonal attention pattern.

5. Data Augmentation Matters More

Ablations show that dual-timestep scheduling remains effective when cross-noise interactions are blocked, pointing to noise-state data augmentation rather than stronger self-supervision as the main source of gains.

  • Ablation evidence: 25.19 to 25.06 FID and 66.75 to 72.94 IS at 800K show Attention Separation can improve dual-timestep training.The separated model matches or exceeds the interacting version despite blocking cross-noise attention.
  • Noise-state augmentation: Dual-timestep training improves FID under full attention from 32.10/28.42/26.34 to 30.20/26.89/25.19 at 400K/600K/800K.The gains persist across the reported training stages.
  • Noise-state augmentation: Under Attention Separation, dual-timestep training improves FID from 32.45/28.30/25.81 to 29.89/26.97/25.06.IS also improves across all reported stages, even though cross-noise token interactions are explicitly blocked.
  • Noise-state augmentation: Each image exposes the model to multiple token subsets at different noise states within one training iteration.This expands the effective training distribution without requiring stronger self-supervision through cross-noise attention.

6. Attention Separation Is Also a Data Aug- mentation

Attention Separation can augment training even when every token shares one timestep by turning a full image into multiple non-interacting partial views.

  • View augmentation: Under single-timestep training, Attention Separation still improves performance by partitioning image tokens into non-interacting groups.Because all tokens share one timestep, these gains cannot depend on cross-noise interactions.
  • View augmentation: One image becomes multiple part-conditioned training views processed by shared parameters and optimized with the same objectives.Each partial view acts as an effective training sample without introducing external data.
  • Mask-ratio ablation: α = 0.25 gives the best IS and comparable FID, while larger ratios substantially degrade FID under Attention Separation.The degradation is attributed to stronger training–inference mismatch as groups become similarly sized and incomplete.
  • Mismatch mitigation: Mixing full-image single-timestep samples with separated dual-timestep samples reduces mismatch from strong separation.The mixed batch preserves augmentation while exposing the model to the global attention pattern used at inference.

7. Putting Things Together

The final design keeps SRA’s internal alignment while treating dual-timestep scheduling and Attention Separation as complementary augmentations. Together, they create more diverse noise-state and part-conditioned training views without external encoders.

  • Design: The training scheme retains SRA’s self-alignment objective to provide representation learning without relying on external encoders.
  • Design: Dual-timestep scheduling augments each image along the noise-state dimension by exposing it to heterogeneous noise conditions.
  • Design: Attention Separation further creates part-conditioned training views by restricting token interactions while sharing model parameters.

8. System-Level Comparison

The proposed method improves diffusion-transformer generation across ImageNet resolutions while remaining competitive with external-encoder alignment. Its strongest gains appear in FID, IS, and training efficiency relative to vanilla SiT-XL/2.

  • ImageNet 256×256: FID improves from 2.06 to 1.44 against vanilla SiT-XL/2 at 256×256, using 4M rather than 7M training steps.The comparison is reported for ImageNet 256×256.
  • ImageNet 256×256: IS improves from 270.3 to 315.3 against vanilla SiT-XL/2 at 256×256.
  • ImageNet 256×256: The method achieves the best IS and second-best FID among compared methods, while remaining comparable to REPA without an external pretrained encoder.
  • ImageNet 512×512: At 512×512, the method matches REPA’s best FID at 2.08, achieves the highest IS of 282.7, and outperforms SRA and Self-Flow in both metrics.
  • ImageNet 512×512: FID decreases from 2.62 to 2.08 against vanilla SiT-XL/2 at 512×512, using 1M rather than 3M training steps.

9. Conclusion

The paper concludes that Self-Flow’s improvement over SRA is better explained by data augmentation than stronger cross-noise self-supervision. Attention Separation reinforces this interpretation by generating multiple effective views from each image.

  • Dual-timestep scheduling preserves performance without cross-noise token interaction, indicating that its main benefit is noise-state data augmentation.
  • Attention Separation splits one image into multiple effective training parts, expanding the effective data through shared-parameter optimization.
  • Combining dual-timestep scheduling and Attention Separation with self-representation alignment yields an effective ImageNet training scheme.The scheme is evaluated at 256×256 and 512×512 resolutions.
Loading 2607.02508v1…