기술 요약: MergeTok

문제 정의

현재 이미지 생성을 위한 시각적 토크나이저는 두 가지 별개의 계열로 나뉘어 있으며, 각각 상호 보완적인 구조적 한계를 겪고 있습니다:

1. 연속형 VAE 기반 토크나이저: 높은 재구성 충실도를 제공하지만, 내재적인 조직화가 결여된 잠재 공간에 의존합니다. 강력한 압축이 가해지면, 밀집된 잠재 표현 내에서 의미론적 요소들이 서로 얽히게 되어, 결과적으로 의미론적 제어력을 저해하는 약한 분리(weakly disentangled) 코드를 생성합니다.
2. 이산형 VQ 기반 토크나이저: 자기회귀(autoregressive) 생성을 가능하게 하지만, 최적화 병목 현상에 직면합니다. 양자화의 비미분적 특성은 희소하고 불균일한 그래디언트 업데이트를 초래하여, 많은 코드북 엔트리가 충분한 감독을 받지 못하게 만듭니다. 이는 코드북 붕괴(codebook collapse)의 위험을 높이고 표현 효율성을 감소시킵니다.

이러한 한계는 재구성 충실도, 의미론적 구조, 효율성, 그리고 학습 안정성 사이의 지속적인 긴장 관계를 형성합니다.

방법론: MergeTok 프레임워크

MergeTok은 단일 인코더-디코더 아키텍처 내에서 연속적(VAE) 및 이산적(VQ) 토큰화를 공동으로 최적화하는 통합 이중 브랜치 토크나이저를 제안합니다. 핵심 혁신은 두 브랜치 사이의 의미론적 가교로서 **토큰 병합(Token Merging, ToMe)**을 사용하는 것입니다.

아키텍처 개요

입력 이미지 $X$가 주어지면, 공유 CNN 인코더 $E_c$가 특징 맵 $Z$를 생성합니다. 이 맵은 두 개의 병렬 브랜치에 의해 처리됩니다:

1. VAE 브랜치 (연속형):

   • 토큰 병합 (Token Merging): ToMe 모듈을 갖춘 어텐션 기반 인코더가 유사한 토큰들을 병합함으로써 $L$개의 토큰 시퀀스를 $K$개의 토큰($K < L$)으로 압축합니다. 이 과정은 응축된 표현 $Z^{(vae)}_K$와 원래의 토큰들을 병합된 그룹에 할당하는 기록인 소스 맵(source map) $S$를 생성합니다.
   • 의미론적 정렬 (Semantic Alignment): 병합된 토근들은 코사인 거리 손실($L_{align}$)을 사용하여 (마찬가지로 ToMe를 갖춘) 프리즈된(frozen) DINO 스타일의 교사 모델과 정렬됩니다. 이는 연속적 잠재 공간을 의미론적으로 인지되고 분리된 표현으로 정규화합니다.
   • 재구성 (Reconstruction): 하이브리드 디코더는 응축된 토큰들과 소스 맵 $S$를 모두 사용하여 그룹을 "언병합(unmerge)"하고 픽셀 공간 이미지를 재구성합니다.

2. VQ 브랜치 (이산형):

   • 그룹 인지 양자화 (Group-Aware Quantization): 이 브랜치는 전체 시퀀스 길이($L$ 토큰)를 보존하기 위해 순전파(forward pass) 동안 ToMe 연산을 우회합니다. 그러나 VAE 브랜치에서 생성된 소스 맵 $S$를 구조적 사전 정보(structural prior)로 활용합니다.
   • 정규화 (Regularization): 소스 맵은 $L$개의 토큰을 $K$개의 의미론적 그룹으로 분할합니다. 두 가지 그룹 인지 손실이 코드북 할당에 적용됩니다:
     • 그룹 내 다양성 ($L_{div}$): 단일 코드로의 붕괴를 방지하기 위해 각 그룹 내 엔트로피를 최대화합니다.
     • 그룹 간 배타성 ($L_{cons}$): 서로 다른 의미론적 그룹이 별도의 코드 서브셋을 사용하도록 보장하기 위해 그룹 간 중첩을 최소화합니다.
   • 재구성 (Reconstruction): 표준 VQ 디코더가 양자화된 토큰으로부터 이미지를 재구성합니다.

학습 전략

• 통합 목적 함수 (Unified Objective): 모델은 재구성, 적대적, 그리고 병합 인지 정규화 항을 결합한 전체 손실 함수로 엔드투엔드(end-to-end) 학습됩니다:
  $$L_{total} = \lambda_{vae}L_{VAE} + \lambda_{vq}L_{VQ} + \lambda_{ali}L_{align} + \lambda_{div}L_{div} + \lambda_{cons}L_{cons}$$
• 입도 인식 샘플링 (Granularity-Aware Sampling): 학습 중에는 병합 비율(따라서 유지되는 토큰 수 $K$)이 하이퍼파라미터 $K^*$를 중심으로 하는 이산 분포에서 동적으로 샘플링됩니다. 이를 통해 인코더가 다양한 병목 너비에 노출되도록 하여, 다양한 의미론적 입도(granularity)에 걸친 강건성을 촉진합니다.
• 추론 (Inference): 추론 시, 토크나이저는 항상 256개의 고정된 토큰 시퀀스를 출력합니다. 다운스트림 생성기(자기회귀 및 확산 모델 모두)는 이 고정된 시퀀스 위에서 동작하며, 내부의 병합 구조에는 무관합니다.

주요 기여

1. 통합 이중 브랜치 토크나이저: MergeTok은 토큰 병합을 의미론적 인터페이스로 사용하여 VAE 학습의 안정성을 강화하고 VQ의 구조적 이점을 활용함으로써, 단일 목적 함수 아래에서 연속적 및 이산적 잠재 공간을 공동으로 최적화하는 최초의 아키텍처입니다.
2. 병합 인지 학습 제약 조건: 저자들은 병합 과정에서 유도된 두 가지 새로운 목적 함수를 도입했습니다:
   • 병합된 토큰 정렬: 연속적 잠재 공간의 의미론적 구조를 개선합니다.
   • 그룹 인지 양자화: 그룹 내 다양성과 그룹 간 배타성을 통해 VQ 학습을 안정화하고 코드북 활용도를 높입니다.
3. 입도 인식 병합 비율 샘플링: 다양한 토큰 입도를 샘플링하는 학습 시 전략을 통해, 다양한 압축률에 걸쳐 강건한 통합 토큰화를 학습할 수 있게 합니다.

실험 결과

실험은 ImageNet-256 (256×256 해상도) 데이터셋과 256개의 고정된 토큰 예산을 사용하여 수행되었습니다.

• 재구성 충실도 (Reconstruction Fidelity):
  • 이산형 (VQ): MergeTok-BL은 0.78의 rFID를 달성하여, GigaTok-BL (0.81) 및 MergeVQ-GR (1.12)과 같은 강력한 베이스라인을 능가했습니다. 프리즈된 DINO 판별기를 사용할 경우, MergeTok-BL*은 0.48 rFID에 도달했습니다.
  • 연속형 (VAE): MergeTok-BL은 0.47의 rFID를 달성하여, RAE (0.49)를 근소하게 앞질렀으며, SoftVQ-L (0.61) 및 DC-AE-f32 (0.69)를 크게 앞질렀습니다.
• 표현 품질 (Representation Quality):
  • MergeTok-BL은 **78.3%**의 선형 프로빙 정확도를 달성하여, 나열된 이산형 토크나이저 중 가장 높은 수치를 기록했습니다. 이는 UniTok (70.8%) 및 MergeVQ-GR (77.9%)을 상회하며, 우수한 의미론적 조직화를 나타냅니다.
• 다운스트림 생성 (Downstream Generation):
  • 자기회귀 (VQ 브랜치): LlamaGen-XXL과 결합했을 때, MergeTok-BL은 (CFG 없이) 1.93, (CFG 적용 시) 2.13의 gFID를 달성하여 GigaTok-BL (2.03)을 능가했습니다.
  • 확산/플로우 (VAE 브랜치): SiT-XL/1과 결합했을 때, MergeTok-BL은 (CFG 없이) 1.79, (CFG 적용 시) 1.29의 gFID를 달성하여 표 내 최상위 수준과 일치하는 성능을 보였습니다.

의의 및 주장

본 논문은 MergeTok이 단일 아키텍처를 통해 시각적 토크나이저에 강건한 의미론적 조직화와 생성기에 친화적인 이산성을 모두 부여할 수 있음을 입증한다고 주장합니다.

• 통합: VAE의 높은 재구성 충실도와 자기회귀 생성에 필요한 이산적 구조 사이의 트레이드오프를 해결합니다.
• 의미론적 가교: 토큰 병합을 구조적 사전 정보로 활용함으로써, 모델이 연속 브랜치로부터 의미론적 그룹화 정보를 성공적으로 전달받아 이산 브랜치를 안정화하고, 역으로도 상호 작용하게 합니다.
• 호환성: MergeTok은 생성기 자체를 수정할 필요 없이, 자기회귀(예: LlamaGen) 및 확산/플로우(예: DiT, SiT) 생성기 모두와 호환되는 256-토큰 시퀀스를 생성합니다.

저자들은 MergeTok을 단순한 재구성 도구가 아니라, 이산적 자기회귀 생성, 연속적 확산 생성, 그리고 의미론적으로 구조화된 표현 학습을 동시에 지원하는 통합 인터페이스로 포지셔닝하고 있습니다.