Diamond Maps for Protein Binder Design: Inference-Time Scaling Survives Stochastic Flow Map Distillation

Motivation

• Known: 생성형 단백질 모델(RFDiffusion, Genie 3, Proteina-Complexa)은 Sequential Monte Carlo나 reward-guided sampling 같은 inference-time scaling 기법과 결합될 때 hallucination 기반 방법(BindCraft, BoltzDesign1)과 경쟁하거나 능가할 수 있다. 이러한 기법들은 AlphaFold-Multimer 같은 구조 예측 oracle을 reward로 사용해 value function을 추정하며, 이 오라클 평가 비용이 전체 계산 비용을 지배한다.
• Gap: Flow map distillation(예: Lagrangian distillation)은 다단계 probability flow ODE를 단일 네트워크 평가로 압축해 생성 속도를 높일 수 있지만, 표준 flow map은 결정론적이어서 주어진 xt에 대해 posterior p1|t(z|xt)를 단일 점으로 붕괴시켜 value function 추정에 편향(Jensen gap)을 유발하고, 결과적으로 inference-time scaling의 이점을 상쇄한다.
• Why: 구조 예측 oracle 평가가 계산 비용을 지배하는 상황에서 생성 모델을 가속화하는 것은 그 절감분을 더 많은 search rollout에 재투자할 때만 의미가 있는데, 결정론적 distillation은 이 재투자의 근거가 되는 정확한 value function 추정 능력 자체를 훼손하므로, stochastic distillation이 inference-time scaling을 보존하는지 여부는 실용적 단백질 binder 설계 파이프라인의 실현 가능성에 직결되는 중요한 문제이다.
• Approach: GLASS Flows 기반의 Posterior Diamond Maps를 Proteina-Complexa의 joint latent 및 Cα flow에 적용하여, posterior p1|t(z|xt)를 단일 forward pass로 샘플링하는 stochastic flow map student를 distill하고, 이를 AlphaFold-Multimer interface confidence를 reward로 사용하는 SMC와 결합해 PD-L1 binder 설계 문제에서 검증한다.

Achievement

1. Posterior Diamond Map distillation: Proteina-Complexa의 joint latent와 Cα space에서 동작하는 Posterior Diamond Map을 distill하고, 이를 이질적(heterogeneous) scheduler를 갖는 GLASS 재매개변수화로 확장했다.
2. AF2M 기반 SMC 절차: AlphaFold-Multimer interface predicted aligned error(ipAE)를 reward로 사용하는 SMC inference 절차를 AF2M 벤치마크 프로토콜 하에서 평가했다.
3. Jensen gap 정량화: 실제 teacher에 대해 GLASS 재매개변수화가 결정론적 denoiser와 GLASS posterior mean 사이에 0.117–0.162의 non-zero relative L2 Jensen gap을 보임을 실증했다.
4. 탐색 전략 비교: 결정론적 고차 적분(Heun)이 동일 계산량에서 stochastic search보다 낮은 성능에서 정체되는 반면, 5K-step student를 사용한 PDM-SMC는 search width 증가에 따라 tail metric이 단조적으로 개선되며, NFE=400에서 single-pass의 best reward를 능가함을 보였다.

How

• Stochastic interpolant/flow matching 프레임워크(Albergo et al.)를 배경으로, marginal velocity field ut와 denoiser Dt를 정의하고 flow map Xs,t를 Lagrangian PDE로 특성화한다.
• GLASS Flows의 inner ODE를 이용해, 고정된 (xt, t)에서 posterior p1|t(z|xt)로 이어지는 inner stochastic interpolant와 inner velocity field ¯us를 사전학습된 teacher denoiser의 닫힌 형태 함수로 유도한다(재학습 불필요).
• Proteina-Complexa가 latent와 Cα 성분에 서로 다른 scheduler scaling을 적용하는 점을 고려해, GLASS 구성을 성분별(componentwise)로 확장하되 conditioned denoiser query는 결합된(joint) forward pass로 처리한다.
• Boffi ansatz를 사용해 경계 조건을 아키텍처적으로 강제하는 Posterior Diamond Map Xθ_{s,s'}를 정의하고, Proteina-Complexa backbone을 재사용하며 teacher 가중치로 초기화, conditioning encoder는 초기 5K step 동안 동결한다.
• Lagrangian loss(forward-mode automatic differentiation으로 ∂s' 계산)를 사용해 GLASS inner velocity에 대해 student를 distill하며, Boffi et al.의 strip-then-full training schedule을 (s, s')에 적용한다.
• Distill된 student로부터 posterior 샘플 zk를 저비용으로 얻어 value function 추정치 V̂t(xt)를 log-mean-exp 형태로 계산하고, 이를 SMC의 particle reweighting/resampling에 결합하여 PD-L1 binder 설계에서 AlphaFold-Multimer interface confidence reward로 평가한다.

Originality

• Deterministic flow map distillation이 inference-time scaling에 필요한 value function 추정을 편향시킨다는 문제를 명확히 정식화하고, 이를 stochastic flow map(Diamond Maps)으로 해결하려는 최초 시도이다.
• Posterior Diamond Maps를 단백질 생성 모델(Proteina-Complexa)에 최초로 적용하며, joint latent와 explicit Cα trace라는 이질적 구조를 갖는 도메인에 GLASS 재매개변수화를 성분별로 확장한 것이 기술적 독창성이다.
• 재학습 없이 frozen teacher의 denoiser로부터 GLASS inner velocity를 closed-form으로 유도하고, 이를 Boffi ansatz 기반 아키텍처와 Lagrangian distillation loss로 단일 forward pass 샘플러에 압축한 파이프라인 설계가 새롭다.

Limitation & Further Study

• 실험이 PD-L1이라는 단일 binder 설계 문제에 국한되어 있어, 일반화 가능성 및 다양한 target에 대한 검증이 필요하다.
• 5K-step student라는 비교적 제한된 학습 규모에서의 결과이며, 더 긴 distillation 학습이나 더 큰 student capacity에서의 성능 변화가 불명확하다.
• Jensen gap 정량화가 특정 노이즈 레벨/조건에서의 relative L2 측정에 의존하므로, 이 지표가 최종 binder 설계 품질(성공률, 실험적 검증)과 얼마나 상관되는지에 대한 추가 분석이 필요하다.
• AlphaFold-Multimer reward 자체의 신뢰성(oracle bias)에 대한 논의가 부족하며, 향후 연구로 다른 reward 모델이나 실험적 검증과의 비교가 필요하다.

Evaluation

총평: Inference-time scaling과 flow map distillation 간의 근본적 긴장을 명확히 짚어내고, stochastic Diamond Maps를 통해 이를 해결하는 참신하고 잘 설계된 워크숍급 연구로, 단일 사례 연구라는 한계가 있지만 향후 생성형 단백질 설계 가속화 연구에 중요한 방향을 제시한다.