MacTok: 이미지 생성을 위한 강건한 연속 토큰화 (Robust Continuous Tokenization for Image Generation)

이 논문은 MacTok이라는 새로운 연속 토큰화 (Continuous Tokenization) 방법을 제안합니다. 이는 기존 KL-VAE 기반의 연속 토큰화 모델이 겪는 후방 붕괴 (Posterior Collapse) 문제를 해결하면서도, 매우 적은 수의 토큰 (64 또는 128 개) 으로 고품질의 이미지 재구성과 생성을 가능하게 합니다.

1. 문제 정의 (Problem)

이미지 생성 모델 (Diffusion, Autoregressive 등) 의 효율성을 높이기 위해 이미지를 압축된 잠재 공간 (Latent Space) 으로 변환하는 토큰화가 필수적입니다. 토큰화는 크게 이산형 (Discrete, 예: VQ-VAE) 과 연속형 (Continuous, 예: KL-VAE) 으로 나뉩니다.

• 연속형 토큰화의 한계: KL-VAE 는 부드러운 잠재 표현을 학습할 수 있지만, 강한 압축 (적은 토큰 수) 을 적용할 때 후방 붕괴 (Posterior Collapse) 현상이 빈번하게 발생합니다.
  • 원인: KL 발산 (KL Divergence) 정규화가 너무 강하게 작용하면, 인코더가 의미 있는 정보를 잠재 공간에 인코딩하는 대신 사전 분포 (Prior) 에만 의존하게 됩니다.
  • 결과: 잠재 변수가 입력 이미지의 정보를 담지 못하게 되어, 재구성 품질이 떨어지고 생성 모델의 성능이 저하됩니다.
• 기존 해결책의 부족: 기존 연구들은 KL 가중치 조정 등을 통해 이를 완화하려 했으나, 이는 민감한 하이퍼파라미터 튜닝이 필요하며 근본적인 해결이 되지 못했습니다.

2. 방법론 (Methodology)

MacTok 은 마스킹 (Masking) 과 표현 정렬 (Representation Alignment) 을 결합하여 후방 붕괴를 방지하고 강건한 잠재 표현을 학습합니다.

2.1. 아키텍처

• ViT 기반: 인코더와 디코더 모두 Vision Transformer (ViT) 를 사용합니다.
• 1D 토큰화: 이미지를 1 차원 시퀀스로 변환하여 유연한 토큰 길이 (64, 128 등) 를 지원합니다.
• KL-VAE 프레임워크: 가우시안 사전 분포를 가진 연속 잠재 공간을 가정합니다.

2.2. 이미지 마스킹 (Image Masking)

후방 붕괴를 막기 위해 인코더가 불완전한 시각 정보를 바탕으로 이미지를 재구성하도록 강제합니다. 두 가지 전략을 50% 확률로 혼합하여 사용합니다.

1. 랜덤 마스킹 (Random Masking): 입력 이미지 패치 중 일부를 무작위로 마스킹합니다. 이는 잠재 변수가 누락된 영역을 추론하기 위해 필수적인 정보를 인코딩하도록 유도합니다.
2. DINO 기반 시맨틱 마스킹 (Semantic Masking):
   • DINOv2 의 분류 토큰 (CLS token) 과 패치 토큰 간의 유사도를 계산하여 가장 정보량이 많은 영역 (정보성 높은 패치) 을 식별합니다.
   • 이러한 중요한 영역을 마스킹함으로써, 모델이 단순한 국소 패턴이 아닌 전체적인 구조와 의미 (Global Structural Semantics) 를 학습하도록 강제합니다.
   • 이는 이미지 공간의 시맨틱 사전 지식을 잠재 공간으로 전이시키는 효과를 가집니다.

2.3. 글로벌 및 로컬 표현 정렬 (Global & Local Representation Alignment)

잠재 공간이 DINOv2 와 같은 사전 훈련된 시각 모델의 특징과 정렬되도록 하여 의미론적 일관성을 강화합니다.

• 로컬 정렬 (Local Alignment): 잠재 토큰을 DINOv2 의 패치 특징과 정렬하여 세부적인 구조와 지역적 의미를 보존합니다.
• 글로벌 정렬 (Global Alignment): 잠재 공간의 평균 풀링된 특징을 DINOv2 의 CLS 토큰과 정렬하여 전체 이미지의 의미적 맥락을 유지합니다.
• 효과: 이는 잠재 공간이 잘 구조화되고 (Well-structured), 구별력 (Discriminative) 이 높도록 하여 재구성과 생성 품질을 동시에 향상시킵니다.

2.4. 학습 목적 함수

재구성 손실, 지각적 손실 (Perceptual Loss), 적대적 손실 (GAN Loss), KL 정규화, 그리고 제안된 표현 정렬 손실 (Representation Alignment Loss) 을 결합하여 최적화합니다.

3. 주요 기여 (Key Contributions)

1. 후방 붕괴의 근본적 해결: 연속 토큰화에서의 후방 붕괴를 "과도한 정규화"로 정의하고, 불완전한 입력으로부터의 재구성을 통해 정보 흐름을 유지함으로써 이를 효과적으로 방지하는 메커니즘을 제시했습니다.
2. MacTok 프레임워크 제안: 랜덤 마스킹과 DINO 기반 시맨틱 마스킹, 그리고 글로벌/로컬 표현 정렬을 통합하여, 강한 압축 환경에서도 강건하고 의미론적으로 구조화된 잠재 공간을 학습합니다.
3. 압도적인 성능 달성: ImageNet 데이터셋에서 64 개 또는 128 개의 토큰만 사용하여 기존 256~1024 개 토큰이 필요한 모델들을 능가하는 성능을 달성했습니다.

4. 실험 결과 (Results)

ImageNet 256x256 및 512x512 해상도에서 SiT-XL 및 LightningDiT 생성 모델과 결합하여 평가되었습니다.

• 생성 품질 (gFID):
  • 256x256: 128 토큰 사용 시 gFID 1.44 (SiT-XL 사용, CFG 적용 시).
  • 512x512: 128 토큰 사용 시 gFID 1.52 (SiT-XL 사용, CFG 적용 시). 이는 해당 해상도에서 State-of-the-Art (SOTA) 성능입니다.
  • 기존 KL-VAE 기반 모델들 (SoftVQ-VAE, MAETok 등) 보다 64 배까지 적은 토큰 수로 더 높은 품질을 달성했습니다.
• 재구성 품질 (rFID): 64 토큰 기준 rFID 0.75 (256x256), 0.89 (512x512) 를 기록하며, 적은 토큰으로도 높은 재구성 정확도를 유지했습니다.
• 잠재 공간 분석: t-SNE 및 UMAP 시각화를 통해 MacTok 은 후방 붕괴가 발생한 모델과 달리, 의미론적으로 잘 분리되고 구조화된 잠재 공간을 형성함을 확인했습니다. 또한 선형 프로빙 (Linear Probing) 정확도가 높아 잠재 표현의 질이 우수함을 입증했습니다.

5. 의의 및 결론 (Significance)

MacTok 은 연속형 토큰화 분야에서 오랫동안 해결되지 않았던 후방 붕괴 문제를 마스킹 전략과 표현 정렬을 통해 성공적으로 해결했습니다.

• 효율성: 생성 모델의 학습 및 추론 비용을 획기적으로 줄일 수 있는 1D 연속 토큰화 솔루션을 제공합니다.
• 고충실도: 적은 토큰 수로도 고해상도 이미지의 세부 사항과 전역 구조를 모두 보존하는 고품질 생성이 가능합니다.
• 미래 전망: 이 연구는 마스킹과 사전 훈련된 시각 모델의 지식을 결합하는 것이 잠재 공간 학습의 핵심임을 보여주며, 향후 더 효율적이고 강력한 이미지 생성 모델 개발의 방향성을 제시합니다.

요약하자면, MacTok 은 적은 토큰으로 더 좋은 이미지를 만들어내는 새로운 표준을 제시하며, 연속형 토큰화의 실용성을 크게 높인 획기적인 연구입니다.