[논문 리뷰] SE(3) Equivariant Augmented Coupling Flows

Laurence I. Midgley, Vincent Stimper외 4명|arXiv (Cornell University)|2023. 08. 20.

Computer Graphics and Visualization TechniquesComputer Science인용 수 3

한 줄 요약

이 논문은 분자 시스템에서 SE(3) 및 치환 불변성을 유지하면서도 빠르고 정확한 샘플링과 밀도 평가가 가능한 SE(3) 불변 증강 결합 흐름(E-ACF)을 제안한다. 학습된 증강 차원을 통해 원자 좌표를 SE(3)-불변 기저로 투영함으로써, 볼츠만 분포 학습에서 최신 기술을 초월하며, CNF 및 확산 모델 대비 샘플링 속도가 10배 이상 빠르고 DW4 및 LJ13 시스템에서 에너지 기반 학습이 가능하다.

ABSTRACT

Coupling normalizing flows allow for fast sampling and density evaluation, making them the tool of choice for probabilistic modeling of physical systems. However, the standard coupling architecture precludes endowing flows that operate on the Cartesian coordinates of atoms with the SE(3) and permutation invariances of physical systems. This work proposes a coupling flow that preserves SE(3) and permutation equivariance by performing coordinate splits along additional augmented dimensions. At each layer, the flow maps atoms' positions into learned SE(3) invariant bases, where we apply standard flow transformations, such as monotonic rational-quadratic splines, before returning to the original basis. Crucially, our flow preserves fast sampling and density evaluation, and may be used to produce unbiased estimates of expectations with respect to the target distribution via importance sampling. When trained on the DW4, LJ13, and QM9-positional datasets, our flow is competitive with equivariant continuous normalizing flows and diffusion models, while allowing sampling more than an order of magnitude faster. Moreover, to the best of our knowledge, we are the first to learn the full Boltzmann distribution of alanine dipeptide by only modeling the Cartesian positions of its atoms. Lastly, we demonstrate that our flow can be trained to approximately sample from the Boltzmann distribution of the DW4 and LJ13 particle systems using only their energy functions.

연구 동기 및 목표

• 카르테시안 좌표에서 작동하고 SE(3) 및 치환 불변성을 강제하며, 빠른 샘플링과 정확한 밀도 평가가 가능한 정규화 흐름이 부족한 문제를 해결하기 위해.
• 정규화 흐름의 타당한 샘플링과 밀도 평가를 보장함으로써 분자 시스템에서 에너지 기반 학습을 가능하게 하기 위해.
• 오직 카르테시안 좌표와 에너지 함수만을 사용하여 알라닌 다펩타이드의 전체 볼츠만 분포를 학습하기 위해.
• 내부 좌표에 의존하지 않고 분자 그래프 및 입자 시스템 전반에 일반화 가능한 확장 가능한 불변 흐름 아키텍처를 제공하기 위해.
• SE(3) 불변 결합 흐름이 연속 정규화 흐름 및 확산 모델과 비슷하거나 이를 능가하는 성능을 달성하면서도 샘플링 속도가 훨씬 빠르다는 것을 입증하기 위해.

제안 방법

• 표준 결합 흐름을 확장하여, SE(3)에 따라 변환되는 보조 변수를 포함한 증강 입력 공간을 도입함으로써 불변성을 달성한다.
• 각 레이어에서 원자 위치는 학습된 변환을 통해 SE(3)-불변 기저로 투영되며, 여기에 표준 정규화 흐름 레이어(예: 단조적 유리적 2차 곡선 스플라인 등)가 적용된다.
• 출력은 다시 원래의 SE(3)-불변 공간으로 투영되며, 이로써 흐름의 가역성과 자코비안 행렬식의 계산 가능성 유지된다.
• CNF나 확산 모델과 달리 반복적 샘플링을 피하기 때문에, 이 방법은 빠른 샘플링과 밀도 평가를 유지한다.
• 최대 우도 학습과 에너지 기반 학습을 모두 지원하며, 비정규화된 볼츠만 밀도만을 사용함으로써 MD 샘플이 필요 없이 학습이 가능하다.
• 불변성과 표현력 최적화를 위해 벡터 투영, 카르테시안 투영, 구면 투영 등의 다양한 투영 전략을 평가한다.

Figure 2: Inter-atomic distances for samples from the train-data and Spherical-proj E-ACF .

실험 결과

연구 질문

• RQ1 카르테시안 좌표에서 작동하고 SE(3) 및 치환 불변성을 유지하면서도 빠른 샘플링과 정확한 밀도 평가가 가능한 결합 흐름을 설계할 수 있는가?
• RQ2 오직 카르테시안 좌표와 에너지 함수만을 사용하여 알라닌 다펩타이드의 전체 볼츠만 분포를 학습할 수 있는가?
• RQ3 CNF 및 확산 모델이 느린 샘플링으로 인해 비가역적인 DW4 및 LJ13 시스템에서 제안된 흐름을 에너지 기반 학습으로 훈련시킬 수 있는가?
• RQ4 동일한 훈련 설정 하에서 제안된 흐름의 샘플링 속도 및 가능도는 SE(3) CNF 및 확산 모델과 비교해 어떻게 되는가?
• RQ5 효율적 샘플 크기가 클 경우, 특히 분포의 저밀도 영역에서 중요도 재가중이 샘플 품질을 향상시키는가?

주요 결과

• 카르테시안 투영 E-ACF는 전체 QM9-좌표 데이터셋에서 -188.6의 최저 음수 로그우도(NLL)를 기록하여 다른 변종과 경쟁 모델을 압도했다.
• 중요도 재가중을 통해 알라닌 다펩타이드 시스템에서 800 K에서 칼루바크-라이블러 발산(KLD)을 1.88×10⁻³에서 1.39×10⁻³으로 감소시켜 저밀도 영역의 샘플 품질을 크게 향상시켰다.
• 300 K에서 알라닌 다펩타이드 시스템에서 Spherical-proj E-ACF는 재가중 후 KLD가 1.97×10⁻³을 기록했고, 내부 좌표 기반 흐름의 (7.57±3.80)×10⁻³보다 뛰어난 성능을 보였다.
• E-ACF 모델은 SE(3) 연속 정규화 흐름 및 확산 모델 대비 10배 이상 빠른 샘플링 속도를 기록했으며, DW4에서 샘플당 6.3–7.5초, LJ13에서 42.7–45.0초의 추론 시간을 기록했다.
• MD 생성 데이터가 필요 없이도, 완전히 불변적인 흐름을 통해 DW4 및 LJ13에서 에너지 기반 학습이 처음으로 가능해졌다.
• 이 방법은 내부 좌표나 분자 그래프 의존성 없이 오직 카르테시안 좌표만을 사용하여 알라닌 다펩타이드의 전체 볼츠만 분포를 학습한 최초의 방법이다.

Figure 3: Ramachandran plots, i.e. marginal distribution of the dihedral angles $\phi$ and $\psi$ (see Sec. C.3.2 ), obtained with MD (ground truth) and various models.