기술 요약: 확산 이중성, 제 II 장: Ψ-샘플러

문제 제기

이산 확산 모델은 언어 및 이미지 생성을 위한 자기회귀 모델의 강력한 대안으로 부상하여, 병렬 디코딩과 양방향 컨텍스트를 제공합니다. 이 중 균일 상태 확산 모델 (USDMs) 은 토큰을 여러 번 수정하며 스스로 오류를 수정하는 능력 덕분에 소수 단계 생성 및 안내 작업에서 탁월한 성능을 발휘합니다. 그러나 USDM 은 두 가지 중요한 한계에 직면해 있습니다:

1. 샘플링 품질 정체: USDM 을 위한 표준 조상 샘플러는 소수 단계 영역에서의 성능과 달리, 샘플링 단계 수 (NFEs) 가 증가함에 따라 생성 품질이 향상되지 않습니다.
2. 훈련 비효율성: 커리큘럼 학습 (예: Duo) 을 통해 USDM 과 마스킹 확산 모델 (MDM) 간의 가능도 격차를 좁히려는 이전 시도들은 계산 비용이 많이 듭니다. 훈련 단계에서 고차원 잠재 벡터를 구체화해야 하므로 높은 메모리 사용량과 긴 훈련 시간이 필요합니다.

반면 마스킹 확산 모델 (MDM) 은 일반적으로 더 높은 가능도를 달성하지만, 한 번 마스킹이 해제된 토큰은 다시 마스킹될 수 없어 추론 중 오류가 누적되는 문제가 있습니다. 예측자 - 수정자 (PC) 방법 (예: ReMDM) 에 대한 최근 연구는 재마스킹을 허용함으로써 MDM 을 개선했지만, USDM 에 대한 유사한 진전은 아직 충분히 탐구되지 않았으며, 기존 USDM 용 PC 방법은 종종 조상 샘플러보다 낮은 성능을 보입니다.

방법론

본 논문은 $\Psi$-샘플러를 통한 추론과 확장 가능한 커리큘럼을 통한 훈련이라는 두 가지 주요 혁신을 통해 이러한 과제를 해결하는 Duo++ 프레임워크를 소개합니다.

1. $\Psi$-사후분포 및 $\Psi$-샘플러

저자들은 $\Psi$-사후분포라고 명명된 비마코프 사후분포 계열을 제안하여, 기존 PC 방법을 임의의 잡음 사전 (마스킹 및 균일 모두) 으로 일반화합니다.

• 정의: $\Psi$-사후분포는 표준 역방향 사후분포 ($q_{s|t}$) 와 전방향 과정 주변분포 ($q_s$) 의 선형 중첩으로 정의됩니다:
  $$\Psi_{s|t}(\cdot|x^\ell, z^\ell_t) = \kappa_t q_{s|t}(\cdot|z^\ell_t, x^\ell) + (1-\kappa_t) q_s(\cdot|x^\ell)$$
  여기서 $\kappa_t \in [0, 1]$은 시간에 의존하는 계수입니다.
• 메커니즘: 이 공식은 생성 과정에서 잡음 ( $q_s$ 항을 통해) 을 주입할 수 있게 합니다.
  • MDM의 경우, 이는 재마스킹을 가능하게 합니다. 디노이저가 오류를 범한 경우 이전에 디코딩된 토큰을 마스킹 상태로 되돌려 누적 오류를 수정할 수 있습니다.
  • USDM의 경우, 디노이저가 토큰에 거의 0 에 가까운 확률을 할당하더라도 모든 토큰이 샘플링될 0 이 아닌 확률을 유지하도록 하여, 샘플러가 지역 최적점에 갇히는 것을 방지합니다.
• 이론적 보장: $\Psi$-사후분포는 표준 마코프 이산 확산 과정과 정확히 동일한 주변분포를 공유하도록 구성됩니다. 이는 충분한 샘플이 주어지면 과정이 올바른 데이터 분포로 수렴함을 보장합니다.
• 일반화: 이 프레임워크는 특정 $\kappa_t$ 일정을 선택한 경우의 특수 사례로서 기존 PC 방법 (예: ReMDM, Campbell 등, Gat 등) 을 포괄합니다.

2. 더 빠른 훈련을 위한 확장 가능한 커리큘럼

Duo(Sahoo 등, 2025a) 에서 제안된 커리큘럼 학습 전략의 계산 병목 현상을 해결하기 위해, 저자들은 메모리 효율적인 구현을 소개합니다.

• 희소성 활용: 기존 커리큘럼은 이산 토큰 입력을 가우시안 잠재의 저온 소프트맥스 완화로 대체합니다. 대규모 어휘 ($K > 100,000$) 의 경우, 전체 $K$ 차원 가중치 벡터를 구체화하는 것은 불가능합니다.
• Top-$k$ 근사: 저자들은 저온에서 소프트맥스 분포가 희소하여 질량이 소수의 항목에 집중된다는 것을 관찰합니다. 이에 따라 전체 선형 결합을 Top-$k$ 항목 ($k \ll K$) 만 사용하여 근사할 것을 제안합니다.
• 효율적 샘플링: 전체 벡터를 생성하고 정렬하는 대신, 이 방법은 순서 통계를 사용하여 전체 벡터를 구체화하지 않고도 기본 가우시안 분포에서 직접 Top-$k$ 값과 해당 인덱스를 샘플링합니다.
• 정규화: 샘플링된 Top-$k$ 항과 샘플링되지 않은 항의 기대 기여도를 고려하여 소프트맥스 정규화 인자 ($\tilde{Z}$) 를 근사하는 폐형식 식을 유도합니다.
• 효율성: 이 접근법은 동일한 가능도 성능을 유지하면서 피크 메모리 사용량을 33% 줄이고 훈련 속도를 25% 가속화합니다.

주요 기여

1. $\Psi$-사후분포: 임의의 잡음 사전 (마스킹 및 균일) 에 대한 예측자 - 수정자 샘플링을 가능하게 하면서 주변분포 일관성을 유지하는 비마코프 중첩 사후분포를 도입한 이산 확산을 위한 통합 프레임워크.
2. $\Psi$-샘플러: 기존 PC 방법 (ReMDM 포함) 을 일반화하는 샘플러 계열. 실험적으로 USDM 이 NFE 에 따라 더 잘 확장되도록 하여 텍스트 생성에서 MDM 과의 성능 격차를 해소하고 이미지 생성에서는 이를 능가합니다.
3. 효율적 커리큘럼: 저온 소프트맥스 완화의 희소성을 활용하여 모델 품질을 희생하지 않으면서 메모리 및 계산 비용을 크게 절감하는 확장 가능한 훈련 전략.

실험 결과

언어 모델링 (OpenWebText)

• 성능: $\Psi$-샘플러는 조상 샘플링에 비해 USDM 의 생성 퍼플렉시티 (Gen. PPL) 를 크게 향상시킵니다.
• 확장: NFE 가 증가함에 따라 정체되는 조상 샘플링과 달리, $\Psi$-샘플러는 더 많은 단계로 계속 개선됩니다. 일치된 유니그램 엔트로피 (다양성) 에서 Duo++ 는 $\Psi$-샘플러를 사용하여 ReMDM 을 사용하는 MDLM 보다 낮은 Gen. PPL 을 달성합니다.
• 하류 작업: Duo++ 는 25% 적은 FLOPs 를 사용하면서도 다중 선택형 질문 답변 작업 (예: ARC, HellaSwag) 에서 원래 Duo 모델의 성능과 일치합니다.

이미지 모델링 (CIFAR-10)

• 품질: CIFAR-10 에서 Duo++ 는 $\Psi$-샘플러를 사용하여 조상 샘플링과 ReMDM 모두보다 우수한 FID 및 Inception Score (IS) 를 달성합니다.
• 확장: 이미지 품질의 개선은 다양한 샘플링 단계 수에 걸쳐 일관되며, Duo++ 는 MDLM 변형들을 능가합니다.

훈련 효율성

• 메모리 및 속도: 제안된 효율적 커리큘럼은 원래 Duo 커리큘럼에 비해 피크 메모리 사용량을 33% 줄이고 총 훈련 시간을 25% 단축합니다.
• 가능도: 효율적 커리큘럼으로 훈련된 모델은 LM1B 및 OpenWebText 에서 원래 Duo 모델과 비교 가능한 퍼플렉시티를 달성합니다.

중요성 및 주장

본 논문은 마스킹 확산이 확산 기반 언어 모델링의 불가피한 미래라는 가정이 의문시된다고 주장합니다. $\Psi$-샘플러를 도입함으로써 저자들은 균일 상태 확산 모델이 다음을 달성할 수 있음을 보여줍니다:

1. 우월한 확장성: 조상 샘플러와 달리 증가된 샘플링 단계에 따라 성능이 계속 향상됩니다.
2. 경쟁력 있거나 우월한 품질: (PC 샘플링과 결합 시) 텍스트 생성에서 MDM 과 일치하거나 능가하며, 이미지 생성에서는 이를 능가합니다.
3. 실용적 효율성: 이전 커리큘럼 학습 접근법보다 메모리 및 시간 효율성이 훨씬 높은 훈련 파이프라인.

이 연구는 가우시안과 이산 확산 사이의 "이중성"과 비마코프 샘플링 전략의 결합이 이산 생성 모델링을 위한 견고하고 확장 가능한 길을 제시하며, 마스킹 접근법의 지배력을 도전한다고 시사합니다.