Likelihood Matching for Diffusion Models

기술 요약: 확산 모델을 위한 가능도 매칭 (Likelihood Matching)

문제 정의
현재 주류인 확산 모델, 특히 점수 기반 생성 모델(Score-Based Generative Models, SGMs)은 주요 학습 목표로 **점수 매칭(Score Matching, SM)**에 의존한다. SM은 경험적으로는 효과적이지만, 매개변수화된 점수 함수와 다양한 노이즈 수준에서의 실제 점수 함수 사이의 차이를 최소화하는 방식이다. 이론적으로 이 접근법은 원래 데이터 분포의 음의 로그 가능도(negative log-likelihood)에 대한 상한만을 최소화할 뿐, 가능도 자체를 직접 최소화하는 것은 아니다. 이러한 간접적인 연결은 MLE(최대 가능도 추정)와 비교했을 때 통계적 효율성 저하를 초래할 수 있는데, MLE는 일관되고 효율적인 추정량을 생성하는 것으로 알려져 있다. 더욱이, 표준 SM은 일차 모멘트 정보(점수/그래디언트)에만 집중하여, 전이 밀도(transition densities)에 포함된 고차 구조 정보를 간과할 가능성이 있다.

방법론
저자들은 타겟 데이터의 가능도와 역방향 확산 과정의 경로 가능도(path likelihood) 사이의 등가성을 활용하여, 데이터 분포의 가능도를 직접적으로 겨냥하는 가능도 매칭(Likelihood Matching, LM) 프레임워크를 제안한다.

1. 이론적 토대:
   본 방법은 역방향 확산 경로의 결합 가능도가 (순방향 과정과 관련된 상수를 제외하면) 원래 데이터 분포의 가능도와 동일하다는 성질에 근거한다. 저자들은 역방향 과정의 정확한 경로 가능도를 최대화하는 것이 데이터 로그 가능도를 최대화하는 것과 동일함을 입증하였다.

2. 준최대 가능도 추정 (Quasi-Maximum Likelihood Estimation, QMLE):
   정확한 역방향 전이 밀도는 다루기 어렵기 때문에, 저자들은 준최대 가능도(Quasi-Maximum Likelihood) 접근법을 사용하여 이를 근사한다. 각 역방향 전이 밀도를 조건부 평균과 공분산이 실제 역방향 과정과 일치하는 가우시안 분포로 모델링한다.

   • 핵심 통찰: 역방향 과정의 조건부 평균과 공분산은 점수 함수($\nabla \log q_t$)뿐만 아니라 헤시안 함수($\nabla^2 \log q_t$)에도 의존한다.
   • 매개변수화: 저자들은 점수와 헤시안을 모두 신경망(구체적으로 U-Net)을 사용하여 매개변수화한다. 계산 복잡도를 관리하기 위해, 헤시안은 종종 대각 행렬에 저계수(low-rank) 구조를 더한 형태로 근사된다.

3. 가능도 매칭 목적 함수:
   제안된 목적 함수인 **가능도 매칭(LM)**은 역방향 전이의 준-로그 가능도(quasi-log-likelihood)를 최소화함으로써 유도된다.

   • 이 목적 함수는 점수 매칭과 공분산 매칭을 통합한다.
   • 이는 매칭된 공분산 행렬 $(I + \sigma^2 H_t)^{-1/2}$를 통해 암묵적인 가능도 가중치 메커니즘을 포함하며, 국소적 곡률(local curvature)에 따라 점수 매칭 항의 가중치를 자동으로 조정한다.
   • 학습에는 순방향 SDE의 진화적 정보를 효율적으로 사용하기 위해, 확산 궤적을 따라 시간 단계를 무작위로 선택하는 확률적 샘플러(stochastic sampler)가 사용된다.

4. 확률적 샘플러 (Stochastic Sampler):
   학습된 점수와 헤시안 정보를 모두 활용하는 새로운 샘플링 절차가 도입된다. 역방향 단계는 다음과 같이 정의된다:
   $$\tilde{Y}_{t-1} = \hat{\mu}_{t-1|t}(\tilde{Y}_t) + \hat{\Sigma}_{t-1|t}^{1/2}(\tilde{Y}_t) Z_t$$
   여기서 평균 $\hat{\mu}$와 공분산 $\hat{\Sigma}$는 각각 추정된 점수와 헤시안의 함수이다.

주요 기여

• 새로운 학습 목적 함수: 본 논문은 역방향 경로 로그 가능도를 근사하는 준-가능도 기반 학습 목적 함수(LM)를 도입하여, 점수 매칭과 공분산 매칭을 암묵적 가능도 가중치와 함께 효과적으로 결합한다.
• 헤시안 강화 샘플링: 점수와 헤시안 정보를 모두 활용하는 확률적 샘플러를 도출하였으며, 이는 표준 DDPM 유형의 샘플러보다 역방향 전이 밀도를 더 정확하게 근사한다.
• 이론적 보장:
  • 수렴성: 총 변동 거리(total variation distance) 측면에서 비점근적 수렴 보장이 제공된다. 역방향 단계 오차는 $O(d^3 \log^{4.5} T / T)$로 스케일링되며, 점수 및 헤시안 추정 오차에 대한 추가 항이 포함된다.
  • 일관성: 저자들은 오라클 매개변수 설정(oracle parametric setting) 하에서 제안된 준-MLE의 일관성을 증명하였으며, 샘플 크기와 확산 단계가 증가함에 따라 추정량이 실제 데이터 파라미터로 수렴함을 보여준다.
• 효율적 구현: 저계수 헤시안 근사와 Sherman-Morrison-Woodbury 공식을 사용하여 가능도 계산 및 샘플링 단계를 효율적으로 수행한다.

실험 결과
저자들은 합성 데이터와 실제 데이터셋 모두에서 방법을 검증하였다:

• 합성 데이터: 가우시안 및 t-분포 혼합 모델에서, LM은 특히 잘 분리된 모드(modes)를 포착하는 데 있어 표준 점수 매칭(SM)보다 일관되게 우수한 성능을 보였다. 곡선형 매니폴드를 나타내는 고도로 꼬인 스위스 롤(Swiss-roll) 분포에서, LM은 헤시안 기반 공분산의 이점을 입증하며 탁월한 성능을 보였다.
• 파라미터 추정: 파라미터 가우시안 혼합 설정에서, LM은 SM에 비해 더 낮은 평균 절대 오차(MAE)와 표준 오차를 달로, 이는 이론적 일관성 주장을 뒷받침한다.
• 이미지 데이터셋 (MNIST, CIFAR-10, CelebA):
  • LM은 SM 베이스라인에 비해 더 낮은 Fréchet Inception Distance (FID)와 더 나은 음의 로그 가능도(NLL) 점수를 달성했다.
  • 성능은 헤시안 근사의 계수($r$)가 높아짐에 따라 향상되었으며, $r=20-30$ 부근에서 정점을 찍었다.
  • 속도-품질 트레이드오프: LM은 유사한 실제 실행 시간(wall-clock time) 대비 더 낮은 FID를 달성했으며, 높은 시각적 품질에 도달하기 위해 더 적은 역방향 샘플링 단계를 필요로 했다. 예를 들어, MNIST에서 LM은 단 20회의 역방향 반복만으로 인식 가능한 숫자를 생성한 반면, SM 샘플은 여전히 흐릿한 상태였다.

의의 및 주장
본 논문은 가능도 매칭이 준-가능도 프레임워크를 통해 데이터 가능도를 직접 겨냥함으로써 점수 매칭에 대한 보다 통계적으로 엄밀한 대안을 제공한다고 주장한다. 헤시안(2차 정보)을 통합함으로써, 이 방법은 더 풍부한 전이 밀도를 포착하고, 특히 낮은 단계의 레짐(low-step regime)과 복잡하고 곡선적인 데이터 매니폴드에서 샘플링 효율성을 개선한다. 저자들은 LM이 헤시안 모델링으로 인해 단계당 더 높은 계산 비용이 발생하지만, 필요한 샘플링 단계의 감소와 생성 품질(FID 및 NLL)의 향상이 고충실도 생성 모델 학습을 위한 효과를 입증한다고 강조한다. 이 연구는 MLE의 통계적 효율성과 확산 모델의 실용적 성공 사이의 간극을 메운다.