고압축 고품질 이미지 재구성을 위한 위톡(WeTok) 토크나이저
초록
위톡은 그룹‑와이즈 룩업‑프리 양자화(GQ)와 생성형 디코더(GD)를 결합해, 기존 이산 토크나이저가 겪던 압축률‑재구성 품질 간의 트레이드오프를 크게 완화한다. 메모리·연산 효율성을 유지하면서 수백 배 압축에서도 rFID 0.12 수준의 초저 오류를 달성했으며, 768배 압축에서도 rFID 3.49를 기록해 경쟁 모델을 크게 앞선다.
상세 분석
위톡의 핵심 혁신은 두 가지이다. 첫째, 그룹‑와이즈 룩업‑프리 양자화(GQ)는 잠재 피처를 채널 차원에서 g 개의 그룹으로 나눈 뒤, 각 그룹에 고정된 {-1, 1} 코드북을 적용한다. 기존 LFQ가 전체 d 차원 공간에 대해 엔트로피 손실을 계산하면서 메모리 폭증을 초래한 것과 달리, GQ는 엔트로피 손실을 g 개의 독립적인 서브스페이스로 분해해 메모리 사용량을 상수 수준으로 유지한다. 또한, 논문에서는 GQ의 근사 오차가 BSQ보다 엄격히 작다는 정리를 제시해, 정확도‑효율성 양면에서 우수함을 이론적으로 뒷받침한다. 둘째, 생성형 디코더(GD)는 기존 이산 토크나이저가 조건 U_Q 만을 입력받아 deterministic하게 복원하던 방식을 탈피한다. 표준 정규분포 z 를 추가 채널로 결합함으로써, 디코더는 z 조건부 확률 모델이 되어 동일 토큰에 대해 다중 가능한 이미지 분포를 샘플링한다. 이는 고압축 상황에서 “평균화된 흐릿함”을 방지하고, 고주파 디테일을 복원하는 데 크게 기여한다. 학습은 두 단계로 진행되는데, 첫 단계에서는 기존 재구성 손실(L2, LPIPS, GAN)과 GQ 손실을 사용해 안정적인 베이스 모델을 만든 뒤, 두 번째 단계에서 디코더에 z 입력을 위한 채널을 zero‑initialization으로 확장해 기존 파라미터와 충돌 없이 조건부 생성 능력을 부여한다. 실험에서는 ImageNet‑50k 검증에서 400% 압축률에 rFID 0.12(FLUX‑VAE 0.18, SD‑VAE 3.5 0.19)와, 768배 압축에서도 rFID 3.49(코스모스 4.57)라는 기록을 세웠다. 메모리 사용량 측면에서도 GQ는 d = 8,16,24,32 모두 10.5~10.6 GB로 OOM 없이 학습이 가능했으며, BSQ와 LFQ는 높은 차원에서 OOM이 발생한다. Ablation 결과는 그룹 수 G 증가가 재구성 품질을 크게 향상시키고, 베이스 채널 C 와 residual 블록 B 의 적절한 조합(C=256, B=4)이 최적 성능을 만든다는 것을 보여준다. 전체적으로 위톡은 코드북 규모를 사실상 무제한으로 확장하면서도 메모리·연산 비용을 억제하고, 생성형 디코더를 통해 고압축에서도 시각적 디테일을 보존하는 새로운 패러다임을 제시한다.
댓글 및 학술 토론
Loading comments...
의견 남기기