Drug-TTA: Test-Time Adaptation for Drug Virtual Screening via Multi-task Meta-Auxiliary Learning
Shen, Yuan, Ma, Du, Huang, Wang | Fudan University | ICML 2025, PMLR 267 | GitHub: ShenAoAO/Drug-TTA
🎯 仮想スクリーニングにテストタイム適応を初導入。5種補助タスク+メタ学習でゼロショット精度を大幅改善
① 背景と課題: 訓練・推論間のドメインギャップ

DrugCLIP 等の contrastive 仮想スクリーニングは、学習バッチ内の「他ポケットのリガンド」をネガティブペアとして使用する。しかし実際の DUD-E 等のベンチマークでは、ネガティブ(decoy)は同ポケットに結合しない類似構造であり、はるかに困難。

訓練時: positive と negative の特徴距離が大きく分離(分類が容易)
推論時: positive と negative の特徴距離が混在(実際のスクリーニングは困難)
ゼロショット問題: テスト標的の分布を事前に知ることができない

→ TTA(テストタイム適応)で各テストインスタンスにモデルを動的適応させる

② 手法: 5補助タスク

Fine-grained(3種):

🔬 座標再構成: masked atom の 3D 座標予測
📏 距離行列再構成: 原子ペア間距離予測
⚗️ 原子タイプ再構成: masked 原子の元素予測

Coarse-grained(2種):

📊 KL divergence 最小化: 潜在分布→一様化(多様な表現学習)
🔗 SimCLR: 元構造 vs masked 版を2ビューとして対照学習
② 手法: メタ学習 + ALBM
【学習時 meta-learning】
① 特徴抽出 → ② 補助 loss で layer norm 更新
→ ③ 更新済みモデルで主タスク loss 計算
→ ④ 全パラメータ更新
【推論時 TTA】
各テスト pocket-mol ペアで
① 補助 loss → layer norm 動的更新
→ ② 類似度スコア計算

ALBM: top/middle/final 層特徴量の MLP で各 loss の重みを動的生成

③ 全体構成
  • Primary branch: pocket-molecule contrastive scoring(DrugCLIP ベース)
  • Auxiliary branch: 5補助タスク + ALBM(pocket/molecule 各々)
  • Meta-learning: 訓練・推論で同一パイプライン
  • 更新対象: layer norm のみ(推論コスト最小化)
+9.86%
TTA なし DrugCLIP 比の平均 AUROC 向上(5ベンチマーク)
④ 主要結果 (a) DUD-E ベンチマーク
DUD-E AUROC (%) Glide-SP Planet DrugCLIP Drug-TTA 50 70 90 100 67.4% 71.5% 80.9% 93.2% +12.2%
④ 主要結果 (b) LIT-PCBA: 最難ベンチマーク
LIT-PCBA BEDROC (%) Glide-SP DrugCLIP Drug-TTA 0 25 50 0.5% 6.2% 45.1% +38.9% (7倍改善!)
④ 主要結果 (c) 5ベンチマーク比較
BenchmarkDrugCLIP AUROCDrug-TTA AUROC改善
DUD-E80.93%93.16%+12.23%
LIT-PCBA57.17%71.24%+14.07%
AD81.19%92.62%+11.43%
DEKOIS 2.077.98%83.61%+5.63%
CASF-201685.92%91.88%+5.96%
+9.86%
5ベンチマーク平均 AUROC 改善
④ 主要結果 (d) アブレーション検証
設定DUD-E AUROC
Drug-TTA(フル)93.16%
w/o meta-learning82.22%
w/o coarse-grained tasks
w/o fine-grained tasks
ALBM → fixed weight83.63%
ALBM → weight=171.13%
DrugCLIP(baseline)80.93%

メタ学習と ALBM がともに重要な貢献

⑤ テイクホームメッセージ
🎯 TTA を創薬 VS に初適用
推論時に各テストインスタンスへオンライン適応。ゼロショット問題に根本的に対応する新パラダイム。
📊 LIT-PCBA で劇的改善
最難ベンチマーク LIT-PCBA で EF(1%) が DrugCLIP 比 約 9 倍向上。真の仮想スクリーニング課題で有効。
⚙️ ALBM が key component
固定 weight より動的 ALBM が優れる(83.63% vs 93.16%)。multi-scale 特徴量から適応的に weight を計算。
🔗 DrugCLIP への drop-in 拡張
DrugCLIP encoder を継承し auxiliary branch を追加するだけで実装可能。OSS 公開済み。
ケムインフォマティクスへの応用
適用先ユースケース
lib/dockingUniDockRunner スコア後の TTA ベース reranker として実装
lib/dockingProLIFCalculator の接触 FP を補助タスク入力として活用

新規標的(ゼロショット)でのスクリーニング精度向上に特に有用。訓練時に見たことのない標的でも動的適応で精度維持。

本研究のインパクト
  • TTA を仮想スクリーニングに初導入した先駆的研究(ICML 2025)
  • 5ベンチマーク全ゼロショット設定で SOTA 達成
  • LIT-PCBA という厳密なベンチマークでの10倍近い EF 改善
  • GitHub 公開: https://github.com/ShenAoAO/Drug-TTA