[논문 리뷰] An optimal control perspective on diffusion-based generative modeling

Julius Berner, Lorenz Richter외 1명|arXiv (Cornell University)|2022. 11. 02.

Bayesian Methods and Mixture Models인용 수 6

한 줄 요약

이 논문은 확률적 최적 제어와 확산 기반 생성 모델 간의 새로운 연결 고리를 설정한다. 확산 모델의 역시간 SDE가 하밀턴–자코비–벨리만(HJB) 방정식을 만족함을 보여주며, 제어 이론의 검증 정리로부터 증거 하한값(ELBO)을 유도하고, 생성 모델링을 경로 공간 상의 KL 발산 최소화 문제로 재해석한다. 또한 비정규화된 밀도 샘플링 작업에서 기존 방법들을 능가하는 새로운 시간 역전 확산 샘플러(DIS)를 제안한다.

ABSTRACT

We establish a connection between stochastic optimal control and generative models based on stochastic differential equations (SDEs), such as recently developed diffusion probabilistic models. In particular, we derive a Hamilton-Jacobi-Bellman equation that governs the evolution of the log-densities of the underlying SDE marginals. This perspective allows to transfer methods from optimal control theory to generative modeling. First, we show that the evidence lower bound is a direct consequence of the well-known verification theorem from control theory. Further, we can formulate diffusion-based generative modeling as a minimization of the Kullback-Leibler divergence between suitable measures in path space. Finally, we develop a novel diffusion-based method for sampling from unnormalized densities -- a problem frequently occurring in statistics and computational sciences. We demonstrate that our time-reversed diffusion sampler (DIS) can outperform other diffusion-based sampling approaches on multiple numerical examples.

연구 동기 및 목표

• 확률적 최적 제어와 확산 기반 생성 모델 간의 공식적 연결 고리를 수립하는 것.
• 제어 이론의 검증 정리로부터 확산 모델의 ELBO가 자연스럽게 유도됨을 보여주는 것.
• 경로 공간 상의 Kullback–Leibler 발산 최소화를 통해 확산 모델링을 재해석하는 것.
• 시간 역전 SDE를 활용해 비정규화된 타깃 밀도에서 샘플링하는 새로운 샘플링 방법을 개발하는 것.
• 제안된 방법(DIS)이 기존의 확산 기반 샘플러들보다 실용적으로 뛰어나다는 것을 보여주는 것.

제안 방법

• 역시간 SDE의 로그 밀도를 지배하는 하밀턴–자코비–벨리만(HJB) 방정식을 유도하여 최적 제어와 연결한다.
• 제어 이론의 검증 정리를 활용해 확산 모델의 ELBO를 공식적으로 유도한다.
• 연속 시간 확률 과정 경로에 대한 조건부 측도 간 KL 발산 최소화 문제로 확산 모델링을 재구성한다.
• 비정규화된 타깃 밀도에서 샘플링하기 위해 제어 과정을 학습하는 새로운 시간 역전 확산 샘플러(DIS)를 제안한다.
• 신경망을 사용해 역시간 SDE의 드리프트를 매개변수화하고, 스코어 함수 및 경계 조건에 기반한 초기화를 적용한다.
• 유연한 초기 분포와 기준 SDE를 적용하여 사전 지식을 샘플링 과정에 통합할 수 있도록 한다.

실험 결과

연구 질문

• RQ1 확산 모델과 확률적 최적 제어 간의 연결 고리가 HJB 방정식을 통해 공식적으로 수립될 수 있는가?
• RQ2 제어 이론의 검증 정리가 확산 기반 생성 모델의 ELBO를 자연스럽게 도출하는가?
• RQ3 경로 공간 상의 KL 발산 최소화가 확산 모델링에 새로운 해석을 제공할 수 있는가?
• RQ4 새로운 시간 역전 확산 샘플러(DIS)가 비정규화된 밀도에서의 샘플링 성능에서 기존 방법들을 능가할 수 있는가?
• RQ5 다중모달성 또는 고에너지 장벽을 가진 분포에서 DIS의 성능은 ULA와 같은 랭지비안 기반 방법과 비교해 어떻게 되는가?

주요 결과

• 확산 과정의 역시간 로그 밀도는 하밀턴–자코비–벨리만(HJB) 방정식을 만족하며, 이는 최적 제어와의 공식적 연결 고리를 확립한다.
• 확산 모델의 증거 하한값(ELBO)은 제어 이론의 검증 정리로부터 직접 유도되는 결과이다.
• 확산 기반 생성 모델링은 경로 공간 상의 측도 간 KL 발산 최소화 문제로 재해석될 수 있다.
• 제안된 시간 역전 확산 샘플러(DIS)는 다중모달성 및 고에너지 장벽을 가진 분포를 포함한 여러 수치 예제에서 다른 확산 기반 샘플링 방법보다 뛰어난 성능을 보였다.
• DIS는 유한 시간 내에 정확한 샘플링을 달성하는 반면, ULA 및 유사한 랭지비안 방법은 크렘러의 법칙에 따라 에너지 장벽으로 인해 지수적 혼합 시간을 겪는다.
• DIS에서 학습된 드리프트는 더 큰 시간 간격에서 최적 드리프트에 더 가까워지며, 스코어 함수를 통합한 매개변수화 방식은 더 안정적이고 정확한 결과를 얻는다.