VFM-VAE: Vision Foundation Models Can Be Good Tokenizers for Latent Diffusion Models

1. 문제 정의 (Problem)

잠재 확산 모델 (Latent Diffusion Models, LDMs) 의 성능은 시각적 토크나이저 (Visual Tokenizer, 주로 VAE) 의 품질에 크게 의존합니다. 최근 연구들은 Vision Foundation Model(VFM) 의 표현 능력을 토크나이저 학습에 활용하기 위해 증류 (Distillation) 기법을 도입했습니다 (예: VA-VAE).
그러나 저자들은 기존 증류 기반 접근법의 근본적인 한계를 발견했습니다:

• 표현의 취약성 (Brittleness): VFM 의 표현을 모방하도록 학습된 VAE 는, VFM 이 원래 가진 강력한 의미론적 (Semantic) 표현력이 증류 과정에서 손실되거나 약화됩니다.
• 반전 변환에 대한 민감성: 의미론적 변환 (노이즈 추가, 회전, 스케일 변경 등) 이 가해졌을 때, 기존 증류 기반 토크나이저들은 VFM 과의 정렬 (Alignment) 이 급격히 떨어지는 '취약성'을 보입니다. 이는 중요한 정보 손실을 의미합니다.

2. 제안 방법 (Methodology)

저자들은 증류를 통한 모방 학습을 버리고, 동결된 (Frozen) VFM 을 직접 LDM 의 인코더로 활용하는 새로운 아키텍처인 VFM-VAE를 제안합니다.

핵심 아키텍처

1. 동결된 VFM 인코더 (Frozen VFM Encoder):

   • VFM 을 재학습시키지 않고 고정된 상태로 사용합니다.
   • VFM 의 얕은 층 (Shallow), 중간 층 (Middle), 깊은 층 (Final) 에서 추출한 다중 스케일 특징 (Multi-scale features) 을 결합하여 풍부한 의미론적 정보를 보존합니다.
   • 이를 경량 프로젝션 네트워크를 통해 잠재 공간 (Latent Space) 으로 매핑합니다.

2. 고품질 재구성을 위한 새로운 디코더:

   • VFM 특징은 의미론적으로 풍부하지만 픽셀 수준의 재구성에는 부적합할 수 있으므로, 이를 해결하기 위해 두 가지 혁신을 도입했습니다.
   • 다중 스케일 잠재 융합 (Multi-Scale Latent Fusion): 전역 (Global) 특징과 공간 (Spatial) 특징을 분리하여 처리합니다. 전역 특징은 전체 스타일을 제어하고, 공간 특징은 초기 저해상도 블록에서 구조를 잡는 데 사용됩니다.
   • 점진적 해상도 재구성 블록 (Progressive Resolution Reconstruction Blocks): 8x8 에서 256x256 까지 해상도를 단계적으로 높이는 구조를 채택했습니다. 각 단계에서 Modulated ConvNeXt Block을 사용하여 전역 스타일 정보를 주입하고, 각 해상도 단계마다 ToRGB 헤드를 통해 직접적인 픽셀 손실 (Pixel-level supervision) 을 적용하여 학습 안정성을 높입니다.

3. 학습 목표 (Training Objectives):

   • 표현 정규화 손실 (Representation Regularization Loss): KL 발산과 VFM 특징과의 유사성 손실 (VF loss) 을 결합하여 잠재 공간이 VFM 의 의미론적 구조를 유지하도록 합니다.
   • 다중 해상도 재구성 손실: 각 디코더 블록마다 픽셀 단위 (L1) 와 지각적 (LPIPS, GAN) 손실을 적용하여 고해상도 재구성을 보장합니다.

4. 새로운 평가 지표: SE-CKNNA

   • 기존 CKNNA (Centered Kernel Nearest-Neighbor Alignment) 는 깨끗한 이미지에서만 유효할 수 있습니다.
   • 저자들은 SE-CKNNA (Semantic-Equivariant CKNNA) 를 제안하여, 의미론적 변환 (노이즈, 회전 등) 을 가한 다양한 입력에 대한 정렬 점수의 평균을 계산함으로써 토크나이저의 강건성 (Robustness) 을 더 정확하게 평가합니다.

3. 주요 기여 (Key Contributions)

1. 증류 없는 직접 활용: VFM 을 동결된 인코더로 직접 사용하여 증류로 인한 표현 열화를 방지하고, 특수한 디코더 설계를 통해 고품질 재구성을 달성했습니다.
2. SE-CKNNA 및 분석: 토크나이저와 VFM 간의 정렬을 평가하는 새로운 지표 (SE-CKNNA) 를 도입하고, 토크나이저의 표현 품질이 확산 모델 내부의 특징 학습에 어떻게 영향을 미치는지 체계적으로 분석했습니다.
3. 이중 측면 정렬 (Dual-side Alignment): 토크나이저 (VFM-VAE) 와 확산 모델 (Diffusion Model) 양쪽에서 VFM 정렬을 강화하는 전략을 통해 효율성과 성능을 극대화했습니다.

4. 실험 결과 (Results)

ImageNet 256x256 기준 실험 결과, VFM-VAE 는 기존 방법론을 압도하는 성능을 보였습니다.

• 학습 효율성 (Speedup):
  • 기존 토크나이저 (VA-VAE 등) 가 800 에포크 이상 학습해야 달성하던 성능을, VFM-VAE 는 80 에포크 만에 달성했습니다 (약 10 배 빠른 수렴).
  • 80 에포크 학습 시 gFID (w/o CFG) 2.22 달성.
  • 640 에포크 학습 시 gFID (w/o CFG) 1.62 달성 (CFG 적용 시 1.31). 이는 기존 SOTA 대비 매우 우수한 수치입니다.
• 표현 강건성:
  • SE-CKNNA 평가에서 기존 VA-VAE 는 변환 시 정렬 점수가 33.2% 감소했으나, VFM-VAE 는 오히려 1.6% 증가하여 VFM 특징과의 정렬이 매우 강건함을 증명했습니다.
• 다양한 VFM 호환성:
  • SigLIP2, DINOv2, EVA-CLIP 등 다양한 VFM 백본과 호환되며, SigLIP2 를 사용할 때 텍스트 - 이미지 생성 (Text-to-Image) 에서도 우수한 성능을 보였습니다.

5. 의의 및 결론 (Significance)

이 논문은 Vision Foundation Model 이 단순히 지식 증류의 원천이 아니라, 잠재 확산 모델의 핵심 토크나이저로서 직접 활용될 수 있음을 입증했습니다.

• 효율성: 동결된 VFM 을 활용함으로써 대규모 VAE 학습 비용을 절감하면서도, 더 강력한 의미론적 표현력을 확보할 수 있습니다.
• 품질: 기존 증류 방식의 한계를 극복하고, VFM 의 원래 의미론적 특성을 보존한 채 고해상도 이미지를 재구성 및 생성할 수 있습니다.
• 미래 방향: VFM 기반 토크나이저는 LDM 의 학습 속도를 획기적으로 높이고 생성 품질을 개선하는 새로운 패러다임을 제시합니다.

요약하자면, VFM-VAE는 "증류 (Distillation) 를 피하고 동결된 VFM 을 직접 활용하여, 더 빠르고 강력하며 견고한 잠재 확산 모델을 만드는 방법"을 제시한 획기적인 연구입니다.