Boltz-Jump: Accelerated Sampling of the Conformational Landscape of Biomolecular Structure Prediction Models

Motivation

• Known: AlphaFold 3, Boltz-2와 같은 diffusion 기반 biomolecular 구조예측 모델은 반복적 diffusion sampling을 통해 구조 ensemble을 생성할 수 있으며, AlphaFlow, Conformix, BioEmu 등 사전학습 모델을 활용하거나 MD 데이터로 추가 학습하여 conformational ensemble을 생성하는 다양한 접근이 존재한다. 또한 walk-jump sampling(JAMUN)은 작은 peptide의 MD trajectory에 대해 학습된 score 기반 생성 모델로서 표준 MD 대비 훨씬 빠른 sampling을 보인 바 있다.
• Gap: 기존 diffusion 기반 구조예측 모델을 이용한 ensemble 생성은 reverse diffusion process에 필요한 많은 function evaluation으로 인해 계산 비용이 매우 크며, JAMUN과 같은 walk-jump 접근은 대규모 biomolecule에 대해 생성하기 비싼 수렴된 MD trajectory로 별도 학습해야 한다는 한계가 있다.
• Why: Boltz-2와 같은 오픈소스 all-atom diffusion 모델에 추가 학습 없이 바로 적용 가능한 가속 sampling 기법을 제공함으로써, 단백질뿐 아니라 소분자·핵산을 포함한 다양한 biomolecular 시스템에 대해 저비용으로 conformational ensemble 및 folding/unfolding trajectory를 얻을 수 있게 되어 신약개발 등 구조 기반 응용에 실용적 가치가 크다.
• Approach: 저자들은 Boltz-2의 AtomDiffusion 모듈에 이미 파라미터화되어 있는 denoiser를 이용해 score function을 유도하고, 이를 walk-jump sampling(SDE 기반 Langevin dynamics로 walk, Tweedie's formula 기반 jump)에 적용하여 별도 학습 없이 ensemble을 빠르게 생성한다.

Achievement

Figure 2. Structural snapshots from the Boltz-Jump trajectory for

1. 10배 가속 sampling: 추가 학습 없이 Boltz-2 기반으로 conformational ensemble 생성 속도를 최대 10배 향상시켰다.
2. Ensemble 특성 재현 향상: ATLAS, mdCATH 데이터셋에서 노출 잔기 예측, 약하고 일시적인 접촉(weak and transient contacts) 등에서 기존 Boltz-2 반복 diffusion sampling 대비 유의미하게 개선된 재현력을 보였다.
3. 범용성 확보: 오픈소스 Boltz-2를 직접 활용함으로써 단백질, 소분자, 핵산 등 Boltz-2가 지원하는 모든 biomolecular 시스템에 대한 ensemble sampling과 user-defined potential을 이용한 sampling 유도(guidance)를 지원한다.
4. 새로운 능력 확장: 작은 단백질에 대한 folding 및 unfolding trajectory sampling을 가능하게 하여 기반 모델(base model)이 갖지 못한 새로운 capability를 제공한다.

How

Figure 1. Overview of the walk-jump sampling process. The

• Boltz-2의 AtomDiffusion 모듈에서 학습된 denoiser $\hat{x}(y;\sigma)$를 Tweedie's formula($\hat{x}(y;\sigma)=y+\sigma^2\nabla_y\log p_Y(y)$)를 통해 score function으로 변환
• 초기 clean 구조 $x(0)$에 노이즈를 더해 초기 sample $y(0)=x(0)+\sigma\epsilon(0)$ 생성
• Langevin dynamics를 나타내는 SDE($dy=v_y dt$, $dv_y=\nabla_y\log p_Y(y)dt-\gamma v_y dt+M^{-\frac12}\sqrt2 dB_t$)를 BAOAB discretized solver로 수치적으로 적분하여 noisy distribution $p_Y$에서 다수의 sample $y^{(1)},...,y^{(N)}$을 walk 단계로 얻음
• 각 $y_i$를 denoiser를 통해 clean distribution $p_X$로 jump back시켜 conformer $\hat{x}_i$ 획득
• denoiser는 Boltz-2의 Pairformer stack이 계산한 sequence/pair representation에 조건화되어 있어 별도 재학습 없이 그대로 사용
• ATLAS(CHARMM36m force field, 1390개 monomeric chain)와 mdCATH(CHARMM22* force field, 320K 온도) 데이터셋으로 벤치마크하고, AlphaFlow가 정의한 구조적 타당성 지표(precision, recall, diversity of lDDT) 등을 사용해 평가

Originality

• JAMUN에서 제안된 walk-jump sampling 프레임워크를 별도 재학습 없이 사전학습된 대규모 all-atom biomolecular 구조예측 모델(Boltz-2)에 직접 적용한 최초의 시도로 보임
• 소규모 peptide MD 데이터로 학습해야 했던 기존 walk-jump 기반 접근(JAMUN)의 확장성 한계를 오픈소스 사전학습 diffusion 모델 재사용으로 우회한 아이디어
• 단백질뿐 아니라 소분자, 핵산 등 Boltz-2가 지원하는 다양한 시스템 전반에 동일한 sampling 가속 기법을 적용 가능하게 한 범용적 설계
• user-defined potential을 통한 sampling 유도(guidance) 및 folding/unfolding trajectory 생성이라는 부가적 capability를 함께 제시

Limitation & Further Study

• 본 방법은 diffusion sampling의 계산 비용 문제만 해결할 뿐, 저자들도 언급하듯 모델이 static 3D 구조 위주로 학습되어 MD 시뮬레이션 분포와 정확히 일치하지 않는다는 근본적 분포 불일치 문제는 해결하지 못함
• 발췌된 본문에서 정량적 벤치마크 결과(수치, 표)가 충분히 제시되지 않아 성능 향상의 구체적 크기와 통계적 유의성을 판단하기 어려움
• BAOAB solver의 하이퍼파라미터(마찰 γ, 질량 M, step size 등) 선택이 결과 품질에 미치는 영향에 대한 민감도 분석이 필요해 보임
• 후속 연구로 학습 기반 보정(fine-tuning) 또는 MD 데이터와의 분포 정합을 위한 추가 guidance 기법 결합을 고려할 수 있음

Evaluation

총평: 추가 학습 없이 기존 오픈소스 diffusion 구조예측 모델의 sampling 효율을 크게 개선하는 실용적이고 일반화 가능한 아이디어로, workshop 논문 수준에서 임팩트가 크지만 제공된 발췌본만으로는 정량적 검증의 충분성을 판단하기 어렵다.