FlexiFlow: decomposable flow matching for generation of flexible molecular ensemble

Motivation

• Known: flow matching과 diffusion model은 3D 분자 생성에서 강력한 성능을 보여왔으며, EDM, MiDi, EQGAT-diff, SemlaFlow와 같은 모델들이 QM9와 GEOM Drugs 벤치마크에서 높은 valid/unique/novel 분자 생성 성능을 달성했다. 또한 CREST와 같은 물리 기반 방법이 분자의 저에너지 conformer를 탐색하는 데 사용되어 왔다.
• Gap: 기존 3D de-novo 분자 생성 모델들은 분자당 단일 conformation만 생성할 수 있어, 실제로 결합 친화도와 생물학적 활성에 큰 영향을 미치는 conformational diversity를 포착하지 못한다는 한계가 있었다. 또한 물리 기반 conformer 탐색 방법(CREST 등)은 정확하지만 계산 비용이 매우 크다는 문제가 있다.
• Why: 약물 발견 과정에서 분자의 conformational landscape는 관측 가능한 물성과 표적 단백질에 대한 결합력을 결정하므로, 대표적인 저에너지 conformer 집합을 함께 생성할 수 있다면 원하는 열역학적 특성을 갖는 분자를 더 효과적으로 설계할 수 있다.
• Approach: 저자들은 flow matching의 조건부 확률 경로를 conditional independence 가정 하에 분해하여, 하나의 참조 conformation을 갖는 분자 그래프와 임의 개수의 추가 conformer 집합을 동시에 샘플링하는 FlexiFlow 아키텍처를 제안한다. 이 아키텍처는 equivariance와 permutation invariance를 동시에 만족하도록 설계되었다.

Achievement

Figure 3. RMSD before and after energy minimization (left), and precision–recall curves for Cov (middle) and AMR (right)

1. State-of-the-art 3D 분자 생성: QM9와 GEOM Drugs 데이터셋에서 valid, unique, novel한 분자를 학습 데이터 분포에 높은 충실도로 생성하며 최신 성능을 달성함.
2. 효율적인 conformational ensemble 생성: CREST와 같은 물리 기반 방법과 유사한 conformer coverage를 훨씬 짧은 inference 시간에 제공함.
3. 단백질 조건부 리간드 생성으로의 전이: conformation 정보가 없는 static pocket 데이터셋에서도 단백질 조건부 리간드 생성 task에 성공적으로 전이됨.
4. 이론적 기여: flow matching objective를 conditional independence를 이용해 분해하는 새로운 프레임워크를 제시하고, 이를 뒷받침하는 수학적 증명(block diagonal 행렬의 행렬식 분해 등)을 제공함.

How

Figure 2. The FlexiFlow architecture takes equivariant and invari-

• 확률 밀도 함수 pt(x, S)를 대표 벡터 x와 conformer 집합 S = {y_i} 간의 conditional independence 가정 하에 곱의 형태로 분해 (pt(x,S) = ∏ pt(x,y))
• 이에 대응하는 벡터장 vt와 flow ψt를 m개의 독립적인 flow의 concatenation으로 정의 (ψt(x,S) = ⊙ ψt(x,y))
• push-forward 방정식과 block diagonal 행렬의 행렬식 성질을 이용해 이 분해가 Lipman et al. (2023)의 flow matching 이론과 일관됨을 증명
• 이를 바탕으로 conditional flow matching (CFM) objective를 확장한 LFM(θ)를 정의하여, 그래프와 conformer 쌍 (x,y) 각각에 대해 손실을 합산
• FlexiFlow 아키텍처는 두 좌표 집합(참조 conformation, 추가 conformer)에 대해 equivariant한 신경망 구조를 사용하며, categorical feature(원자/결합 타입)에는 permutation invariance를 유지
• QM9, GEOM Drugs에서 무조건적 분자 생성 실험을, CREST와 비교하는 conformer ensemble 실험을, 그리고 단백질 pocket 조건부 리간드 생성 실험을 수행

Originality

• flow matching의 conditional probability path를 conditional independence 가정으로 분해하여 하나의 분자 그래프와 다중 conformer를 동시에 생성하는 새로운 이론적 프레임워크 제안
• 두 좌표 집합(참조 conformation과 conformer 집합)에 대해 동시에 equivariance를 보장하는 새로운 아키텍처 설계
• Boltzmann generator와 달리 추가 MD 데이터로 재학습할 필요 없이 정확한 에너지 추정 없이도 고품질 conformer를 생성할 수 있음을 보임
• conformation 데이터가 없는 static pocket 데이터셋에서도 학습된 프레임워크를 단백질 조건부 생성 task로 전이 가능함을 시연

Limitation & Further Study

• conformer 개수 m을 고정하거나 사전에 정해야 하는지, 가변적인 conformer 수를 다루는 방식에 대한 명확한 설명이 부족할 수 있음
• conditional independence 가정이 실제 conformer들 간의 상관관계(예: 에너지 장벽에 따른 전이 확률)를 완전히 포착하지 못할 가능성이 있음
• CREST와의 비교가 특정 분자 크기/유형에 국한되어 있을 수 있어, 더 크고 유연한 약물유사 분자에 대한 일반화 검증이 추가로 필요함
• 단백질 조건부 생성에서 conformation 정답 데이터 없이 학습된 모델의 실제 결합 친화도 예측 정확도에 대한 실험적 검증(예: 실제 실험값과의 비교)이 제한적일 수 있음
• 후속 연구로 다양한 단백질 표적에 대한 대규모 검증, 실제 실험적 결합 친화도와의 상관성 분석, conformer 수를 적응적으로 결정하는 방법 등이 필요함

Evaluation

총평: flow matching을 conditional independence로 분해하여 분자 그래프와 다중 conformer를 동시에 생성하는 아이디어는 이론적으로 탄탄하고 실용적 가치가 높으며, 약물 발견 분야에 실질적으로 기여할 수 있는 참신한 연구이다.