SNCE: 확장 가능한 이산 이미지 생성을 위한 기하학 인지 지도 학습

이산 이미지 생성 모델에서 고화질 구현을 위해 코드북 크기를 키우면 학습이 매우 어려워지는 문제를 해결했다. 기존의 '정답 아니면 오답' 식의 학습 대신 정답과 유사한 의미를 가진 주변 토큰들에게도 점수를 주는 방식을 도입해 학습 효율을 높였다. 이를 통해 더 적은 데이터와 시간으로도 정교한 이미지 생성과 편집이 가능해졌다.

이산 이미지 생성은 이미지를 격자 형태의 토큰으로 변환하여 학습한다. 이때 VQ(Vector Quantization)를 통해 연속적인 특징을 가장 가까운 코드북 토큰으로 매핑하는데, 화질을 높이려면 코드북 크기를 키워야 한다. 하지만 코드북이 커질수록 각 토큰이 학습 데이터에 등장하는 빈도가 급격히 낮아지는 '희소성' 문제가 발생하며, 기존의 Cross-Entropy 손실 함수는 정답 토큰 하나에만 모든 확률을 부여하므로 학습 신호가 매우 약해지는 한계가 있다.

SNCE는 t-SNE 시각화 기법에서 영감을 받아, 정답 토큰뿐만 아니라 임베딩 공간에서 정답과 가까운 주변 토큰들에게도 거리에 비례한 확률을 배분한다. 이는 모델이 단순히 특정 인덱스를 맞추는 것을 넘어, 이미지의 의미론적 구조가 반영된 임베딩 공간의 기하학적 관계를 학습하도록 유도한다.

결과적으로 모델은 정답과 아주 유사한 '오답'을 내놓았을 때 강한 벌칙을 받는 대신 긍정적인 피드백을 받게 된다. 이는 학습 초기 단계에서 모델이 유의미한 방향으로 빠르게 수렴하게 하며, 데이터가 부족한 대규모 코드북 환경에서도 안정적인 성능을 보장한다.

SNCE는 연속적인 잠재 벡터 $z_i$와 코드북 벡터 $v_k$ 사이의 거리를 기반으로 이웃 분포 $q_k(z_i)$를 정의한다. [입력 잠재 벡터와 모든 코드북 벡터 사이의 거리 $d$를 입력으로] → [지수 함수 $\exp(-d/2\tau^2)$를 계산하고 전체 합으로 나누는 연산을 수행해] → [각 토큰에 대한 확률 값 $q_k$를 얻고] → [이 값은 해당 토큰이 정답과 얼마나 의미적으로 가까운지를 나타내는 가중치가 된다].

기존의 Cross-Entropy 손실 함수 $J_{CE}$에서 원-핫 인디케이터 함수를 위에서 계산한 $q_k(z_i)$로 대체하여 $J_{SNCE}$를 구성한다. [이웃 분포 $q_k$와 모델 예측 로그 확률 $\log p_\theta$를 입력으로] → [두 값을 곱하여 모든 토큰에 대해 합산하는 연산을 수행해] → [최종 손실 값을 산출하며] → [이 값이 작아질수록 모델은 기하학적으로 타당한 주변 토큰들을 정답 후보로 더 잘 인식하게 된다].

그라디언트 분석을 통해 SNCE의 동작 원리를 명확히 한다. [타겟 확률 $w_k$와 모델의 예측 확률 $p_\theta$를 입력으로] → [두 값의 차이($w_k - p_\theta$)를 구하는 연산을 수행해] → [각 로짓에 대한 그라디언트를 얻고] → [이 값이 양수면 해당 토큰의 생성 확률을 높이고 음수면 낮추는 방향으로 가중치를 갱신한다]. 이를 통해 정답 근처의 토큰들도 긍정적인 업데이트를 받게 된다.

ImageNet-256 클래스 조건부 생성 실험에서 131,072 크기의 대규모 코드북을 사용했을 때, SNCE는 100 에포크 만에 FID 3.62를 기록하여 기존 CE(7.53) 대비 2배 이상의 성능 향상을 보였다. 300 에포크 학습 시에도 FID 3.42로 CE(5.44)를 크게 앞질렀으며, 이는 코드북이 커질수록 SNCE의 최적화 효율이 더 극대화됨을 입증한다.

1024x1024 고해상도 텍스트-이미지 합성 태스크(LaViDa-O 기반)에서 SNCE는 FID를 3.67 개선하고 HPSv3 점수를 0.12 높였다. GenEval 벤치마크에서도 0.78점을 기록하며 기존 모델들보다 뛰어난 텍스트 정렬 능력을 보였다. 특히 단일 객체뿐만 아니라 여러 객체의 관계나 색상 속성을 정확하게 표현하는 능력이 향상되었다.

이미지 편집 벤치마크인 ImgEdit에서도 SNCE는 전체 점수 3.89를 기록하여 CE(3.76) 및 기존 소규모 코드북 모델(3.71)보다 우수한 성능을 보였다. 정성적 분석 결과, SNCE로 학습된 모델은 얼굴 특징이나 눈과 같은 미세한 디테일을 더 정교하게 보존하면서도 편집 지시사항을 충실히 따르는 것으로 나타났다.

SNCE는 이산 토큰 생성을 연속적인 임베딩 공간에서의 회귀 문제와 분류 문제 사이의 가교 역할을 하도록 설계되었다. 기존 CE가 토큰 간의 독립성을 가정하는 것과 달리, SNCE는 VQ 임베딩 공간의 기하학적 연속성을 귀납적 편향(Inductive Bias)으로 주입한다. 이는 수학적으로 모델의 예측 분포와 기하학적 이웃 분포 사이의 KL Divergence를 최소화하는 것과 동일하다.

이론적으로 SNCE는 범주형 변분 오토인코더(Categorical VAE)의 확률적 양자화 과정으로 해석될 수 있다. 또한 지식 증류(Knowledge Distillation) 관점에서는 토크나이저를 '약한 교사'로 활용하여 생성 모델인 '강한 학생'에게 임베딩 공간의 연속성 정보를 전달하는 구조를 가진다. 이는 모델이 단순히 레이블을 외우는 것이 아니라 데이터의 내재적 구조를 이해하게 만든다.

하드웨어 효율성을 위해 t-SNE의 적응형 대역폭 대신 고정된 온도 파라미터 $\tau$를 사용한다. 실험 결과 $\tau=0.71$에서 최적의 성능을 보였으며, 이는 임베딩 공간의 밀도 변화에 유연하게 대응하면서도 학습 신호의 잡음을 최소화하는 균형점임이 확인됐다. 또한 이 방식은 별도의 추가 파라미터 없이도 대규모 코드북의 선형 투영 계층이 유의미한 양(+)의 신호를 받도록 유도한다.

생성된 이미지가 여전히 픽셀 단위에서 완벽하지 않으며 미세한 아티팩트가 포함될 수 있다. 또한 기반 모델인 LaViDa-O로부터 환각(Hallucination) 현상이나 사회적 편향성을 그대로 상속받을 가능성이 존재한다.