[논문리뷰] Adaptive Length Image Tokenization via Recurrent Allocation

ICLR 2025. [Paper] [Github]
Shivam Duggal, Phillip Isola, Antonio Torralba, William T. Freeman
MIT CSAIL
4 Nov 2024

Introduction

이미지 임베딩 또는 표현 학습을 위한 일반적인 프레임워크는 인코더-디코더 방식이다. Reconstruction loss가 다양한 수준의 디테일을 포착하여 더 나은 재구성을 가능하게 하는 적응형 표현 학습을 촉진하기 때문이다. 현재 SOTA 인코더-디코더 접근 방식은 discrete한 토큰 공간에서 작동하며, 이미지를 학습된 토큰으로 인코딩한 다음 다시 이미지 픽셀로 디코딩한다. 이러한 접근 방식들은 입력 패치 수준에서 약간 압축한 다음 인코더-디코더 네트워크 전체 깊이에 걸쳐 토큰 수(패치 수)를 유지한다. 따라서 모든 이미지의 표현 길이는 토큰 수로 고정되며, 이는 인간 엔지니어가 결정한 고정된 패치 크기와 같다. 또한, 토큰 수와 패치 수가 같기 때문에 다양한 이미지에 대한 적응형 표현이나 압축을 불가능하게 한다.

모달리티에 구애받지 않는 아키텍처를 목표로, Google DeepMind는 입력 데이터 토큰을 고정된 1D 토큰 집합으로 추출하는 transformer 기반 아키텍처인 Perceiver를 제안했다. 이러한 latent-token distillation 과정은 2D 이미지 입력을 가장 관련성이 높은 feature를 포착하는 더 압축된 1D latent 토큰으로 압축한다. Perceiver와 마찬가지로, 본 논문에서도 self-supervised reconstruction loss를 통해 2D 이미지 토큰을 훨씬 적은 수의 1D latent 토큰으로 인코딩하는 latent-token distillation 방식을 사용하였다.

본 논문에서는 입력 이미지를 점진적으로 증가하는 수의 1D latent 토큰으로 autoregressive하게 정제함으로써 적응형 또는 가변 길이 표현 학습 문제를 해결하였다. 다양한 길이의 여러 표현을 동시에 한 번에 학습하는 마트료시카 방식과는 달리, 본 논문에서는 공유 인코더-디코더 아키텍처를 통해 입력 이미지 또는 2D 이미지 토큰을 1D latent 토큰으로 재귀적으로 정제하여 각 이미지 토큰이 latent 토큰으로 충분히 처리/정제될 때까지 반복한다. 이러한 순환 과정의 각 iteration마다 새로운 학습 가능한 latent 토큰 형태의 추가적인 연산 자원을 제공하여 모델이 여러 iteration에 걸쳐 적응형 및 가변 길이 표현을 학습할 수 있도록 한다.

본 논문에서는 Adaptive Length Image Tokenizer (ALIT)를 제안하고, self-supervised reconstruction loss를 사용하여 학습시켰다. 점진적으로 증가하는 표현 용량 덕분에, 각 반복 업데이트를 통해 latent 토큰은 특정 영역에 attention한다. 또한 이미지당 토큰 개수를 유연하게 설정할 수 있기 때문에 하위 task/모델에 대한 지식 수준에 따라 필요한 용량을 결정할 수 있다.

Method

본 논문에서는 각 이미지가 고유하며 서로 다른 토큰 수를 가져야 한다고 주장한다. 이 문제를 해결하기 위해, 본 논문에서는 이미지에 더 많은 표현 용량, 즉 토큰을 autoregressive하게 할당하는 새로운 프레임워크를 제안하였다. 이 프레임워크는 테스트 시점에 각 이미지에 대해 가변적인 토큰 수를 허용한다.

2D 이미지 토큰에서 1D 토큰으로의 latent distillation

입력 이미지를 1D latent 토큰으로 매핑하고자 한다. 압축/적응형 표현 학습의 핵심 문제에 초점을 맞춰, 먼저 기존 VQGAN 이미지 tokenizer를 활용하여 입력 이미지를 2D 이미지 토큰 집합 \(\textbf{K}_\textrm{2D}^{t=0}\)로 매핑한다. 사전 학습된 VQGAN 모델은 256$\times$256 크기의 이미지를 16$\times$16 크기의 2D 공간 토큰으로 매핑할 때 디테일 손실을 최소화한다. 각 16$\times$16 토큰은 학습된 VQGAN codebook에 있는 quantized code 중 하나를 가리키는 포인터이다. 본 논문에서는 256개의 공간 토큰을 최소 32개의 1D 토큰으로 distillation하였다.

2D→1D→2D Distillation

공간 이미지 토큰 \(\textbf{K}_\textrm{2D}^{t=0}\)가 주어졌을 때, 토큰 축을 따라 latent 토큰 \(\bar{\textbf{K}}_\textrm{1D}\)와 concat하고, 이를 distillation 인코더 $\textrm{Enc}$에 통과시킨다. \(\bar{\textbf{K}}_\textrm{1D}\)는 학습된 임베딩으로 초기화된다. 인코더는 모든 토큰에 대해 joint self-attention을 수행하고 \(\textbf{K}_\textrm{1D}\)로 distillation한다. \(\textbf{K}_\textrm{1D}\)는 distillation 디코더 $\textrm{Dec}$로 전달되어 마스킹된 토큰 \(\bar{\textbf{M}}_\textrm{2D}\)와 concat되고, latent 토큰을 다시 2D 공간 토큰으로 distillation한다. 2D 이미지 토큰, 1D latent 토큰, 2D 마스크 토큰은 서로 별도의 위치 인코딩이 더해진다. 인코더 출력은 디코더에 입력하기 전에 학습된 1D codebook에서 가장 가까운 code를 샘플링함으로써 quantization된다.

\[\begin{aligned} \textbf{K}_\textrm{2D}^{t=1}, \textbf{K}_\textrm{1D}^{t=1} &= \textrm{Enc}([\; \textbf{K}_\textrm{2D}^{t=0}; \; \bar{\textbf{K}}_{1D} \;]) \\ \textbf{M}_\textrm{2D}^{t=1} &= \textrm{Enc}([\; \bar{\textbf{M}}_\textrm{2D}; \; \textbf{K}_\textrm{1D}^{t=1} \;]) \end{aligned}\]

주요 학습 loss는 \(\textbf{M}_\textrm{2D}^{t=1}\)와 \(\textbf{K}_\textrm{2D}^{t=0}\) 사이의 reconstruction loss이다.

가변 tokenization를 위한 autoregressive 프레임워크

본 논문에서는 가변 길이 압축 표현을 학습하기 위해 입력 이미지에 대한 순환 처리를 수행한다. 각 roll-out iteration을 통해 인코더-디코더 아키텍처를 재귀적으로 확장하여 처리 용량을 늘릴 뿐만 아니라, 새로운 기록 가능한 토큰 측면에서 추가적인 계산 메모리를 제공하여 이미지 토큰을 더 많은 1D latent로 효율적으로 distillation한다. \(\textbf{K}_\textrm{1D}^{t=1}\)는 이미지당 최소한의 개수의 토큰이다.

Latent distillation의 각 iteration에서, 이전 iteration의 latent 토큰 \(\textbf{K}_\textrm{1D}^{t=T}\)에 학습된 임베딩으로 초기화된 새로운 토큰 \(\bar{\textbf{K}}_{1D}\)를 concat한다. 선택적으로, 이전 iteration에서 완벽하게 distillation되지 않은 이미지 토큰에 인코더가 집중할 수 있도록, 마지막 iteration의 처리된 이미지 토큰 \(\textbf{K}_\textrm{2D}^{t=T}\)에 마스킹/동적 정지 연산 $\textrm{Mask}$를 적용한다. 이 마스크는 재구성된 출력 \(\textbf{M}_\textrm{2D}^{t=T}\)와 원본 이미지 토큰 $\textbf{K}_\textrm{2D}^{t=0}$$ 간의 정렬에 따라 결정된다. 마스킹된 이미지 토큰은 latent 토큰과 concat되어 인코더-디코더를 통과한다.

\[\begin{aligned} \textbf{K}_\textrm{1D}^{t=T} &= [\; \textbf{K}_\textrm{1D}^{t=T} \; \bar{\textbf{K}}_\textrm{1D} \;] \\ \textbf{K}_\textrm{2D}^{t=T} &= \textrm{Mask}(\; \textbf{K}_\textrm{2D}^{t=T} \mid \textbf{M}_\textrm{2D}^{t=T}, \textbf{K}_\textrm{2D}^{t=0} \;) \\ \textbf{K}_\textrm{2D}^{t=T+1}, \textbf{K}_\textrm{1D}^{t=T+1} &= \textrm{Enc}([\; \textbf{K}_\textrm{2D}^{t=T}; \; \textbf{K}_\textrm{1D}^{t=T} \;]) \\ \textbf{M}_\textrm{2D}^{t=T+1} &= \textrm{Enc}([\; \bar{\textbf{M}}_\textrm{2D}; \; \textbf{K}_\textrm{1D}^{t=T+1} \;]) \end{aligned}\]

이 과정은 여러 iteration에 걸쳐 수행된다. 주요 학습 loss는 새로운 재구성 토큰 \(\textbf{M}_\textrm{2D}^{t=T+1}\)와 원본 이미지 토큰 \(\textbf{K}_\textrm{2D}^{t=0}\) 사이의 reconstruction loss이다. 각 iteration에서 distillation된 latent 토큰은 공유 1D codebook을 사용하여 quantization된다. 따라서 서로 다른 iteration에서 학습된 토큰은 동일한 임베딩 공간에 속한다.

요약하자면, 각 iteration에서 이전 단계의 latent 토큰은 잔여 업데이트를 받는 동시에 새로운 계산 메모리가 도입된다. 이러한 새로운 토큰은 기존 latent 토큰이 특정 영역에 집중할 수 있도록 해주어 더욱 정밀하고 sparse한 attention을 구현한다.

학습 절차

본 논문에서는 다단계 접근 방식을 사용하여 학습을 진행하였다. 먼저, 사전 학습된 VQGAN 이미지 토큰을 distillation된 latent 토큰으로부터 재구성한다. 이때, discrete VQGAN 토큰에는 smooth cross-entropy loss를 적용한다. 다음으로, 이미지 픽셀 재구성을 목표로 VQGAN 인코더-디코더와 distillation 인코더-디코더를 동시에 최적화한다. 이후, 이미지 재구성 품질이 여러 표현에서 만족스러운 수준에 도달하면, 사실성을 더욱 향상시키기 위해 학습 과정 후반에 GAN loss를 도입한다. 1D codebook 학습을 촉진하기 위해 학습 과정 전반에 걸쳐 quantization loss를 적용한다.

Experiments

1. Analysis

다음은 다양한 복잡도의 out-of-distribution 이미지에 대한 reconstruction 결과이다.

다음은 (왼쪽) classification과 (오른쪽) depth를 기준으로 최소 토큰 수를 선택한 결과이다.

다음은 reconstruction loss를 기준으로 최소 토큰 수를 선택한 결과이다.

다음은 모델 강도에 따른 성능을 비교한 결과이다.

다음은 latent 토큰 attention map을 시각화한 것이다.

2. Comparison

다음은 reconstruction FID를 비교한 결과이다.

다음은 reconstruction 결과와 reconstruction loss 분포를 비교한 것이다.

3. Ablations

다음은 ablation study 결과이다.