[2025-2] 전연주 - MaeTok: Masked Autoencoders Are Effective Tokenizers for Diffusion Models

논문 링크: arxiv.org/pdf/2502.03444

깃허브 링크: Hhhhhhao/continuous_tokenizer

허깅페이스 링크: MAETok/maetok-b-128 · Hugging Face

---

Abstract

최근 latent diffusion model의 발전은 고해상도 이미지 생성에서 뛰어난 성능을 보였다. 그러나 diffusion model 학습과 생성에 있어 좋은 latent 공간의 성질에 대해서는 아직 충분히 연구되지 않았다. 본 논문에서는 이론적·실험적 분석을 통해, discriminative한 특성과 단순한 구조(적은 Gaussian Mixture 모드)를 갖는 latent 분포가 더 우수한 생성 품질로 이어진다는 사실을 밝혔다.

이에 따라 reconstruction fidelity를 유지하면서도 의미 있는(latent space가 semantically rich한) 표현을 학습할 수 있도록 mask modeling 기반 Autoencoder인 MAETok을 제안한다. 기존의 VAE와 달리, variational한 제약(KL constraint) 없이도 성능을 극대화할 수 있음을 보인다. 특히, MAETok은 단 128개의 token만으로도 ImageNet 512×512 이미지 생성에서 gFID 1.69, 31배 빠른 추론, 76배 빠른 학습 성능을 기록했다.

이 연구는 diffusion model의 성능에 있어 중요한 요소는 latent 공간의 구조이며, variational 형태(VAE)는 필수적이지 않다는 점을 강조한다. 

---

1. Introduction

Diffusion Model은 최근 다양한 이미지 생성 과제에서 SOTA 성능을 달성하며 주목받고 있다. 초기에는 pixel space 상에서 작동하였으나, 이후 연구들은 tokenizer를 통해 압축된 latent space에서 diffusion을 수행하는 것이 훨씬 효율적임을 보여주었다.

Problems

대부분의 latent diffusion model은 VAE(Variational Autoencoder)를 tokenizer로 활용한다. 이는 latent 분포가 부드럽고(smooth) 연속적인 구조를 가지도록 유도하지만, 그 대가로 재구성 이미지의 품질이 저하되는 문제가 있다.

반면 plain Autoencoder(AE)는 고품질 재구성이 가능하지만, 그 latent space는 너무 entangled(얽혀있고 복잡)하거나 구조화되지 않아 diffusion model 학습에 부적절할 수 있다.

diffusion model 학습에 효과적인 "좋은 latent space"란 어떤 것인가?

기존의 smooth한 분포가 아니라도, 클래스 간 분리가 잘된 구조적(discriminative + structured) latent space가 diffusion에 더 적합할 수 있다.

MAETok

이 질문에 답하기 위해, AE, VAE, 그리고 representation-aligned VAE의 latent space를 Gaussian Mixture Model (GMM)으로 분석하였다. 그 결과:

• GMM 모드 수가 적고
• 특징(feature)이 더 뚜렷하게 분리된 latent space가

→ diffusion model의 학습 손실이 낮고, 생성 품질이 우수하다는 사실을 발견하였다.

이에 따라 기존 AE를 확장한 Masked Autoencoder Tokenizer (MAETok)를 제안한다. 이 구조는 다음과 같은 특징을 갖는다:

• Transformer 기반 AE 구조
• 입력 이미지의 일부 patch를 mask하고, 그 복원을 통해 학습
• 다중 목표(Multi-target):
  • 픽셀 복원
  • HOG feature
  • DINOv2, CLIP 등 의미 기반 representation 복원
• Auxiliary Decoder를 통해 다양한 수준의 정보를 동시에 학습

이러한 설계를 통해, encoder는 더 풍부하고 구별력 있는 latent 표현을 학습하고, decoder는 고해상도 복원 성능을 확보한다.

Contributions

• 이론 및 실험 분석을 통해, latent space의 구조가 diffusion 성능에 미치는 영향을 규명하였다.
• MAETok은 단순 AE 구조임에도 불구하고 mask modeling과 multi-target reconstruction을 통해 SOTA 성능을 달성하였다.
• ImageNet 512×512에서 128 tokens만으로 기존 2B 파라미터 모델보다 우수한 gFID 1.69, IS 304.2 성능을 달성하였다.

---

2. On the Latent Space and Diffusion Models

좋은 latent space가 diffusion model의 학습과 성능에 어떤 영향을 주는지를 분석하기 위해, 다음 세 가지 tokenizer를 비교하였다:

• AE (Autoencoder)
• VAE (Variational Autoencoder)
• VAVAE (Representation-Aligned VAE)

이들은 동일한 학습 설정과 latent 차원을 갖도록 설정해, 비교의 공정성을 확보하였다. 각 tokenizer의 latent space에 대해 diffusion model을 학습시킨 후, 이미지 생성 품질(gFID)과 학습 난이도(diffusion loss)를 비교했다.

2.1 Empirical Analysis(실증적 분석)

기존 이론에 따르면, diffusion 모델의 학습 품질은 latent 분포의 구조 복잡도에 영향을 받는다. 특히 latent 분포가:

• 너무 복잡하거나 multi-modal 구조일 경우 → 학습 난이도 증가
• 더 간단하고 구조화된 경우 → 학습 효율 증가

즉, latent 분포가 얼마나 separable(구분 가능한)지, 즉 얼마나 discriminative한지가 중요하다는 것이다.

GMM(Gaussian Mixture Model)을 이용한 분석

latent 공간의 구조를 정량적으로 평가하기 위해, 각 tokenizer의 latent 분포에 GMM을 fitting하였다.

• GMM의 모드 수(mode)는 분포의 복잡도를 나타낸다.
• 적은 모드 수 = 간단한 구조 → diffusion loss가 작음

[a] GMM 모드 수에 따른 NLL(Negative Log-Likelihood) 비교
[b] Diffusion 모델 학습 시의 손실(loss) 비교

→ MAETok과 VAVAE는 적은 GMM 모드 수로도 충분히 분포를 표현함
→ 이에 따라 diffusion loss가 낮고, gFID도 작아짐

즉,

→ latent 분포가 더 단순하고 구별력이 있을수록, diffusion model 학습이 더 쉽고 성능이 높다.

2.2 Theoretical Analysis(이론적 분석)

• 데이터 분포는 평균이 다른 K개의 Gaussian 분포의 혼합(GMM)이라고 가정
• 데이터 차원은 d이며, 각 mode의 평균(norm)은 상한 B를 가짐

이러한 GMM 구조에서 DDPM(Denoising Diffusion Probabilistic Models)이 주어진 수준의 생성 품질을 달성하기 위해 필요한 샘플 수(n')를 수학적으로 분석

• GMM 기반 데이터 분포$p_0 = \frac{1}{K} \sum_{i=1}^{K} \mathcal{N}(\mu^*_i, I)$
• DDPM의 score matching loss:$$\min_w \mathbb{E}\left[\|s_w(x, t) - \nabla_x \log p_t(x)\|^2\right]$$
• 주요 결과 (Theorem 2.1):$$n' = \Theta\left( \frac{K^4 d^5 B^6}{\epsilon^2} \right)$$
• Diffusion model이 오류 $O(T\epsilon^2)$이하의 KL divergence를 달성하기 위해 필요한 샘플 수:

→ K(모드 수)가 많아질수록, 필요한 샘플 수가 급격히 증가한다는 점을 수학적으로 증명

즉,

• diffusion 모델의 학습 성능은 latent space의 구조, 특히 GMM 모드 수 K에 따라 큰 영향을 받음
• 실험적으로도, 모드 수가 적은 MAETok이 가장 낮은 diffusion loss와 gFID를 달성
• 이론적으로도, 모드 수가 많을수록 학습이 어려워진다는 점을 KL divergence 기준으로 입증

---

3. Method

앞선 이론 및 실험 결과로부터, diffusion model을 잘 학습시키기 위해 반드시 VAE처럼 variational constraint(KL loss)를 적용할 필요는 없다는 결론을 도출하였다. 단순한 Autoencoder(AE) 구조라도 discriminative한 latent space를 학습하면 충분하다.

이에 따라 제안하는 모델이 바로 MAETok (Masked AutoEncoder Tokenizer)이다.

3.1 Architecture

MAETok은 최신 1D tokenizer 아키텍처를 기반으로 하며, 핵심 구성은 다음과 같다:

Encoder

• Vision Transformer (ViT) 기반
• 입력 이미지 $I \in \mathbb{R}^{H \times W \times 3}$를 patch 단위로 분할해 토큰화함
• 각 패치는 D 차원으로 임베딩됨 → $\mathbf{x} \in \mathbb{R}^{N \times D}$
• 추가적으로 학습 가능한 latent tokens $\mathbf{z} \in \mathbb{R}^{L \times D}$를 사용
• 최종 입력: patch tokens + latent tokens → $[\mathbf{x}; \mathbf{z}] \in \mathbb{R}^{(N+L) \times D}$
• Encoder 출력: 오직 latent token에 대한 output → $\mathbf{h} = \mathcal{E}([\mathbf{x}; \mathbf{z}]) \in \mathbb{R}^{L \times H}$

Decoder

• reconstruction을 위한 이미지 토큰 $\mathbf{e} \in \mathbb{R}^{N \times H}$를 learnable하게 설정
• decoder 입력: $[\mathbf{e}; \mathbf{h}]$ → 이미지 복원
• 최종 복원 결과: linear layer를 거쳐 $\hat{\mathbf{x}} \to \hat{I}$

Position Encoding

• patch tokens: 2D RoPE 사용
• latent tokens: 1D absolute positional embedding 사용

Training Loss

$$\mathcal{L} = \mathcal{L}_{\text{recon}} + \lambda_1 \mathcal{L}_{\text{percep}} + \lambda_2 \mathcal{L}_{\text{adv}}$$​

• $\mathcal{L}_{\text{recon}}$​: 픽셀 단위 MSE
• $\mathcal{L}_{\text{percep}}$​: perceptual loss (e.g., VGG 기반)
• $\mathcal{L}_{\text{adv}}$​: GAN 기반 adversarial loss
• VAE처럼 prior/posterior KL loss는 사용하지 않음

3.2 Mask Modeling

MAETok의 핵심은 MAE처럼 self-supervised mask modeling을 적용한다는 점이다.

Masking Strategy

• encoder 입력 이미지의 40–60% patch를 무작위로 마스킹
• 마스크된 위치는 learnable한 mask token $\mathbf{m} \in \mathbb{R}^{D}$으로 대체

Auxiliary Decoders

• pixel 복원 외에도 추가적인 feature target들을 예측
  • HOG (edge 정보)
  • DINOv2 (semantic features)
  • CLIP (vision-language features)
  • (LAION-COCO에선 text BPE index 예측도 추가됨)
• 각 auxiliary decoder는 얕은(3 layer) 구조의 transformer로 구성되며, loss는 다음과 같이 계산:

$$\mathcal{L}_{\text{mask}} = \sum_{j} \| \mathbf{M} \otimes (\hat{\mathbf{y}}_j - \mathbf{y}_j) \|^2$$

• 여기서 $\mathbf{M}$: 마스크 인디케이터, $\otimes$: elementwise product

목적: 다양한 수준의 정보를 학습시켜 더 구조화된(discriminative) latent space 획득

3.3 Pixel Decoder Fine-Tuning

mask modeling을 적용하면 encoder는 좋은 latent를 학습하지만, decoder는 마스크된 입력 때문에 reconstruction 성능이 저하될 수 있음.

해결책: Decoder Fine-Tuning

1. 학습 완료된 encoder를 freeze
2. decoder만 full image로 추가 학습 (mask 없음)
3. 복원 성능이 향상됨

이 과정에서도 loss는 동일하게:

$$\mathcal{L} = \mathcal{L}_{\text{recon}} + \lambda_1 \mathcal{L}_{\text{percep}} + \lambda_2 \mathcal{L}_{\text{adv}}$$​

→ semantic 정보는 유지하면서도, 복원 품질까지 확보 가능

즉,

• MAETok은 plain AE에 MAE 스타일의 mask modeling과 auxiliary decoder를 추가하여 VAE 없이도 SOTA 성능을 달성
• 핵심 설계 요소:
  • 학습 가능한 latent token: 고정된 128개
  • mask modeling: encoder에 자기지도 학습 신호 제공
  • multi-target 학습: pixel + feature 복원
  • decoder fine-tuning: fidelity 향상
• 결과적으로, 구조화된 latent space + 고해상도 복원이 모두 가능함

---

4. Experiments

본 실험은 MAETok의 구성 요소들이 diffusion 모델 학습 및 이미지 생성에 어떤 영향을 미치는지를 체계적으로 분석하고, 다양한 벤치마크에서의 성능을 비교한다.

4.1 실험 설정

Tokenizer 학습 설정

• 모델 구조: ViT-Base (총 176M 파라미터)
• latent token 수: L = 128
• latent 차원: H = 32
• 학습 데이터: ImageNet 256×256 / 512×512, LAION-COCO
• mask ratio: 40–60%
• auxiliary targets:
  • HOG (low-level texture)
  • DINOv2 (high-level semantic)
  • CLIP
  • LAION-COCO에서는 text caption의 BPE index도 포함
• loss weight: λ₁ = 1.0, λ₂ = 0.4
• pixel decoder fine-tuning: mask ratio를 60% → 0%로 줄이며 50K iteration 학습

Diffusion model 설정

• 모델 종류: SiT (Scalable interpolant Transformer), LightningDiT
• 모델 크기: SiT-XL (675M), LightningDiT (더 경량이지만 고성능)
• 학습 스텝: SiT-XL은 4M, LightningDiT는 400K
• 평가 지표:
  • rFID: reconstruction FID
  • PSNR, SSIM: 재구성 품질
  • gFID, IS: 생성 품질
  • Precision, Recall (Appendix C.1 참조)

4.2 Ablation Study

1) Mask Modeling 효과

→ AE는 mask modeling으로 latent가 구조화되며 gFID가 크게 향상됨. decoder fine-tuning을 통해 rFID까지 회복 가능

2) 재구성 타겟 종류

→ DINO, CLIP과 같은 semantic teacher가 gFID에 유리. 다양한 타겟을 조합하면 성능 균형 확보

3) Mask Ratio 변화

→ mask ratio가 높아질수록 latent는 더 discriminative, 하지만 pixel fidelity는 낮아짐 → decoder fine-tuning으로 보완

4) Auxiliary Decoder 깊이

→ 너무 얕거나 깊은 구조는 오히려 latent 품질 저해. 적절한 depth (3 layer)가 가장 효과적

4.3 Latent Space 분석

1) UMAP 시각화

AE, VAE, MAETok의 latent 분포를 2D로 투영한 결과

• AE/VAE: 클래스 간 혼합되어 구조가 불분명
• MAETok: 명확하게 분리된 클러스터 구조 → 더 discriminative

2) Linear Probing Accuracy vs. gFID

[a] latent feature로 linear classifier 학습
→ linear probing accuracy가 높을수록 gFID 낮음 (성능 좋음)

[b] diffusion training curve
→ MAETok이 더 빠르게 학습 수렴 (낮은 loss)

4.4 Main Results

ImageNet 256×256 기준

→ MAETok이 VAE 기반 토크나이저를 압도. 토큰 수는 절반 이하이면서도 성능 우수

ImageNet 512×512 기준

→ MAETok이 2B 파라미터의 모델보다도 높은 성능 기록

---

4.5 Discussion

1. Training & Inference Efficiency

• 일반적으로 256×256 또는 512×512 이미지는 VAE 기반 tokenizer에서 256~1024 tokens로 인코딩된다.
• MAETok은 단 128 tokens만 사용하면서도 높은 성능을 달성함.

→ 이로 인해 diffusion model의 계산 비용(Flops)과 추론 시간이 크게 감소:

• 예: 512×512 이미지 기준
  • 기존: 1024 tokens → 373.3 GFlops, 0.1 images/sec
  • MAETok: 128 tokens → 48.5 GFlops, 3.12 images/sec

→ MAETok은 REPA보다 76배 빠른 학습, 31배 빠른 추론 달성

2. Unconditional Generation 성능

• 일반적으로 diffusion model은 CFG(Classifier-Free Guidance)를 사용해 조건부(class-conditioned) 생성을 한다.
• 그러나 MAETok은 CFG 없이도 생성 품질이 상당히 높음

→ 이는 latent space 자체가 매우 의미 있고 구조화되어 있어, 조건 없이도 좋은 이미지를 생성할 수 있음을 의미함

→ 결과적으로 CFG의 효과는 상대적으로 줄어들며, 향후 더 정교한 guidance 방식이 필요함

3. 설계 요소별 효과 분석

아래는 MAETok의 핵심 구성 요소 3가지를 조합해보며 어떤 조합이 가장 성능이 좋은지를 비교한 표이다.

 

→ 세 가지 요소가 모두 함께 사용될 때 가장 좋은 결과를 보임
→ 특히 learnable latent token의 부재는 성능을 심각하게 저해함

---

5. Related Work

Image Tokenization

• 초기에는 AE(Hinton & Salakhutdinov, 2006), VAE 기반 tokenization이 주류
• 이후 VQ-VAE, VQ-GAN 등의 discrete tokenizer가 발전
• 최근에는 tokenization의 fidelity + semantics를 동시에 추구하는 방향으로 발전
  • 예: SoftVQ, TexTok, REPA 등

MAETok의 차별점:
→ Variational regularization 없이도 높은 fidelity와 semantic 표현 모두 달성

Image Generation

• Autoregressive 모델:
  • CNN 기반 PixelCNN → Transformer 기반 MaskGIT, MAGE 등으로 발전
• Diffusion 모델:
  • DDPM → LDM → DiT → SiT, LightningDiT 등
• Latent diffusion을 통해 효율성과 성능을 동시에 확보하는 방향이 주류가 됨

MAETok의 기여:
→ 기존 VAE 기반 tokenizer를 대체할 수 있는 plain AE 기반 구조로 diffusion 기반 생성 성능을 더 빠르고 가볍게 향상시킴

---

6. Conclusion

• 본 논문은 diffusion model 학습에 있어 latent space의 구조가 핵심임을 이론적·실험적으로 입증함
• VAE의 variational constraint 없이도,
  • Mask modeling + Auxiliary decoding + Learnable latent tokens 조합을 통해
  • 더 구조화되고 구별력 있는 latent space를 구성할 수 있음
• MAETok은 단 128 tokens만으로도,
  • SOTA 성능 (ImageNet 512×512 기준 gFID 1.69)
  • 76배 빠른 학습, 31배 빠른 추론 달성

"Latent space 구조 > Variational regularization"
→ MAETok은 diffusion model을 위한 새로운 tokenizer 패러다임