1D 시맨틱 토크나이저를 활용한 엔드투엔드 자기회귀 이미지 생성

기존의 이미지 생성 모델은 토크나이저와 생성 모델을 별도로 학습시켜 최적의 성능을 내기 어려웠으나, 이 논문은 이를 동시에 학습시키는 엔드투엔드 파이프라인을 제안한다. 특히 2D 구조에 얽매이지 않는 1D 시맨틱 토크나이저를 통해 자기회귀 모델의 효율성을 극대화하고 ImageNet 256x256 벤치마크에서 뛰어난 FID 점수를 기록했다.

기존의 자기회귀(Autoregressive) 이미지 생성은 이미지를 2D 격자 형태의 토큰으로 나누어 처리한다. 하지만 Transformer와 같은 모델이 다음 토큰을 예측할 때, 2D 격자 구조는 왼쪽에서 오른쪽, 위에서 아래로 흐르는 단방향 예측 방식과 근본적으로 어긋나는 다방향 의존성을 가진다. 이로 인해 모델이 이미지의 전체적인 맥락을 파악하는 데 효율성이 떨어진다.

이 논문은 이미지를 처음부터 순차적인 1D 시퀀스로 압축하는 토크나이저를 사용한다. 마치 문장에서 다음 단어를 예측하듯 이미지 토큰을 예측하게 함으로써 자기회귀 모델의 본질적인 동작 방식에 최적화했다. 특히 학습 과정에서 토크나이저를 고정하지 않고 생성 모델과 함께 학습시켜, 생성 모델이 예측하기 더 쉬운 형태의 토큰 표현을 토크나이저가 스스로 찾아내도록 유도한다.

결과적으로 이러한 엔드투엔드 접근법은 토큰 공간에서의 예측 정확도뿐만 아니라 실제 사람이 보기에 자연스러운 픽셀 단위의 이미지 품질을 동시에 개선한다. 이는 복잡한 2D 구조적 제약 없이도 고해상도 이미지를 효율적으로 생성할 수 있는 새로운 방향을 제시한다.

전체 시스템은 1D Causal ViT Encoder, Quantizer, Autoregressive Transformer, 그리고 1D Causal ViT Decoder로 구성된다. 이미지는 패치로 분할된 후 학습 가능한 쿼리 토큰과 함께 인코더에 입력되어 1D 잠재 표현(Latent Representation)으로 압축된다. 이 과정에서 IBQ(Index Backpropagation Quantization)를 사용하여 이산적인 토큰으로 변환된다.

핵심 메커니즘인 APR(Autoregressive Prediction Reconstruction) 손실은 다음과 같이 계산된다. AR 모델이 예측한 다음 토큰의 확률 분포 [입력: 이전 토큰들 → 연산: Transformer 추론 → 출력: 다음 토큰 확률]를 기반으로, Teacher-forcing 방식으로 예측된 토큰들을 디코더에 통과시켜 픽셀 이미지로 복원한다. 이후 원본 이미지와의 MSE 및 지각적 손실(Perceptual Loss)을 계산하여 [입력: 예측 토큰 시퀀스 → 연산: 디코더 복원 → 출력: 복원된 이미지와 원본의 차이] 토크나이저와 생성 모델을 동시에 업데이트한다.

시맨틱 정렬을 위해 DINOv2의 특징 맵을 활용한다. 인코더의 중간 계층인 Hidden Patch Embedding과 VFM의 임베딩 사이의 코사인 유사도를 계산하여 [입력: 두 임베딩 벡터 → 연산: 내적 및 정규화 → 출력: 유사도 점수] 이를 최대화하는 방향으로 학습한다. 이는 1D 토큰이 2D 공간 구조에 종속되지 않으면서도 풍부한 시각적 의미 정보를 담을 수 있게 한다.

ImageNet 256x256 벤치마크에서 Classifier Guidance 없이 FID 1.48을 달성하며 기존 SOTA 모델들을 능가했다. 특히 644M 파라미터 규모의 EOSTok-H 모델은 기존의 2D 기반 모델들보다 훨씬 적은 파라미터로도 더 정교한 이미지를 생성함을 입증했다.

Ablation Study를 통해 APR 손실의 중요성을 확인했다. APR 손실이 없을 경우 토크나이저가 특정 토큰만 반복해서 사용하는 코드북 붕괴(Codebook Collapse) 현상이 발생하여 gFID가 8.01까지 치솟았으나, APR 손실 적용 시 3.32로 급격히 개선되었다. 또한 1D 토큰의 순서를 뒤섞거나 반전시켰을 때 성능이 크게 하락하는 실험을 통해, 제안된 모델이 이미지의 순차적 의존성을 성공적으로 학습했음을 보여주었다.

EOSTok 아키텍처는 TiTok의 구조를 계승하되, 인코더와 디코더 모두에 Causal Masking을 적용하여 1D 시퀀스 특성을 강화했다. 인코더는 2D 패치 토큰과 L개의 학습 가능한 쿼리 토큰을 입력받아 최종적으로 쿼리 토큰들만을 1D 잠재 표현으로 출력한다. 디코더는 이 토큰들을 다시 2D 격자 구조의 마스크 토큰과 결합하여 이미지를 복원한다.

학습 목적 함수는 LVQVAE(재구성), LNTP(다음 토큰 예측), LAPR(예측 기반 재구성), Lalign(시맨틱 정렬)의 가중합으로 정의된다. 특히 Lalign은 'Implicit Alignment' 방식을 채택하여, 최종 잠재 토큰 z가 아닌 인코더의 중간 출력 h_Enc를 VFM의 특징과 정렬시킨다. 이는 z가 생성 작업에 더 유연하게 최적화될 수 있도록 자유도를 부여하면서도 VFM의 지식을 전이받을 수 있게 한다.

구현 측면에서는 8개의 H100 GPU를 사용하여 400 에포크 동안 학습되었으며, Adam 옵티마이저와 Cosine Learning Rate 스케줄러를 적용했다. 추론 시에는 KV 캐시를 활용한 자기회귀 샘플링을 통해 속도를 최적화했으며, Classifier-free Guidance 대신 AutoGuidance 기법을 사용하여 가이드 성능을 극대화했다.

토큰 시퀀스의 길이가 길어질수록 재구성 품질(rFID)은 향상되지만, 자기회귀 모델의 예측 난이도가 상승하여 생성 품질(gFID)이 특정 지점(192~256 토큰) 이후로는 오히려 하락하는 트레이드오프 관계가 존재함을 명시했다.