REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training
Created by
Haebom
저자
Ziqiao Wang, Wangbo Zhao, Yuhao Zhou, Zekai Li, Zhiyuan Liang, Mingjia Shi, Xuanlei Zhao, Pengfei Zhou, Kaipeng Zhang, Zhangyang Wang, Kai Wang, Yang You
개요
본 논문은 확산 트랜스포머(DiT)의 느린 학습 속도 문제를 해결하기 위해 새로운 방법인 HASTE(Holistic Alignment with Stage-wise Termination for Efficient training)를 제안합니다. 기존의 표현 정렬(REPA) 방법은 초기 학습 속도를 높이지만, 이후 성능이 정체되거나 저하되는 한계를 보입니다. 이는 학습 과정에서 생성 모델이 데이터 분포를 학습하기 시작하면, 교사 모델의 저차원 임베딩과 어텐션 패턴이 오히려 학습을 방해하기 때문입니다. HASTE는 이러한 문제를 해결하기 위해 두 단계 학습 전략을 제시합니다. 1단계에서는 교사 모델의 어텐션 맵과 특징 벡터를 DiT의 중간 레이어에 동시에 주입하여 빠른 수렴을 유도하고, 2단계에서는 특정 반복 횟수에 도달하면 정렬 손실을 비활성화하여 DiT가 생성 능력을 최대한 활용하도록 합니다. ImageNet 256x256 데이터셋에서 HASTE는 기존 방법보다 훨씬 적은 학습 단계로 동등하거나 더 나은 성능을 달성하며, MS-COCO 데이터셋의 텍스트-이미지 생성 작업에서도 성능 향상을 보였습니다.
시사점, 한계점
•
시사점:
◦
DiT의 학습 속도를 획기적으로 향상시키는 효율적인 학습 방법을 제시합니다. (28배 감소)
◦
다양한 DiT 아키텍처와 작업에 적용 가능한 일반적인 방법론을 제공합니다.
◦
ImageNet과 MS-COCO 데이터셋에서 성능 향상을 실험적으로 검증했습니다.
◦
단순하면서도 원리에 기반한 방법으로 구현이 용이합니다.
•
한계점:
◦
제안된 방법의 효과는 특정 트리거(고정된 반복 횟수)에 의존하며, 최적의 트리거 설정은 데이터셋이나 아키텍처에 따라 달라질 수 있습니다.
