[논문 리뷰] Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks

David Alvarez-Melis, Yair Schiff외 1명|Digital Access to Scholarship at Harvard (DASH) (Harvard University)|2021. 06. 01.

Topological and Geometric Data Analysis참고 문헌 49인용 수 9

한 줄 요약

이 논문은 입력-볼록 신경망(ICNNs)을 사용하여 JKO 스킴 내 볼록 함수를 매개변수화함으로써 확률 측도 공간 위의 함수를 최적화하는 확장 가능한 방법인 JKO-ICNN을 제안한다. 이는 수렴 보장과 함께 안정적이고 고차원 최적화를 가능하게 하며, 해석적 해가 있는 저차원 PDE와 고차원 분자 생성에 대해 검증되었으며, 약물 유사도(QED)가 향상됨을 보여준다.

ABSTRACT

Gradient flows are a powerful tool for optimizing functionals in general metric spaces, including the space of probabilities endowed with the Wasserstein metric. A typical approach to solving this optimization problem relies on its connection to the dynamic formulation of optimal transport and the celebrated Jordan-Kinderlehrer-Otto (JKO) scheme. However, this formulation involves optimization over convex functions, which is challenging, especially in high dimensions. In this work, we propose an approach that relies on the recently introduced input-convex neural networks (ICNN) to parametrize the space of convex functions in order to approximate the JKO scheme, as well as in designing functionals over measures that enjoy convergence guarantees. We derive a computationally efficient implementation of this JKO-ICNN framework and experimentally demonstrate its feasibility and validity in approximating solutions of low-dimensional partial differential equations with known solutions. We also demonstrate its viability in high-dimensional applications through an experiment in controlled generation for molecular discovery.

연구 동기 및 목표

• 고차원에서 확률 측도 위의 함수 최적화의 계산적 과제를 해결하기 위해.
• 특히 고차원 공간에서 JKO 스킴 내 볼록 함수 최적화의 어려움을 극복하기 위해.
• 딥러닝을 활용한 워샤르 스페이스 위의 경량 유동에 대한 확장 가능하고 미분 가능한 프레임워크를 개발하기 위해.
• 이론적 수렴 보장을 갖는 복잡한 확률 함수 최적화(예: 분자 생성에서의 함수 최적화)를 실용적으로 가능하게 하기 위해.
• 기본 해가 있는 저차원 PDE와 실제 고차원 응용 분야 모두에서 방법의 실현 가능성을 입증하기 위해.

제안 방법

• 입력-볼록 신경망(ICNNs)을 사용하여 JKO 스킴 내 볼록 잠재 함수를 매개변수화함으로써 입력에 대한 볼록성을 보장한다.
• 브레니에의 정리에 따라 JKO 변분 문제를 볼록 함수에 대한 최적화로 재구성함으로써, 미분 가능한 최적화를 가능하게 한다.
• 유한한 샘플과 자동 미분를 활용한 적응형 경량 하강법을 사용하여 JKO-ICNN을 구현하고 엔드 투 엔드 학습을 수행한다.
• 잠재 함수 기반의 볼록 대체 모델(Residual ICNN)을 사용하여 분자 임베딩에서 QED 점수를 예측함으로써 잠재 함수 기반의 볼록성을 보장한다.
• 학습 가능한 가중치를 갖는 분포 거리 측도(Sinkhorn 또는 MMD)를 통합하여 유동이 목표 분포로 향하도록 이끈다.
• 재훈련 없이도 효율적인 추론이 가능하도록, 미리 학습된 운반 지도를 통해 계산 비용을 약간으로 줄인다.

실험 결과

연구 질문

• RQ1 ICNNs는 고차원 확률 최적화를 위한 JKO 스킴 내 볼록 함수 공간을 효과적으로 매개변수화할 수 있는가?
• RQ2 JKO-ICNN 프레임워크는 고차원에서 확장 가능하고 미분 가능한 최적화를 가능하게 하면서도 수렴 보장을 유지하는가?
• RQ3 JKO-ICNN은 해석적 해가 있는 저차원 PDE의 해를 얼마나 잘 근사하는가?
• RQ4 JKO-ICNN은 고차원 분자 생성에서 약물 유사도(QED)와 같은 복잡한 함수 최적화를 성공적으로 수행할 수 있는가?
• RQ5 다양한 분포 거리 측도(Sinkhorn 대비 MMD)와 가중치 설정은 모드 붕괴와 최적화 안정성에 어떤 영향을 미치는가?

주요 결과

• JKO-ICNN는 해석적 해가 있는 저차원 PDE의 해를 성공적으로 근사하여 정확성과 수렴성을 입증하였다.
• QM9 데이터셋에서 MMD를 사용하고 λ₂ = 1,000일 경우 중앙값 QED가 0.315에서 0.452로 상승하여 상대적 향상률 44.8%를 기록하였다.
• Sinkhorn 거리와 λ₂ = 10,000을 사용한 경우 중앙값 QED는 0.419에 도달하여 초기 분포 대비 33.0% 상승하였다.
• 분자 생성에서 유효성 92.7%와 유일성 81.9%를 달성하여 효과적인 모드 커버리지와 구조적 품질을 보였다.
• 실험 결과 고 λ₂ 값(10,000) 또는 MMD 기반 거리 측도 사용이 모드 붕괴를 유도하는 것으로 나타나 하이퍼파rameter 선택에 민감함을 확인하였다.
• 학습된 운반 지도가 미리 보지 않은 샘플로 일반화되어 재훈련 없이도 효율적인 추론이 가능함을 보여주었으며, 계산의 약간화를 입증하였다.