Protein Language Model Embeddings Improve Generalization of Implicit Transfer Operators

Motivation

• Known: Generative molecular dynamics (GenMD) 방법론은 Boltzmann Generators/Emulators와 Implicit Transfer Operators (ITO) 두 갈래로 발전해왔으며, 이들은 MD 대비 효율적인 샘플링을 가능하게 하지만 학습에 대량의 장시간 MD trajectory 데이터를 필요로 한다.
• Gap: 기존 TITO 및 Boltzmann Emulator 계열 모델들은 새로운(미학습) 단백질 시스템에 대한 transferability와 데이터 효율성이 제한적이며, 서열·구조 정보와 같은 보조 정보를 활용해 일반화를 향상시키는 방법이 충분히 탐구되지 않았다.
• Why: MD 시뮬레이션은 물리·화학·생물학 전반에서 실험 관측값을 정량적으로 예측하는 핵심 도구이지만 계산비용이 매우 크므로, 미학습 단백질계에도 일반화되면서 데이터 효율적인 대체 샘플링 모델을 개발하는 것은 fast-folding protein을 포함한 다양한 생물물리학적 응용에 중요하다.
• Approach: Coarse-grained TITO 프레임워크에 ESM 기반 protein language model 임베딩, Proteina 구조 임베딩, 온도, LLM 기반 주석 등 다양한 보조 조건화 신호를 결합하여 PLaTITO를 구성하고, 이를 다중 온도의 다양한 off-equilibrium MD trajectory로 학습하여 out-of-distribution 단백질계에 대한 일반화 성능을 검증한다.

Achievement

Figure 3. Scalability of PLaTITO-19M with training compute.

1. Coarse-grained transferable TITO 개발: 다양한 다중 온도 off-equilibrium MD trajectory로 처음부터 학습하여 시스템별 fine-tuning 없이 out-of-distribution 단백질계에 일반화되는 coarse-grained TITO를 제시했다.
2. PLaTITO를 통한 데이터 효율성 개선: ESM pLM 임베딩 등 보조 정보를 조건화에 활용한 PLaTITO가 base TITO 대비 적은 학습 데이터·연산 자원으로도 equilibrium sampling 벤치마크에서 state-of-the-art 성능을 달성함을 보였다.
3. Non-Arrhenius 접힘 동역학 재현: 온도 조건부 PLaTITO 모델이 fast-folding protein의 non-Arrhenius 접힘/풀림 속도를 재현하여 복잡한 rugged free energy landscape와 일치하는 결과를 얻었다.

How

Figure 1. PLaTITO generalizes to unseen protein systems while improving data efficiency. Given the molecular state of a

• Molecular dynamics를 discrete-time Markov process로 정식화하고, Boltzmann 분포 및 장시간 transition density 학습을 목표로 함.
• Conditional Flow Matching (CFM) 프레임워크를 채택해 rectified flow 형태의 linear conditional path로 velocity field를 회귀 학습.
• Implicit Transfer Operator (ITO) 프레임워크를 확장하여 backbone coordinate xt, 아미노산 서열 S, 온도 T, time step ∆t를 조건으로 장시간 transition density p(xt+∆t | xt, S, T, ∆t)를 근사.
• Protein Backbone Transformer 구조에 ESM 서열 임베딩, Proteina 구조 임베딩, LLM 기반 주석(ALLM)을 residue/pair representation 형태로 추가 조건화.
• 3D Gaussian noise에서 시작해 iterative rollout을 통해 다양한 시간 스케일(10ns~1μs)의 conformational dynamics를 샘플링.
• Coarse-grained TITO를 Boltzmann Emulator와 데이터 효율성 측면에서 비교하고, PLaTITO-19M으로 모델 용량을 스케일링하여 equilibrium sampling 벤치마크(fast-folding protein 포함) 및 training compute에 따른 확장성을 평가.

Originality

• Coarse-grained TITO에 protein language model(ESM) 임베딩을 조건화 신호로 도입하여 out-of-distribution 일반화를 개선한 최초의 시도.
• 구조 임베딩(Proteina), 온도, LLM 기반 주석 등 다양한 보조 정보 소스를 통합적으로 비교·분석하는 체계적 ablation 설계.
• 온도 조건부 모델을 통해 non-Arrhenius 접힘/풀림 속도라는 물리적으로 의미 있는 현상을 학습된 모델에서 재현함으로써 단순 샘플링 정확도를 넘어선 동역학적 타당성을 검증.

Limitation & Further Study

• Coarse-grained 표현을 사용하므로 원자 수준의 세부 정보(atomistic detail)가 손실될 수 있으며 all-atom 복원 과정의 정확도에 대한 추가 검증이 필요하다.
• 벤치마크가 주로 fast-folding protein 등 특정 단백질군에 집중되어 있어, 훨씬 크거나 다중 도메인 단백질, 복합체(단백질-단백질, 단백질-리간드)로의 일반화는 추가 검증이 요구된다.
• pLM 및 구조 임베딩에 의존하므로 이러한 사전학습 모델의 편향이나 한계가 TITO 성능에 전이될 가능성이 있다.
• 본문 발췌에서는 정량적 비교 수치와 baseline 대비 개선폭이 구체적으로 제시되지 않아, 성능 향상의 통계적 유의성 및 다양한 시스템에서의 강건성에 대한 추가 논의가 필요하다.

Evaluation

총평: Protein language model 임베딩과 같은 보조 정보를 활용해 transferable implicit transfer operator의 데이터 효율성과 일반화를 크게 향상시킨 실질적이고 체계적인 연구로, GenMD 분야에서 의미 있는 진전을 보여준다.