[논문 리뷰] Efficient Neural Network Approaches for Conditional Optimal Transport with Applications in Bayesian Inference

Zheyu Oliver Wang, Ricardo Baptista외 3명|arXiv (Cornell University)|2023. 10. 25.

Gaussian Processes and Bayesian Inference인용 수 5

한 줄 요약

이 논문은 베이지안 추론에서 효율적인 조건부 최적 운반(COT)을 위한 두 가지 딥러닝 기반 방법—PCP-Map와 COT-Flow—을 제안한다. 부분적으로 볼록인 네트워크와 정규화된 플로우를 활용해 단조성과 유일성을 보장하는 모놀로닉 운반 맵을 신경망으로 파arameter화함으로써, 복잡한 조건부 사후분포로부터 빠른 샘플링과 밀도 추정이 가능해지며, 고차원 역문제에서 최신 기법들보다 수렴 속도와 수치 정확도에서 뛰어난 성능을 보인다.

ABSTRACT

We present two neural network approaches that approximate the solutions of static and dynamic $\unicode{x1D450}\unicode{x1D45C}\unicode{x1D45B}\unicode{x1D451}\unicode{x1D456}\unicode{x1D461}\unicode{x1D456}\unicode{x1D45C}\unicode{x1D45B}\unicode{x1D44E}\unicode{x1D459}\unicode{x0020}\unicode{x1D45C}\unicode{x1D45D}\unicode{x1D461}\unicode{x1D456}\unicode{x1D45A}\unicode{x1D44E}\unicode{x1D459}\unicode{x0020}\unicode{x1D461}\unicode{x1D45F}\unicode{x1D44E}\unicode{x1D45B}\unicode{x1D460}\unicode{x1D45D}\unicode{x1D45C}\unicode{x1D45F}\unicode{x1D461}$ (COT) problems. Both approaches enable conditional sampling and conditional density estimation, which are core tasks in Bayesian inference$\unicode{x2013}$particularly in the simulation-based ($\unicode{x201C}$likelihood-free$\unicode{x201D}$) setting. Our methods represent the target conditional distribution as a transformation of a tractable reference distribution. Obtaining such a transformation, chosen here to be an approximation of the COT map, is computationally challenging even in moderate dimensions. To improve scalability, our numerical algorithms use neural networks to parameterize candidate maps and further exploit the structure of the COT problem. Our static approach approximates the map as the gradient of a partially input-convex neural network. It uses a novel numerical implementation to increase computational efficiency compared to state-of-the-art alternatives. Our dynamic approach approximates the conditional optimal transport via the flow map of a regularized neural ODE; compared to the static approach, it is slower to train but offers more modeling choices and can lead to faster sampling. We demonstrate both algorithms numerically, comparing them with competing state-of-the-art approaches, using benchmark datasets and simulation-based Bayesian inverse problems.

연구 동기 및 목표

• 고차원 베이지안 추론에서 조건부 최적 운반(COT) 문제를 해결하는 데 있어 계산적으로 비가능한 문제를 해결한다.
• 빠른 샘플링과 밀도 추정을 가능하게 하는 스케일러블하고 미분 가능한 신경망 파arameterization을 통해 COT 맵을 개발한다.
• 측도 운반에서의 적대적 학습의 한계를 극복하기 위해 학습을 위해 우도 최대화를 사용한다.
• 다중 조건 변수에 걸쳐 암시적 조건부 샘플링을 가능하게 하여 각 추론 비용을 감소시킨다.
• 우도나 사전확률이 계산이 어려운 상황에서 조건부 사후분포 근사화를 위한 실용적이고 재현 가능한 프레임워크를 제공한다.

제안 방법

• PCP-Map는 부분적으로 입력-볼록 신경망(PICNN)의 그래디언트로 조건부 운반 맵을 파arameterization하여 볼록 잠재함수를 통해 단조성과 유일성을 보장한다.
• COT-Flow는 정규화된 신경미분방정식(Neural ODE)의 플로우를 통해 조건부 운반을 모델링하며, 신경망으로 파arameter화된 속도장에 의해 보존 동역학을 강제한다.
• 두 방법 모두 관측된 타겟 샘플의 우도 최대화를 통해 학습되며, 적대적 최적화를 피하고 엔드 투 엔드 백프로파게이션을 가능하게 한다.
• PCP-Map는 샘플링을 위해 볼록 최적화를 사용하고, COT-Flow는 ODE 통합을 사용하여 훈련 후 더 빠른 샘플링을 제공한다.
• 다양한 데이터셋과 문제 설정에서 강건한 구성 요소를 식별하기 위해 이중 단계 히ュ리스틱을 사용한 하이퍼파ram터 탐색 전략을 적용한다.
• 이 방법들은 스토케스틱 로트카-볼테라 및 1D 얕은 수면 방정식을 포함한 벤치마크 역문제에 적용되었으며, 고차원 조건부 변수를 포함한다.

Figure 2: Posterior samples in log scale and MAP point quality comparisons between proposed approaches and ABC with ${\bf x}^{*}=(0.01,0.5,1,0.01)^{\top}$ . Left : posterior samples generated by proposed approaches trained on 50k samples. Middle : posterior samples generated by proposed approaches t

실험 결과

연구 질문

• RQ1 신경망이 조건부 최적 운반 맵을 효율적으로 파arameterization하면서도 단조성과 유일성과 같은 바람직한 이론적 성질을 유지할 수 있는가?
• RQ2 베이지안 추론에서 우도 기반 학습이 적대적 학습과 비교해 수렴 속도와 샘플 품질 측면에서 어떻게 성능을 내는가?
• RQ3 정적(PCP-Map)과 동적(COT-Flow) 접근 방식 간에 훈련 시간과 샘플링 속도 사이의 상호 보완적 관계는 어떠한가?
• RQ4 伝통적인 격자 기반 또는 다항식 근사가 실패하는 고차원 문제에 대해 이러한 방법들이 얼마나 스케일러블한가?
• RQ5 하이퍼파라미터 설정에 대해 이러한 방법들이 얼마나 강건한가? 그리고 단순한 히ュ리스틱이 다양한 응용 분야에서 구성 선택을 향상시킬 수 있는가?

주요 결과

• PCP-Map는 14개의 유효 파라미터와 3,500개의 측정값을 가진 고차원 역문제에서 약 715초 만에 수렴하였다.
• PCP-Map는 정적 형태임에도 불구하고 암시적 CP-Flow 방법보다 더 빠른 수렴을 보였다.
• COT-Flow는 ODE 통합 덕분에 PCP-Map보다 더 빠른 샘플링 속도를 제공했으며, 특히 잘 훈련된 속도장이 학습된 후에 두드러졌다.
• 1D 얕은 수면 방정식 문제에서 SMC-ABC보다 두 방법 모두 더 높은 수치 정확도를 보였다. 이는 ABC가 강력한 이론적 보장을 지녔음에도 불구하고.
• PCP-Map는 특성 너비, 조건 너비, 깊이 등 다양한 하이퍼파라미터에 대해 일관된 성능를 보였으며, 핵심 하이퍼파라미터는 총 3개 뿐이었다.
• 우도나 사전확률 평가가 불가능한 경우에도 이 방법들은 효과적이었으며, SDE나 레거시 코드를 포함한 실제 역문제에서의 유용성을 입증하였다.

Figure 3: Posterior samples log scale and MAP point quality comparisons between proposed approaches and ABC with ${\bf x}^{*}=(0.02,0.02,0.02,0.02)^{\top}$ . Left : posterior samples generated by proposed approaches trained on 50k samples. Middle : posterior samples generated by proposed approaches