큐빅 이산 확산: 고차원 표현 토큰 기반의 이산적 시각 생성

기존 이미지 생성 AI는 용량이 작은 토큰을 사용해 세밀한 의미 파악에 한계가 있었으나, 이 논문은 768차원 이상의 고차원 토큰을 직접 생성하는 기술을 구현했다. 이를 통해 이미지의 이해와 생성을 동일한 고성능 토큰으로 처리할 수 있어 진정한 의미의 통합 멀티모달 AI 구축을 가능하게 한다.

기존의 이산 생성 모델은 이미지를 아주 작은 차원의 벡터로 압축한 뒤 이를 예측한다. 하지만 이런 압축 과정에서 이미지의 풍부한 의미 정보가 손실되어, 생성된 토큰을 이미지 분류나 추론 같은 이해 작업에 그대로 쓰기에는 성능이 부족했다. CubiD는 이 문제를 해결하기 위해 고성능 시각 모델이 사용하는 768차원 이상의 고차원 특징 벡터를 직접 생성 대상으로 삼는다.

핵심 아이디어는 '세밀한 마스킹'이다. 기존 모델들이 이미지의 특정 위치 전체를 가리고 예측했다면, CubiD는 특정 위치의 수많은 차원 중 일부만 가리고 나머지 차원들과의 관계를 통해 이를 복원하는 법을 배운다. 이는 마치 퍼즐 조각의 모양뿐만 아니라 그 조각 내부의 미세한 색상 정보까지 부분적으로 관찰하며 전체를 맞추는 것과 같다.

결과적으로 모델은 공간적인 배치 관계뿐만 아니라, 하나의 위치 안에서 수백 개의 차원 정보가 서로 어떻게 얽혀 있는지를 정교하게 학습한다. 이러한 접근은 고차원 데이터가 가진 정보의 중복성을 역으로 활용하여, 훨씬 적은 단계만으로도 고품질의 이미지를 생성할 수 있게 만든다.

Dimension-wise Quantization (DQ) 과정을 통해 고차원 특징 맵 $z \in \mathbb{R}^{h \times w \times d}$의 각 차원을 독립적으로 $L$개의 레벨로 양자화한다. [연속적인 스칼라 값 $z_{x,y,i}$ 입력 → $L$단계 중 가장 가까운 인덱스로 매핑 → 이산 토큰 $q_{x,y,i}$ 출력] 과정을 거치며, 이는 고차원 공간에서 클러스터링이 어려운 기존 Vector Quantization의 한계를 극복한다.

Cubic Masking Strategy는 $h \times w \times d$ 크기의 3차원 텐서에서 임의의 위치와 차원을 독립적으로 마스킹한다. [이진 마스크 $M$ 생성 → 마스킹 비율 $r$을 Truncated Gaussian 분포에서 샘플링 → $M$이 1인 위치를 [MASK] 토큰으로 치환]하여 모델이 부분적인 관측값으로부터 전체를 추론하도록 유도한다.

Transformer 아키텍처는 각 공간 위치($h \times w$)를 하나의 시퀀스 토큰으로 취급하되 내부적으로는 $d$차원의 정보를 모두 포함한다. [마스킹된 $d$차원 벡터 입력 → Bidirectional Attention을 통한 공간 및 차원 간 관계 계산 → MLP 헤드를 통해 $d \times L$개의 로짓 출력] 구조를 통해 차원 수 $d$가 늘어나도 Transformer의 시퀀스 길이는 고정되어 연산 효율을 유지한다.

ImageNet-256 벤치마크에서 CubiD-XXL(3.7B) 모델은 FID 1.88을 기록하며 기존 이산 생성 모델들을 압도하는 성능을 보였다. 특히 Classifier-free guidance 없이도 2.02 gFID를 달성하여 토큰 자체의 표현력이 매우 우수함을 입증했다.

Ablation Study 결과, 공간 위치 전체를 마스킹하는 기존 방식(gFID 22.22)보다 차원 단위로 세밀하게 마스킹하는 본 방식(gFID 5.33)이 월등히 우수한 성능을 보였다. 이는 고차원 토큰 생성 시 차원 간의 상관관계를 학습하는 것이 필수적임을 시사한다.

모델 스케일링 분석에서 파라미터 수가 9억 개에서 37억 개로 증가함에 따라 gFID가 5.25에서 4.68로 꾸준히 개선되는 현상을 확인했다. 또한 DINOv2와 SigLIP2 등 서로 다른 사전 학습 인코더에서도 일관되게 높은 성능을 보여 범용성을 증명했다.

CubiD 아키텍처는 표준 Transformer를 기반으로 하며, Bidirectional Attention을 사용하여 마스킹된 텐서 전체의 문맥을 파악한다. 특징적인 점은 $d$차원의 이산 토큰을 다시 연속적인 스칼라로 역양자화하여 Transformer의 입력 임베딩으로 사용함으로써 사전 학습된 인코더의 특징 공간을 최대한 보존한다는 것이다.

학습 시에는 Cross-Entropy Loss를 사용하며, 마스킹된 위치($i \in M$)에 대해서만 손실을 계산한다. [정답 토큰 $q_i$와 예측 확률 $p$ 입력 → $-\log p$ 계산 → 마스킹된 모든 위치에 대해 합산 및 평균] 과정을 통해 모델이 데이터의 결합 분포를 학습하도록 한다. 마스킹 비율은 $\mu=1.0$인 Truncated Gaussian 분포에서 샘플링하여 모델이 아주 적은 정보만으로도 예측할 수 있도록 강하게 훈련시킨다.

추론 시에는 Cosine Schedule에 따라 점진적으로 토큰을 공개하는 Iterative Refinement 방식을 채택한다. 이는 $O(hwd)$의 복잡도를 갖는 자기회귀 방식과 달리, 차원 수에 관계없이 고정된 $T$ 단계(보통 256~512회)만으로 생성을 완료할 수 있게 하여 고차원 데이터 생성의 계산적 난제를 해결했다.

생성 품질이 사용된 사전 학습 인코더의 재구성 능력에 종속되며, 현재 재구성 PSNR이 약 18dB 수준으로 제한되어 아주 미세한 디테일 표현에 한계가 있다. 또한 연속형 확산 모델에 비해 여전히 수백 단계의 많은 생성 과정이 필요하다.