Learning Compressed Shape-Aware Molecular Representations for Virtual Screening

Motivation

• Known: 3D shape 기반 virtual screening은 2D fingerprint 방식보다 구조적으로 다양하지만 pharmacophore가 유사한 분자를 발굴할 수 있음이 알려져 있으며, ROCS/FastROCS 같은 도구가 이를 구현하고 있다. 최근에는 contrastive learning으로 2D 그래프와 3D conformer 간의 관계를 학습하여 geometric awareness를 주입하는 연구들(3D-Infomax, GraphMVP 등)과, shape similarity를 예측하는 학습 기반 표현(Atsango et al., Sellner et al.)도 제안된 바 있다.
• Gap: 기존 shape 기반 스크리닝은 분자당 수백 개의 conformer 샘플링과 쌍별 정렬이 필요해 십억 개 규모 라이브러리에는 계산적으로 감당하기 어렵고, 기존 shape-aware 학습 표현 연구들은 billion-scale 확장성을 다루지 않거나 post-hoc quantization에 의존해 encoder와 discretization이 분리 최적화되는 한계가 있다.
• Why: 십억~조 단위로 성장하는 Enamine REAL Space 같은 가상 라이브러리를 대상으로 shape 기반 스크리닝을 실용적으로 수행하려면, conformer 생성·정렬 없이 빠르고 압축된 검색이 가능한 방법이 필수적이며, 이는 생성 화학(generative chemistry)과 구매 가능한 화합물(purchasable chemical matter)을 연결하는 새로운 응용도 가능하게 한다.
• Approach: 차별화 가능한 Spearman correlation 기반 rank-preserving contrastive learning으로 2D 그래프 임베딩이 3D shape overlap 순위를 보존하도록 학습하고, encoder와 two-level IVF-PQ discretization을 end-to-end로 공동 최적화하는 quantization-aware training을 결합한다.

Achievement

Figure 2. Compression performance comparison between our quantization aware training (QAT) method and a FAISS IndexIVFPQ

1. Shape 상관성 확보: 학습된 임베딩의 cosine similarity가 실제 3D shape overlap 점수와 R=0.86의 높은 상관관계를 보여 2D 그래프만으로 3D shape 유사도를 근사함을 입증했다.
2. 압축 효율 개선: end-to-end quantization-aware training이 동일한 검색 품질에서 post-hoc k-means quantization 대비 약 4배 더 나은 압축률을 달성했다.
3. 초고속 대규모 검색: 단일 GPU 노드에서 100억(10 billion) 개 이상의 분자를 1초 이내에 검색할 수 있어, 전통적 방법 대비 10^8배 이상의 속도 향상을 달성했다.
4. 실용적 응용 제시: AI로 생성된 리간드(Flowr-generated ligands)에 대해 구매 가능한 shape-analog 화합물을 SAND로 검색해 in silico 설계에서 실험 검증으로 이어지는 실용적 경로를 제시했다.

How

Figure 3. Shape overlay scores vs cosine similarities of the learned SAND representations between query molecules (Left:

• Graph neural network encoder로 2D 분자 그래프를 연속 벡터 표현으로 인코딩
• ROCS/ROSHAMBO2 등 shape alignment 도구로 계산된 pairwise 3D shape overlap 점수를 ground-truth로 사용
• 임베딩 간 cosine similarity와 ground-truth shape overlap 간의 순위를 보존하도록 differentiable Spearman correlation loss를 활용한 contrastive learning 수행
• 두 단계(inverted file index + product quantization)로 구성된 IVF-PQ discretization을 neural discrete representation learning 기법으로 encoder와 jointly 학습 (quantization-aware training)
• 학습된 모델을 FAISS 등 vector database framework에 통합하여 ANN lookup 기반의 대규모 검색 파이프라인 구축
• 압축 성능은 post-hoc quantization과 비교, shape 상관성은 3D overlay score와의 상관계수로 평가

Originality

• Spearman correlation을 differentiable하게 만들어 contrastive learning의 loss로 직접 사용함으로써 rank-preserving 표현 학습을 구현한 점
• encoder와 IVF-PQ quantization codebook을 분리하지 않고 end-to-end로 공동 최적화하는 quantization-aware training scheme을 shape-aware 분자 표현 학습에 최초로 결합
• 학습된 표현을 실제 FAISS 기반 vector database에 통합해 100억 개 규모의 실질적 검색 시스템으로 구현하고 검증한 점
• 생성 화학(generative chemistry) 결과물에 대한 구매 가능한 shape-analog 탐색이라는 새로운 실용적 응용 시나리오 제시

Limitation & Further Study

• 3D shape overlap ground-truth 자체가 ROCS/ROSHAMBO2와 같은 conformer 샘플링 기반 도구에 의존하므로, 이 도구들의 conformer 생성 편향이나 부정확성이 SAND 학습 신호에 그대로 전이될 수 있음
• R=0.86 상관관계가 모든 화학 공간(예: 매우 유연한 분자, 특이 골격)에 걸쳐 일관되게 유지되는지에 대한 세부 분석이 제한적일 수 있음
• 후속 연구로는 pharmacophore 특징이나 정전기적 특성 등 shape 외 추가 3D 특성을 포함한 확장, 그리고 다양한 target class에 대한 실제 hit rate 검증(생물학적 활성 검증)이 필요
• quantization-aware training의 계산 비용 및 학습 안정성에 대한 상세한 분석이 추가로 요구됨

Evaluation

총평: 3D shape 기반 virtual screening의 확장성 문제를 differentiable rank-preserving contrastive learning과 quantization-aware training의 결합으로 실질적으로 해결한 인상적인 연구로, 실제 십억 규모 검색 성능과 오픈소스 공개까지 갖춘 실용성이 돋보인다.