[논문 리뷰] Path Integral Sampler: a stochastic control approach for sampling

Qinsheng Zhang, Yongxin Chen|arXiv (Cornell University)|2021. 11. 30.

Model Reduction and Neural Networks참고 문헌 55인용 수 6

한 줄 요약

Path Integral Sampler (PIS)는 신경망으로 제어되는 SDE를 통해 슈뢰딩거 브리지 문제를 해결함으로써 비정규화된 대상 밀도에서 샘플을 생성하는 새로운 확률적 제어 기반 샘플링 알고리즘입니다. 자유형 아키텍처를 허용하는 엔드 투 엔드 훈련이 가능하며, 최적화되지 않은 제어 조건 하에서도 이론적으로 워샤르슈타인 거리 보장을 제공하고, 경로 적분 기반 중요도 가중치를 사용해 편향을 보정하여 기존 방법들보다 샘플 품질과 효율성에서 뛰어난 성능을 발휘합니다.

ABSTRACT

We present Path Integral Sampler~(PIS), a novel algorithm to draw samples from unnormalized probability density functions. The PIS is built on the Schrödinger bridge problem which aims to recover the most likely evolution of a diffusion process given its initial distribution and terminal distribution. The PIS draws samples from the initial distribution and then propagates the samples through the Schrödinger bridge to reach the terminal distribution. Applying the Girsanov theorem, with a simple prior diffusion, we formulate the PIS as a stochastic optimal control problem whose running cost is the control energy and terminal cost is chosen according to the target distribution. By modeling the control as a neural network, we establish a sampling algorithm that can be trained end-to-end. We provide theoretical justification of the sampling quality of PIS in terms of Wasserstein distance when sub-optimal control is used. Moreover, the path integrals theory is used to compute importance weights of the samples to compensate for the bias induced by the sub-optimality of the controller and time-discretization. We experimentally demonstrate the advantages of PIS compared with other start-of-the-art sampling methods on a variety of tasks.

연구 동기 및 목표

• 정규화된 밀도 모델이 필요 없이 비정규화된 확률 밀도를 위한 유연하고 엔드 투 엔드 훈련 가능한 샘플링 방법을 개발하는 것.
• 스체로딩거 브리지 프레임워크를 활용하여 목표 분포에 연결된 종료 비용을 갖는 확률적 최적 제어 문제로 샘플링을 공식화하는 것.
• 기존 MCMC보다 더 적은 단계로 높은 품질의 샘플링을 가능하게 하고 이론적 성능 한계를 제공하는 것.
• 경로 적분 이론을 활용해 최적화되지 않은 제어 및 시간 이산화 오차를 보완하기 위해 중요도 가중치를 계산하는 것.
• 명시적 밀도 제약 조건이 필요한 변분 추론과 천천히 혼합되며 학습이 어려운 제안 분포를 갖는 MCMC의 한계를 극복하기 위해 암시적이고 확장 가능하며 미분 가능한 샘플링 프레임워크를 제공하는 것.

제안 방법

• 실행 비용으로 제어 에너지를, 종료 비용으로 목표 밀도에서 유도된 비용을 갖는 확률적 최적 제어 문제로 샘플링 문제를 공식화하는 것.
• 자유형 아키텍처 설계와 엔드 투 엔드 훈련을 가능하게 하기 위해 제어 정책을 신경망으로 모델링하는 것.
• 기르산وف 정리를 사용해 SDE 동역학을 재가중하고, 제어된 확산을 통해 초기 유클리드 프로세스를 목표 분포로 연결하는 것.
• 최적화되지 않은 제어 정책에 기인한 편향과 이산화 오차를 보완하기 위해 경로 적분 이론을 적용해 중요도 가중치를 계산하는 것.
• 제어된 SDE의 종료 분포와 목표 측도 사이의 KL 발산 목표 함수를 통해 컨트롤러를 훈련하는 것.
• 특히 저차원에서 안정성을 확보하기 위해 애드조인을 사용하지 않은 옐러-마루야마 SDE 통합을 사용하고, 메모리가 제한된 경우 역전형 헤운 방법을 사용하는 것.

실험 결과

연구 질문

• RQ1 스체로딩거 브리지 문제 기반의 확률적 제어 프레임워크를 사용해 비정규화된 밀도를 위한 자유형, 엔드 투 엔드 훈련 가능한 샘플링 알고리즘을 설계할 수 있는가?
• RQ2 최적화되지 않은 제어 정책은 어떻게 수정하여 渐진적으로 편향이 없는 샘플링을 보장할 수 있는가?
• RQ3 최적화되지 않은 제어기 조건을 사용할 경우 샘플링 품질에 대해 이론적으로 어떤 성능 보장을 제공할 수 있는가?
• RQ4 시간 영역 T와 시간 단계 크기 Δt의 선택은 훈련 안정성과 샘플 품질에 어떤 영향을 미치는가?
• RQ5 PIS는 MCMC와 노멀라이징 플로우와 같은 최첨단 방법보다 다양한 작업에서 샘플 품질과 효율성 면에서 뛰어나게 성능을 발휘할 수 있는가?

주요 결과

• PIS는 최적화되지 않은 제어 정책을 사용하더라도 워샤르슈타인 거리 기반으로 샘플링 품질에 대한 이론적 보장을 제공한다.
• 경로 적분 기반 중요도 가중치의 사용은 최적화되지 않은 제어 및 시간 이산화로 인한 편향을 효과적으로 보정한다.
• 더 큰 시간 영역 T는 훈련 안정성과 모드 커버리지를 향상시키지만, 더 작은 T는 더 높은 리프시츠 조건을 갖는 제어 정책을 초래하여 학습이 더 어려워진다.
• 더 작은 시간 단계 Δt는 이산화 오차를 감소시키고 샘플 품질을 향상시키며, 더 높은 통합 단계와 함께 사용할 경우 尤히 효과적이다.
• 가우시안 혼합 모델, 펌넬 분포, 그리고 MNIST 및 CIFAR-10과 같은 실제 데이터셋을 포함한 다양한 작업에서 PIS는 기준 방법들보다 샘플 품질과 효율성 면에서 뛰어난 성능을 발휘한다.
• 배치 정규화와 기울기 클리핑을 사용하면 안정성이 향상되며, 실무에서는 애드조인 기반 방법보다 애드조인 없이 옐러-마루야마 통합을 사용하는 것이 더 신뢰할 수 있다.