Demystifying Multimodal Biomolecular Co-design With Intrinsic Geodesic Coupling

Motivation

• Known: 단백질/소분자 리간드 같은 biomolecule의 sequence와 3D structure 간 상호작용을 포착하기 위해 diffusion/flow 기반 생성모델들이 multimodal co-design(예: MultiFlow, DPLM-2, La-Proteina, MolCRAFT, MolPilot)을 통해 두 modality를 공동 모델링해왔다.
• Gap: 기존 co-design 방법들은 대부분 marginal generative process를 병렬 실행하며 고정된 synchronous coupling을 암묵적으로 강제하거나, random coupling으로 완화를 시도하지만 이는 high-variance supervision과 training-inference mismatch를 유발하여 intermediate state의 modality consistency를 저해한다는 점이 간과되어 왔다.
• Why: 서로 다른 특성(연속 3D 좌표 vs. 이산 서열)을 가진 heterogeneous modality 간 temporal coupling을 최적화하면, 학습 안정성과 생성물의 물리적 타당성 및 다양성을 동시에 개선할 수 있어 구조 기반 신약 설계와 단백질 설계 등 실제 응용에 직접적인 영향을 미친다.
• Approach: GeoCoupling은 multimodal generation을 Temporal Optimal Transport 문제로 재정식화하여, prior에서 data로 이동하는 데 필요한 에너지를 최소화하는 temporal coupling schedule을 Gaussian Process surrogate 기반 bi-level optimization(Bayesian optimization)으로 탐색한다.

Achievement

Figure 2. Co-designability vs. diversity

1. 동적 불일치(dynamic mismatch) 규명: heterogeneous modality 간 synchronous coupling은 high-bias, random coupling은 high-variance 문제를 유발함을 규명.
2. 모델-비종속적 학습 프레임워크 제안: Gaussian Process 기반 Bayesian optimization으로 학습 중 low-energy temporal coupling을 자동 발견하는 model-agnostic 프레임워크(GeoCoupling)를 제시.
3. 실험적 우수성 입증: structure-based drug design 및 unconditional protein design 벤치마크에서 synchronous·random coupling 베이스라인 대비 학습된 coupling이 일관되게 우수한 physical validity와 diversity를 보임을 확인.

How

Figure 5. Visualization of the learned coupling against baselines.

• 각 modality(연속 구조 좌표, 이산 서열)에 대해 독립적인 diffusion/flow 시간 축을 설정하고, 이를 결합하는 temporal coupling schedule을 정의.
• 이 문제를 Temporal Optimal Transport로 재구성하여, prior에서 data로 가는 transport cost(에너지)를 최소화하는 schedule을 탐색 목표로 설정.
• schedule 최적화 목표(최종 validation 성능)가 비미분가능하고 평가 비용이 높다는 점을 고려해, 전체 학습 궤적을 통한 hypergradient 역전파 대신 Gaussian Process surrogate를 활용한 bi-level Bayesian optimization을 적용.
• structure-based drug design(MolCRAFT/MolPilot 계열)과 unconditional protein design(MultiFlow 계열) 두 과제에 프레임워크를 적용하여 학습된 coupling과 synchronous, random coupling 베이스라인을 비교.
• scRMSD, pLDDT, co-designability, diversity, trajectory 상의 structure MSE 및 sequence identity 변화를 통해 학습된 geodesic coupling의 효과를 정량적으로 분석.

Originality

• 기존 co-design 연구들이 대부분 marginal process의 병렬 실행(synchronous) 또는 무작위 조합(random coupling)에 머물렀던 것과 달리, temporal coupling 자체를 학습 가능한 자유도로 명시적으로 정식화한 최초의 체계적 프레임워크.
• 문제를 Temporal Optimal Transport로 재정의하고, 비미분가능·고비용 목적함수를 다루기 위해 Gaussian Process surrogate 기반 bi-level Bayesian optimization을 도입한 방법론적 결합이 독창적.
• Qiu et al.(2025)의 test-time schedule 최적화(post-hoc, 반복 평가 필요)와 달리, 학습 단계에서부터 coupling을 최적화하여 근본적인 training dynamics 불안정성을 해결하려는 접근.

Limitation & Further Study

• Bayesian optimization 기반 bi-level 최적화가 실제로 얼마나 many-modality(3개 이상)로 확장 가능한지, 차원이 커질수록 Gaussian Process surrogate의 샘플 효율성이 어떻게 변하는지에 대한 분석이 부족해 보임.
• 학습된 coupling이 특정 데이터셋/아키텍처에 과적합될 가능성이 있으며, 다양한 모델 아키텍처(예: latent 기반 La-Proteina, tokenizer 기반 DPLM-2)에 대한 일반화 검증이 제한적일 수 있음.
• 후속 연구로 coupling schedule을 conditional하게(예: 입력 pocket이나 서열 길이에 따라 적응적으로) 학습하는 방향, 그리고 더 많은 modality(예: side-chain, ligand pose 등)로의 확장이 필요.

Evaluation

총평: 멀티모달 biomolecular 생성에서 간과되어온 temporal coupling 문제를 정확히 짚어내고 이를 Temporal Optimal Transport 및 Bayesian optimization으로 풀어낸 참신하고 실용적인 프레임워크로, SBDD와 단백질 설계 양쪽에서 일관된 성능 향상을 보여 학계에 의미 있는 기여를 한다.