기술 요약: SSDD: 효율적인 이미지 토큰화를 위한 단일 단계 확산 디코더 (Single-Step Diffusion Decoder)

문제 정의

최첨단 생성 이미지 모델은 고차원 시각 신호를 압축된 잠재 표현(latent representations)으로 압축하기 위해 토크나이저에 의존합니다. KL-정규화된 변이형 오토인코더(KL-VAE)와 같은 현재의 표준 토크나이저들은 효과적이긴 하지만, 왜곡(pixel-level similarity)과 분포 변화(perceptual realism and diversity) 사이의 근본적인 트레이드오프에 직면해 있습니다. 결정론적 디코더(Deterministic decoders)는 왜곡을 최소화하지만 실제 이미지의 복잡한 분포를 모델링하는 데 실패하여 결과물이 흐릿해지는 경우가 많습니다. 반대로, 생성적 디코더(Generative decoders)는 이 분포를 모델링할 수 있지만, 일반적으로 반복적인 샘플링(느린 추론)이 필요하거나 스케일링 및 안정화가 어려운 적대적 손실(GANs)에 의존해야 합니다.

기존의 확산 기반 디코더(Diffusion-based decoders)는 잠재값으로부터 이미지의 조건부 분포를 모델링하는 원칙적인 접근 방식을 제공하지만, 두 가지 주요 한계점을 가집니다:

1. 느린 추론 속도: 여러 번의 디노이징 단계가 필요하여 다운스트림 생성 모델의 병목 현상을 일으킵니다.
2. 학습 불안정성: 최첨단 결과를 얻기 위해서는 GAN 손실이 필수적인 경우가 많으며, 이는 스케일링이 어렵고 학습 과정이 발산(diverging)할 위험이 있습니다.

방법론

저자들은 확장 가능한 하이브리드 아키텍처, GAN-free 학습, 그리고 단일 단계 증류(single-step distillation) 전략을 결합하여 이러한 한계를 극복하도록 설계된 새로운 아키텍처인 **SSDD (Single-Step Diffusion Decoder)**를 소개합니다.

1. 확장 가능한 하이브리드 아키텍처

픽셀 공간 확산에서 흔히 쓰이는 순수 컨볼루션 U-Net이나 잠재 모델링에서 흔히 쓰이는 순수 트랜스포머에만 의존하는 대신, SSDD는 U-ViT 기반 아키텍처를 활용합니다.

• 구조: 4단계의 컨볼루션 ResNet 블록을 가진 U-Net 백본을 사용하되, 중간 어텐션 블록을 $8 \times 8$ 픽셀 패치를 나타내는 토큰을 사용하는 **트랜스포머(Transformer)**로 대체합니다.
• 컨디셔닝(Conditioning): 잠재 그리드를 업샘플링하고 이를 노이즈가 섞인 입력과 결합함으로써 잠재 표현 $z$에 조건을 부여합니다. 또한, ResNet 및 트랜스포머 블록 내에서 Adaptive Group Normalization (AdaGN) 및 Adaptive Layer Normalization (AdaLN)을 통해 컨디셔닝을 더욱 강화합니다.
• 스케일링: 채널 너비와 트랜스포머 깊이를 늘림으로써 13.4M (SSDD-S)에서 153.8M (SSDD-XL) 파라미터까지 효율적으로 확장됩니다.

2. GAN-Free 학습 체계

SSDD는 세 가지 특정 손실 함수를 결합하여 적대적 학습의 필요성을 제거합니다:

• Flow Matching Loss: 기본 목적 함수는 노이즈와 데이터 사이의 속도장(velocity field, $v = x - \epsilon$)을 예측하는 최적 운송 플로우 매칭(optimal transport flow matching)을 사용하며, $L_2$ 손실로 학습됩니다.
• Perceptual Loss (LPIPS): 생성 과정이 지각적으로 올바른 영역을 향하도록 유도하고 고수준 특징을 정렬하기 위해 추가되었습니다.
• REPA Loss: 트랜스포머 표현을 안정화하고 학습을 가속화하기 위해 DINOv2-B 특징을 사용하는 특징 정규화 손실입니다.
• 다단계 학습: 모델은 먼저 $128 \times 128$ 크롭과 무작위 크기 조정 증강(random resizing augmentation)으로 학습된 후, 타겟 해상도($256 \times 256$)에서 미세 조정(fine-tuning)됩니다.

3. 단일 단계 증류 (Single-Step Distillation)

추론 속도의 병목 현상을 해결하기 위해, SSDD는 증류 전략을 채택합니다:

• 교사-학생 설정 (Teacher-Student Setup): 다단계 교사 디코더(8단계로 학습됨)가 고품질의 다양한 재구성 이미지를 생성합니다. 이후 학생 디코더는 이러한 출력들을 단일 단계로 재현하도록 미세 조정됩니다.
• 손실 보존: 오직 $L_2$ 회귀에만 의존하는 표준 증류와 달리, SSDD 증류는 전체 학습 목적 함수를 보존합니다. 학생 모델은 교사의 다단계 출력을 대상으로 계산된 flow-matching 및 LPIPS 손실을 최소화합니다.
• 결과: 이를 통해 모델은 다단계 교사의 지각적 및 생성적 특성을 상속받으면서도 단일 단계 모델의 처리량(throughput)으로 동작할 수 있습니다.

4. 공유 인코더 (Shared Encoders)

저자들은 SSDD 디코더가 공유 인코더(구체적으로 Rombach et al., 2022의 표준 KL-정규화된 인코더)와 함께 학습될 수 있음을 입증했습니다. 이를 통해 인코더를 재학습하거나 각 디코더 크기마다 새로운 잠재 공간을 생성할 필요 없이, 기존 VAE 디코더의 드롭인 교체(drop-in replacement)로 SSDD를 사용할 수 있습니다.

주요 기여

1. 최초의 GAN-Free 단일 단계 확산 디코더: SSDD는 적대적 손실 없이 최첨단 연속 토큰화를 달实现하는 최초의 아키텍처로, 모델 규모에 따른 학습 안정성을 보장합니다.
2. 하이브리드 아키텍처: 픽셀 공간 디코더에서 컨볼루션의 귀납적 편향(inductive bias)과 트랜스포머의 생성 능력을 통합함으로써, 순수 U-Net보다 우수한 스케일링 특성을 보여줍니다.
3. 효율적인 증류: 제안된 증류 방법은 품질 저하를 최소화하면서 다단계 확산 동작을 단일 단계로 성공적으로 압축하여 처리량을 크게 개선합니다.
4. 공유 잠재 공간 호환성: 표준 사전 학습된 인코더를 사용할 수 있는 능력 덕분에 SSDD는 기존 생성 파이프라인에 즉시 적용 가능합니다.

실험 결과

ImageNet-1k (256$\times$256) 및 COCO에서 평가했을 때, SSDD는 KL-VAE 및 다른 확산 디코더들에 비해 유의미한 개선을 보여줍니다:

• 재구성 품질 (f8c4 설정):
  • SSDD-S (13.4M 파라미터)는 재구성 FID (rFID)를 KL-VAE의 0.87에서 0.46으로 개선했습니다.
  • SSDD-M (48M 파라미터, KL-VAE와 유사한 크기)은 0.39의 rFID를 달성했습니다.
  • SSDD는 픽셀 단위의 왜곡(PSNR, SSIM)에서는 경쟁력 있는 수치를 유지하면서도, 지각적 지표(LPIPS, DreamSim)에서 모든 베이스라인을 능가합니다. 단, 픽셀 단위의 완벽한 왜곡보다는 지각적 품질을 우선시합니다.
• 처리량 (Throughput):
  • SSDD-S는 더 나은 품질을 제공하면서도 KL-VAE보다 1.4배 높은 처리량을 달성했습니다.
  • 다단계 확산 디코더(예: $\epsilon$-VAE)와 비교했을 때, SSDD는 유사한 rFID 수준에서 최대 10배 빠른 디코딩 속도를 제공합니다.
• 다운스트림 생성:
  • SSDD를 사용하여 Diffusion Transformer(DiT)를 학습했을 때, SSDD는 KL-VAE 대비 더 나은 생성 FID (gFID)를 보여줍니다.
  • 더 높은 압축률(f16c4)을 사용할 때, SSDD는 생성 품질을 유지하면서 DiT 모델의 샘플링 속도를 3.8배 빠르게 만듭니다.
• 일반화 성능:
  • SSDD 모델은 처음부터 다시 학습하지 않고도 COCO 및 더 높은 해상도(최대 1024$\times$1024)에 잘 일반화되며, DiTo 및 DC-AE와 같은 프리즈(frozen)된 인코더로부터의 잠재값을 디코딩할 수 있습니다.

의의

본 논문은 속도, 안정성, 생성적 충실도 사이의 긴장을 해결함으로써 SSDD가 이미지 토큰화의 새로운 표준을 세운다고 주장합니다. GAN에 대한 의존성과 반복적인 샘플링을 제거함으로써, SSDD는 KL-VAE를 대체할 수 있는 확장 가능하고 즉각 적용 가능한 대안을 제공하며, 더 높은 품질과 더 빠른 생성 모델의 학습을 가능하게 합니다. 이 연구는 확산 기반 토큰화가 다단계 추론의 계산적 오버헤드나 적대적 학습의 불안정성 없이도 효율적이고 안정적일 수 있음을 입증하며, 미래의 대규모 생성 모델링을 위한 길을 열어줍니다.