스냅샷을 활용한 일반화된 이산 확산 모델 (Generalized Discrete Diffusion from Snapshots)

기존의 텍스트 생성 AI는 단어를 하나씩 순차적으로 생성하는 방식이 주류였으나, 이 논문은 모든 단어를 동시에 생성하는 이산 확산 방식의 한계를 극복했다. 임의의 노이즈 과정을 지원하는 통합 프레임워크를 통해 학습 효율과 생성 품질을 획기적으로 높여, 대규모 언어 모델에서도 확산 모델이 자기회귀 모델보다 우수한 성능을 낼 수 있음을 입증했다.

기존의 텍스트 생성은 주로 단어를 하나씩 이어 붙이는 자기회귀(Autoregressive) 방식을 사용하거나, 토큰을 마스킹했다가 복구하는 이산 확산 방식을 사용했다. 하지만 확산 방식은 노이즈를 섞는 과정이 단순 마스킹이나 균등 분포 교체로 제한되어 있어, 언어의 복잡한 의미적 구조를 충분히 반영하기 어려웠다.

GDDS는 연속 시간 마르코프 연쇄(CTMC) 이론을 바탕으로, 어떤 형태의 노이즈 과정도 수용할 수 있는 일반화된 구조를 제안한다. 특히 'Uniformization' 기법을 도입하여, 수만 개의 단어가 있는 어휘집에서 계산이 불가능에 가까웠던 복잡한 확률 전이 행렬을 Poisson 분포 기반의 점프 횟수 샘플링으로 치환함으로써 연산 효율성을 극대화했다.

또한, 학습 시 전체 노이즈 경로를 다 추적하는 대신 특정 시점의 데이터 상태인 '스냅샷'만을 보고 원래 데이터를 예측하도록 설계했다. 이는 딥러닝 모델이 노이즈 섞인 입력에서 깨끗한 입력을 복원하는 표준적인 Denoising 방식과 완벽히 일치하게 만들어, Transformer와 같은 기존 아키텍처를 그대로 활용하면서도 더 강력한 생성 능력을 갖추게 한다.

전방 노이징 과정은 $K_t = \alpha_t I_m + (1 - \alpha_t) \Pi_t$ 형태의 보간 행렬로 정의된다. [입력 토큰 $x_0$에 대해 $\alpha_t$ 확률로 원래 값을 유지하고 $1-\alpha_t$ 확률로 $\Pi_t$에 따른 새로운 토큰을 샘플링하여 → 시간 $t$에서의 노이즈 섞인 토큰 $x_t$를 생성한다]. 여기서 $\alpha_t$는 노이즈 강도를 조절하는 스케줄이며, $\Pi_t$는 확률 질량이 토큰 간에 어떻게 재분배되는지를 결정하는 믹싱 행렬이다.

대규모 어휘집에서의 효율적 계산을 위해 Uniformization 기법을 사용한다. [전이율 행렬 $Q_t$와 시간 $t$를 입력으로 받아 → Poisson 분포에서 점프 횟수 $N_t$를 추출하고 해당 횟수만큼 상태 전이를 수행하여 → 행렬 지수 함수 계산 없이 정확한 $x_t$ 샘플을 얻는다]. 이 과정은 각 토큰별로 독립적으로 병렬 수행이 가능하여 대규모 시퀀스 처리에 유리하다.

역방향 학습은 Snapshot ELBO를 최적화한다. 모델 $\mu_\theta(x_t, t)$는 노이즈 섞인 스냅샷 $x_t$로부터 원본 토큰 $x_0$의 분포를 직접 예측한다. [노이즈 토큰 $x_t$와 시간 $t$를 신경망에 입력하여 → 각 단어별 확률 분포를 출력하고 → 실제 정답 $x_0$와의 Cross-Entropy 손실을 계산하여 → 모델 가중치를 갱신한다]. 이는 기존의 복잡한 경로 기반 손실 함수보다 분산이 낮고 학습이 안정적이다.

Text8 캐릭터 단위 언어 모델링 벤치마크에서 GDDS Absorb 모델은 1.16 BPC(Bits Per Character)를 기록했다. 이는 기존 MDM(1.58 BPC)이나 AR 모델(1.35 BPC)보다 현저히 낮은 수치로, 이산 확산 모델이 동일 연산량 조건에서 AR 모델을 처음으로 앞선 사례이다.

OpenWebText(OWT) 데이터셋 실험에서 GDDS Gauss(의미론적 정보를 반영한 노이즈 과정)는 7.65 PPL(Perplexity)을 달성했다. 이는 기존 SOTA 모델인 HDLM(23.25 PPL)이나 재학습된 AR 모델(20.49 PPL) 대비 압도적인 성능 향상을 보여주며, 확산 모델의 강력한 모델링 능력을 입증했다.

제로샷 전이(Zero-shot transfer) 성능에서도 PTB, Wikitext, LM1B 등 다양한 데이터셋에서 일관되게 가장 낮은 Perplexity를 기록했다. 특히 의미론적으로 구조화된 노이징 과정이 단순 마스킹보다 더 강력한 일반화 능력을 유도함을 확인했으며, 이는 데이터의 구조적 특징을 노이징 과정에 반영하는 것이 중요함을 시사한다.

GDDS는 연속 시간 마르코프 연쇄(CTMC)의 전방 및 역방향 진화를 통합적으로 다룬다. 전방 과정은 Kolmogorov 전방 방정식을 따르며, 역방향 과정은 Bayes 정리를 통해 유도된 역방향 전이 확률로 정의된다. 핵심 차별점은 'Jump-states parametrization'으로, 기존 방식이 역방향 전이율 전체를 학습하려 했던 것과 달리 GDDS는 탈출율(exit rate)을 고정하고 점프 목적지(jump destination) 분포만을 학습함으로써 최적화 난이도를 낮추고 표준적인 Cross-Entropy 손실 함수를 사용할 수 있게 했다.

'Snapshot ELBO'는 경로 기반 ELBO(Path-wise ELBO)의 변동성 문제를 해결한다. 경로 기반 방식은 전체 궤적을 고려해야 하므로 분산이 크지만, 스냅샷 방식은 특정 시점의 데이터만으로도 충분한 학습 신호를 제공하며 이는 이론적으로 정보-교정 분해(Information-calibration decomposition)를 통해 정당화된다. 구현 측면에서는 Semantic-Informed Kernel(SIK)을 도입하여 토큰 간의 임베딩 거리에 기반한 노이징을 수행하며, 이를 위해 KNN 또는 KeOps 라이브러리를 사용해 수만 개의 토큰 사이의 유사도 계산을 효율적으로 처리한다.

SIK(Semantic-Informed Kernel)를 사용한 조상 샘플링(Ancestral sampling)은 여전히 계산 비용이 높다. 각 단계마다 전방 연산자의 근사가 필요하며, 이 과정에서 발생하는 오차가 누적되어 샘플링 단계가 많아질수록 생성 품질이 저하되는 현상이 관찰되었다.