From SRA to Self-Flow: Attention Separation Reveals the Real Mechanism Is Data Augmentation

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 help 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. Surprisingly, blocking the interaction does not hurt — and often helps — showing the improvement over SRA is really data augmentation along the noise dimension, not cross-noise reasoning. They further show 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 ImageNet 256²/512² CFG results reported.

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 [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 [Table 2].
The augmentation effect is masking-ratio-sensitive: mild ratios preserve the gain, larger ratios degrade due to train/inference mismatch [Table 3].
Mixing full-image single-timestep samples with Attention-Separated dual-timestep samples reduces the train/inference mismatch caused by strong separation [Fig. 5].
Combining self-representation alignment + dual-timestep scheduling + Attention Separation is the best-performing configuration on ImageNet 256² and 512² with CFG [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 does. 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

Ablation under dual-timestep scheduling: blocking cross-noise attention matches or improves quality vs full attention, disproving Self-Flow’s cross-noise-interaction explanation [Table 1].
Under single-timestep training (isolating the augmentation from the heterogeneous-noise effect), Attention Separation still improves FID, confirming it acts as data augmentation [Table 2].
Sweep over mask ratio at 800K iterations shows mild ratios preserve the augmentation benefit; larger ratios hurt due to stronger training–inference mismatch [Table 3].
Adding full-image single-timestep samples alongside separated dual-timestep samples substantially 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].
Result reproduces on ImageNet 512² [Table 5].

Why it’s interesting

This is the first paper to isolate which piece of Self-Flow’s mechanism is actually doing the work, 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 can beat 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 implication is that the dual-timestep scheduler can be replaced or extended with cheaper augmentation operators — an easier engineering target than the two-forward-passes-per-step cost Self-Flow incurs for its EMA-teacher setup.

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. The Representation Autoencoders page currently treats Self-Flow as a partial supersession of external-DINO recipes; this paper adds a caveat that even Self-Flow’s proposed mechanism was under-examined and its true lever may be a more mundane (and cheaper) augmentation.