Self-transcendence: Is External Feature Guidance Indispensable for Accelerating Diffusion Transformer Training?

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 확산 모델 (Diffusion Models), 특히 확산 트랜스포머 (DiT, Diffusion Transformer) 는 이미지 생성 분야에서 뛰어난 성능을 보이고 있지만, 학습에 많은 계산 자원이 소요되고 수렴 속도가 느린 문제가 있습니다.
• 기존 접근법의 한계:
  • 최근 연구 (REPA 등) 는 DINO 와 같은 외부 사전 학습된 비전 인코더 (External Vision Encoders) 를 사용하여 중간 특징 (Intermediate Features) 을 정렬함으로써 DiT 학습을 가속화하고 생성 품질을 향상시켰습니다.
  • 그러나 이러한 방법은 외부 모델에 대한 의존성을 초래하며, 추가적인 데이터 전처리와 모델 유지 비용이 발생합니다.
• 핵심 질문: "DiT 학습을 가속화하기 위해 외부 특징 지도 (External Feature Guidance) 가 필수적인가?"
• 문제점: 외부 의존성을 제거하려는 시도가 있었으나 (LayerSync, SRA 등), 내부에서 생성된 특징이 초기 학습 단계에서 구조적으로 불완전하고 의미론적으로 미숙하여 외부 지도 방식 (REPA) 보다 성능이 떨어지는 경향이 있었습니다.

2. 제안 방법: Self-Transcendence (Methodology)

저자들은 DiT 모델이 외부 지도 없이 자기 자신의 내부 특징을 활용하여 학습을 가속화할 수 있다고 주장하며, 이를 실현하는 Self-Transcendence라는 두 단계 (Two-stage) 학습 전략을 제안합니다.

핵심 아이디어

효과적인 내부 지도 특징은 다음 두 가지 조건을 충족해야 합니다:

1. 구조적 정결함 (Structurally Clean): 얕은 레이어 (Shallow blocks) 가 노이즈와 신호를 명확히 분리할 수 있도록 돕는 깨끗한 구조.
2. 의미론적 구분력 (Semantically Discriminative): 얕은 레이어가 유익한 표현을 학습할 수 있도록 돕는 의미론적 정보.

학습 파이프라인 (2 단계)

1. 1 단계: VAE 구조 지도 (VAE Structure Guidance)

   • 목적: 초기 학습 단계에서 얕은 레이어가 노이즈와 신호를 구분할 수 있도록 구조적 정결함을 제공.
   • 방식: 표준 잠재 확산 (Latent Diffusion) 프레임워크의 구성 요소인 VAE 잠재 특징 (VAE Latent Features) 을 사용하여 DiT 의 중간 특징을 정렬합니다.
   • 손실 함수: DiT 의 중간 특징을 경량 MLP 를 통해 변환한 후, Ground-truth VAE 잠재 벡터 $z$와 L2 손실로 정렬합니다.
   • 효과: 외부 데이터 없이도 구조적으로 깨끗한 표현을 빠르게 학습하게 하여 수렴을 가속화합니다.

2. 2 단계: 자기 주도 표현 정렬 (Self-guided Representation Alignment)

   • 목적: 1 단계 학습이 완료된 모델의 내부 특징을 활용하여 의미론적 구분력을 강화.
   • 방식:
     • 부분적으로 학습된 모델 (Teacher) 의 깊은 레이어 특징을 추출합니다.
     • 클래스프리 가이드 (Classifier-Free Guidance, CFG) 개념을 특징 공간에 적용합니다. 조건부 (Conditional) 와 무조건부 (Unconditional) 특징을 추출하여 가이드 스케일 ($\omega$) 로 결합함으로써, 내부 특징의 의미론적 표현력을 증폭시킵니다.
     • 이렇게 강화된 내부 특징 ($f_g$) 을 새로운 DiT 모델 (Student) 의 얕은 레이어 학습을 위한 지도 신호로 사용합니다.
   • 특징: 외부 인코더 없이 모델 자체의 내부 표현을 정제하여 지도 신호로 활용합니다.

3. 주요 기여 (Key Contributions)

1. 완전 자기 주도 학습 프레임워크: 외부 특징 (DINO 등) 에 의존하지 않고 DiT 모델의 내부 특징만으로 학습을 가속화하는 Self-Transcendence를 제안했습니다.
2. 효율적인 2 단계 파이프라인:
   • VAE 잠재 특징을 통한 구조적 지도 (1 단계) 와 CFG 기반의 의미론적 강화 (2 단계) 를 결합하여, 내부 특징의 품질을 극대화했습니다.
3. 성능 입증:
   • 외부 지도 방식인 REPA 를 능가하거나 동등한 성능을 달성하면서도, 외부 데이터와 모델 의존성을 제거했습니다.
   • ImageNet 및 텍스트 - 이미지 생성 태스크에서 수렴 속도와 생성 품질 (FID, IS) 모두에서 기존 자기 주도 방법 (SRA, LayerSync) 보다 월등히 우수한 결과를 보였습니다.

4. 실험 결과 (Results)

• ImageNet 256x256 (SiT-XL/2):
  • 80 에포크 학습 시, 제안된 방법은 FID 7.51을 기록하여, 200 에포크 학습한 LayerSync (FID 8.80) 보다 우수하고, 외부 DINO 를 사용하는 REPA (FID 7.90) 보다도 더 좋은 성능을 보였습니다.
  • LightningDiT-XL/1 모델에서는 64 에포크 학습 시 FID 3.55를 달성했습니다.
• ImageNet 512x512 및 텍스트 - 이미지 생성:
  • 고해상도 및 텍스트 조건부 생성 태스크에서도 REPA 를 능가하는 성능을 보여주었습니다 (예: Text-to-Image FID 4.56 vs REPA 4.90).
• 학습 비용 및 효율성:
  • 외부 모델 (DINOv2) 의 사전 학습 비용이 전혀 들지 않습니다.
  • 메모리 오버헤드는 초기 단계에만 일시적으로 증가하며, 전체 학습 시간은 베이스라인과 유사하거나 REPA 보다 빠릅니다.
• 시각화 분석:
  • PCA 시각화를 통해 제안된 방법이 얕은 레이어에서도 더 구조화되고 의미론적으로 분리된 (Discriminative) 특징을 학습함을 확인했습니다.

5. 의의 및 결론 (Significance)

• 내부 표현의 잠재력 재발견: 확산 모델이 외부 지도 없이도 자체적인 내부 표현을 정제하여 효과적인 지도 신호로 활용할 수 있음을 입증했습니다.
• 실용성 향상: 외부 의존성을 제거함으로써 학습 파이프라인을 단순화하고, 계산 비용을 절감하며, 다양한 도메인으로의 확장성을 높였습니다.
• 미래 방향: DiT 학습 가속화를 위한 새로운 자기지도 (Self-supervised) 패러다임을 제시하며, 향후 비디오나 3D 생성 등 다른 생성 모달리티로의 확장을 위한 기초를 마련했습니다.

요약하자면, Self-Transcendence는 "외부 지도가 필수적인가?"라는 질문에 **"아니오, 모델 내부의 정제된 특징을 올바르게 활용하면 외부 지도보다 더 빠르고 우수한 학습이 가능하다"**는 것을 증명하는 획기적인 연구입니다.