Scaling Pocket Docking with Data Augmentation and Heterogeneous Equivariant Graph Attention

Motivation

• Known: DiffDock 계열의 diffusion 기반 docking 모델들은 SE(3) manifold 상의 reverse-diffusion을 통해 pose를 생성하며, DiffDock-Pocket, ArtiDock, DTMol, DiffDock-RL++, SigmaDock 등 pocket-specific docking으로 발전해왔다. 이들 모델은 all-atom 모델링과 confidence 기반 ranking을 통해 pose 품질을 평가한다.
• Gap: 기존 방법들은 좁은 학습 분포로 인한 data-scarcity, all-atom 모델링의 높은 연산 비용, 그리고 RMSD 기반 ranking이 물리적으로 타당한 pose를 제대로 선택하지 못해 Top-1과 Oracle 성공률 간 큰 격차가 존재하는 문제를 해결하지 못했다.
• Why: Top-1과 Oracle 성능 간 gap을 줄이는 것은 실제 신약 개발에서 다수의 후보 pose 중 최적의 pose를 자동으로 선택하는 능력을 의미하며, 이는 고비용의 실험적 검증 없이도 신뢰할 수 있는 virtual screening을 가능하게 하므로 실용적으로 매우 중요하다.
• Approach: PLINDER와 SAIR 데이터를 leakage 제거, 구조 품질 필터링, dataset-specific pocket cutoff로 curation하여 diffusion 모델을 확장 학습시키고, confidence 모델의 fully connected tensor-product convolution을 HeteroEGA로 대체하여 ligand·receptor-residue·receptor-atom 그래프 간 interaction-specific attention을 수행함으로써 pose ranking 정확도를 개선한다.

Achievement

Figure 1. Architecture overview. (1) DiffDock’s TPConv backbone uses 9 tensor product convolutions per layer across thre

1. 대규모 데이터 증강: PLINDER(약 10만 개 시스템)와 SAIR(약 87만 개로 정제된 시스템, 524만 샘플)를 결합하여 학습 분포를 크게 확장하고 leakage를 제거함으로써 일반화 성능을 개선했다.
2. HeteroEGA confidence 모델: 기존 TPConv 기반 confidence 모델을 attention 기반 HeteroEGA로 대체하여 ligand, receptor-residue, receptor-atom 간 이질적 상호작용을 명시적으로 모델링했다.
3. SOTA 성능 달성: PoseBusters-308 벤치마크에서 Top-1 success rate 81.85%, Oracle success rate 94.51%를 달성하여 SigmaDock과 DiffDock-RL++를 능가했다.
4. 효율적 후처리 대비 우위: Vina minimization과 GNINA reranking을 결합한 물리 기반 refinement track과 비교했을 때, HeteroEGA confidence 모델이 근소하게 더 우수하면서도 훨씬 빠른 ranking 속도를 보였다.
5. 연산 효율성 개선: cuEquivariance 하드웨어 가속 커널을 활용해 all-atom pocket docking 모델의 추론 속도를 크게 향상시켰다.

How

Figure 1. Architecture overview. (1) DiffDock’s TPConv backbone uses 9 tensor product convolutions per layer across thre

• PLINDER-TIME split과 SAIR를 결합한 multi-source dataset을 구성하고, PDB/CCD 코드 및 상동 서열 매칭을 통한 two-tier leakage 필터링을 적용함
• RDKit으로 생성한 초기 ligand conformer와 ground truth 간 RMSD가 2Å를 초과하는 샘플을 제거하는 structural quality filtering을 수행함
• 데이터 소스별로 pocket extraction radius를 개별 튜닝하는 dataset-specific pocket cutoff 최적화를 도입함
• DiffDock/DiffDock-Pocket의 diffusion 기반 백본을 유지하되, confidence 모델의 fully connected tensor-product convolution을 HeteroEGA로 교체하여 ligand, receptor-residue, receptor-atom 그래프 간 interaction-specific attention을 수행함
• cuEquivariance 커널을 도입해 all-atom equivariant 모델의 추론 효율을 높임
• Vina minimization 후 GNINA reranking을 적용하는 독립적 physics-based refinement track을 별도로 평가하여 HeteroEGA 기반 confidence 모델과 비교함

Originality

• 기존에 별도로 다뤄지던 data augmentation(PLINDER+SAIR), 연산 효율화(cuEquivariance), physics-based refinement(Vina+GNINA), attention 기반 ranking(HeteroEGA)을 하나의 통합 pipeline으로 결합한 점이 새로움
• confidence 모델에서 fully connected tensor-product convolution을 대체하는 Heterogeneous Equivariant Graph Attention을 도입하여 ligand, receptor-residue, receptor-atom이라는 이질적 그래프 유형 간 interaction-specific attention을 명시적으로 모델링한 것이 핵심 기여
• 단일 고정 radius가 아닌 dataset-specific pocket cutoff 최적화라는 실용적이지만 기존에 간과되었던 전처리 전략을 제시함

Limitation & Further Study

• 발췌된 본문에서는 HeteroEGA의 구체적인 attention 메커니즘 수식과 세부 아키텍처, ablation 결과가 충분히 제시되지 않아 기술적 세부사항 검증이 제한적임
• PoseBusters-308이라는 단일 벤치마크에 대한 평가만 제시되어 있어, 더 다양한 벤치마크나 실제 실험적 검증(wet-lab)과의 상관성에 대한 추가 검증이 필요함
• confidence 모델이 물리 기반 refinement(Vina+GNINA)를 근소하게 앞서는 수준이며, 두 접근을 결합했을 때의 상승효과나 실패 사례에 대한 분석이 부족해 보임
• 후속 연구로 다양한 transformer 기반 attention 아키텍처의 탐색과, SAIR와 같은 synthetic structure의 품질이 최종 성능에 미치는 영향에 대한 심층 분석이 필요함

Evaluation

총평: 데이터 증강과 attention 기반 equivariant ranking을 결합해 pocket docking의 Top-1/Oracle 성능 격차를 실질적으로 줄인 견고한 엔지니어링 및 방법론적 기여이며, 실용적 가치가 높은 연구이다.