Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think

Recent studies have shown that the denoising process in (generative) diffusion models can induce meaningful (discriminative) representations inside the model, though the quality of these representations still lags behind those learned through recent self-supervised learning methods. We argue that one main bottleneck in training large-scale diffusion models for generation lies in effectively learning these representations. Moreover, training can be made easier by incorporating high-quality external visual representations, rather than relying solely on the diffusion models to learn them independently. We study this by introducing a straightforward regularization called REPresentation Alignment (REPA), which aligns the projections of noisy input hidden states in denoising networks with clean image representations obtained from external, pretrained visual encoders. The results are striking: our simple strategy yields significant improvements in both training efficiency and generation quality when applied to popular diffusion and flow-based transformers, such as DiTs and SiTs. For instance, our method can speed up SiT training by over 17.5$\times$, matching the performance (without classifier-free guidance) of a SiT-XL model trained for 7M steps in less than 400K steps. In terms of final generation quality, our approach achieves state-of-the-art results of FID=1.42 using classifier-free guidance with the guidance interval.

한 줄 요약

REPA 정규화를 통해 디퓨전 트랜스포머의 생성 성능과 학습 효율성이 크게 향상된다.

핵심 기여도

REPA 정규화를 제안하여 디퓨전 트랜스포머의 내부 표현과 외부 사전학습 시각 표현을 정렬.
SiT 모델 학습을 17.5× 가속화하며, 400K 스텝에서 7M 스텝 모델과 동일한 성능 달성.
FID 1.42 달성으로, 클래스-조건부 ImageNet 생성에서 최신 기술 수준의 성능 보여.

핵심 아이디어

기존 디퓨전 모델은 생성 과정에서 의미 있는 내부 표현을 학습하지만, 이 표현의 품질이 최근의 자기감독 학습 방법에 비해 떨어진다. 이 연구는 디퓨전 트랜스포머가 생성 성능을 향상시키기 위해서는 고질량 표현을 학습하는 것이 필수적이라는 가설을 제시한다. 이를 위해 외부 사전학습된 시각 표현(예: DINOv2)을 활용한 정규화 기법인 REPA를 제안한다. REPA는 노이즈 입력의 히든 상태와 외부 표현 사이의 정렬을 강화함으로써, 디퓨전 모델의 학습을 용이하게 하고 생성 품질을 향상시킨다. 특히, REPA는 트랜스포머 블록의 앞부분에서만 정렬이 충분히 이루어질 수 있음을 보여주며, 이는 후반 블록이 고주파 세부 정보를 집중적으로 학습할 수 있도록 돕는다.

기술적 접근법

**REPA 정규화**: 노이즈 입력의 히든 상태와 외부 사전학습 시각 표현(예: DINOv2)을 정렬하는 정규화 기법.
**모델**: DiT, SiT와 같은 디퓨전 트랜스포머에 적용.
**데이터셋**: ImageNet 사용.
**하이퍼파라미터**: 정규화 계수 λ=0.5에서 성능이 포화됨.
**학습 과정**: SiT-XL/2 모델을 400K 스텝 학습하여 7M 스텝 모델과 동일한 성능 달성.

주요 결과

**SiT 학습 가속**: 17.5× 빠르게 학습하여 400K 스텝에서 7M 스텝 모델과 동일한 성능.
**FID 향상**: 클래스-조건부 ImageNet 생성에서 FID=1.42 달성 (기존 최고 기록).
**정렬 효과**: 다양한 노이즈 수준에서 표현 간의 semantic gap 감소 (CKNNA 값 증가).

의의 및 한계

REPA는 디퓨전 트랜스포머의 학습 과정을 단순화하고, 생성 성능을 향상시키는 간단한 정규화 기법으로, 기존 자기감독 학습과 생성 모델 간의 통합 가능성을 제시한다. 특히, 외부 표현을 활용함으로써 학습 효율성을 극대화할 수 있다는 점에서 실용적 가치가 높다. 그러나 REPA는 외부 사전학습 모델에 의존하므로, 이러한 모델의 품질에 따라 성능이 변동할 수 있는 한계가 있다. 또한, REPA가 다양한 종류의 시각 표현(예: MAE, MoCov3)에도 효과적임을 보여주지만, 다른 도메인(예: 텍스트-이미지 생성)에서의 일반화 가능성은 추가 연구가 필요하다.

실용적 활용

REPA는 대규모 디퓨전 모델의 학습 시간을 줄이고, 생성 품질을 향상시키는 데 활용될 수 있다. 특히, 이미지 생성, 영상 생성, 텍스트-이미지 생성 등 고해상도 시각 콘텐츠를 요구하는 산업 분야에서 유용하게 사용될 수 있다. 또한, 자기감독 학습과 생성 모델 간의 표현 통합을 위한 기초 연구로도 활용 가능하다.