[논문 리뷰] Learning Interpolations between Boltzmann Densities

Bálint Máté, François Fleuret|arXiv (Cornell University)|2023. 01. 18.

Gaussian Processes and Bayesian Inference인용 수 5

한 줄 요약

이 논문은 연속 정규화 흐름에 대한 새로운 학습 목표를 제안하며, 역 KL 발산 대신 시간에 따라 변화하는 벡터장과 에너지 함수의 보간을 통해 보츠만 밀도 간의 보간을 학습함으로써 이를 대체한다. 이 방법은 연속성 방정식 위반을 최소화하여 다모드 목표, 예를 들어 가우시안 혼합 및 양자역학적 双함수 터널 잠금 위치에서의 모드 붕괴를 감소시키고 학습 효율성을 향상시킨다.

ABSTRACT

We introduce a training objective for continuous normalizing flows that can be used in the absence of samples but in the presence of an energy function. Our method relies on either a prescribed or a learnt interpolation $f_t$ of energy functions between the target energy $f_1$ and the energy function of a generalized Gaussian $f_0(x) = ||x/σ||_p^p$. The interpolation of energy functions induces an interpolation of Boltzmann densities $p_t \propto e^{-f_t}$ and we aim to find a time-dependent vector field $V_t$ that transports samples along the family $p_t$ of densities. The condition of transporting samples along the family $p_t$ is equivalent to satisfying the continuity equation with $V_t$ and $p_t = Z_t^{-1}e^{-f_t}$. Consequently, we optimize $V_t$ and $f_t$ to satisfy this partial differential equation. We experimentally compare the proposed training objective to the reverse KL-divergence on Gaussian mixtures and on the Boltzmann density of a quantum mechanical particle in a double-well potential.

연구 동기 및 목표

• 다모드 목표 밀도에서 기대값을 추정할 때 역 KL 발산으로 학습하는 정규화 흐름에서 발생하는 모드 붕괴 문제를 해결한다.
• 역 KL 발산의 한계를 극복하기 위해 연속된 에너지 함수 보간을 따라 운반되는 물리학에 영감을 받은 학습 목표를 도입한다.
• 진짜 샘플에 접근할 수 없더라도 에너지 함수 평가와 그 기울기만으로도 학습이 가능하도록 한다.
• 학습된 에너지 보간을 따라 시간에 따라 변화하는 벡터장을 최적화하여 연속성 방정식을 만족함으로써 샘플 품질과 학습 속도를 동시에 향상시킨다.
• 다양한 도전적인 목표, 예를 들어 다모드 가우시안과 양자역학적 보츠만 밀도에서의 방법의 효과성을 입증한다.

제안 방법

• 기본 에너지 함수 $ f_0 $ 는 일반화된 가우시안이고, 목표 에너지 함수 $ f_1 $ 이며, 시간에 따라 변화하는 에너지 함수 $ f_t = (1-t)f_0 + t f_1 + t(1-t)f^\theta(t) $ 를 정의한다.
• 보츠만 밀도의 가족 $ p_t \propto e^{-f_t} $ 를 유도하고, 이 가족을 따라 샘플을 운반하기 위해 시간에 따라 변화하는 벡터장 $ V_t $ 를 학습한다.
• 적절한 운반 조건을 연속성 방정식 $ \partial_t p_t + \nabla \cdot (p_t V_t) = 0 $ 으로 기술하며, 이는 $ V_t $ 와 $ p_t $ 를 연결한다.
• 연속성 방정식의 국소적 위반 정도를 최소화함으로써 $ V_t $ 와 $ f_t $ 를 최적화하며, 손실 함수 $ \mathcal{L} = \int \mathcal{E} + \mathcal{E}^2 \, dt $ 를 사용한다. 여기서 $ \mathcal{E} $ 는 잔차 오차이다.
• 시간에 따라 공통된 아키텍처를 사용하는 신경망(MLP)으로 $ f_t $ 와 $ V_t $ 를 매개변수화하고, 자동 미분을 통해 기울기를 계산한다.
• Adam 옵timizer를 사용하고 코즈인 학습률 감쇠를 적용하며, 4차 룬게-쿠타 방법을 사용해 50단계로 흐름을 통합한다.

실험 결과

연구 질문

• RQ1 에너지 함수 평가 외에 진짜 샘플에 접근할 수 없는 조건에서, 연속성 방정식 기반의 학습 목표가 정규화 흐름에서 모드 붕괴를 줄일 수 있는가?
• RQ2 역 KL 발산 최소화와 비교했을 때, 제안된 방법의 성능은 KL 발산, 유효 샘플 크기, 모드 커버리지 측면에서 어떻게 다른가?
• RQ3 기본과 목표 에너지 함수 사이의 비선형 보간 $ f_t $ 를 학습하는 것이 선형 보간보다 성능을 향상시키는가?
• RQ4 고차원 다모드 목표에서 학습 속도와 샘플 품질 측면에서 이 방법의 확장성은 어느 정도인가?
• RQ5 이 방법은 양자 입자가 이중함수 터널 잠금 위치에 있을 때의 보츠만 밀도와 같은 복잡한 물리계에 일반화될 수 있는가?

주요 결과

• 제안된 방법은 가우시안 혼합 및 양자 이중함수 터널 잠금 목표 모두에서 역 KL 학습보다 유의미하게 낮은 KL 발산을 달성한다.
• 가우시안 혼합 벤치마크에서 역 KL 대비 유효 샘플 크기를 최대 2.5배 향상시킨다.
• 모드 커버리지가 크게 향상되어 2차원 가우시안 혼합의 네 개의 모드를 모두 학습한 흐름이 붕괴 없이 학습된다.
• 학습 수렴 속도가 더 빠르고 안정적이며, 10,000 반복만으로도 역 KL 모델이 250,000 반복 학습한 결과보다 더 좋은 성능을 낸다.
• 학습된 비선형 보간 $ f_t $ 는 선형 보간보다 성능이 뛰어나며, 특히 에너지 장의 고곡률 영역에서 두각을 나타낸다.
• 이 방법은 PINN 프레임워크에 자연스럽게 통합되며, 연속성 오차가 잔차 손실로 기능함으로써 PDE 제약 학습에 대한 더 넓은 적용 가능성을 시사한다.