TC-AE: 딥 압축 오토인코더를 위한 토큰 용량 최적화 연구

이미지 생성 모델의 효율성을 높이기 위해 토큰 압축률을 극단적으로 높이면 이미지 품질이 급격히 저하되는 '잠재 표현 붕괴' 현상이 발생합니다. 이 논문은 복잡한 구조 변경 없이 토큰 압축 단계를 나누고 자기지도학습을 결합하는 것만으로도 압축 효율과 생성 품질을 동시에 잡을 수 있음을 증명했습니다.

기존의 이미지 토크나이저는 고해상도 이미지를 작은 격자(Patch)로 나누어 임베딩한 뒤, 이를 Transformer 블록을 거쳐 매우 작은 잠재 공간(Latent Space)으로 압축한다. 이때 압축률을 높이기 위해 토큰의 채널 수는 늘리고 공간 해상도는 줄이는데, 이 과정에서 토큰 간의 관계 정보가 소실되는 '잠재 표현 붕괴'가 발생하여 생성된 이미지의 질이 떨어진다.

TC-AE는 이 문제를 '토큰 수 확장'과 '단계적 압축'으로 해결한다. 먼저 Patch Size를 줄여 초기 토큰 수를 대폭 늘려 세부 정보를 확보한다. 그다음, 한 번에 압축하는 대신 Transformer 레이어 중간에 압축 계층을 두어 점진적으로 해상도를 줄인다. 이는 정보가 병목 구간에서 한꺼번에 사라지는 것을 막는 완충 작용을 한다.

결과적으로 모델은 더 많은 시각적 세부 사항을 유지하면서도 계산 효율적인 작은 잠재 공간을 생성할 수 있게 된다. 여기에 자기지도학습을 더해 각 토큰이 이미지의 의미적 특징을 더 잘 담아내도록 함으로써, 생성 모델이 학습하기 훨씬 수월한 '친화적인' 잠재 공간을 제공한다.

전체 아키텍처는 Vision Transformer(ViT) 기반의 인코더와 대칭 구조의 디코더로 구성된다. 인코더는 입력 이미지 X를 p x p 크기의 패치로 나누어 N개의 토큰 T를 생성한다. [입력 이미지 HxWx3 → 패치 임베딩 → N x d 크기의 토큰 시퀀스 생성]

핵심인 Staged Token Compression(STC)은 인코더의 12개 레이어를 M개와 (12-M)개로 나눈다. 첫 M개의 레이어는 고해상도 토큰을 처리하고, 중간 압축 계층이 이를 1/4로 줄인 뒤 나머지 레이어가 처리한다. [N개 토큰 → M개 레이어 연산 → 1/4 축소 → 나머지 레이어 연산 → 최종 잠재 표현 z 생성]

학습 시에는 iBOT의 Student-Teacher 구조를 차용하여 자기지도학습 손실 함수(LiBOT)를 추가한다. 이는 마스킹된 이미지 패치를 복구하는 과정과 전체 이미지의 클래스 토큰을 일치시키는 과정을 포함한다. [이미지 증강 및 마스킹 → 패치 단위 예측 및 클래스 토큰 정렬 → 의미론적 구조 강화]

최종 손실 함수 LTC-AE는 픽셀 단위 복원 손실(Lpix), 지각 손실(Lp), 적대적 손실(Lg)에 LiBOT를 가중치 α=0.1로 결합하여 정의된다. [복원 오차 + 의미론적 정렬 오차 → 가중 합산 → 전체 모델 최적화]

ImageNet-1K 데이터셋 실험 결과, TC-AE는 기존 방식보다 적은 연산량(164 GFLOPs vs DC-AE 607 GFLOPs)으로 더 우수한 성능을 보였다. 특히 64개의 토큰만 사용하는 극한의 압축 환경에서 gFID 점수를 기존 26.44에서 7.16으로 대폭 낮추었다.

확산 모델(DiT) 학습 시 수렴 속도가 비약적으로 향상되었다. 동일한 gFID 수준에 도달하는 데 필요한 학습 반복 횟수가 베이스라인 대비 4.7배 적었으며, Inception Score(IS) 기준으로는 3.5배 빠른 수렴을 기록했다.

Ablation Study를 통해 STC와 SSL의 시너지 효과를 확인했다. STC는 토큰-잠재 공간 간의 정보 보존율(Linear Probing Accuracy)을 0.08에서 0.31로 높였으며, SSL은 재구성 품질을 약간 희생하는 대신 생성 성능(gFID)을 32.92에서 16.39로 크게 개선했다.

TC-AE는 ViT 기반 토크나이저에서 발생하는 'Aggressive Token-to-Latent Compression' 문제를 최초로 식별하고 이를 구조적으로 해결했다. 기존 연구들이 잠재 공간의 채널 수를 늘리는 데 집중한 것과 달리, 본 연구는 토큰 공간(Token Space)의 해상도 유지와 단계적 축소가 생성 성능에 더 결정적임을 밝혔다.

아키텍처 설계에서 M=6 설정이 재구성 품질과 생성 성능 사이의 최적의 균형점임을 실험적으로 도출했다. 이는 초기 레이어가 충분한 해상도에서 시각적 특징을 추출하고, 후반부 레이어가 압축된 토큰의 의미를 정제하는 역할을 분담하기 때문이다.

자기지도학습(SSL)의 통합 방식 또한 차별화된다. 외부의 거대한 사전학습 모델(DINOv2 등)을 Teacher로 사용하는 대신, 토크나이저 학습 과정에서 자체적인 Student-Teacher 구조를 통해 의미 구조를 형성함으로써 데이터 도메인에 특화된 잠재 공간을 구축할 수 있게 했다.

SSL을 도입할 경우 이미지의 세부적인 픽셀 복원 품질(rFID, PSNR)이 순수 복원 목적의 모델보다 다소 낮아지는 경향이 있다. 이는 모델이 픽셀 단위의 정확도보다 이미지의 전체적인 의미 구조 학습에 더 많은 용량을 할당하기 때문으로 분석된다.