CausalBind: Causal Modeling and Learning for Protein-Molecule Virtual Screening
Organizations: Mohamed bin Zayed University of Artificial Intelligence · Carnegie Mellon University
Abstract
Protein-molecule virtual screening is increasingly cast as a problem of representation learning in a shared embedding space. Existing methods rely on dense holistic alignment, entangling invariant binding determinants with nuisance correlations and limiting transfer to new targets. It has been noted that binding in protein-molecule systems involves sparse cross-modality interactions: binding is governed by a small contact interface and a few decisive local interactions (e.g., hydrogen bonds, hydrophobic contacts, and salt bridges) rather than the global structures of the protein and molecule. We hypothesize that uncovering and leveraging sparse interaction patterns is critical for generalization beyond the training data, as these patterns are reusable and expected to improve performance across different scenarios. In this paper, we aim to identify and leverage sparse interaction patterns, and verify our hypothesis. Since the training data contain only observed binding pairs, we formalize this prior via a V-structure causal model under Heckman-style selection, and establish three theoretical results: (i) the latent concepts of interacting proteins and molecules are not identifiable without appropriate sparsity constraints; (ii) these concepts and their sparse interactions are component-wise identifiable under structural sparsity conditions; and (iii) a low-rank relaxation of these conditions yields subspace identifiability of the concepts and interactions. Inspired by these principles, we propose CausalBind with three implementation variants. Extensive experiments on DUD-E and LIT-PCBA benchmarks show that all variants consistently outperform strong retrieval baselines, with the largest gains on LIT-PCBA early enrichment, and further generalize to target- and scaffold-level out-of-distribution splits. Code is available at https://github.com/lokali/CausalBind.
Figures & tables
| Method | DUD-E ( ) | LIT-PCBA ( ) | ||||
|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF 1% | AUROC | BEDROC 80.5 | EF 1% | |
| Glide-SP ( Friesner et al., 2004 ) | 0.767 | 0.407 | 16.18 | 0.532 | 0.040 | 3.41 |
| DeepDTA ( Öztürk et al., 2018 ) | 0.584 | 0.051 | 2.28 | 0.563 | 0.025 | 1.47 |
| Gnina ( McNutt et al., 2021 ) | 0.782 | 0.299 | 17.73 | 0.609 | 0.054 | 4.63 |
| RTMScore ( Shen et al., 2022 ) | 0.753 | 0.434 | 27.10 | 0.525 | 0.039 | 2.94 |
| TankBind ( Lu et al., 2022 ) | 0.751 | 0.330 | 13.00 | 0.597 | 0.039 | 2.90 |
| Method | DUD-E ( ) | LIT-PCBA ( ) | ||||
|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF 1% | AUROC | BEDROC 80.5 | EF 1% | |
| LigUnity ( Feng et al., 2025 ) | 0.896 0.005 | 0.660 0.018 | 43.14 1.42 | 0.586 0.010 | 0.077 0.001 | 6.46 0.18 |
| HypSeek ( Wang et al., 2026 ) | 0.914 0.013 | 0.690 0.081 | 44.12 6.14 | 0.603 0.009 | 0.067 0.007 | 5.62 0.71 |
| CausalBind SP | 0.939 0.004 | 0.746 0.004 | 47.99 0.39 | 0.635 0.006 | 0.088 0.007 | 7.55 0.73 |
| CausalBind SP | 0.935 0.004 | 0.732 0.014 | 46.93 1.07 | 0.623 0.005 | 0.084 0.004 | 7.17 0.36 |
| CausalBind LR | 0.937 0.004 | 0.732 0.011 | 46.87 0.91 | 0.630 0.006 | 0.085 0.006 | 7.22 0.31 |
| DUD-E | LIT-PCBA | Cross-modal mask statistics | |||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|
| AUC | BEDROC 80.5 | EF@1% | AUC | BEDROC 80.5 | EF@1% | avg | %zeros | Anti. viol. | |||
| 0.938 | 0.746 | 48.14 | 0.617 | 0.080 | 6.83 | 0.93 | 0.0 | 1 | 25.0 | 32,512 | |
| 0.940 | 0.749 | 47.94 | 0.635 | 0.088 | 7.26 | 0.92 | 0.0 | 1 | 25.0 | 32,512 | |
| 0.938 | 0.735 | 47.09 | 0.607 | 0.078 | 6.29 | 0.88 | 0.0 | 1 | 26.0 | 32,512 | |
| 0.935 | 0.744 | 47.69 | 0.639 | 0.097 | 8.56 | 0.61 | 0.0 | 1 | 34.0 | 32,512 | |
| 0.943 | 0.768 | 49.73 | 0.614 | 0.073 | 5.49 | 0.21 | 20.5 | 45.5 | 79.5 | 0 | |
Appendix figures & tables25 assets
Supplementary material from the paper’s appendix.
Appendix
| Physics-based docking | Structure-aware ML scoring | Retrieval-based learning | |
|---|---|---|---|
| Inference input | (protein, ligand) pair + pose search | (protein, ligand) pair, often pose-conditioned | protein and ligand encoded independently |
| Binding score | physics + empirical scoring on the pair | one forward pass through a joint network on the pair | dot product (or cosine) of two precomputed vectors |
| Pre-computable embeddings? | No (pose-coupled) | No (pair-coupled) | Yes (dual-tower) |
| Cost per (P, L) at inference | seconds–minutes (pose search dominates) | 10–100 ms (one full forward) | sub-millisecond (one dot product) |
| Practical library scale | – | – | |
| Where the inductive bias lives | force field + sampling protocol | complex graph / 3D-voxel architecture | shape of the shared embedding space and the similarity used on it |
| Method | Core Design | Relevance Here |
| Physics-based docking systems | ||
| Glide-SP [ Friesner et al., 2004 ] | Hierarchical sampling with empirical scoring function | Industry standard; accurate on small candidate sets but expensive at library scale |
| AutoDock Vina [ Trott and Olson, 2010 ] | Empirical scoring + iterated local search | Open-source, widely used; not designed for ultra-large libraries |
| Structure-aware machine-learning scoring models | ||
| DeepDTA [ Öztürk et al., 2018 ] | Sequence/graph-based affinity prediction without pairwise retrieval design | Useful DTA baseline; not optimized for target-wise early enrichment |
| Gnina [ McNutt et al., 2021 ] | 3D CNN scoring on docking-grid voxels | Strong pose-rescoring; depends on docked poses |
| Variant | Mask object | Forward bottleneck | Closest theory link | Role in the paper |
| CausalBind SP | Dense concept-pair mask with sparsity and hard thresholding. | Concept axis before pooling; each active entry selects a protein-molecule concept pair. | Direct sparse-support form of Theorem 2 ; learned masks can also enter the low-rank regime of Theorem 3 . | Primary method in the main text. |
| CausalBind LR | Rank- concept-basis mask with . | Concept axis before pooling; rank controls the number of interaction subspaces. | Direct implementation of the low-rank subspace-identifiability structure in Theorem 3 . | Appendix variant for explicit low-rank implementation. |
| CausalBind EMB | Dense mask in pooled embedding space. | After mean-pooling; the mask acts on pooled features rather than concept pairs. | Compatible with a compressed interaction view, but not a direct instantiation of concept-pair sparse support. | Appendix alternative and robustness check. |
| Method | Protocol | DUD-E | LIT-PCBA | ||||
|---|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF@1% | AUROC | BEDROC 80.5 | EF@1% | ||
| LigUnity | last | 0.897 | 0.674 | 44.20 | 0.599 | 0.075 | 6.50 |
| best_auc | 0.915 | 0.716 | 46.42 | 0.581 | 0.089 | 8.15 | |
| best_bedroc | 0.896 | 0.671 | 44.00 | 0.599 | 0.076 | 6.38 | |
| HypSeek | last | 0.909 | 0.605 | 37.83 | 0.603 | 0.061 | 4.62 |
| best_auc | 0.913 | 0.629 | 39.63 | 0.608 | 0.063 | 5.13 | |
| Protocol | DUD-E | LIT-PCBA | |||||
|---|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF@1% | AUROC | BEDROC 80.5 | EF@1% | ||
| 32 | last | 0.938 | 0.749 | 48.41 | 0.620 | 0.081 | 6.90 |
| best_auc | 0.939 | 0.750 | 48.38 | 0.625 | 0.083 | 6.85 | |
| best_bedroc | 0.938 | 0.750 | 48.57 | 0.621 | 0.081 | 6.90 | |
| 64 | last | 0.930 | 0.724 | 46.55 | 0.637 | 0.085 | 6.41 |
| best_auc | 0.933 | 0.715 | 45.90 | 0.638 | 0.088 | 7.88 | |
| Protocol | DUD-E | LIT-PCBA | |||||
|---|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF@1% | AUROC | BEDROC 80.5 | EF@1% | ||
| 128 | last | 0.919 | 0.589 | 35.75 | 0.615 | 0.062 | 5.25 |
| best_auc | 0.922 | 0.629 | 39.43 | 0.630 | 0.060 | 4.75 | |
| best_bedroc | 0.920 | 0.592 | 36.03 | 0.620 | 0.062 | 5.35 | |
| 256 | last | 0.912 | 0.600 | 36.93 | 0.614 | 0.062 | 5.68 |
| best_auc | 0.922 | 0.683 | 43.07 | 0.620 | 0.067 | 4.73 | |
| Protocol | DUD-E | LIT-PCBA | |||||
|---|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF@1% | AUROC | BEDROC 80.5 | EF@1% | ||
| last | 0.938 | 0.746 | 48.14 | 0.617 | 0.080 | 6.83 | |
| best_auc | 0.943 | 0.744 | 47.68 | 0.629 | 0.073 | 6.26 | |
| best_bedroc | 0.934 | 0.732 | 47.03 | 0.613 | 0.075 | 6.10 | |
| last | 0.940 | 0.749 | 47.94 | 0.635 | 0.088 | 7.26 | |
| best_auc | 0.941 | 0.750 | 47.75 | 0.628 | 0.086 | 7.41 | |
| Protocol | DUD-E | LIT-PCBA | |||||
|---|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF@1% | AUROC | BEDROC 80.5 | EF@1% | ||
| last | 0.940 | 0.763 | 49.20 | 0.625 | 0.087 | 6.98 | |
| best_auc | 0.946 | 0.776 | 49.99 | 0.616 | 0.082 | 6.53 | |
| best_bedroc | 0.942 | 0.757 | 48.48 | 0.622 | 0.085 | 7.95 | |
| (default) | last | 0.935 | 0.744 | 47.69 | 0.639 | 0.097 | 8.56 |
| best_auc | 0.934 | 0.722 | 46.41 | 0.628 | 0.074 | 5.84 | |
| branch | max | (sum) | zeros | %zeros | ||||
|---|---|---|---|---|---|---|---|---|
| mol | 0.920 | 1.311 | 15068 | 000 0 | 0 0.0% | 1 | 5 | |
| poc | 0.993 | 1.398 | 16267 | 000 0 | 0 0.0% | 1 | 1 | |
| seq | 1.023 | 1.374 | 16758 | 000 0 | 0 0.0% | 1 | 1 | |
| mol | 0.910 | 1.301 | 14914 | 000 0 | 0 0.0% | 1 | 5 | |
| poc | 0.991 | 1.395 | 16228 | 000 0 | 0 0.0% | 1 | 1 | |
| seq | 1.023 | 1.374 | 16755 | 000 0 | 0 0.0% | 1 | 1 |
| Configuration | Primary metrics | ROC-enrichment (RE) | |||||
|---|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF@1% | @0.5% | @1% | @2% | @5% | |
| DrugCLIP [ Gao et al., 2023 ] | 0.809 | 0.505 | 31.89 | 73.97 | 41.79 | 23.68 | 11.16 |
| LigUnity [ Feng et al., 2025 ] | 0.897 | 0.674 | 44.20 | 104.69 | 57.47 | 33.76 | 13.88 |
| HypSeek [ Wang et al., 2026 ] | 0.909 | 0.605 | 37.83 | 95.20 | 54.49 | 30.74 | 14.01 |
| CausalBind SP ( ) | 0.943 | 0.752 | 47.93 | 127.56 | 69.94 | 37.83 | 16.39 |
| CausalBind SP ( ) | 0.935 | 0.744 | 47.69 | 128.06 | 69.26 | 37.07 | 15.99 |
| Method | Pearson | Spearman |
|---|---|---|
| LigUnity | 0.532 | 0.499 |
| HypSeek | 0.423 | 0.390 |
| CausalBind SP | 0.598 | 0.548 |
| Method | EF@0.5% | EF@1% |
|---|---|---|
| S 2 Drug (paper-reported) | 11.44 | 7.38 |
| CausalBind SP | 10.88 | 8.56 |
| CausalBind LR | 11.59 | 8.84 |
| CausalBind EMB | 12.20 | 9.60 |
| Method | Target-level OOD (9 targets) | Scaffold-level OOD (6 targets) | ||||
|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF@1% | AUROC | BEDROC 80.5 | EF@1% | |
| HypSeek [ Wang et al., 2026 ] | 0.819 | 0.593 | 20.06 | 0.880 | 0.630 | 44.47 |
| LigUnity [ Feng et al., 2025 ] | 0.859 | 0.570 | 19.82 | 0.915 | 0.612 | 43.11 |
| CausalBind SP | 0.860 | 0.658 | 23.00 | 0.945 | 0.712 | 43.66 |
| CausalBind LR | 0.867 | 0.653 | 22.71 | 0.949 | 0.716 | 44.70 |
| CausalBind ATOM | 0.837 | 0.666 | 22.94 | 0.941 | 0.744 | 50.77 |
| Stage | Evidence | Observation | Purpose |
|---|---|---|---|
| Local structure | Heavy-atom contacts ( Å) | 98 contacts involving 27 ligand and 47 pocket atoms | Identify the local binding interface |
| Learned concept | Contact concentration in the most contact-aligned concept pair | 0.0505 vs. 0.00881 uniform baseline (5.73 enrichment) | Test whether contacts concentrate in concepts |
| Predicted score | Mask the associated 4 ligand and 8 pocket atoms (encoder fixed) | Score decreases ( ) | Test contribution to the prediction |
| Random control | Mask 64 matched random sets of 4 ligand and 8 pocket atoms | Mean decrease ; guided decrease (99.2nd percentile) | Exclude an arbitrary masking effect |
| Params/branch | Protocol | DUD-E | LIT-PCBA | |||||
|---|---|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF@1% | AUROC | BEDROC 80.5 | EF@1% | |||
| 1 | 256 | last | 0.935 | 0.720 | 45.86 | 0.630 | 0.087 | 8.41 |
| best_auc | 0.946 | 0.756 | 48.78 | 0.636 | 0.079 | 6.53 | ||
| best_bedroc | 0.935 | 0.724 | 46.28 | 0.632 | 0.089 | 8.56 | ||
| 4 | 1,024 | last | 0.934 | 0.731 | 47.14 | 0.632 | 0.085 | 7.45 |
| best_auc | 0.933 | 0.710 | 45.14 | 0.633 | 0.080 | 7.67 | ||
| Method | DUD-E | LIT-PCBA | |||||
|---|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF@1% | AUROC | BEDROC 80.5 | EF@1% | ||
| CausalBind LR ( ) | 3 | ||||||
| CausalBind LR ( ) | 3 | ||||||
| CausalBind LR ( ) | 3 | ||||||
| CausalBind SP main | 4 | ||||||
| Params/branch | Protocol | DUD-E | LIT-PCBA | |||||
|---|---|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF@1% | AUROC | BEDROC 80.5 | EF@1% | |||
| 64 | 128 | last | 0.935 | 0.704 | 44.87 | 0.632 | 0.103 | 8.84 |
| best_auc | 0.934 | 0.702 | 44.63 | 0.638 | 0.099 | 9.03 | ||
| best_bedroc | 0.936 | 0.716 | 45.92 | 0.631 | 0.095 | 6.54 | ||
| 128 | 256 | last | 0.935 | 0.720 | 45.86 | 0.630 | 0.087 | 8.41 |
| best_auc | 0.946 | 0.756 | 48.78 | 0.636 | 0.079 | 6.53 | ||
| capacity | %zeros | Interpretation | |||
|---|---|---|---|---|---|
| 1 | 0.86 | 0.0% | 1 | 1 | rank-one by construction |
| 4 | 0.63 | 0.0% | 1 | 1 | trained solution collapses to rank one |
| 8 | 0.54 | 8.9% | 1 | 1 | some thresholded zeros, still rank one |
| 16 | 0.51 | 23.3% | 1 | 1 | more zeros, same effective rank |
| 32 | 0.51 | 34.3% | 1 | 1 | unused rank capacity |
| 64 | 0.50 | 39.7% | 1 | 1 | unused rank capacity |
| Protocol | DUD-E | LIT-PCBA | |||||
|---|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF@1% | AUROC | BEDROC 80.5 | EF@1% | ||
| 32 | last | 0.942 | 0.743 | 47.35 | 0.638 | 0.090 | 7.48 |
| best_auc | 0.942 | 0.747 | 47.84 | 0.640 | 0.085 | 6.86 | |
| best_bedroc | 0.937 | 0.716 | 45.79 | 0.643 | 0.066 | 4.38 | |
| 64 | last | 0.939 | 0.754 | 48.56 | 0.636 | 0.102 | 9.60 |
| best_auc | 0.941 | 0.765 | 49.47 | 0.635 | 0.097 | 8.18 | |
| Protocol | DUD-E | LIT-PCBA | |||||
|---|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF@1% | AUROC | BEDROC 80.5 | EF@1% | ||
| 128 | last | 0.943 | 0.732 | 47.09 | 0.606 | 0.074 | 5.88 |
| best_auc | 0.944 | 0.726 | 46.59 | 0.617 | 0.075 | 6.97 | |
| best_bedroc | 0.943 | 0.738 | 47.49 | 0.602 | 0.071 | 6.03 | |
| 256 | last | 0.939 | 0.754 | 48.56 | 0.636 | 0.102 | 9.60 |
| best_auc | 0.941 | 0.765 | 49.47 | 0.635 | 0.097 | 8.18 | |
| Model | Protocol | DUD-E | LIT-PCBA | ||||
|---|---|---|---|---|---|---|---|
| AUROC | BEDROC 80.5 | EF@1% | AUROC | BEDROC 80.5 | EF@1% | ||
| HypSeek [ Wang et al., 2026 ] | last | 0.909 | 0.605 | 37.83 | 0.603 | 0.061 | 4.62 |
| best_auc | 0.913 | 0.629 | 39.63 | 0.608 | 0.063 | 5.13 | |
| best_bedroc | 0.923 | 0.676 | 42.71 | 0.607 | 0.063 | 4.56 | |
| CausalBind SP | last | 0.935 | 0.744 | 47.69 | 0.639 | 0.097 | 8.56 |
| best_auc | 0.934 | 0.722 | 46.41 | 0.628 | 0.074 | 5.84 | |
| Branch | %zeros | ||
|---|---|---|---|
| mol–pocket | 0.0% | 1 | 46 |
| mol–sequence | 0.0% | 1 | 46 |