VibeToken: Scaling 1D Image Tokenizers and Autoregressive Models for Dynamic Resolution Generations

다음은 "VibeToken: 동적 해상도 생성을 위한 1 차원 이미지 토크나이저 및 자기회귀 모델 확장" 논문에 대한 상세한 기술 요약입니다.

1. 문제 제기

자기회귀 (AR) 이미지 생성 모델 (예: LlamaGen) 은 경쟁력 있는 성능을 보여주었으나, 확산 모델 (diffusion models) 에 비해 생산 환경 도입에 있어 상당한 장벽에 직면해 있습니다. 주요 병목 현상은 해상도 및 종횡비에 따른 확장성입니다.

• 토큰 폭주: 기존의 2 차원 토크나이저 (예: VQGAN) 는 이미지 해상도에 비례하여 토큰 시퀀스가 선형적으로 증가합니다. 1024×1024 이미지의 경우 수천 개의 토큰 (예: 4,096 개 토큰) 이 생성됩니다.
• 2 차 계산 비용: AR 모델은 자기 어텐션 (self-attention) 에 의존하며, 이는 시퀀스 길이 $T$에 대해 $O(T^2)$의 복잡도를 가집니다. 해상도가 증가함에 따라 추론 FLOPs 가 2 차적으로 증가하여 고해상도 생성이 계산적으로 불가능해집니다.
• 경직성: 대부분의 AR 모델은 고정된 해상도 (예: 256×256) 로 훈련되며, 재훈련이나 복잡한 초해상도 파이프라인 없이는 임의의 종횡비나 더 높은 해상도로 일반화하는 데 어려움을 겪습니다.
• 확산 모델의 우세: 확산 모델은 자연스럽게 임의의 해상도를 처리하므로, 유연하고 생산 수준의 사용 사례에서 AR 모델을 뒤처지게 만듭니다.

2. 방법론

저자들은 VibeToken(새로운 토크나이저) 과 VibeToken-Gen(AR 생성기) 의 두 부분으로 구성된 솔루션을 제안합니다.

A. VibeToken: 해상도 무관 1 차원 토크나이저

VibeToken 은 임의의 해상도의 이미지를 고정된, 사용자가 제어 가능한 토큰 수(예: 32~256 개) 로 인코딩하도록 설계된 1 차원 Transformer 기반 이미지 토크나이저로, 토큰 수를 공간 해상도에서 분리합니다.

주요 아키텍처 혁신은 다음과 같습니다.

1. 동적 그리드 위치 임베딩: 고정된 절대 위치 임베딩 대신, VibeToken 은 입력 패치 그리드 ($T_H \times T_W$) 에 맞추기 위해 차등적으로 크기 조정 (이중 선형/이중 3 차 보간을 통해) 되는 학습된 임베딩 그리드를 사용합니다. 이를 통해 모델은 품질 저하 없이 임의의 종횡비를 처리할 수 있습니다.
2. 적응형 패치 임베딩: FlexiViT 에서 영감을 받아, 모델은 훈련 및 추론 중 가변적인 패치 크기 ($k$) 를 지원합니다. 가중치 크기 조정 연산자가 기본 투영 가중치를 현재 패치 크기에 매핑하여, 재파라미터화 없이 모델이 다양한 해상도에 적응할 수 있게 합니다.
3. 적응형 디코더 해상도: 디코더는 고정된 고해상도 (패치 크기의 4 배) 에서 픽셀을 예측하고, 학습 가능한 2 차원 합성곱 다운샘플링 계층을 사용하여 목표 해상도로 매핑합니다. 이는 별도의 업스케일러 모듈 없이 네이티브 초해상도를 가능하게 합니다.
4. 동적 길이 토크나이징: 모델은 잠재 길이 $L \in [32, 256]$의 균일한 샘플링으로 훈련됩니다. 인코더는 정확히 $L$개의 토큰을 생성하고 디코더는 정확히 $L$개의 토큰을 소비하므로 패딩이나 고정 길이 제약이 필요 없습니다.
5. 양자화: 압축과 재구성 품질의 균형을 맞추기 위해 8 개의 코드북을 사용하는 다중 벡터 양자화 (MVQ) 를 사용하여 어휘 크기를 32,768 로 달성합니다.

B. VibeToken-Gen: 효율적인 AR 생성기

VibeToken-Gen 은 VibeToken 이 생성한 짧고 해상도 무관한 토큰 시퀀스에서 작동하는 클래스 조건부 AR 모델 (LlamaGen 아키텍처 기반) 입니다.

• 일정한 계산 예산: 토큰 시퀀스 길이 $L$이 출력 해상도와 관계없이 고정되어 있으므로 (예: 64 개 토큰), 추론 FLOPs 는 일정하게 유지됩니다.
• 조건부 설정: 기하학적 늘어남 아티팩트를 방지하기 위해 모델은 클래스 레이블과 함께 목표 해상도 $(H, W)$ 및 종횡비를 조건으로 받습니다.
• 훈련: 고정 해상도 베이스라인과 비교 가능한 토큰 예산을 사용하여 다양한 해상도 (256×256~512×512) 와 종횡비가 혼합된 ImageNet-1k 로 훈련되었습니다.

3. 주요 기여

1. 해상도 무관 토크나이징: 짧은 동적 토큰 예산 (32~256 개 토큰) 을 유지하면서 임의의 해상도와 종횡비를 네이티브로 지원하는 최초의 1 차원 Transformer 토크나이저입니다.
2. 계산 효율성: AR 추론 비용을 이미지 해상도에서 분리합니다. VibeToken-Gen 은 모든 해상도에서 179 GFLOPs의 일정한 비용을 유지하는 반면, LlamaGen 과 같은 베이스라인은 1024×1024 이미지에서 약 11 TFLOPs로 확장됩니다 (이는 63.4 배의 효율성 향상입니다).
3. 네이티브 초해상도: 이 아키텍처는 별도의 업스케일링 모듈 없이 저해상도 입력 (또는 임의의 목표) 에서 고해상도 이미지를 생성할 수 있습니다.
4. 최첨단 성능: 2 차원 토크나이저 및 확산 모델과 비교하여 훨씬 적은 수의 토큰을 사용하면서도 경쟁력 있는 재구성 (rFID) 및 생성 (gFID) 품질을 달성합니다.

4. 실험 결과

이 논문은 다양한 해상도 (256×256~1024×1024) 와 종횡비에 걸쳐 ImageNet-1k 및 FFHQ 데이터셋에서 VibeToken 과 VibeToken-Gen 을 평가했습니다.

• 재구성 품질:
  • VibeToken-LL 은 256×256 에서 0.40 rFID를, 1024×1024 에서 2.40 rFID를 달성합니다.
  • 해상도 일반화 측면에서 다른 1 차원 토크나이저 (예: TiTok, One-D-Piece) 보다 우수하며, 압축 효율성 측면에서 2 차원 토크나이저 (예: IBQ, MAGViT) 와 맞먹거나 능가합니다.
• 생성 품질:
  • VibeToken-Gen (XXL) 은 64 개 토큰만을 사용하여 1024×1024 이미지를 합성하며 gFID 3.94를 기록합니다.
  • 비교:
    • 확산 모델 (NiT) 대비: VibeToken-Gen 은 1024×1024 에서 2.35 배 더 빠릅니다 (0.46 초 vs 1.08 초) 그리고 더 나은 품질을 달성합니다 (3.94 gFID vs 5.87 gFID).
    • AR 베이스라인 (LlamaGen) 대비: LlamaGen 은 1024×1024 에 대해 약 11 TFLOPs 를 필요로 하는 반면, VibeToken-Gen 은 일정한 179 GFLOPs 를 사용합니다.
• 효율성:
  • 1024×1024 이미지의 토큰 수는 2 차원 토크나이저의 약 4,096 개에서 64 개로 감소했습니다.
  • 고해상도 생성에 대한 추론 속도는 베이스라인 AR 모델이 유사한 해상도에서 32 초 이상 소요되는 것과 비교해 1 초 미만 (0.46 초) 입니다.

5. 의의

이 연구는 유연성과 효율성 측면에서 자기회귀 모델과 확산 모델 간의 격차를 해소합니다.

• 생산 타당성: AR 모델이 일정한 계산으로 임의의 해상도를 처리할 수 있게 함으로써, 입력/출력 차원이 변하는 실제 응용 프로그램 (예: 모바일 앱, 다양한 디스플레이 형식) 에서 AR 생성을 실현 가능하게 만듭니다.
• 확장성: AR 모델이 2 차 계산 폭주 없이 고해상도로 확장될 수 있음을 보여주어, 고해상도 생성형 AI 로 가는 유일한 경로가 확산 모델이라는 관념에 도전합니다.
• 미래 방향: 이 접근법은 생성 콘텐츠의 공간적 복잡성과 관계없이 고정된 사용자 제어 계산 예산 하에서 작동할 수 있는 통합 멀티모달 모델 (텍스트, 이미지, 비디오) 의 문을 엽니다.

요약하자면, VibeToken은 생성 모델의 부담을 유연하고 해상도 무관한 토크나이저로 이동시킴으로써 AR 생성의 "해상도 확장" 문제를 해결하며, 현재 최첨단 대안들보다 더 빠르고 저렴하며 유연한 시스템을 제공합니다.