CaTok: 1차원 인과적 이미지 토큰화를 위한 Mean Flows 제어

기존 이미지 생성 AI는 2D 데이터를 1차원으로 펼치면서 정보의 순서와 인과관계가 깨지는 문제가 있었다. 이 논문은 이미지를 언어처럼 순차적으로 이해하고 생성할 수 있는 1D 인과적 토큰화 기술을 제안하여, 고품질 이미지 복원과 효율적인 생성을 동시에 가능하게 한다.

기존의 자기회귀(Autoregressive) 모델은 텍스트처럼 순서가 있는 데이터를 처리하는 데 최적화되어 있다. 하지만 이미지는 2차원 구조이기에 이를 1차원 시퀀스로 변환할 때 토큰 간의 인과관계(앞의 토큰이 뒤의 토큰에 영향을 주는 구조)를 정의하기 어렵다. 기존 방식은 이미지를 격자 형태로 잘라 단순히 나열하거나 무작위 순서를 부여했는데, 이는 언어 모델의 '다음 토큰 예측' 패턴과 맞지 않아 학습 효율이 떨어지는 한계가 있었다.

CaTok은 이미지를 1차원 인과적 토큰으로 변환하기 위해 MeanFlow 디코더를 도입한다. 이는 확산 모델의 원리를 활용하되, 특정 시간 간격 내의 평균 속도 필드(velocity field)를 직접 모델링하는 방식이다. 인코더가 추출한 1D 토큰들을 시간 간격에 따라 디코더의 조건으로 입력함으로써, 토큰들이 이미지 생성 과정에서 논리적인 순서와 인과성을 갖도록 강제한다.

결과적으로 CaTok은 단 한 번의 연산으로 이미지를 복원하는 속도와 여러 단계에 걸쳐 정교하게 생성하는 품질 사이의 균형을 잡았다. 특히 REPA-A라는 정규화 기법을 통해 인코더가 시각 기초 모델(VFM)의 풍부한 지식을 빠르게 습득하게 함으로써, 기존 모델들보다 훨씬 적은 학습량으로도 뛰어난 복원 및 생성 성능을 보여준다.

CaTok은 Causal ViT 인코더와 MeanFlow DiT 디코더로 구성된 확산 오토인코더 구조를 채택한다. 인코더는 이미지 패치와 레지스터(Register) 토큰을 입력받아 1D 토큰 시퀀스 $V_K$를 생성하며, 이때 Causal Attention Mask를 적용하여 각 토큰이 이전 토큰의 정보에만 의존하도록 제한한다.

디코더는 MeanFlow 목적 함수를 사용하여 학습된다. 두 타임스텝 $r, t$ 사이의 평균 속도 $u_\theta$를 예측하도록 설계되었으며, 수식으로는 $u(z_t, r, t) = \frac{1}{t-r} \int_r^t v(z_\tau, \tau) d\tau$로 정의된다. [노이즈가 섞인 이미지 $z_t$와 시간 간격 $[r, t]$에 해당하는 토큰 $V_{r:t}$를 입력으로] → [신경망 $D_\theta$를 통해 해당 구간의 평균 변화율을 계산하여] → [예측된 평균 속도 $u_\theta$를 출력하고] → [이 값이 실제 노이즈와 원본 이미지 사이의 직선 경로를 얼마나 잘 따르는지 측정하여] 가중치를 업데이트한다.

학습 가속화를 위해 REPA-A(Representation Alignment for Autoencoder)를 제안한다. 이는 인코더의 특징 $H_e$와 사전 학습된 DINOv2 같은 VFM의 특징 $H_{vfm}$ 사이의 코사인 유사도를 극대화하는 손실 함수를 추가하는 방식이다. [인코더의 중간 출력값과 VFM의 출력값을 입력으로] → [두 벡터 사이의 각도를 계산하는 코사인 유사도 연산을 수행해] → [유사도 점수를 결과로 얻고] → [이 점수가 높을수록 인코더가 이미 시각적으로 검증된 풍부한 표현력을 갖게 됨을 의미한다].

ImageNet-1K 256x256 벤치마크에서 CaTok-L-256 모델은 0.75 FID, 22.53 PSNR, 0.674 SSIM을 기록했다. 이는 기존의 FlowMo(0.95 FID)나 Semanticist(0.78 FID) 같은 최신 확산 기반 토큰화 모델들을 능가하는 수치이다. 특히 Semanticist 대비 절반 이하의 학습 에포크(160 vs 400)만으로도 더 우수한 성능을 달성하며 학습 효율성을 증명했다.

자기회귀 생성 실험에서도 CaTok은 우수한 성능을 보였다. LlamaGen 아키텍처를 기반으로 학습했을 때, 160 에포크의 짧은 학습만으로도 300 에포크 이상 학습된 기존 모델들과 대등하거나 더 낮은 gFID(2.95)를 기록했다. 이는 CaTok이 생성한 1D 토큰이 언어 모델의 예측 구조에 매우 적합함을 시사한다.

Ablation Study를 통해 MeanFlow 목적 함수와 REPA-A 정규화의 효과를 검증했다. REPA-A를 적용했을 때 학습 초기 손실 값의 급격한 튐(spike) 현상이 억제되고 수렴 속도가 빨라졌으며, 토큰 선택 전략에서 특정 구간 $[r, t]$를 활용하는 것이 전체 토큰을 사용하는 것보다 인과관계 학습에 더 유리함을 확인했다.

CaTok의 핵심은 확산 오토인코더의 디코더에 MeanFlow 공식을 적용하여 1D 토큰의 인과성을 확보한 점이다. 기존의 확산 모델은 모든 토큰에 동시에 의존하는 비인과적(non-causal) 구조를 가졌으나, CaTok은 시간 간격 $[r, t]$에 따라 토큰 시퀀스의 부분 집합 $V_{r:t}$를 조건으로 주어 토큰 간의 순차적 의존성을 모델링한다.

아키텍처 측면에서 인코더는 ViT-B/8 구조를 기반으로 하며, 16개의 레지스터 토큰을 사용하여 풍부한 전역 정보를 캡처한다. 디코더는 DiT-B/4 또는 DiT-L/2 구조를 사용하며, 고정된 VAE(MAR-VAE)의 잠재 공간 위에서 동작하여 연산 효율을 높였다.

REPA-A 정규화는 인코더의 특징 맵을 DINOv2-B/16과 정렬시킨다. 이는 기존 REPA-E가 VAE를 직접 정규화하던 방식과 달리, 조건부 확산 오토인코더의 인코더에 특화된 정렬 방식을 제안한 것이다. 이를 통해 인코더가 학습 초기부터 유의미한 시각적 특징을 추출하게 되어 전체 시스템의 안정성을 획기적으로 높였다.

학습 시 MeanFlow 손실($L_{MF}$)과 Rectified Flow 손실($L_{RF}$)을 함께 사용한다. $L_{RF}$는 순간 속도 필드를 모델링하여 학습 안정성을 돕고, $L_{MF}$는 구간 평균 속도를 학습하여 적은 수의 샘플링 단계에서도 정확한 복원을 가능하게 한다.

VQGAN 같은 전통적인 2D 토큰화 방식에 비해 복원 FID 수치에서 여전히 약간의 격차가 존재한다. 이는 2D 방식이 GAN 손실이나 복잡한 학습 레시피를 더 많이 활용하기 때문이며, CaTok은 이러한 최적화 없이도 경쟁력 있는 결과를 냈으나 절대적인 수치 면에서는 향후 개선의 여지가 있다.