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

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.

한 줄 요약

Self-Flow의 성능 향상이 주로 데이터 증강 효과에서 비롯됨을 밝히고, Attention Separation이라는 새로운 증강 기법을 제안한다.

핵심 기여도

Self-Flow의 dual-timestep scheduling이 self-supervision이 아닌 noise-state data augmentation 효과에서 비롯됨을 밝힘.
Attention Separation이라는 새로운 기법을 제안하여, cross-noise-level 토큰 상호작용을 차단하면서도 성능 향상 가능함을 보임.
Attention Separation은 단일 이미지를 여러 훈련 샘플로 분할하여 훈련 데이터를 확장하는 효과를 가짐.
ImageNet 256×256 및 512×512에서 기존 self-alignment 및 external-encoder 기반 방법을 능가함.

핵심 아이디어

Self-Flow는 SRA에 비해 dual-timestep scheduling을 도입하여 성능 향상을 보인다. 이전 연구는 이 향상이 'cleaner token이 noisier token을 도와주는 self-supervision'에서 비롯된다고 주장했으나, 본 연구는 이를 재검토한다. 대신, dual-timestep scheduling이 noise level 차원에서 데이터 증강 효과를 제공한다고 주장한다. 즉, 동일 이미지라도 다양한 노이즈 상태에서 훈련되게 하여 더 다양한 조건을 학습하게 만드는 것이다. 이를 검증하기 위해 Attention Separation이라는 기법을 도입한다. 이 기법은 cross-noise-level 토큰 간의 어텐션을 차단하면서도 dual-timestep 입력을 유지한다. 결과적으로, 토큰 간 상호작용이 없어도 성능이 유지되거나 향상됨으로써, 성능 향상이 self-supervision이 아닌 데이터 증강에서 비롯됨을 입증한다.

기술적 접근법

**Dual-timestep scheduling**: Self-Flow에서 사용된 입력 방식으로, 하나의 샘플이 두 가지 다른 노이즈 레벨의 토큰을 포함함.
**Attention Separation**: cross-noise-level 토큰 간의 어텐션을 차단하는 기법. 동일 노이즈 레벨 내의 토큰은 어텐션 가능.
**Self-Representation Alignment (SRA)**: 이전 레이어의 높은 노이즈 조건에서의 잠재 표현을, 이후 레이어의 낮은 노이즈 조건 표현과 정렬.
**ImageNet 256×256, 512×512 데이터셋**에서 실험 수행.
**FID, IS (Inception Score)**를 주요 평가 지표로 사용.

주요 결과

ImageNet 256×256에서 FID 1.44, IS 315.3 달성 (기존 SiT-XL/2 대비 FID -0.62, IS +45.0).
ImageNet 512×512에서 FID 2.08, IS 282.7 달성 (기존 SiT-XL/2 대비 FID -0.54, IS +20.0).
SRA 및 Self-Flow 대비 FID, IS 모두 상승.
REPA (external-encoder 기반)와 유사하거나 더 우수한 성능 보임.

의의 및 한계

본 연구는 self-alignment 기반 디퓨전 트랜스포머 훈련에서 성능 향상의 근본 원인을 재정의한다. 기존의 self-supervision 중심 설명 대신, noise-state data augmentation이라는 새로운 관점에서 해석함으로써, 디퓨전 모델 훈련 전략을 재조정할 수 있는 가능성을 제시한다. 또한, Attention Separation은 훈련 데이터를 효과적으로 확장하는 새로운 증강 기법으로, 기존의 외부 인코더에 의존하지 않는 훈련 방식을 강화한다. 다만, 본 연구는 ImageNet 데이터셋에만 적용되었으며, 다른 모델 아키텍처나 다중 모달성 데이터에 대한 일반화 가능성은 추가 연구가 필요하다.

실용적 활용

이 연구는 디퓨전 트랜스포머의 훈련 효율성을 높이고 외부 인코더에 대한 의존성을 줄이는 데 기여할 수 있다. 특히, 대규모 이미지 생성, 텍스트-이미지 생성, 텍스트-비디오 생성 등 다양한 생성 모델 개발에 활용 가능하다. Attention Separation은 훈련 데이터가 제한된 상황에서도 효과적인 증강 기법으로 활용될 수 있다.