MolFORM: Multi-modal Flow Matching for Structure-Based Drug Design

MolFORM: 구조 기반 약물 설계(SBDD)를 위한 멀티모달 플로우 매칭 기술 요약

1. 문제 정의

구조 기반 약물 설계(SBDD)는 3D 구조 정보를 활용하여 특정 단백질 타겟에 효과적으로 결합하는 리간드 분자를 생성하는 것을 목표로 한다. 핵심 과제는 단백질 포켓을 조건으로 하여 이산 변수(원자 유형)와 연속 변수(3D 원자 좌표)의 결합 분포를 모델링하는 것이다. 확산 기반 생성 모델이 SBDD의 지배적인 패러다임이 되었으나, 대안적인 비자기회적(non-autoregressive) 프레임워크는 여전히 미개척 영역으로 남아 있다. 또한, 표준 생성 모델은 단일 목적 함수에 과적합되지 않으면서 결합 친화도를 최적화하는 동시에 약물 유사성 및 합성 가능성을 유지하는 것과 같은 하위 사용자 선호도와 정렬하는 데 어려움을 겪는 경우가 많다.

2. 방법론

본 논문은 MolFORM을 제안하며, 이는 이산적 및 연속적 분자 양상(modality)을 공동으로 모델링하기 위해 **멀티모달 플로우 매칭(Multi-modal Flow Matching)**을 사용하는 새로운 생성 프레임워크이다. 방법론은 두 가지 주요 단계로 구성된다:

A. 멀티모달 플로우 매칭 (기본 모델)
MolFORM은 생성 작업을 상미분 방정식(ODE)에 의해 제어되는 수송 매핑(transport mapping) 문제로 취급한다. 이 모델은 두 개의 별도 브랜치를 가진 조건부 플로 매칭(Conditional Flow Matching, CFM) 프레임워크를 채택한다:

• 연속적 양상 (3D 좌표): 모델은 가우시안 사전 분포(Gaussian prior)를 실제 데이터 분포의 원자 위치로 수송하기 위한 시간 의존적 벡터 필드를 학습한다. 훈련 목적 함수는 예측된 좌표와 타겟 좌표 사이의 평균 제곱 오차(MSE)를 최소화한다.
• 이산적 양상 (원자 유형): 모델은 **이산 플로우 매칭(Discrete Flow Matching, DFM)**을 채택한다. 모델은 시간 $t \in [0, 1]$ 동안 원자 유형이 균등 노이즈 분포에서 데이터 분포로 진화하는 균등 훼손(uniform corruption) 과정을 사용한다. 모델은 훼손된 시퀀스로부터 원래의 원자 유형을 예측하며, 이때 교차 엔트로피 손실을 사용한다.
• 기하학적 정렬: 정확한 공간적 충실도를 보장하기 위해, 전체 목적 함수에 Chamfer 손실이 통합된다. 이 양방향 거리 메트릭은 예측된 원자 포인트 클라우드가 실제 구조와 밀접하게 정렬되도록 유도한다.
• 아키텍처: 모델은 리간드 원자와 단백질 포켓 원자 간의 상호작용을 처리하기 위해 SE(3)-등변 그래프 신경망(GNN)을 사용하여 은닉 임베딩과 좌표를 번갈아 업데이트한다.

B. 선호도 가이드 미세 조정 (DPO)
생성 품질을 높이고 모델을 화학적 선호도에 정렬하기 위해, 저자들은 멀티플로우 직접 선호 최적화(Multi-flow Direct Preference Optimization, DPO) 단계를 도입한다.

• 보상 신호: Vina 점수(결합 친화도)가 각 단백질 포켓에 대한 선호 쌍(승리한 리간드 vs 패배한 리간드)을 구축하기 위한 화학 정보 기반 보상으로 사용된다.
• 공동 모델링 전략: 표준 DPO와 달리, MolFORM은 이산 및 연속 양상 모두에 DPO를 동시에 적용한다.
  • 연속 좌표의 경우, DPO 손실은 승리한 샘플과 패배한 샘플에 대한 모델의 예측값과 참조 모델의 예측값 사이의 Chamfer 거리 차이를 사용하여 공식화된다.
  • 이산 원자 유형의 경우, 균등 노이징 스킴 하에서 유도된 이산 플로우 매칭 변형 DPO를 사용하여, 원자 유형 공간에서의 선호도를 정렬하도록 속도 행렬(rate matrix)을 최적화한다.
• 훈련 전략: 저자들은 샘플링 중에 비균등 시간 스케줄링 전략을 사용하여, 최종 정밀화 단계에서 미세한 구조적 세부 사항을 포착하기 위해 구간 $[0.8, 1]$에 더 많은 단계를 할당한다.

3. 주요 기여

• 새로운 프레임워크: SBDD를 위해 멀티플로우 매칭을 적용하여 이산 원자 유형과 연속 3D 좌표의 공동 생성을 효과적으로 처리하는 최초의 프레임워크인 MolFORM을 소개한다.
• 멀티플로우 DPO: 명시적인 보상 함수 추정 없이도 미세한 선호도 정렬을 가능하게 하는, 이산 및 연속 양상 모두에 직접 선호 최적화를 동시에 적용하는 새로운 공동 모델링 전략을 제안한다.
• 성능 향상: 플로우 기반 모델이 SOTA 확산 모델과 대등한 성능을 달성할 수 있음을 입증하였으며, DPO를 추가함으로써 추가적인 개선을 이끌어냈다.
• 효율성: 모델은 확산 기반 베이스라인(예: TargetDiff의 경우 3428초)에 비해 현저히 빠른 샘플링 시간(100개 샘플에 69초)을 달-성하며, 1000단계 대신 100단계의 적분만 필요로 한다.

4. 실험 결과

실험은 CrossDocked2020 데이터셋에서 수행되었으며, LiGAN, Pocket2Mol, TargetDiff, DiffBP를 포함한 12개의 베이스라인과 MolFORM을 비교하였다.

• 결합 친화도: MolFORM-DPO는 Vina Score (-6.16), Vina Min (-7.18), Vina Dock (-8.13)에서 SOTA 성능을 달성하여, 베이스라인인 TargetDiff 및 기타 확산 모델을 능가했다.
• 화학적 특성: 모델은 약물 유사성 지수(QED) 및 합성 가능성(SA)에서 개선을 보였다. 특히, TargetDiff를 재학습하여 적용했을 때의 작은 개선(각각 2%, 2%)과 비교하여, DPO 미세 조정을 거친 MolFORM은 QED를 4%, SA를 8% 개선했다.
• 구조적 품질: 모델은 베이스라인에 비해 더 낮은 Chamfer 거리와 더 나은 기하학적 충실도(더 낮은 JSDBL 및 JSDBA)를 달성했다. DPO 미세 조정 후에도 생성된 분자의 다양성은 0.77로 높게 유지되어 모드 붕괴(mode collapse)가 없음을 보여주었다.
• 절제 연구(Ablation Studies): DPO 목적 함수에서 Chamfer 손실을 제거하면 Vina 점수와 기하학적 지표에서 상당한 하락이 발생하여, 기하학적 정렬 항의 중요성을 확인시켜 주었다.

5. 의의 및 주장

본 논문은 정교하게 설계된 플로우 기반 모델이 확산 기반 접근 방식의 성능을 따라잡거나 능가할 수 있음을 보여줌으로써, MolFORM이 SBDD의 유망한 방향임을 주장한다. 주요 의의는 멀티플로우 DPO 공동 모델링 전략에 있으며, 저자들은 이것이 멀티플로우 설정에서 선호도 정렬을 적용한 첫 사례라고 밝히고 있다. 이 접근 방식은 단일 목적 함수에 과적합되지 않고 여러 지표에 걸쳐 일관된 개선을 이끌어내면서, 원자 유형과 3D 위치 모두에 대해 선호도를 동시에 정렬할 수 있게 한다.

저자들은 이 프레임워크가 (TargetDiff와의 비교에서 나타난 것처럼) 단일 모달 미세 조정보다 DPO 미세 조정에 더 큰 잠재력을 제공한다고 강조한다. 이는 이산적 및 연속적 양상의 구조적 분리(structural disentanglement) 덕분이며, 이를 통해 기하학을 방해하지 않고 원자 유형을 정밀하게 수정하는 것과 같이 더 표적화된 업데이트를 가능하게 하여 그래디언트 간섭을 줄이고 안정성을 높인다. 이 연구는 이러한 멀티모달 플로우 프레임워크가 재료 발견 및 칩 레이아웃 최적화를 포함한 더 넓은 컴퓨터 보조 설계 도메인으로 확장될 가능성이 있음을 시사한다.