MixTok이 1D 전역 토큰과 희소 2D 토큰을 섞어 이미지 토큰을 절약한다
MixTok: Adaptive Visual Tokenization with Global 1D Tokens and Sparse 2D Refinement
무엇인가
시각 토크나이저는 이미지를 압축된 잠재 표현으로 바꿔 복원과 하위 생성 모델의 입력 공간을 제공한다. 문제는 표현 용량을 어디에 얼마나 쓸 것인가다. 기존 2D 그리드 토크나이저는 토큰과 공간 위치의 대응을 유지하지만, 매끄러운 하늘 같은 단순 영역에도 복잡한 물체·텍스처 영역과 똑같은 토큰 밀도를 배정해 용량을 낭비한다. 반대로 TiTok 계열의 1D 토크나이저는 이미지 전체를 짧은 시퀀스로 집계해 압축률은 뛰어나지만, 짧은 시퀀스 안에서 미세한 국소 디테일을 보존하기 어렵다. 여기에 더해 대부분의 토크나이저는 이미지 내용과 무관하게 같은 해상도면 같은 토큰 수를 모든 이미지에 배정한다. MixTok은 이 두 축, 즉 어디에 국소 토큰을 쓸지와 이미지마다 몇 개를 줄지를 동시에 다룬다.
어떻게 동작하나
구조는 하이브리드 표현이다. 1D 분기가 길이 L의 토큰 시퀀스 T_1D를, 2D 분기가 h×w 그리드 T_2D를 만든다. 2D 그리드는 N=hw개의 후보 정제 위치를 제공한다. 1D 토큰은 공간 대응이 없으므로, 양자화된 특징 F_1D를 공간 질의 그리드 S로 확장하는 G_1D를 거쳐 학습 가능한 어댑터 P_1D로 디코더 입력 공간에 투영해 F̄_1D를 만든다. 2D 분기의 F_2D는 이미 같은 디코더 차원과 공간 구조를 갖고 있어 두 표현이 같은 특징 공간에서 결합된다. 선택된 위치 집합 R에서는 1D에서 복원한 특징을 2D 특징으로 교체하고, 나머지 위치는 F̄_1D를 그대로 둔 뒤 공유 디코더 D가 이미지를 복원한다. 최종 토큰 수는 B = L + K이며, 선택 위치는 디코딩 때 위치를 되찾기 위한 사이드 정보로 저장될 뿐 콘텐츠 토큰으로 계산되지 않는다.
무엇과 다른가
핵심은 콘텐츠 적응형 라우터다. Reconstruction Grid Scorer는 정렬된 두 분기 특징 F̄_1D, F_2D와 1D만으로 만든 베이스 복원 x̂_base를 입력받아 위치별 정제 우선순위 점수 ρ를 예측한다. 이 점수를 내림차순으로 정렬한 뒤 온도 T를 가진 softmax로 정규화하고, 누적 확률이 임계값 p에 도달하는 최단 접두 길이를 K(x)로 삼는 top-p 선택을 적용한다. 확률 질량이 소수 위치에 몰려 있으면 정제 토큰을 적게, 넓게 퍼져 있으면 많이 배정하는 셈이다. p는 학습 이미지 일부에서 목표 평균 토큰 예산에 맞춰 보정한 뒤 평가에서는 고정하되, 이미지별 토큰 수는 달라지게 둔다. 점수 예측은 1D 베이스 복원과 완전한 2D 복원 사이의 국소 복원 오차 감소량으로 보조 지도를 받고, 미분 가능한 선택 대리(surrogate)를 통해 혼합 복원 손실로도 최적화된다. 토크나이저 학습 목표는 L1 복원 손실, LPIPS 지각 손실, GAN 손실의 가중 합이다.
어떻게 쓰나
생성은 2단계 MaskGIT 파이프라인으로 확장된다. 1단계에서 사전학습된 TiTok MaskGIT이 길이 L의 1D 토큰을 만들고, 이를 코드북 조회·공간 확장·특징 정렬로 베이스 그리드 F̄_1D로 바꾼다. 생성 시점에는 2D 특징을 쓸 수 없으므로, 별도의 Generation Grid Scorer가 1D 베이스 그리드만으로 정제 우선순위를 예측한다. 이 scorer는 고정된 Reconstruction Grid Scorer를 증류해 학습한다. Budget Controller가 같은 top-p 규칙으로 정제 집합 R과 크기 K를 정한다. 베이스 그리드는 이미지 디코더의 특징 공간에 있는 반면 2D 생성기는 토큰 임베딩 공간에서 동작하므로, Embedding Adapter A가 베이스 특징을 공간 컨텍스트 임베딩 H_ctx로 변환한다. 초기 입력 임베딩은 선택 위치에서 마스크 토큰, 나머지 위치에서 H_ctx가 되고, 미세조정된 Halton-MaskGIT이 R 안에서만 2D 토큰을 생성한다. 샘플링은 R로 제한된 Halton 순서를 따른다.
전제와 한계
복원 실험은 ImageNet 256×256에서 TiTok-L-32(1D 토큰 32개)와 LlamaGen VQ-16(16×16 그리드, 공유 이미지 디코더)을 조합해 수행했다. 평균 약 128개 콘텐츠 토큰으로 rFID 1.1950, PSNR 19.51 dB를 기록했다. 논문은 이 수치가 비슷한 토큰 예산에서 TiTok 변형들과 One-D-Piece보다 rFID를 개선한 것이며 FlexTok, DetailFlow와는 경쟁력 있는 수준이라고 서술한다. 백본 전이성도 확인했는데, ALIT과 VQGAN 조합에서는 1D 64개와 2D 64개로 총 128토큰을 맞춰 ALIT-128 대비 rFID를 2.9453에서 2.4720으로, PSNR을 18.71 dB에서 18.76 dB로 개선했다.
생성에서는 약 131토큰으로 gFID 1.95를 냈다. 비슷한 토큰 예산의 TiTok-S-128이 gFID 2.00에 샘플링 라운드 64회를 쓰는 것과 비교해 gFID를 1.95로 낮추면서 라운드를 16회로 줄였다. 비슷한 파라미터 수의 DetailFlow-32와 비교하면 gFID 1.95 대 2.63으로 더 낮으면서 토큰은 약 절반, 라운드는 16 대 39로 적다. 단일 이미지 지연 시간 기준으로 각각 3.77배, 2.48배 빠르다고 논문은 밝힌다. 같은 1D 토큰에서 출발해 희소 2D 정제를 추가하면 8회의 추가 샘플링 라운드로 gFID가 4.33에서 1.95로 개선된다는 점도 정제 단계의 기여를 보여준다.
절제 실험은 두 구성 요소를 각각 검증한다. 위치 선정에서는 Reconstruction Grid Scorer가 rFID 1.1950으로 가장 좋았고, 2D 특징 없이 1D 베이스 특징만 쓰는 Generation Grid Scorer도 rFID 1.3479, PSNR 18.4776으로 무작위·균일·그래디언트 엔트로피 기반 순위보다 우수했다. 예산 배정에서는 동일한 scorer를 쓰고 top-p와 고정 K=96, 그리고 이미지별 예산을 무작위로 뒤섞은 대조군을 비교했는데, top-p가 비슷한 평균 예산에서 rFID와 PSNR 모두 가장 좋았다. 예산을 섞으면 두 지표가 나빠진다는 것은 단순히 토큰 수를 다양하게 하는 것 이상으로, 예산을 이미지 내용에 맞추는 것 자체가 이득이라는 주장의 근거다.
실무 관점에서 MixTok은 토큰 예산이 곧 비용인 이미지 생성·압축 파이프라인에 관심 있는 개발자에게 유용하다. 이미지마다 다른 토큰 수를 쓰면서도 평균 예산을 목표치에 맞추는 top-p 방식과, 1D 베이스만 보고 정제 위치를 예측하도록 scorer를 증류하는 설계는 생성 단계에서 별도 정보 없이도 적응적 할당을 구현하는 패턴으로 참고할 만하다. 다만 저자들이 밝힌 한계도 분명하다. 콘텐츠 토큰 수는 줄였지만 복원 과정은 여전히 조밀한 2D 인코딩을 필요로 하고, 2단계 생성기는 전체 공간 그리드 위에서 동작한다. 평가도 ImageNet 256×256과 좁은 범위의 평균 토큰 예산에 국한되어 있다. 향후 과제로 저자들은 희소 연산 도입과 함께 더 넓은 예산 범위, 더 높은 해상도, 더 다양한 시각 데이터로 적응적 전역-국소 할당을 확장하는 것을 제시한다.