Adapting Self-Supervised Representations as a Latent Space for Efficient Generation

1. 문제 정의 (Problem)

최근 확산 모델 (Diffusion Models) 및 흐름 기반 모델 (Flow-based Models) 은 고품질 이미지 생성에 탁월한 성능을 보이지만, 고차원 픽셀 공간에서 벡터장을 회귀하거나 2D 잠재 공간 (Latent Space) 을 사용함에 따라 막대한 계산 비용이 요구된다는 한계가 있습니다.

• 기존 잠재 공간의 비효율성: 기존 잠재 확산 모델 (LDM) 은 VAE 를 통해 이미지를 2D 그리드 형태의 잠재 공간으로 압축합니다. 이는 이미지의 높은 공간적 중복성 (Spatial Redundancy) 을 완전히 제거하지 못하며, 여전히 많은 토큰을 필요로 합니다.
• 생성 효율성 부족: 2D 그리드 구조는 토큰 간의 상호작용 (Attention) 을 필요로 하여 훈련 및 추론 비용이 높습니다.
• SSL 표현의 활용 부족: 사전 학습된 자기 지도 학습 (Self-Supervised Learning, SSL) 모델 (예: DINO, CLIP) 은 강력한 의미론적 표현을 가지지만, 기존 연구들은 이를 주로 생성 모델의 가이드 (Guidance) 로만 사용하거나, 2D 그리드 구조를 유지하며 특징 정렬 (Feature Alignment) 에만 활용했습니다.

2. 방법론 (Methodology)

저자들은 **RepTok (Representation Tokenizer)**을 제안합니다. 이는 SSL 인코더의 표현을 단일 연속 잠재 토큰 (Single Continuous Latent Token) 으로 변환하여 생성 모델링에 활용하는 프레임워크입니다.

핵심 구성 요소 및 학습 전략

1. 단일 토큰 표현 (Single Token Formulation):

   • 사전 학습된 SSL 비전 트랜스포머 (예: DINOv2) 의 [cls] 토큰을 기반으로 합니다.
   • 기존 2D 그리드 토큰 대신, 이미지 전체를 단 하나의 연속 벡터로 압축합니다. 이는 공간적 중복성을 극도로 제거하여 효율성을 높입니다.

2. 경량화된 파인튜닝 (Minimal Fine-tuning):

   • SSL 인코더의 대부분 가중치는 동결 (Freeze) 시킵니다.
   • 오직 **[cls] 토큰 임베딩 (Semantic Token Embedding)**만 미세 조정 (Fine-tuning) 합니다.
   • 이를 통해 고수준 의미론적 정보에 저수준의 재구성 관련 세부 정보 (Low-level details) 를 주입하여, 단일 토큰으로도 고품질 이미지 재구성이 가능하도록 합니다.

3. 공동 학습 (Joint Training) 및 흐름 정합 (Flow Matching):

   • **인코더 (E) 와 생성 디코더 (D)**를 함께 학습합니다.
   • Flow Matching Objective: 표준 흐름 정합 손실 함수를 사용하여, 가우시안 노이즈에서 타겟 이미지로의 매핑을 학습합니다.
   • Cosine-Similarity Loss: 파인튜닝된 토큰이 원래 SSL 공간의 매끄러운 기하학적 구조 (Smooth Geometry) 를 유지하도록 정규화합니다.
     • 수식: $L_{cos} = \lambda (1 - \cos(z, z_{frozen}))$
     • 이 손실 함수는 토큰이 너무 멀리 벗어나지 않도록 제한하면서도, 디코더가 필요한 세부 정보를 학습할 수 있는 균형을 맞춥니다.

4. Attention-Free 생성 아키텍처:

   • 입력이 단일 토큰으로 압축되었기 때문에, 토큰 간 상호작용 (Self-Attention) 이 불필요합니다.
   • 따라서 MLP-Mixer와 같은 Attention 이 없는 경량 아키텍처를 생성 모델 (Generator) 로 사용하여 계산 비용을 획기적으로 줄입니다.

5. 텍스트 - 이미지 생성 (Text-to-Image) 확장:

   • 크로스 어텐션 (Cross-Attention) 을 도입하여 텍스트 조건을 단일 잠재 토큰에 통합함으로써, 텍스트 - 이미지 합성으로 자연스럽게 확장됩니다.

3. 주요 기여 (Key Contributions)

• SSL 표현의 직접적 활용: SSL 인코더가 생성 모델의 가이드가 아닌, 직접적인 인코더로 작용할 수 있음을 증명했습니다. 최소한의 적응 (클래스 토큰 파인튜닝) 만으로 매끄럽고 의미론적으로 구조화된 잠재 공간을 생성할 수 있습니다.
• 단일 연속 토큰 압축: 2D 그리드 기반의 공간적 중복성을 제거하고, 단일 연속 벡터로 이미지를 표현하여 재구성과 생성을 동시에 수행합니다.
• 극도의 효율성: Attention 메커니즘과 보조 손실 함수 (Adversarial, Perceptual loss 등) 없이도, 단순한 MLP 기반 아키텍처로 고품질 생성을 가능하게 하여 훈련 비용을 대폭 절감합니다.
• 범용성: ImageNet 클래스 조건 생성뿐만 아니라, 제한된 훈련 예산 (20 시간 미만, 4 개 A100 GPU) 으로도 MS-COCO 에서 경쟁력 있는 제로샷 (Zero-shot) 텍스트 - 이미지 생성 성능을 달성합니다.

4. 실험 결과 (Results)

ImageNet 클래스 조건 생성 (Class-Conditional Generation)

• 성능: RepTok 은 DiT-XL/2, SiT-XL/2 등 최신 트랜스포머 기반 확산 모델과 경쟁력 있는 FID 점수를 기록합니다.
  • RepTok-L 모델: gFID 1.88 (CFG 사용 시).
• 효율성:
  • 계산 비용: SiT-XL/2 대비 90% 이상의 훈련 비용 절감 (Total Compute 기준).
  • FLOPs: SiT 대비 약 1.7% 수준의 FLOPs 만 소모하여 동일한 파라미터 수 대비 훨씬 효율적입니다.
  • 재구성: 단일 토큰임에도 불구하고 rFID 1.85, PSNR 14.94 를 기록하여 기존 2D/1D 토크나이저 (TiTok, FlexTok 등) 와 비교해도 우수한 재구성 품질을 보입니다.

텍스트 - 이미지 생성 (Text-to-Image Synthesis)

• 성능: MS-COCO 벤치마크에서 200k 스텝 (약 20 시간 훈련) 만으로도 경쟁력 있는 제로샷 FID 점수를 달성했습니다.
• 확장성: 언어 백본 (CLIP, InternVL, Gemma) 의 크기를 늘리면 성능이 선형적으로 향상됨을 확인했습니다.

기타 분석

• 잠재 공간 보간 (Interpolation): SSL 인코더의 매끄러운 기하학적 특성을 유지하여, 잠재 공간에서의 선형 보간이 의미론적 및 시각적 세부 사항의 자연스러운 전환을 가능하게 합니다.
• SSL 백본 일반화: DINOv2 외에도 MAE, CLIP 등 다양한 SSL 모델에서 동일한 효과를 입증했습니다.

5. 의의 및 결론 (Significance)

이 논문은 자기 지도 학습 (SSL) 표현 자체가 생성 모델링을 위한 이상적인 잠재 공간이 될 수 있음을 보여줍니다.

• 효율성의 패러다임 전환: 고비용의 Attention 기반 확산 모델 대신, 사전 학습된 SSL 인코더의 강력한 표현력을 활용하고 이를 경량 MLP 로 연결하는 새로운 효율적 생성 패러다임을 제시했습니다.
• 자원 제약 환경에서의 활용: 제한된 컴퓨팅 자원 (저비용 GPU, 짧은 훈련 시간) 으로도 고품질 생성이 가능해짐으로써, 생성 AI 의 접근성을 높이고 환경 부담을 줄이는 데 기여합니다.
• 간결함과 성능의 조화: 복잡한 아키텍처와 보조 손실 함수 없이 단일 토큰과 흐름 정합 (Flow Matching) 만으로 높은 성능을 달성함으로써, 생성 모델의 설계 원리에 대한 새로운 통찰을 제공합니다.

요약하자면, RepTok은 "SSL 인코더의 [cls] 토큰을 미세 조정하여 단일 잠재 벡터로 만들고, 이를 경량 디코더와 결합함으로써 기존 모델들의 계산 비용을 90% 이상 절감하면서도 경쟁력 있는 생성 품질을 달성한 혁신적인 프레임워크"입니다.