DrugRx — multi-task ADMET / CYP450 inhibition ensembles

Deep ensembles over fixed molecular descriptors for ADMET endpoints and CYP450 (1A2 / 2C9 / 2D6 / 3A4) pIC50 prediction. Trained on Colab A100s and L4s.

Built for Talindrew — a computational biology workspace with molecular cloning, an electronic lab notebook, protein structure prediction and the FoldRx drug-discovery pipeline. These models are served inside Talindrew's FoldRx pipeline as a per-endpoint router with an applicability-domain check.

This card documents what was measured, including what failed. A large part of this programme's output is null results, and they are recorded here deliberately: the reason these models are small, fixed-feature ensembles rather than anything more fashionable is that the fashionable alternatives were tried and did not win.

The headline result is a negative one

Across four separate attempts — a regularisation sweep, recursive weight sharing, a 16-member ensemble, and simply more compute — no model-class change moved the score outside the noise band. What did move it was features and metric exploitation:

Architecture

A plain MLP: in_dim → d_model with GELU, n_layers residual blocks with LayerNorm and dropout, then a linear head per task. d_model=1024, n_layers=3, dropout=0.15 for the shipped configuration.

Members are trained simultaneously, not sequentially. torch.func's stack_module_state + vmap(grad_and_value(...)) trains 16–32 members as one batched program, which is what makes a deep ensemble affordable; randomness="different" is required or every member draws the same dropout mask and the ensemble degenerates into one model.

Features (~2,500 dims): count-based ECFP4, RDKit 2D descriptors, Avalon, and ErG — the MapLight recipe. Descriptors are standardised using training-fold statistics only, and only over the descriptor span: an early version standardised everything past index 2048, which silently included fingerprint columns.

Loss

The CYP models are trained against the challenge metric rather than against MAE. The metric (ST-RAE) is zero when the prediction falls inside the label's credible interval and otherwise measures distance to the nearer bound, so the loss is a hinge to the interval bounds composed with a Huber transition:

d = relu(lo - pred) + relu(pred - hi)      # 0 inside the interval
q = min(d, 1)
loss = mean(0.5 * q^2 + (d - q))           # quadratic near the bound, linear far

Auxiliary heads (TDI conditions, Emax, log2fc, external Veith labels) are plain MSE at weight 0.3.

Honest caveats — read these before quoting a number

1. "Best epoch" is probably selection on noise. At 3000 epochs the validation score stops improving around epoch 20 and then oscillates in a band of mean 0.4857, sd 0.0095, for another 800 epochs. Two seeds scored 0.4441 and 0.4448 with best epochs 360 and 25 — an optimum that relocates by 335 epochs without changing the score is not an optimum. Any score reported as a minimum-over-epochs on the same split used to select it (which includes our challenge submission) is optimistically biased by an unknown amount. A three-way-split test of band-averaging vs best-epoch selection is in scripts/cyp_curse.py.
2. There is no benefit to training past ~100 epochs on this data. We measured 861 epochs to be sure.
3. A 7x apparent improvement was leakage. Feeding the auxiliary assay columns as features scored 0.0617 against 0.4432, because the TDI-condition pIC50 correlates r=0.90–0.99 with the direct-inhibition target — it is the target measured under a slightly different condition — and those columns do not exist at inference. That arm was deleted rather than kept as an "upper bound".
4. Identity must be chemical, not textual. Matching on raw SMILES hid 425 train/test duplicate molecules (3.2%) behind differing salt forms. Dedup is on a standardised InChIKey (salt stripped, charge neutralised, tautomer canonicalised) with stereochemistry preserved — an earlier version merged L- and D-alanine.
5. Uncertainty is from ensemble disagreement and is not calibrated against a held-out conformal procedure.

Training mechanisms published alongside

code/ in this repo carries the training package and scripts, including:

• drugrx/cyp/resume.py — full-state checkpointing (stacked params, optimizer moments, scheduler position, patience counter, RNG, curve, and the best snapshot separately), mirrored to the Hub. Verified by killing a run at epoch 50 and confirming it resumed at epoch 50 with its peak intact, because the A100 runtime has died mid-experiment twice.
• scripts/gpu_review.py — audits every Colab session for live processes, GPU fill and artifact staleness; proposes shutdown only for sessions with no process and no writes, and reports an unreachable session as UNKNOWN rather than dead.
• scripts/cyp_curse.py — the three-way-split test of best-epoch selection.
• drugrx/metrics.py, drugrx/cyp/metric.py — the metric and the paired bootstrap used for every verdict above.

Intended use and limits

Research triage and hypothesis generation. Not for clinical or regulatory use. Predictions outside the training applicability domain are unreliable and the router artifact carries an explicit domain check for that reason.

The number worth holding alongside any virtual-screening claim: in blind prospective screening (CACHE #1, LRRK2), 23 teams produced a 3.7% hit rate at K_D < 150 µM — and most teams found nothing (7 of 23 on CACHE #1, 5 of 23 on CACHE #2 at an overall 0.7%).

Where these are used

In Talindrew's FoldRx pipeline and AI workspace, as interchangeable model implementations behind a per-endpoint router. The router refuses to score molecules outside the training applicability domain rather than returning a confident number for a compound unlike anything it was trained on.

License

MIT for the code and weights. The underlying datasets keep their own terms (TDC/MoleculeNet sources, ChEMBL CC-BY-SA 3.0, PubChem public domain, and the OpenADMET CYP challenge data under its own terms).