CaTok: Taming Mean Flows for One-Dimensional Causal Image Tokenization

이 논문은 Vision Foundation 모델의 특징 정렬을 통해 학습을 안정화하고, MeanFlow 디코더를 활용한 1 차원 인과적 이미지 토크나이저 'CaTok'을 제안하여 ImageNet 재구성 및 생성 성능에서 최첨단 결과를 달성했다고 요약할 수 있습니다.

핵심

🎨 그림을 그리는 두 가지 방식: "한 번에 다 그리기" vs "한 줄씩 쓰기"

지금까지 AI 가 그림을 그릴 때는 주로 두 가지 방식 중 하나를 썼습니다.

1. VQGAN 같은 방식 (2D 토큰): 그림을 작은 퍼즐 조각 (패치) 으로 잘게 나누고, 이를 2 차원 격자에 채워 넣습니다. 하지만 이 조각들을 1 차원 줄로 만들 때, 순서가 중요하지 않습니다. 마치 "퍼즐 조각을 아무렇게나 섞어서 옆에 놓는 것"과 비슷해서, AI 가 "다음 조각은 무엇이어야 할까?"라고 순서대로 생각하며 그리는 것이 어렵습니다.
2. Diffusion Autoencoder (확산 모델): 그림을 1 차원 줄 (문장처럼) 로 만들려고 시도했지만, 기존 방식들은 모든 조각을 한 번에 다 보여주고 그림을 완성했습니다. 이는 "글을 쓸 때, 앞뒤 문장을 모두 다 보고 한 줄을 쓰는 것"과 같아서, AI 가 순서대로 예측하는 능력 (다음 단어 예측) 을 키우기 어렵게 만들었습니다.

🚀 CaTok 의 등장: "그림을 문장처럼 순서대로 쓰기"

CaTok 은 이 문제를 해결하기 위해 그림을 '문장'처럼 1 차원 줄로 만들고, 순서대로 읽을 수 있게 만들었습니다.

1. 시간의 흐름을 이용한 '순서' (인과성)

CaTok 은 그림을 그릴 때 시간의 흐름을 이용합니다.

• 비유: 그림을 그리는 과정을 '시간이 흐르며 그림이 완성되어 가는 과정'으로 봅니다.
• 작동 원리: AI 는 그림의 **처음 부분 (0~10% 시간)**에 해당하는 조각을 먼저 보고, 그다음 중간 부분, 그리고 나중 부분을 순서대로 학습합니다.
• 결과: 마치 사람이 글을 쓸 때 "첫 문장 → 두 번째 문장 → 세 번째 문장" 순서로 이어가듯, AI 도 "그림의 앞부분 → 중간 → 뒷부분" 순서로 자연스럽게 이어지는 **인과성 (Causality)**을 갖게 됩니다.

2. '평균 흐름'을 보는 눈 (MeanFlow)

기존 방식들은 '순간적인 속도'만 보거나, '무작위로 일부만' 보는 문제가 있었습니다.

• CaTok 의 해결책: CaTok 은 **시간 구간 [r, t] 동안의 '평균 흐름'**을 봅니다.
• 비유: 차를 운전할 때, "지금 이 순간의 속도"만 보는 게 아니라, "지난 10 초 동안의 평균 속도"를 보고 다음 방향을 결정하는 것과 같습니다.
• 장점: 이렇게 하면 AI 는 한 번에 그림을 완성할 수도 있고 (1 스텝), 단계별로 정교하게 다듬을 수도 있습니다 (25 스텝). 마치 빠른 스케치와 정밀한 드로잉을 모두 가능하게 하는 것입니다.

3. REPA-A: "명화 감수성을 가진 선생님"

학습을 더 빠르고 안정적으로 하기 위해, CaTok 은 REPA-A라는 기술을 썼습니다.

• 비유: 초보 화가가 그림을 그릴 때, 이미 유명한 대가 (Vision Foundation Model) 가 그린 명화를 옆에 두고 "이 화가라면 이 부분을 어떻게 그렸을까?"라고 비교하며 배웁니다.
• 효과: AI 가 그림의 의미 (개념, 사물의 특징) 를 더 잘 이해하게 되어, 훨씬 적은 시간으로 더 높은 퀄리티의 그림을 그릴 수 있게 됩니다.

---

🏆 CaTok 의 성과: 왜 이것이 중요한가요?

1. 빠르고 정확한 그림: 256 개의 조각 (토큰) 만으로도 아주 선명한 그림을 만들 수 있습니다.
2. 문장 같은 순서: 그림의 토큰들이 앞뒤로 자연스럽게 연결되어, AI 가 "다음에 어떤 그림 조각이 올지" 예측하는 능력이 뛰어납니다.
3. 유연성: 그림을 한 번에 뚝딱 만들 수도 있고, 천천히 정교하게 다듬을 수도 있습니다.
4. 최고의 기록: ImageNet(유명한 이미지 데이터셋) 에서 다른 최신 기술들보다 더 좋은 화질 (PSNR, SSIM) 을 보여주면서, 학습 시간도 절반 이하로 줄였습니다.

💡 한 줄 요약

CaTok 은 그림을 '퍼즐 조각'이 아니라 '문장'처럼 순서대로 이어 붙여 그릴 수 있게 만든 기술입니다. 이를 통해 AI 는 글을 쓰듯 그림을 순서대로 예측할 수 있게 되었고, 더 빠르고 더 똑똑하게 그림을 그릴 수 있게 되었습니다.

기술 요약

1. 문제 정의 (Problem)

최근 대규모 언어 모델 (LLM) 의 성공은 '자기회귀 (Autoregressive, AR)' 패러다임과 '인과적 토큰화 (Causal Tokenization)'에 기반합니다. 텍스트는 자연스러운 순서가 존재하지만, 이미지를 1 차원 시퀀스로 변환하여 AR 모델에 적용하는 것은 여전히 난제입니다. 기존 시각 토크나이저들의 한계는 다음과 같습니다:

• 2D 토큰화 및 평탄화 (Flattening): VQGAN 등 기존 모델은 2D 패치를 1D 시퀀스로 평탄화하지만, 래스터 순서나 무작위 순서를 따르므로 토큰 간의 인과성 (causality) 이 부족합니다.
• VAR 모델의 한계: VAR 은 다중 스케일 2D 토큰을 사용하여 coarse-to-fine 순서를 강제하지만, 이는 LLM 의 '다음 토큰 예측 (Next-token Prediction)' 패턴과 불일치합니다.
• 확산 오토인코더 (Diffusion Autoencoders) 의 문제:
  • Naïve Flow Decoder: 모든 토큰을 동시에 조건으로 사용하여 인과성이 결여됩니다.
  • Consistency Decoder (Nested Dropout): 무작위 샘플링이나 시간 단계 바인딩을 통해 앞선 토큰들만 조건으로 사용하지만, 초기 토큰이 과도하게 선택되어 **불균형 (Imbalance)**이 발생하고 AR 생성 성능이 저하됩니다.

2. 제안 방법 (Methodology)

저자들은 CaTok을 제안합니다. 이는 MeanFlow 객체와 결합된 1 차원 인과적 이미지 토크나이저로, 확산 오토인코더 아키텍처를 기반으로 합니다.

A. 아키텍처

• 인과적 ViT 인코더: 이미지 패치와 레지스터 (Registers) 를 입력받아 1D 토큰을 추출합니다. 인과적 어텐션 마스크를 적용하여 토큰 간의 의존성을 보장합니다.
• MeanFlow 디코더: 기존 확산 모델이 순간 속도 (instantaneous velocity) 를 학습하는 것과 달리, MeanFlow는 시간 구간 $[r, t]$에 걸친 **평균 속도장 (Average Velocity Field)**을 학습합니다.

B. 핵심 기법: MeanFlow 기반 토큰 선택

• 균형 잡힌 인과성: 디코더는 시간 구간 $[r, t]$에서 샘플링된 1D 토큰 세그먼트 ($V_{r:t}$) 만을 조건으로 사용합니다.
• 작동 원리: 토큰을 무작위로 선택하거나 앞쪽 토큰만 고르는 대신, 시간 구간을 매핑하여 토큰을 선택함으로써 토큰 간의 인과성을 유지하면서도 불균형 문제를 해결합니다. 이는 노이즈에서 이미지로의 생성 과정에서 인과성을 자연스럽게 포착하게 합니다.

C. REPA-A (Representation Alignment - A)

• 정규화 기법: 인코더의 특징을 고품질의 외부 비전 파운데이션 모델 (VFM, 예: DINOv2) 의 표현과 정렬 (Align) 하는 새로운 정규화 방법입니다.
• 효과: 인코더가 더 의미 있고 구별력 있는 시각적 콘텐츠를 추출하도록 하여, 오토인코더의 학습을 안정화하고 수렴 속도를 가속화합니다.

D. 학습 및 생성

• 학습: MeanFlow 손실 (평균 속도 예측) 과 Rectified Flow 손실 (순간 속도 예측) 을 결합하여 학습하며, REPA-A 를 통해 특징 정렬을 수행합니다.
• 생성: 학습된 AR 모델이 1D 토큰 시퀀스를 생성하면, MeanFlow 디코더를 통해 1 스텝 (One-step) 또는 다중 스텝으로 이미지를 재구성할 수 있습니다.

3. 주요 기여 (Key Contributions)

1. 새로운 아키텍처: 확산 오토인코더와 MeanFlow 객체를 기반으로 한 최초의 1D 인과적 이미지 토크나이저 제안.
2. 균형 잡힌 인과성: 시간 구간 기반 토큰 선택 메커니즘을 통해 AR 생성에 적합한 인과성을 유지하면서도 토큰 간 불균형을 해결.
3. REPA-A: VFM 을 활용한 새로운 정규화 기법으로 학습 안정성과 속도를 획기적으로 개선.
4. 유연한 샘플링: 1 스텝 (고속) 및 다중 스텝 (고품질) 샘플링을 모두 지원하며, 토큰 수를 줄여도 coarse-to-fine 재구성이 가능합니다.

4. 실험 결과 (Results)

ImageNet-1K (256x256) 에서의 실험 결과는 다음과 같습니다:

• 재구성 성능 (Reconstruction):
  • CaTok-L-256은 0.75 rFID, 22.53 PSNR, 0.674 SSIM을 기록하여 기존 확산 기반 토크나이저 (Semanticist, FlowMo 등) 보다 우수한 성능을 보였습니다.
  • 특히, 1 스텝 샘플링을 지원하는 토크나이저 중에서도 가장 높은 PSNR 과 SSIM 을 달성했습니다.
  • 학습 에포크가 타 모델 (300+ 에포크) 에 비해 절반 미만 (160 에포크) 으로 적게 소요되어 효율성이 입증되었습니다.
• 자기회귀 생성 (AR Generation):
  • CaTok 으로 토큰화한 후 LlamaGen 기반 AR 모델로 생성 시, 2.95 gFID를 기록하여 최첨단 방법론들과 경쟁력 있는 성능을 보였습니다.
  • 토큰 수를 줄여도 (예: 32 개) 성능 저하가 적어 효율적인 생성이 가능합니다.
• Ablation Study:
  • MeanFlow 객체와 토큰 선택 메커니즘 ($[r, t]$ 구간) 이 AR 생성 성능 (gFID) 향상에 결정적임을 확인했습니다.
  • REPA-A 를 적용했을 때 학습 곡선이 안정화되고 재구성 품질이 향상됨을 시각화 (PCA, Loss Curve) 하였습니다.

5. 의의 및 결론 (Significance)

CaTok 은 언어 모델의 자기회귀 패러다임을 시각 생성 영역으로 성공적으로 확장하는 중요한 디딤돌이 됩니다.

• 인과성의 정립: 이미지 토큰화에서 '인과성'과 '균형'을 동시에 해결하여, AR 모델이 시각 데이터를 효과적으로 학습할 수 있는 토대를 마련했습니다.
• 효율성: 1 스텝 샘플링을 지원하면서도 고화질 재구성을 가능하게 하여, 생성 속도와 품질 간의 트레이드오프를 줄였습니다.
• 미래 지향성: 이 연구는 시각 생성 모델이 언어 모델과 유사한 확장성 (Scalability) 과 일반화 능력을 갖추기 위한 핵심 기술로 평가받으며, 향후 대규모 비전 - 언어 모델 및 효율적인 생성 모델 개발에 중요한 기여를 할 것으로 기대됩니다.