[논문 리뷰] Smooth Normalizing Flows

Jonas Köhler, Andreas Krämer외 1명|arXiv (Cornell University)|2021. 10. 01.

Protein Structure and Dynamics인용 수 15

한 줄 요약

이 논문은 내부 좌표를 갖는 분자 시스템에 필수적인 컴팩트한 간격과 초구면에서 확률 분포를 모델링하기 위해 C∞-스무스 혼합 변환을 사용하는 스무스 정규화 플로우를 소개한다. 역함수 정리에 기반한 미분 가능 역변환을 통해 이중 방향 학습, 힘 매칭, 분자 동역학 시뮬레이션에서의 응용이 가능해지며, 잘 조율된 힘과 높은 샘플링 효율을 보이는 정확하고 스무스한 잠재 에너지 함수를 달성한다.

ABSTRACT

Normalizing flows are a promising tool for modeling probability distributions in physical systems. While state-of-the-art flows accurately approximate distributions and energies, applications in physics additionally require smooth energies to compute forces and higher-order derivatives. Furthermore, such densities are often defined on non-trivial topologies. A recent example are Boltzmann Generators for generating 3D-structures of peptides and small proteins. These generative models leverage the space of internal coordinates (dihedrals, angles, and bonds), which is a product of hypertori and compact intervals. In this work, we introduce a class of smooth mixture transformations working on both compact intervals and hypertori. Mixture transformations employ root-finding methods to invert them in practice, which has so far prevented bi-directional flow training. To this end, we show that parameter gradients and forces of such inverses can be computed from forward evaluations via the inverse function theorem. We demonstrate two advantages of such smooth flows: they allow training by force matching to simulation data and can be used as potentials in molecular dynamics simulations.

연구 동기 및 목표

• 분자 내부 좌표를 모델링하는 데 필수적인 컴팩트한 간격과 초구면에서 C∞-스무스 정규화 플로우의 부족을 보완한다.
• 역함수 정리에 기반한 기울기 계산을 통해 비해석적 역함수를 갖는 플로우의 이중 방향 학습을 가능하게 한다.
• 분자 동역학 시뮬레이션에 적합한 스무스하고 미분 가능한 잠재 에너지 함수를 개발한다.
• 기존 방법이 지지역을 사전에 알 필요가 있거나 불연속적인 힘을 생성하는 등의 한계를 극복한다.

제안 방법

• 컴팩트한 간격과 초구면에서 사용 가능한 학습 가능한 형상 매개변수를 가진 버블 함수를 사용하는 C∞-스무스 혼합 변환의 클래스를 제안한다.
• 플로우의 역방향에서 역을 구하기 위해 블랙박스 근사 해법(예: 다이어그램 기반 검색)을 사용한다.
• 역함수 정리에 기반해 역함수의 기울기를 계산하여 근사 해법 단계를 통해 역전파가 가능하게 한다.
• 최대우도추정(MLE), 힘 매칭(FM), 역KL 발산을 조합한 하이브리드 손실을 사용해 플로우를 학습한다.
• 속도 버웰 적분기를 사용해 학습된 플로우를 분자 동역학 시뮬레이션에서 잠재 에너지로 적용한다.
• OpenMM을 사용해 에너지와 힘을 평가하고, 샘플링 효율성과 에너지 보존 정도를 통해 검증한다.

실험 결과

연구 질문

• RQ1 C∞-스무스 정규화 플로우는 컴팩트한 간격과 초구면 모두를 지원하면서 분자 내부 좌표의 스무스한 모델링이 가능한가?
• RQ2 역함수 정리에 기반해 비해석적 역함수 플로우를 통해 기울기를 역전파할 수 있는가?
• RQ3 밀도 추정과 결합된 힘 매칭은 분자 시스템에서 학습된 잠재 에너지의 정확도를 향상시킬 수 있는가?
• RQ4 스무스 플로우는 전통적인 힘 장력과 비교해 안정적이고 에너지 보존이 잘 되는 분자 동역학 시뮬레이션을 생성하는가?
• RQ5 스무스 플로우는 스퍼링 기반 또는 애프린 플로우보다 샘플링 효율성과 힘 정확도에서 뛰어나지 않는가?

주요 결과

• 스무스 플로우 모델은 알라닌 다펩타이드에서 42%의 샘플링 효율을 달성하여 스퍼링 기반 플로우(38%)와 애프린 플로우(<1%)를 능가했다.
• 힘 매칭과 MLE로 학습한 스무스 플로우는 역KL 발산이 0.0003으로 나타나 밀도 매칭이 거의 완벽한 수준임을 보였다.
• 스무스 플로우를 사용한 MD 시뮬레이션에서의 에너지 변동은 Amber ff99SB-ILDN 힘 장력과 유사하여 수치적 안정성을 나타냈다.
• 스퍼링 기반 플로우보다 스무스 플로우에서 계산된 힘은 훨씬 더 정확했으며, 테스트 데이터에서 더 낮은 힘 잔여치를 보였다.
• 역함수 정리에 기반해 비해석적 역함수 플로우의 기울기를 계산함으로써 기존 연구의 핵심 한계를 극복하고 이중 방향 학습을 가능하게 하였다.
• 스무스 플로우 아키텍처는 잘 조율된 힘과 에너지를 생성하여 힘의 불연속성이 없는 상태로 직접 분자 동역학 시뮬레이션에 활용할 수 있었다.