[논문 리뷰] Rigid Body Flows for Sampling Molecular Crystal Structures

Jonas Köhler, Michele Invernizzi외 2명|arXiv (Cornell University)|2023. 01. 26.

Computer Graphics and Visualization Techniques인용 수 10

한 줄 요약

이 논문은 3D에서 다수의 강체(분자)의 위치와 방향 샘플링에 맞춘 노멀라이징 흐름을 도입하고, 이중 커버 쿼터니언 표현을 사용해 SO(3) 밀도를 매끄럽고 정확하게 모델링합니다. 또한 외부장에 의한 정사면체 시스템과 TIP4P 물에서의 얼음 XI에 대한 볼츠만 제너레이터 스타일 샘플링을 시연하여 분자 결정에서 자유에너지 차이 ΔF 추정을 효과적으로 가능하게 합니다.

ABSTRACT

Normalizing flows (NF) are a class of powerful generative models that have gained popularity in recent years due to their ability to model complex distributions with high flexibility and expressiveness. In this work, we introduce a new type of normalizing flow that is tailored for modeling positions and orientations of multiple objects in three-dimensional space, such as molecules in a crystal. Our approach is based on two key ideas: first, we define smooth and expressive flows on the group of unit quaternions, which allows us to capture the continuous rotational motion of rigid bodies; second, we use the double cover property of unit quaternions to define a proper density on the rotation group. This ensures that our model can be trained using standard likelihood-based methods or variational inference with respect to a thermodynamic target density. We evaluate the method by training Boltzmann generators for two molecular examples, namely the multi-modal density of a tetrahedral system in an external field and the ice XI phase in the TIP4P water model. Our flows can be combined with flows operating on the internal degrees of freedom of molecules and constitute an important step towards the modeling of distributions of many interacting molecules.

연구 동기 및 목표

• 다수의 상호 작용하는 강체의 공동 위치와 방향 샘플링을 효율적이고 가능성 있는 우도 기반으로 촉진하기.
• 쿼터니언 이중 커버를 통해 회전 도메인에서 매끄럽고 가역적인 흐름을 개발하여 물리적 밀도를 보존하고 정확한 밀도 계산을 지원하기.
• 외부 포즈와 내부 자유도를 독립적으로 다룰 수 있도록 하여 내부 자유도 흐름 모델과의 통합을 가능하게 하기.
• 다중 모드 회전 밀도와 얼음 XI 상의 서로 다른 크기와 온도에서의 Boltzmann 제너레이터를 구성하여 접근 방식을 시연하기.

제안 방법

• 각 강체를 (x0, R, Psi)로 표현하는데, 여기서 x0는 평행이동, R은 SO(3) 회전, Psi는 내부 자유도.
• 이중 커버를 사용해 단위 쿼터니언 구 S3로 SO(3) 회전을 승격시키고, S3에서 매끄럽고 플립-대칭인 흐름을 가능하게 함.
• 두 클래스의 플립-대칭 커플링 흐름을 S3에서 구성: 대칭화된 모비우스 변환과 투영 기하학적 볼록 그래디언트 맵(해석적 역함수와 부피 변화 공식을 포함).
• 커플링 기반 흐름 F를 정의해 위치와 쿼터니언을 업데이트하면서 플립-대칭성을 보존하고 이중 커버 투영하에서 정확한 밀도를 보장.
• 학습을 볼츠만 제너레이터로 설명: (편향된) MD 데이터에 대한 최대가능도 값을 역 KL 최소화나 LFEP와 조합하여 자유에너지 차이(ΔF)를 추정.
• 실험 설정 두 가지를 시연: (i) 외부 전장 속 다모드 회전 밀도를 갖는 다면체, (ii) TIP4P 물에서의 얼음 XI를 다양한 크기와 온도에서 샘플링.

Figure 1: Flow on SO(3) via the $S^{3}$ double cover: we first transform a system of Cartesian coordinates $\bm{x}$ via $T$ into the triple $(\bm{x}_{0},\bm{R},\bm{\Psi})$ , containing the global translation $\bm{x}_{0}$ , the global rotation $\bm{R}$ and its fixed inner degrees of freedom $\bm{\Psi

실험 결과

연구 질문

• RQ1 쿼터니언 이중 커버에서의 매끄럽고 플립-대칭인 흐름이 강체의 다모드 회전 밀도를 정확히 표현할 수 있는가?
• RQ2 강체 노멀라이징 흐름이 분자 결정의 샘플링을 얼마나 잘 수행하고 MD 기반 참조와 비교하여 자유에너지 차이(ΔF)를 정확하게 추정할 수 있는가?
• RQ3 프레임워크를 내부 자유도 흐름과 결합하여 결정 속에서 상호 작용하는 분자를 모델링할 수 있는가?
• RQ4 S3를 통한 SO(3) 표현이 결정 시스템의 학습 효율성과 밀도 정확도에 미치는 영향은 무엇인가?

주요 결과

• 투사형 볼록 그래디언트 맵과 플립-대칭 대칭화된 모비우스 변환은 다모드 회전 밀도를 신뢰성 있게 재구성할 수 있으며 외부 전장 테스트의 사강체에서 선형 쿼터니언 계층보다 우수한 성능을 보임.
• TIP4P 물을 사용한 얼음 XI 샘플링에서 흐름 기반 LFEP는 분자당 ΔF를 다양한 시스템 크기와 목표 온도에서 MBAR 참조와 근접하게 일치시킴.
• 기저 흐름으로 매핑된 에너지와 목표 MD 분포 간의 큰 중첩을 얻고 재가중은 반지름 분포 함수를 높은 충실도로 복원함.
• 단일 MD 궤적을 학습과 LFEP 평가에 사용하여 자유에너지 차이의 효율적 추정이 가능해지며, 계층적 MD 시뮬레이션의 필요성을 줄임.
• 이 구조는 외부 포즈와 내부 자유도 간 독립성을 지원하며, 대형 결정 시스템에 대한 치환-불변 및 확장 가능한 흐름 설계와도 호환됨.

Figure 2: Tetrahedron in an external field: a) each colored bead is attracted with the same force towards the external point $\bm{c}$ according to the potential defined in Eq. ( 15 ). b) Density of rotational degree of freedom. Rotations are represented as unit-quaternions $\bm{q}=(x,y,z,w)$ . First