Uniform Diffusion Models 재고: Leave-One-Out Denoiser 및 Absorbing State 재정식화

이산 확산 모델 (Discrete diffusion models)은 종종 깨끗한 데이터 예측 (clean-data prediction)을 통해 학습되지만, 이 예측은 역동학 (reverse dynamics)을 정의하는 데 다양한 방식으로 사용될 수 있습니다. Masked Diffusion Models (MDM)에서는 이러한 선택들이 대체로 일치하는 반면, Uniform Diffusion Models (UDM)에서는 일치하지 않습니다. 본 연구에서는 UDM을 위한 표준 플러그인 브리지 파라미터화 (plug-in bridge parameterization)가 디노이징 사후 확률 (denoising posterior)에 의해 최적화되는 것이 아니라, 자신의 노이즈가 섞인 관측치를 사용하지 않고 각 깨끗한 토큰을 예측하는 Leave-One-Out 사후 확률 (leave-one-out posterior)에 의해 최적화됨을 보여줍니다. 이는 플러그인 ELBO와 일반적인 교차 엔트로피 (cross-entropy) 디노이징 목적 함수 사이의 불일치를 식별합니다. 우리는 Leave-One-Out 타겟을 특성화하고 디노이저 (denoiser), Leave-One-Out 사후 확률, 그리고 스코어 (score) 사이의 정확한 변환을 도출합니다. 이러한 변환을 통해 파라미터화와 학습 목적 함수를 분리할 수 있습니다. 또한 우리의 결과는 정보 기반의 예측기-교정기 샘플러 (predictor-corrector sampler)와 Leave-One-Out 예측기에 기반한 개선된 온도 샘플링 (temperature sampling)을 통해 추가 학습 없이도 추론 성능을 향상시킵니다. 나아가, 우리는 UDM의 결합 법칙 (joint law)을 유지하면서도 이를 Masked Diffusion과 유사한 샘플링 연산으로 분해하는 Uniform Diffusion의 Absorbing State 재정식화 (absorbing-state reformulation)를 도입하며, 이는 더 단순한 디노이징 사후 확률, carry-over unmasking, 그리고 자연스러운 remasking 메커니즘을 제공합니다. 언어 모델링 (language modeling)에서 Leave-One-Out 파라미터화는 UDM 생성을 일관되게 향상시키는 반면, Absorbing 구조는 Masked Diffusion과 대등하거나 이를 능가합니다. 이러한 결과는 Masked Diffusion과 Uniform Diffusion 사이의 경험적 격차가 주변 확률 분포 (marginals) 자체의 선택보다는 파라미터화와 샘플링 설계에 의해 더 크게 좌우됨을 시사합니다. 코드와 모델은 https://github.com/samsongourevitch/rev_udm 에서 확인할 수 있습니다.