단일 토큰을 넘어: 이산 모멘트 매칭(Discrete MMD)을 통한 이산 확산 모델의 증류

기존의 이산 확산 모델은 고품질 결과를 얻기 위해 수백 번의 반복 계산이 필요하여 비용이 매우 높았습니다. 이 논문은 연속형 확산 모델에서 성공적이었던 모멘트 매칭 기법을 이산 영역으로 확장하여, 아주 적은 단계만으로도 원래 모델보다 더 나은 품질의 이미지와 텍스트를 생성하는 효율적인 D-MMD 알고리즘을 제시하여 생성 AI의 실용성을 크게 높였습니다.

확산 모델(Diffusion Model)은 데이터에 노이즈를 점진적으로 추가하고 이를 다시 복원하는 과정을 학습한다. 하지만 텍스트 토큰과 같은 이산 데이터를 다루는 이산 확산 모델은 각 단계에서 토큰을 독립적으로 예측한다고 가정하기 때문에, 생성 단계를 줄이면 토큰 간의 복잡한 상관관계를 놓쳐 품질이 급격히 저하되는 '모드 붕괴' 현상이 발생한다.

D-MMD는 이 문제를 해결하기 위해 '모멘트 매칭(Moment Matching)' 개념을 도입한다. 이는 단순히 정답 토큰 하나를 맞추는 것이 아니라, 학생 모델이 생성한 데이터의 확률 분포 통계(모멘트)가 교사 모델의 분포와 일치하도록 강제하는 방식이다. 특히 이산 영역에서는 불연속적인 토큰 대신 소프트 확률 벡터를 사용하여 미분 가능한 형태로 최적화를 수행함으로써 학습의 안정성을 확보한다.

결과적으로 학생 모델은 교사 모델이 수백 단계에 걸쳐 수행하던 복잡한 상관관계 계산을 단 몇 단계의 연산 안에 압축하여 학습하게 된다. 이를 통해 연산량(FLOPs)은 획기적으로 줄이면서도, 교사 모델이 가진 분포의 특성을 더 정확하게 재현하거나 오히려 노이즈가 정제된 더 고품질의 결과를 내놓는 것이 가능해진다.

D-MMD의 핵심은 교사 모델 $\hat{x}_\theta$와 학생 모델 $\hat{x}_\eta$ 사이의 조건부 기대치를 일치시키는 것이다. 이를 위해 보조 모델(Auxiliary model) $\hat{x}_\phi$를 도입하여 학생 모델의 기댓값을 추정하고, 학생 모델은 이 보조 모델과 교사 모델 사이의 차이를 최소화하도록 적대적으로 학습한다.

이산 데이터 처리를 위해 소프트 확률 벡터를 직접 매칭하는 손실 함수를 설계했다. [입력값: 현재 노이즈 상태 $z_t$ → 연산: 학생, 교사, 보조 모델의 로그 확률 차이를 교차 엔트로피(CE)로 계산 → 출력: $L_{GEN}$ 및 $L_{AUX}$ 손실값 → 의미: 학생 모델이 교사의 분포를 따르도록 가중치를 갱신]. 특히 Masked Diffusion에서는 마스킹되지 않은 토큰 정보를 활용해 보조 모델의 편향을 수정하는 메커니즘을 포함한다.

텍스트 생성 시 주로 사용되는 Top-p 샘플링을 증류 과정에 반영하기 위해 동적 로짓 조정 방식을 사용했다. [입력값: 교사 모델의 로짓 → 연산: 상위 p%에 속하지 않는 카테고리의 로짓을 상수 $\Delta$만큼 감산 → 출력: 조정된 확률 분포 → 의미: 그래디언트 폭주를 방지하면서 교사의 샘플링 특성을 학생에게 전이]. 이를 통해 추론 시의 기법들이 학습 단계에서 자연스럽게 내재화된다.

CIFAR-10 이미지 생성 실험에서 Uniform D-MMD는 단 32단계 만에 FID 3.7을 기록했다. 이는 1024단계의 교사 모델이 기록한 FID 7.5보다 두 배 이상 우수한 수치이며, 연산량은 약 32배 절감한 결과이다. Masked D-MMD 역시 64단계에서 FID 3.5를 달성하며 교사 모델의 성능을 뛰어넘었다.

텍스트 생성(Open Web Text)에서는 새로 제안된 GPT-2 Gradient Moment 지표를 기준으로 평가했다. Masked D-MMD 16단계 모델은 교사 모델의 256단계 성능과 대등하거나 더 나은 품질을 보여주었다. 특히 기존 증류 기법인 SDTT가 반복적인 증류 과정에서 품질이 저하되는 것과 달리, D-MMD는 안정적인 성능 향상을 유지하며 더 낮은 Perplexity와 높은 샘플 엔트로피를 기록했다.

D-MMD는 연속형 MMD의 미분 가능한 수식을 이산 확률 공간으로 확장한 min-max 최적화 프레임워크이다. 학생 모델은 교사 모델 하에서의 손실을 최소화하고 보조 모델 하에서의 손실을 최대화하는 방향으로 학습하며, 보조 모델은 학생의 기댓값을 추적하도록 훈련되어 고정점(Fixed point)에서 학생과 교사의 분포가 일치하게 된다.

이산 확산의 두 가지 주요 유형인 Masked Diffusion과 Uniform Diffusion을 모두 지원하도록 설계되었다. Masked Diffusion의 경우 마스크 토큰이 정보를 제공하지 않는 특성을 활용해 보조 모델 학습 시 소프트 타겟과 하드 타겟을 선택적으로 사용할 수 있는 이론적 근거를 제시하며, Uniform Diffusion에서는 하드 샘플 학습의 필요성을 강조한다.

새로운 평가 지표인 GPT-2 GM은 참조 모델(GPT-2)의 로그 확률 그래디언트 노름을 측정한다. [입력: 생성된 샘플 $x$ → 연산: $abla_\theta \log p(x)$ 계산 및 실제 데이터 분포에서의 그래디언트 평균과의 차이 제곱 → 출력: GM 점수 → 의미: 샘플이 실제 데이터 분포의 통계적 특성과 얼마나 일치하는지 판단]. 이는 Perplexity가 문법적으로는 맞지만 의미 없는 반복 문구에 높은 점수를 주는 취약점을 보완하여 실제 생성 품질을 더 정확히 반영한다.

적대적 최적화(Adversarial optimization) 과정의 특성상 학습 역학이 복잡하여 특정 하이퍼파라미터 설정에 민감하게 반응할 수 있다. 또한 샘플링 단계가 매우 많아질 경우 학생 모델의 성능이 교사 모델로 수렴하게 되어, 증류를 통한 성능 이득이 점차 줄어드는 한계가 존재한다.