Inference-time optimization for experiment-grounded protein ensemble generation

Motivation

• Known: AF3와 같은 sequence-conditioned diffusion model은 강력한 구조 예측 prior를 제공하지만, NOE 기반 NMR 제약이나 X-ray 결정학 밀도와 같은 실험적 우도(likelihood) 항을 reverse diffusion 과정 중 좌표(coordinate)에 gradient guidance로 주입하는 방법들이 최근 제안되었다.
• Gap: 기존 coordinate-space guidance 방법은 고정된 sampling horizon(diffusion length)에 의존하고 초기화(initialization)에 민감하여 수렴이 불안정하며, ensemble을 열역학적으로(thermodynamically) 타당한 상태 가중치로 변환하는 메커니즘이 없어 물리적으로 비합리적인 결과를 산출한다.
• Why: 단백질 기능은 단일 구조가 아닌 동적 conformational ensemble에 의해 결정되므로, binder engineering·효소 활성 부위 최적화·단백질-단백질 인터페이스 설계 등 실제 설계 과제에서 실험과 부합하고 물리적으로 타당한 ensemble을 생성하는 능력은 핵심적으로 중요하다. 또한 ipTM과 같은 설계용 confidence metric의 취약점을 드러내는 것은 binder 설계에서 false discovery를 줄이는 데 실질적 함의를 갖는다.
• Approach: 좌표를 직접 perturb하는 대신 AF3의 Pairformer conditioning embedding(latent representation)을 실험적 우도의 gradient로 업데이트하여 ensemble log-likelihood를 극대화하고, 여기에 force-field 기반 prior와의 product distribution에서 Boltzmann-weighted sampling을 결합하는 이중 전략을 취한다.

Achievement

Figure 4. Summary statistics for NMR and X-ray crystallography benchmarks. Violin plots compare AlphaFold3, Guidance, an

1. IT-Optimization 프레임워크 제안: diffusion length에 독립적이고 초기화 편향이 없으며 외부 제약을 쉽게 결합할 수 있는 representation-space conditioning 기법을 도입해 coordinate-space guidance 대비 일관되게 우수한 성능을 보였다.
2. Energy-weighted Boltzmann sampling: AF3 structural prior와 AMBER99SB 같은 force-field prior의 product distribution에서 샘플링하여 실험 우도와 열역학적 타당성을 동시에 만족하는 ensemble을 생성했다.
3. NMR 벤치마크 개선: NMRDB 데이터셋에서 NOE restraint violation을 NOE-guided AF3 대비 크게 감소시키고, 동시에 더 낮은 effective energy를 갖는 ensemble을 산출했다.
4. X-ray crystallography 개선: altloc 및 protein-bound peptide 벤치마크에서 Rwork/Rfree 감소, 국소 밀도 정합 향상, 시드 간 재현성 향상을 달성했으며, 말단 restraint 없이도 peptide 모델링과 분리된 altloc 복원이 가능함을 보였다.
5. ipTM 취약점 발견: MSA embedding을 perturb하여 ipTM 점수를 인위적으로 부풀릴 수 있음을 밝혀, 현재 design metric의 취약성과 false discovery 위험을 노출시켰다.

How

Figure 2. Nested inference-time optimization of AlphaFold3 conditioning embeddings. (Left) The outer loop runs K diffusi

• AF3를 학습된 구조적 prior로 취급하고, 미분 가능한 목적함수(실험 유래 우도, confidence metric 등)의 gradient로 Pairformer trunk embedding Z를 업데이트하는 nested loop 구조(outer loop: K개의 독립 reverse diffusion 실행, inner loop: 각 diffusion 내 embedding gradient 업데이트)를 설계
• 이전 diffusion 실행에서 최적화된 conditioning embedding을 다음 실행의 초기값으로 이어받아 diffusion trajectory 전반에 걸쳐 embedding을 지속적으로 정제
• force-field 기반 에너지(AMBER99SB)를 이용한 Boltzmann weight w를 도입하여 product-of-experts 방식으로 구조적 prior와 물리적 prior를 결합, ensemble 통계를 열역학적으로 타당하게 재가중
• NMRDB(NMR), X-ray 결정학(altloc, protein-bound peptide) 벤치마크와 ipTM 기반 protein-protein complex 벤치마크(낮은 baseline ipTM 케이스, p53 N-terminal domain 등)에서 정량 평가 수행

Originality

• 기존의 좌표 공간(coordinate-space) gradient guidance 패러다임에서 벗어나, diffusion model의 conditioning latent(Pairformer embedding)를 직접 최적화 대상으로 삼는 새로운 inference-time 전략을 제시
• 구조적 prior(AF3)와 force-field 기반 물리적 prior를 product distribution 형태로 결합하여 Boltzmann-weighted ensemble sampling을 구현한 점이 독창적
• 내부 diffusion iteration 및 기반 생성 모델에 대해 model-agnostic한 meta-guidance layer로 설계되어 기존 coordinate-level steering과 조합 가능
• ipTM 같은 model-internal confidence metric을 최적화 목표로 사용할 때 발생하는 취약점(embedding perturbation을 통한 confidence inflation)을 실험적으로 규명, design metric 신뢰성 문제를 새롭게 제기

Limitation & Further Study

• 초록과 발췌된 본문 기준으로 large-scale 단백질 및 diverse fold에 대한 일반화 가능성, 계산 비용(예: nested loop로 인한 inference 시간 증가)에 대한 구체적 분석이 충분히 드러나지 않음
• ipTM 취약점 발견은 문제 제기 수준이며, 이를 방지하기 위한 구체적 mitigation 전략이나 robust metric 제안은 후속 연구 과제로 남아있음
• force-field 기반 reweighting이 AMBER99SB에 의존하므로 force-field 선택에 따른 편향 가능성 및 다른 force-field와의 비교 검증이 필요
• 실험 데이터(NMR, X-ray)가 제한된 벤치마크 세트에 기반하므로 cryo-EM 등 다른 실험 modality나 더 다양한 단백질 패밀리에 대한 추가 검증이 필요

Evaluation

총평: Latent space에서의 inference-time optimization과 Boltzmann-weighted sampling을 결합해 실험 부합성과 물리적 타당성을 동시에 개선한 견고한 프레임워크이며, ipTM 취약점 발견은 단백질 설계 커뮤니티에 중요한 경고를 제공하는 의미 있는 부가 기여이다.