[논문 리뷰] Maximum Likelihood Training of Score-Based Diffusion Models

Yang Song, Conor Durkan외 2명|arXiv (Cornell University)|2021. 01. 22.

Advanced Neuroimaging Techniques and Applications참고 문헌 54인용 수 46

한 줄 요약

본 연구는 점수 기반 확산 모델의 음의 대수 우도(negative log-likelihood)를 상한하는 확률 가중 목적함수를 도출하여 근사 최대우도 학습을 가능하게 하고, 데이터셋과 SDE 전반에 걸쳐 모델 우도를 향상시키며 CIFAR-10 및 ImageNet32x32에서 증강 없이도 경쟁적인 비트/차원(bit/dim) 성능을 달성합니다.

ABSTRACT

Score-based diffusion models synthesize samples by reversing a stochastic process that diffuses data to noise, and are trained by minimizing a weighted combination of score matching losses. The log-likelihood of score-based diffusion models can be tractably computed through a connection to continuous normalizing flows, but log-likelihood is not directly optimized by the weighted combination of score matching losses. We show that for a specific weighting scheme, the objective upper bounds the negative log-likelihood, thus enabling approximate maximum likelihood training of score-based diffusion models. We empirically observe that maximum likelihood training consistently improves the likelihood of score-based diffusion models across multiple datasets, stochastic processes, and model architectures. Our best models achieve negative log-likelihoods of 2.83 and 3.76 bits/dim on CIFAR-10 and ImageNet 32x32 without any data augmentation, on a par with state-of-the-art autoregressive models on these tasks.

연구 동기 및 목표

• 가중 스코어 매칭 손실을 음의 로그-가능도에 대한 상한으로 연결하여 점수 기반 확산 모델(SBDMs)의 더 높은 가능도 학습을 동기부여하고 가능하게 한다.
• bound를 촘촘하게 하고 모델의 가능도를 향상시키는 확률 가중치(lambda(t) = g(t)^2)를 도입한다.
• SDE/ODE 형식과 가능도 간의 이론적 보장을 제공하고, 실제 학습에 적용하기 위한 분산 감소 및 변분 기법을 제안한다.
• 데이터셋, SDE 패밀리(VE/VP/subVP) 및 모델 아키텍처 전반에서 향상된 로그 가능도를 시연하고, CIFAR-10 및 ImageNet-32x32에서 경쟁력 있는 결과를 제시한다.

제안 방법

• 정방향 SDE와 역방향-SDE를 갖는 점수 기반 확산 모델을 정식화하고, 신경망이 학습할 시간 의존 스코어를 정의한다.
• 확률 가중치 λ(t) = g(t)^2를 사용하면 가중 스코어 매칭 목적함수가 모델 p_theta^SDE에 대한 KL 발산의 상한을 형성하여 음의 로그-가능도 최적화와 연결된다.
• 스코어가 참 시간 의존 스코어와 일치하는 특정 조건에서 p_theta^SDE와 p_theta^ODE(CNF) 사이의 등가성을 보인다.
• 개별 데이터 포인트에 대한 L^SM_theta(x) 및 L^DSM_theta(x)를 통해 효율적 추정과 학습을 가능하게 하는 실용적 경계(bound)를 제공한다.
• 수치적 안정성을 위해 작은 ε > 0를 사용하고 편향 보정을 Jensen의 부등식으로 논의한다.
• 가중 확률 가중치로 인한 높은 분산을 완화하기 위한 중요도 샘플링으로 분산 감소를 도입하고, 이를 변분 디퀀타이제이션과 결합하여 가능도를 향상시킨다.

실험 결과

연구 질문

• RQ1 확률 가중치가 점수 기반 확산 모델의 음의 대수 가능도에 대한 원칙적 상한을 제공하는가?
• RQ2 확률 가중 스코어 매칭 목적함수를 최적화하면 p_theta^SDE 및 p_theta^ODE의 실제 로그 가능도가 데이터셋과 SDE 유형 전반에서 향상되는가?
• RQ3 가능도 상한이 촘촘해지거나 참 데이터 스코어 역학과 등가가 되는 조건은 무엇인가?
• RQ4 학습에서 가능가중을 사용할 때 분산을 어떻게 제어하고, 중요도 샘플링 및 변분 디퀀타이제이션이 가능도를 더 향상시킬 수 있는가?

주요 결과

• 확률 가중치(lambda(t) = g(t)^2)가 데이터에서 p_theta^SDE로의 KL 발산을 상한하도록 하여 근사 최대우도 학습을 가능하게 한다.
• 확률 가중 학습으로 여러 데이터셋, 다양한 SDE 유형(VE/VP/subVP) 및 모델 아키텍처에서 p_theta^SDE의 가능도가 일관되게 향상된다.
• 실험에서 데이터 증강 없이 CIFAR-10에서 2.83 bits/dim, ImageNet-32x32에서 3.76 bits/dim의 경쟁력 있는 로그 가능도를 보여주며, 이는 최첨단 자가회귀 모델과 비슷한 수준이다.
• SDE 기반 SBDMs와 CNF를 확률 흐름 ODE를 통해 연결하고, 정확한 스코어 매칭 하에서 CNF 가능도와 동등함을 보인다.
• 확률 가중치의 분산은 중요도 샘플링으로 효과적으로 감소되어 학습 안정성과 성능을 향상시킨다.
• 이 방법은 표준 Fréchet Inception Distance 지표에서 큰 손실 없이 더 높은 품질의 가능도 추정을 가능하게 한다.