Towards Sub-Second Molecular Docking as a Structural Primitive: A Quantized Consistency Diffusion Framework

Motivation

• Known: co-folding 모델(AlphaFold3, Chai-1, Boltz, Protenix 등)은 protein-ligand complex를 공동으로 모델링하여 rigid-body docking 대비 global binding topology와 atomic geometry를 더 잘 포착하며 최신 분자 도킹의 high-fidelity 방향을 대표한다.
• Gap: co-folding 모델은 큰 pair representation, deep attention block, iterative generative refinement에 의존하여 추론이 순차적이고 비용이 크며, 이는 정확도는 높지만 screening-scale의 실시간 인프라로 배포하기에는 부족한 accuracy-latency gap을 만든다. 또한 blind docking과 interface-informed docking처럼 서로 다른 수준의 structural prior 하에서 평가가 이루어져 모델 성능과 입력 정보량이 혼동되는 문제가 존재한다.
• Why: AI agent가 도구를 조율하고 고빈도 연구 루프를 구동하는 agent-centered scientific discovery 패러다임에서는 분자 도킹이 실시간으로 반복 호출되는 구조적 skill이 되어야 하므로, latency는 단순한 공학적 디테일이 아니라 도킹이 능동적 추론 루프 안에 남을 수 있는지를 결정하는 핵심 요소이다.
• Approach: diffusion dynamics를 few-step inference로 압축하는 reconstruction-anchored consistency tuning(PCR)과, geometry-sensitive 연산은 BF16으로 보존하면서 compute-intensive linear transformation만 선택적으로 quantization하는 Residual-Safe Quantization을 결합하여, 정확도를 유지하면서 sub-second 추론을 달성한다.

Achievement

1. 정확도 저하 없는 sub-second 추론: interface-informed protocol 하에서 state-of-the-art 도킹 정확도를 유지하면서, blind docking에서도 경쟁력 있는 성능을 보인다.
2. 극적인 속도 향상: 256-token 규모의 대표 복합체에 대해 5개 conformation을 단일 NVIDIA H20 GPU에서 0.17초에 생성하여, 동일 벤치마크 설정에서 AlphaFold3 대비 300배 이상의 speedup을 달성했다.
3. information-level-matched 평가 프로토콜 정립: blind와 interface-informed 설정을 분리하여 비교함으로써 모델 능력과 입력 prior 강도를 혼동하지 않는 평가 체계를 제시했다.

How

• 통제된 비교를 통해 Linear parameterization이 cosine/EDM-style denoising parameterization보다 discretization-friendly한 transport target과 update rule을 제공함을 확인하여, few-step 분자 생성이 trajectory-geometry 문제임을 규명
• Progressive Consistency Regularization(PCR): reconstruction 기반 anchor를 유지하는 consistency distillation 전략(Song et al., 2023 기반)으로 diffusion sampling을 소수 step으로 압축
• Residual-Safe Quantization: Pairformer와 token-level Diffusion Transformer(DiT) 블록 내 compute-intensive linear transformation을 선택적으로 quantize하되, 고정밀도가 요구되는 residual stream과 geometry-sensitive 연산은 BF16으로 유지
• 두 구성 요소를 결합하여 co-folding 프레임워크 전체를 재설계, 동일 input-shape 배포 프로토콜 하에서 latency 벤치마크 수행

Originality

• 분자 도킹을 "정확도" 관점이 아니라 agent-centered discovery 인프라의 "실시간 callability" 관점에서 재정의한 문제 설정의 참신성
• few-step 생성 문제를 단순 solver step 축소가 아닌 trajectory geometry(parameterization 선택) 문제로 재해석한 관점
• diffusion 기반 co-folding에 특화된 mixed-precision quantization 전략(Residual-Safe Quantization)으로 geometry-sensitive 연산을 선별적으로 보호하는 설계
• blind/interface-informed 설정을 분리하는 information-level-matched 평가 프로토콜 제안

Limitation & Further Study

• 발췌된 내용만으로는 PCR의 구체적 손실 함수, consistency distillation 학습 절차, quantization 적용 레이어 선택 기준 등 세부 구현이 충분히 검증되지 않음
• 단일 GPU(H20) 및 특정 256-token 규모 복합체에서의 speedup 수치이므로 다양한 크기의 복합체, 다른 하드웨어 환경에서의 일반화 가능성에 대한 추가 검증 필요
• blind docking에서는 "경쟁력 있는" 수준에 그친다고 언급되어 있어, interface-informed 대비 blind 설정에서의 성능 저하 원인과 개선 방향에 대한 후속 연구가 필요
• quantization으로 인한 정확도 손실이 다양한 target class(예: 큰 리간드, flexible loop 등)에서 어떻게 나타나는지에 대한 세밀한 분석이 부족해 보임

Evaluation

총평: 분자 도킹을 정확도뿐 아니라 실시간 배포 가능성 관점에서 재조명하며, consistency regularization과 mixed-precision quantization을 결합해 300배 이상의 speedup과 경쟁력 있는 정확도를 동시에 달성한 실용적이고 시의적절한 연구이다. Agent 기반 과학 발견 인프라 구축에 실질적으로 기여할 수 있는 잠재력이 크다.