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

2026-07-02

AI Summary

This work investigates the source of performance gains in Self-Flow, disentangling whether improvements stem from cross-noise-level token interactions or data augmentation along the noise dimension. To this end, the authors propose an Attention Separation mechanism that preserves dual timesteps as input while explicitly blocking attention-based interactions between distinct noise levels, thereby decoupling self-supervised alignment from data augmentation effects. Experimental results demonstrate that the primary benefit of Self-Flow arises from noise-dimension augmentation rather than cross-noise interactions. Moreover, Attention Separation itself emerges as an effective augmentation strategy. Integrated within a framework combining dual-timestep scheduling, self-representation alignment, and a diffusion Transformer architecture, the approach validates its efficacy on ImageNet—achieving maintained or even improved performance despite the removal of cross-noise attention.

Abstract

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.

Problem

• representation alignment
• data augmentation
• self-supervision
• diffusion transformer
• noise scheduling

Innovation

• Attention Separation
• Data Augmentation
• Self-Supervision
• Diffusion Transformer
• Dual-Timestep Training