디퓨전 기반 범주형 재파라미터화와 그라디언트 추정

초록

본 논문은 범주형 변수에 대한 확률적 최적화 문제에서 연속적인 확산 모델을 이용해 부드러운 재파라미터화를 제안한다. 제안된 ReDGE는 확산 과정에서 얻은 닝스팅된 샘플을 이용해 미분 가능한 근사 샘플링 맵을 구성하고, 이를 통해 Straight‑Through와 REINMAX와 같은 기존 추정기를 특수 경우로 포함하는 새로운 그라디언트 추정기를 제공한다. 실험 결과, 다양한 잠재 변수 모델과 이산 디퓨전 모델의 보상 가이드에서 기존 방법들을 일관되게 능가한다.

상세 분석

논문은 범주형 확률분포 πθ 를 직접 미분하기 어려운 문제를 확산 모델(Denoising Diffusion Model, DDM)과 결합해 해결한다. 기존의 Gumbel‑Softmax와 같은 연속 완화는 온도 τ 를 낮출수록 편향이 감소하지만, 그 과정에서 그래디언트가 소실되거나 수치적으로 불안정해지는 단점이 있다. 저자들은 이러한 문제를 “소음 수준”을 온도와 동일시하는 확산 기반 접근법으로 재구성한다. 핵심 아이디어는 πθ 가 단순히 K‑차원 단위벡터들의 곱으로 표현된다는 점을 이용해, 확산 과정의 역전파 단계에서 각 시점 t 에서의 조건부 분포 πθ,0|t 를 정확히 계산할 수 있다는 것이다. 이는 각 카테고리 i 에 대해 Gaussian 노이즈와 결합된 one‑hot 벡터의 형태이므로, 사후 평균 denoiser ˆxθ,0 (x_t , t) 가 단순히 소프트맥스 형태의 확률 행렬로 닫힌 형태를 가진다. 따라서 별도의 신경망 학습 없이도 샘플링 맵 T_tk (·) 를 정의할 수 있다.

저자들은 작은 소음(즉, τ →0) 영역을 분석해, 이때 샘플링 맵이 거의 계단식(step‑like) 구조를 띠며, 경계 근처에서 급격한 변화를 보이는 것을 보였다. 이는 기존 연속 완화에서 온도를 낮출 때 발생하는 “경사 소실” 현상과 동일하지만, 여기서는 명시적인 확산 스케줄을 통해 제어 가능하다. 또한, 단일 확산 단계만 사용할 경우 ReDGE는 Straight‑Through estimator와 동일한 형태가 되며, 두 단계 이상을 사용하면 REINMAX와 유사한 편향‑분산 절충을 제공한다. 이와 더불어 파라미터‑의존 초기화 전략을 도입해, 초기 단계에서의 샘플링 편향을 최소화하고 전체 학습 비용을 크게 증가시키지 않는다.

실험에서는 (1) 다항식 프로그래밍 문제, (2) 변분 오토인코더(VAE)와 같은 잠재 변수 모델, (3) 이산 디퓨전 모델에서의 보상 가이드 작업을 대상으로 기존 Score‑Function, Straight‑Through, Gumbel‑Softmax, REINMAX와 비교하였다. 전반적으로 ReDGE는 동일하거나 더 낮은 분산을 보이며, 특히 높은 차원의 카테고리 L·K 조합에서 기존 방법보다 빠른 수렴과 더 좋은 최적값을 달성했다. 코드가 공개될 예정이며, 구현은 PyTorch 기반으로 제공된다.

이 논문은 확산 모델을 “훈련‑프리” 재파라미터화 도구로 활용한다는 점에서 혁신적이다. 기존 연속 완화가 반드시 복잡한 신경망 기반 변환에 의존해야 했던 반면, 여기서는 확산 과정 자체가 미분 가능한 변환을 제공한다. 따라서 범주형 변수 최적화에 있어 편향‑분산 트레이드오프를 보다 체계적으로 조절할 수 있는 새로운 패러다임을 제시한다.