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

TL;DR

This paper revisits Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis‘s explanation for why its dual-timestep scheduler improves over No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves (SRA) (SRA). Self-Flow claimed the gain comes from cross-noise-level token interaction — cleaner tokens helping infer noisier ones. The authors introduce Attention Separation, which keeps the same dual-timestep input but blocks attention between tokens at different noise levels via a block-diagonal mask. Blocking the interaction does not hurt and often helps, evidence the improvement over SRA is really data augmentation along the noise dimension, not cross-noise reasoning. Attention Separation is itself an augmentation, splitting one image into multiple part-conditioned training views. Combining self-representation alignment with dual-timestep + attention-separation augmentation yields the strongest reported ImageNet 256²/512² CFG results.

Key claims

• Self-Flow attributes its gain over SRA to interactions between tokens at different noise levels; this paper argues the gain is actually data augmentation along the noise dimension [Abstract].
• Under dual-timestep scheduling, replacing full attention with a block-diagonal mask that blocks cross-noise-level interaction gives comparable or better FID — direct evidence the interaction is not the load-bearing mechanism [§4, Table 1].
• Attention Separation is itself a data-augmentation operator: under single-timestep training (no noise heterogeneity), it still improves FID by partitioning image tokens into multiple non-interacting part-conditioned views [§4, Table 2].
• The augmentation effect is masking-ratio-sensitive: mild ratios preserve the gain, larger ratios degrade due to train/inference mismatch [§4, Table 3].
• Mixing full-image single-timestep samples with Attention-Separated dual-timestep samples reduces train/inference mismatch [§4, Fig. 5].
• Combining self-representation alignment + dual-timestep scheduling + Attention Separation is the best-performing configuration on ImageNet 256² and 512² with CFG [§5, Table 4, Table 5].

Method

Given a Self-Flow-style dual-timestep input where some tokens are noised at t₁ and others at t₂, Attention Separation applies a block-diagonal attention mask so tokens assigned to t₁ can only attend to other t₁ tokens and likewise for t₂. This preserves the heterogeneous-noise input (the augmentation) while removing the cross-noise token interaction Self-Flow credited for its gain. The authors then interpret Attention Separation as a further augmentation: it effectively creates multiple part-conditioned views of one image, expanding the training distribution beyond what dual-timestep scheduling alone achieves. Their final training recipe stacks (a) self-representation alignment from SRA/Self-Flow, (b) dual-timestep noise scheduling from Self-Flow, and (c) Attention Separation as an augmentation with a mild mask ratio and full-image single-timestep samples mixed in.

Results

• Blocking cross-noise attention under dual-timestep scheduling matches or improves quality vs full attention, disproving Self-Flow’s cross-noise-interaction explanation [Table 1].
• Attention Separation under single-timestep training (isolating augmentation from heterogeneous-noise effect) still improves FID, confirming it acts as data augmentation [Table 2].
• Mask ratio sweep at 800K iterations shows mild ratios preserve the augmentation benefit; larger ratios hurt from stronger train–inference mismatch [Table 3].
• Adding full-image single-timestep samples alongside separated dual-timestep samples reduces the mismatch and yields the best training trajectory [Fig. 5].
• On ImageNet 256² with CFG, the combined design (self-alignment + dual-timestep + Attention Separation) achieves the best headline metrics against SRA, Self-Flow, and REPA baselines [Table 4]. Reproduces on ImageNet 512² [Table 5].

Why it’s interesting

This is the first paper to isolate which piece of Self-Flow’s mechanism actually drives its improvement, and the answer overturns Self-Flow’s own explanation. That matters for Diffusion training efficiency, where Self-Flow was the strongest recent argument that self-alignment beats external DINO alignment (Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis vs REPA-style methods). If the improvement is data augmentation rather than cross-noise reasoning, the dual-timestep scheduler can be replaced or extended with cheaper augmentation operators — an easier engineering target than Self-Flow’s two-forward-passes-per-step EMA-teacher cost.

It also connects to REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training‘s HASTE result that REPA-style external alignment silently stops helping mid-training: both papers suggest the training signal in the “representation alignment” family is more subtle and augmentation-adjacent than the original REPA/SRA/Self-Flow framings claim.

See also

• No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves (SRA) — SRA, the single-timestep baseline this work builds on.
• Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis — the paper being re-examined; this work keeps its dual-timestep scheduler but reinterprets why it works.
• From SRA to Self-Flow: Attention Separation Reveals the Real Mechanism Is Data Augmentation — the scholar-inbox blog summary of this paper filed earlier as a separate pointer.
• REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training — HASTE’s parallel finding that external representation alignment quietly stops helping.
• Diffusion training efficiency — sharpens the “self-alignment vs external alignment” fork by clarifying the actual lever inside self-alignment recipes.
• Representation Autoencoders — same stacking is available for any RAE-style latent encoder.