[논문 리뷰] Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think

Sihyun Yu, Sangkyung Kwak외 5명|arXiv (Cornell University)|2024. 10. 09.

Reinforcement Learning in Robotics인용 수 4

한 줄 요약

이 논문은 DINOv2와 같은 사전 훈련된 자기지도 학습 인코더에서 유도된 고품질 외부 시각적 특징과 노이즈가 많은 잠재 표현을 정렬하는 간단한 정규화 기법인 Representation Alignment (REPA)를 소개한다. 이 외부 특징들을 확산 모델의 초기 레이어에 흡수시킴으로써 REPA는 훈련을 17.5배 이상 가속화하고, 분류기 자유 지도 학습을 사용할 경우 SOTA FID=1.42를 달성하여 훈련 효율성과 생성 품질을 모두 크게 향상시킨다.

ABSTRACT

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$ imes$, 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.

연구 동기 및 목표

• 내부 표현 학습이 열악하여 대규모 확산 트랜스포머의 훈련이 느리고 비효율적인 문제를 해결하기 위해.
• 고품질 외부 시각적 표현을 통합하면 확산 모델의 훈련 효율성과 생성 품질을 크게 향상시킬 수 있는지 조사하기 위해.
• 노이즈가 많은 확산 입력과 깨끗한 이미지 훈련을 기반으로 한 시각 인코더 간의 입력 불일치 문제를 단순하고 효과적인 정규화 전략으로 극복하기 위해.
• DINOv2와 같은 강력한 자기지도 학습 인코더와의 표현 정렬이 더 빠른 수렴과 뛰어난 생성 성능을 이끌 수 있음을 입증하기 위해.
• 아키텍처 변경 없이도 적용 가능한 실용적이고 즉시 사용 가능한 방법을 제공하기 위해.

제안 방법

• REPA는 확산 트랜스포머의 노이즈가 많은 잠재 히든 상태의 투영을, 깨끗한 이미지에 적용된 사전 훈련된 시각 인코더(예: DINOv2)의 해당 특징과 정렬하는 정규화 손실을 도입한다.
• 이 방법은 확산 모델의 히든 상태를 외부 시각 인코더의 특징 공간과 동일한 공간으로 매핑하기 위해 선형 투영 헤드를 사용하여 특징 매칭을 가능하게 한다.
• 정렬 손실은 훈련 중에 계산되며, 확산 트랜스포머의 초기 레이어에만 적용되어 최소한의 계산 오버헤드를 유도한다.
• 이 방법은 DiT 및 SiT 아키텍처를 포함한 확산 및 플로우 기반 트랜스포머와 호환되며, VAE 임bedded 입력을 사용하는 잠재 확산 프레임워크와도 작동한다.
• 외부 인코더의 미세조정이 필요로 하거나 확산 모델 아키텍처를 수정할 필요 없이, 쉽게 구현이 가능하다.
• 훈련 과정은 분류기 자유 지도 학습을 활용하며, FID 및 사용자 연구를 포함한 표준 텍스트-이미지 벤치마크에서 평가된다.

Figure 1: Representation alignment makes diffusion transformer training significantly easier. Our framework, REPA, explicitly aligns the diffusion model representation with powerful pretrained visual representation through a simple regularization. Notably, model training becomes significantly more e

실험 결과

연구 질문

• RQ1 고품질 외부 시각적 특징과의 정렬이 확산 트랜스포머 표현을 통해 훈련 수렴을 크게 가속화할 수 있는가?
• RQ2 DINOv2와 같은 사전 훈련된 자기지도 학습 인코더와의 표현 정렬이 기존 확산 훈련 대비 생성된 이미지의 품질을 향상시키는가?
• RQ3 확산 모델 특징과 외부 시각적 특징 간의 정렬이 훈련 중에 어떻게 변화하는가? 그리고 명시적 정규화를 통해 가속화될 수 있는가?
• RQ4 제안된 REPA 방법은 DiT 및 SiT와 같은 다양한 확산 트랜스포머 아키텍처에서 효과적인가?
• RQ5 REPA는 기준 모델보다 훨씬 적은 훈련 스텝 수로 SOTA 생성 품질을 달성할 수 있는가?

주요 결과

• REPA는 SiT-XL/2의 성능을 7M 스텝 기준 모델과 동일하게 맞추기 위해 필요한 훈련 스텝 수를 40만 스텝 이하로 줄여 17.5배의 수렴 속도 향상을 달성한다.
• 분류기 자유 지도 학습을 사용할 경우 지도 스케일 4.0로 SOTA FID 점수 1.42를 달성하여 이전 방법들을 능가한다.
• REPA를 사용할 경우 다운스트림 분류 작업에서 선형 프로빙 정확도가 크게 향상되어 더 나은 의미적 표현 학습을 의미한다.
• DINOv2와의 표현 정렬은 장기적인 훈련과 더 큰 모델 스케일에서 점진적으로 향상되며, 외부 지도 학습의 가치를 확인한다.
• 노이즈가 있는 입력 조건에서도 REPA 정렬 손실은 확산 모델 특징과 강력한 자기지도 학습 인코더의 특징 간의 의미적 갭을 효과적으로 줄인다.
• 정성적 결과는 '소방차', '노트북', '우주 셔틀', '산호초'와 같은 다양한 클래스에서 고해상도, 다양성 있고 프롬프트에 부합하는 생성 결과를 보여준다.

(a) Semantic gap: Linear probing