End-to-End Autoregressive Image Generation with 1D Semantic Tokenizer

"End-to-End Autoregressive Image Generation with 1D Semantic Tokenizer"(EOSTok) 논문에 대한 상세한 기술적 요약입니다.

1. 문제 제기

자기회귀 (AR) 이미지 생성은 이산적인 시각 토큰을 사용하여 이미지를 잠재 표현으로 압축하는 확산 모델의 강력한 대안으로 부상했습니다. 그러나 기존 접근 방식은 두 가지 주요 과제를 안고 있습니다:

• 공간적 불일치: 대부분의 이미지 토크나이저는 2D 그리드 구조 (공간적 레이아웃 보존) 를 사용합니다. 이는 토큰 간의 양방향 의존성을 유발하여, 표준 래스터 순서 AR 모델링에 필요한 단방향 분해와 충돌합니다.
• 비최적의 훈련 패러다임: 현재 방법들은 일반적으로 이단계 훈련 전략을 사용합니다. 먼저 재구성을 위해 토크나이저를 훈련시킨 후 (동결), 결과 토큰으로 AR 모델을 훈련시킵니다. 이는 토크나이저가 생성 작업의 특정 요구 사항에 적응하는 것을 방해합니다.
• 손실 함수 간극: AR 모델에서 사용되는 다음 토큰 예측 (NTP) 손실은 이산 토큰 공간에서 작동하며 최종 픽셀 공간 생성 품질과 직접적인 상관관계가 없습니다. 오직 NTP 만을 최적화하면 "잠재 공간 붕괴"가 발생할 수 있는데, 이는 토크나이저가 예측 손실을 최소화하기 위해 소수의 토큰만 사용하여 이미지 충실도를 저하시키는 현상입니다.

2. 방법론: EOSTok

저자들은 1D 비전 토크나이저와 AR 생성 모델을 단일 단계 파이프라인에서 공동 최적화하는 EOSTok(End-to-end One-dimensional Semantic Tokenizer) 프레임워크를 제안합니다.

A. 1D 비전 트랜스포머 토크나이저

• 아키텍처: TiTok와 유사한 ViT 오토인코더 구조를 기반으로 합니다.
• 프로세스: 2D 이미지 패치는 평탄화되어 학습 가능한 쿼리 토큰과 연결됩니다. 인과적 ViT 인코더가 이 시퀀스를 처리하여 공간적 패치 임베딩을 폐기하고 1D 잠재 토큰만 출력합니다.
• 양자화: 연속적인 잠재 벡터를 이산 코드북에 매핑하여 AR 모델링을 가능하게 하는 **IBQ(Improved Bottleneck Quantization)**를 사용합니다.
• 디코더: 1D ViT 디코더가 양자화된 토큰과 학습 가능한 마스크 토큰을 연결하여 이미지를 재구성합니다.

B. 엔드투엔드 공동 훈련

토크나이저를 동결하는 대신, 인코더, 디코더, AR 모델을 동시에 훈련합니다. 전체 손실 함수는 다음을 결합합니다:

1. 재구성 손실 ($L_{recon}$): 이미지 충실도를 보장하기 위한 표준 VQ-VAE 손실 (L1/L2, 지각적, GAN).
2. 다음 토큰 예측 손실 ($L_{NTP}$): AR 모델을 위한 표준 교차 엔트로피 손실.
3. 자기회귀 예측 재구성 (APR) 손실 ($L_{APR}$):
   • 혁신: 토큰 예측과 픽셀 품질 간의 간극을 메우기 위해, 토크나이저의 디코더가 AR 모델의 예측 토큰 (교사 강제 모드) 을 다시 픽셀 공간으로 디코딩합니다.
   • 메커니즘: 이 손실은 디코딩된 픽셀과 실제 이미지 (ground-truth) 를 비교합니다.
   • 효과: 이는 토크나이저에 엔드투엔드 생성 감독을 제공하여, 토크나이저가 잠재 공간을 붕괴시킴으로써 (너무 적은 토큰 사용) NTP 손실을 "해킹"하는 것을 방지합니다.

C. 의미적 표현 정렬 (VFM)

AR 성능을 저하시킬 수 있는 2D 공간적 사전 지식을 부과하지 않으면서 글로벌 의미 정보를 주입하기 위해, 저자들은 암시적 정렬을 제안합니다:

• 직접 정렬 (거부됨): 1D 잠재 토큰을 2D VFM 특징에 직접 정렬하면 1D 공간에 2D 구조가 강요되어 성능이 저하됩니다.
• 직접 대체 (거부됨): 이미지 패치를 VFM 특징으로 입력을 대체하는 방식.
• 암시적 정렬 (제안됨): 1D 인코더의 숨겨진 패치 임베딩(쿼리 토큰 추출 전) 을 VFM 특징 (예: DINOv2) 과 정렬합니다. 이는 명시적인 2D 공간적 제약을 부과하지 않고도 의미 지식을 순차적 잠재 공간으로 증류합니다.
• 디코더 정렬: 디코더도 VFM 특징과 정렬되어 재구성을 조건부 생성 작업으로 처리하여 수렴을 돕습니다.

3. 주요 기여

1. 엔드투엔드 훈련 프레임워크: 재구성과 생성을 공동 최적화하는 새로운 단일 단계 파이프라인으로, 토크나이저가 AR 모델링에 특화된 표현을 학습하도록 합니다.
2. APR 손실: AR 예측을 픽셀 공간으로 디코딩하는 핵심 손실 함수로, 토크나이저가 재구성 가능하고 생성적으로 예측 가능한 잠재 공간을 학습하도록 하여 잠재 붕괴를 방지합니다.
3. 암시적 의미 정렬: 2D 공간적 편향을 재도입하지 않으면서 1D 토크나이저의 생성 품질을 향상시키는 비전 파운데이션 모델 (VFMs) 활용 전략입니다.
4. 최고 수준의 성능: 분류기 가이드 없이 ImageNet 256×256 생성에서 SOTA 결과를 달성했습니다.

4. 실험 결과

모델은 ImageNet-1K 256×256 벤치마크에서 평가되었습니다.

• 생성 품질:
  • EOSTok-H(644M 파라미터): 가이드 없이 gFID 1.48과 IS 239.5 를 달성했습니다. 이는 가이드 없는 AR 모델의 새로운 SOTA로, 가이드가 있는 REPA-XL/2(gFID 1.42) 및 기타 AR 베이스라인을 능가합니다.
  • EOSTok-L(312M 파라미터): 가이드 없이 gFID 1.74 를 달성했습니다.
• 재구성 품질:
  • EOSTok-L은 rFID 0.73을 달성하여 고품질 재구성 능력을 입증했습니다.
• 애블레이션 연구:
  • APR 손실: APR 없이는 코드북 사용률이 급격히 떨어집니다 (~99% 에서 ~51% 로), gFID 또한 크게 악화됩니다 (8.01 대 3.32).
  • 의미 정렬: 암시적 정렬은 gFID 를 12.27(베이스라인) 에서 3.32 로 개선합니다. 직접 정렬은 성능을 악화시켜 2D 사전 지식이 1D AR 모델링에 해롭다는 것을 확인시켜 줍니다.
  • 시퀀스 순서: 엔드투엔드 감독으로 훈련된 모델은 역순 또는 무작위 시퀀스에 비해 원래 토큰 시퀀스에서 훨씬 더 나은 성능을 보여주어, 잠재 공간이 인과적 예측을 위해 최적화되었음을 증명합니다.
  • 확장성: 모델 크기 (S 에서 H) 와 코드북 크기가 커질수록 성능이 일관되게 향상되며, 더 큰 모델은 재구성 - 생성 간 상충 관계를 완화합니다.

5. 의의

이 논문은 자기회귀 이미지 생성의 패러다임을 근본적으로 전환합니다:

• 이단계에서 원단계로: 토크나이저와 생성기를 공동 훈련하는 것이 전통적인 동결 토크나이저 접근법보다 우수한 결과를 산출함을 입증합니다.
• 1D 토크나이저의 타당성: 의미적 가이드와 엔드투엔드 피드백으로 적절히 훈련될 경우, 공간적 의존성 충돌을 제거함으로써 2D 그리드 기반 토크나이저보다 AR 생성에서 더 우수한 성능을 낼 수 있음을 증명합니다.
• 손실 간극 해소: 토큰 수준의 최적화가 픽셀 수준의 품질로 이어지지 않는다는 디스크리트 생성 모델링의 일반적인 병목 현상을 APR 손실 도입으로 해결합니다.
• 확장성: 이 방법은 효과적으로 확장되어 644M 파라미터 모델로 SOTA 결과를 달성함으로써, 더 큰 크기로 AR 비전 모델을 확장할 수 있는 유망한 길을 제시합니다.