기술 요약: 표현 오토인코더 (RAEv2) 를 통한 개선된 베이스라인

문제 정의

표현 오토인코더 (RAE) 는 시각적 이해와 생성을 위한 토큰화를 통합하는 것을 목표로 하여, 확산 트랜스포머 훈련에서 기존 변분 오토인코더 (VAE) 를 대체할 수 있는 유망한 프레임워크로 부상했습니다. 그러나 원래 RAE 공식은 실제 채택을 방해하는 세 가지 중요한 병목 현상에 직면해 있습니다:

1. 비최적의 재구성: 전문화된 데이터 (예: 텍스트 또는 얼굴) 없이 훈련될 때 특히, 재구성 성능은 전문화된 VAE 보다 뒤처집니다.
2. 가이드런스와 비호환성: RAE 는 기존 클래스프리 가이드런 (CFG) 과 호환되지 않아, 오토가이드런 (AG) 을 위해 2 차의 더 약한 확산 모델을 훈련해야 합니다. 이는 상당한 계산 오버헤드와 복잡성을 추가합니다.
3. 활용되지 않는 인코더 표현: 이전 연구들은 잠재 표현으로 사전 훈련된 비전 인코더의 최종 레이어만 사용하는 것을 기본으로 삼아, 중간 레이어에 존재하는 풍부하고 다양한 추상화를 무시했습니다.

방법론

이 논문은 인코더에 대한 훈련 없는 접근 방식을 유지하면서 세 가지 주요 통찰력과 설계 변경을 통해 이러한 문제를 해결하는 개선된 베이스라인인 RAEv2를 제안합니다.

1. 일반화된 표현 인코더

저자들은 사전 훈련된 비전 인코더의 최종 레이어 ($K=1$) 만 의존하는 대신, 잠재 표현이 마지막 $K$개 인코더 레이어의 합인 일반화된 공식을 제안합니다.

• 메커니즘: 출력은 $\mathbf{x} = \sum_{\ell=L-K+1}^{L} \mathbf{z}_\ell$로 정의되며, 여기서 $\mathbf{z}_\ell$은 레이어 $\ell$의 특징 맵을 나타냅니다.
• 근거: 서로 다른 레이어는 상호 보완적인 특징을 포착합니다. 초기부터 중간 레이어는 더 정교한 공간 구조를 유지하는 반면, 후기 레이어는 글로벌 의미를 강조합니다. 이를 집계하면 잠재 차원성을 증가시키거나 추가 학습 파라미터를 요구하지 않으면서도 기하학적 구조를 보존합니다.
• 구현: 저자들은 단순 덧셈 (MLS) 과 무작위 행렬 투영 (MLR) 을 비교하여, 단순 덧셈이 재구성을 개선하면서도 일관되게 더 나은 생성 성능을 낸다는 것을 발견했습니다.

2. RAE 와 REPA 의 상호 보완적 메커니즘

이 논문은 RAE 가 표현 정렬 (REPA) 의 필요성을 제거한다는 지배적인 가정을 도전합니다. REPA 는 동일한 사전 훈련된 표현을 확산 중간 레이어로 증류합니다.

• 발견: 27 개 비전 인코더에 대한 실증 분석은 RAE 와 REPA 가 상호 보완적 메커니즘을 통해 작동함을 보여줍니다.
  • RAE는 글로벌 의미를 활용하여 (선형 프로빙 정확도로 측정) 의미적으로 풍부한 잠재 공간을 제공합니다.
  • REPA는 중간 확산 특징의 공간 구조를 정규화합니다 (국소 거리 유사성, LDS 로 측정).
• 결과: 동일한 표현을 인코더 (RAE) 로 그리고 중간 레이어의 목표 (REPA) 로 사용하는 것은 일관되게 생성 성능을 향상시킵니다. 이 시너지는 글로벌 및 공간 지표 모두에서 뛰어난 성능을 보이는 더 강력한 인코더 (예: DINOv3-L) 가 우수한 생성 결과를 달성할 수 있게 합니다.

3. 자기 가이드런 (x-예측) 으로서의 REPA

원래 RAE 는 가이드런을 위해 별도의 약한 모델을 필요로 합니다. 저자들은 REPA 헤드를 재구성하여 내부 가이드런을 가능하게 합니다.

• 통찰: RAE 잠재 공간에서 REPA 프로젝션 헤드 (초기 레이어 특징을 깨끗한 잠재로 매핑) 는 효과적으로 x-예측을 수행합니다.
• 메커니즘: 전체 모델 출력을 x-예측이 되도록 재매개변수화함으로써, REPA 헤드의 예측 ($\hat{\mathbf{x}}_{repa}$) 은 전체 모델의 예측 ($\hat{\mathbf{x}}_{full}$) 에 비해 "더 약한" 베이스라인 역할을 합니다.
• 가이드런 공식: 가이드런은 단일 순전파 내에서 내부적으로 적용됩니다: $\hat{\mathbf{x}}_{guided} = \hat{\mathbf{x}}_{full} + w \cdot (\hat{\mathbf{x}}_{full} - \hat{\mathbf{x}}_{repa})$.
• 이점: 이는 별도의 오토가이드런 모델의 필요성을 제거하고 CFG 가 요구하는 추가 순전파를 피하여, 가이드런에 대한 함수 평가 횟수 (NFEs) 를 효과적으로 절반으로 줄입니다.

주요 결과

RAEv2 는 ImageNet, 텍스트-이미지, 및 내비게이션 월드 모델 전반에서 수렴 속도, 생성 품질, 및 재구성에서 상당한 개선을 보여줍니다.

• 훈련 효율성: RAEv2 는 원래 RAE 대비 10 배 빠른 수렴을 달성합니다. ImageNet-256 에서 80 에포크 만에 최첨단 가이드런 gFID 1.06에 도달하는 반면, 원래 RAE 는 3.26 에 도달하는 데 800 에포크가 필요했습니다.
• 새로운 효율성 지표 (EPFID): 이 논문은 절대 gFID 보다 더 실용적인 지표인 EPFID@k (유도되지 않은 gFID $\le k$에 도달하는 에포크) 를 도입합니다. RAEv2 는 EPFID@2 를 35 에포크로 달성하여, 원래 RAE(177 에포크) 및 이전 작업들 (>480 에포크) 보다 획기적인 개선을 보여줍니다.
• 생성 품질: 최근 제안된 표현 프리체트 거리 ($FD_r$) 에서 RAEv2 는 80 에포크 만에 2.17 점수를 달성하여, 후처리를 하지 않고도 800 에포크 동안 훈련된 이전 최선인 3.26 을 능가합니다.
• 재구성: 일반화된 공식 (특히 $K=23$) 은 재구성 (rFID 0.18) 과 생성 사이의 파레토 최적 균형을 달성하여, ImageNet 만으로 훈련되었음에도 불구하고 Flux-VAE, SD3.5 와 같은 독점 VAE 를 능가합니다.
• 일반화:
  • 텍스트-이미지: RAEv2 는 GenEval, DPG-Bench, GenAI-Bench 에서 Flux-VAE 및 원래 RAE 보다 일관된 개선을 보여줍니다.
  • 내비게이션 월드 모델: 미래 프레임 예측 작업 (RECON 데이터셋) 에서 RAEv2-NWM 은 105.61의 FVD 를 달성하여 NWM(200.97) 과 RAE(312.01) 를 크게 능가하며, 자기회귀 롤아웃에서 깜빡임이 감소합니다.

의의 및 주장

이 논문은 RAEv2 가 표현 오토인코더를 위한 단순화되고 매우 효율적인 베이스라인을 제공한다고 주장합니다. 주요 기여점은 다음과 같습니다:

1. 단순화: 인코더 미세 조정, 전문화된 훈련 데이터, 또는 별도의 가이드런 모델의 필요성을 제거합니다.
2. 효율성: 최첨단 성능에 도달하는 데 필요한 훈련 시간을 획기적으로 단축하여, 절대 지표 개선에서 훈련 효율성 (EPFID) 으로 초점을 이동시킵니다.
3. 통합 토큰화: 고정된 인코더의 글로벌 의미 이해 능력을 보존하면서 (선형 프로빙을 통해 검증됨) 생성과 재구성을 강화함으로써, 사전 훈련된 비전 인코더를 이해 및 생성 작업 모두에 대한 통합 토큰화 방식으로 사용하는 타당성을 강화합니다.

저자들은 이러한 통찰력이 다양한 생성 모델링 시나리오에서 표현 오토인코더의 실제 채택을 촉진하기를 바랍니다.