workbench 0.8.330__tar.gz → 0.8.332__tar.gz

{workbench-0.8.330 → workbench-0.8.332}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: workbench
-Version: 0.8.330
+Version: 0.8.332
 Summary: Workbench: A Dashboard and Python API for creating and deploying AWS SageMaker Model Pipelines
 Author-email: SuperCowPowers LLC <support@supercowpowers.com>
 License: MIT License
@@ -53,7 +53,7 @@ Requires-Dist: joblib>=1.3.2
 Requires-Dist: requests>=2.26.0
 Requires-Dist: rdkit>=2026.3.1
 Requires-Dist: mordredcommunity>=2.0.6
-Requires-Dist: workbench-bridges>=0.2.8
+Requires-Dist: workbench-bridges>=0.2.9
 Requires-Dist: cleanlab>=2.9.0
 Requires-Dist: pyyaml>=6.0
 Provides-Extra: ui

{workbench-0.8.330 → workbench-0.8.332}/data/public_data/build_logp_logd_overlap.py

@@ -25,9 +25,9 @@ Per-tier construction
 * Band selections capped at --max-aux (default 1,558) via deterministic
   random sample (--seed) when more candidates exist than the cap.
 
-Approach: run DatasetComparison with logd as reference and logp as query, so
-each LogP compound gets a per-row Tanimoto-to-LogD score. Filter by band, cap,
-merge with the LogD primary, write CSV.
+Approach: build a MultiTaskAlignment with logd as primary and logp as aux. Each
+LogP-only row in the unified results gets a `tanimoto_to_primary` score. Filter
+by band, cap, merge with the LogD primary, write CSV.
 
 Output: output/experiments/logp_logd_overlap_<suffix>.csv
 """
@@ -38,7 +38,7 @@ from pathlib import Path
 
 import pandas as pd
 
-from workbench.algorithms.dataframe import DatasetComparison
+from workbench.algorithms.dataframe import MultiTaskAlignment
 
 log = logging.getLogger("workbench")
 
@@ -58,7 +58,7 @@ TIERS = [
 
 
 def build_one(
-    dc: DatasetComparison,
+    mta: MultiTaskAlignment,
     logp: pd.DataFrame,
     logd: pd.DataFrame,
     suffix: str,
@@ -67,15 +67,18 @@ def build_one(
     max_aux: int,
     seed: int,
 ) -> Path:
-    """Build a single tier CSV. Reuses the already-built DatasetComparison model."""
-    results = dc.results()
-    logp_scored = results[results["dataset"] == "query"].copy()
+    """Build a single tier CSV. Reuses the already-built MultiTaskAlignment model."""
+    results = mta.results()
+    # LogP-only rows in the unified results — `tanimoto_to_primary` is best Tanimoto to any LogD compound.
+    logp_scored = results[results["logp"].notna()].copy()
 
     # Inclusive band [min, max]. The high tier ([0.7, 1.0]) naturally captures
     # exact-SMILES compounds (Tanimoto == 1.0); their merged rows show both
     # logp and logd populated. Lower tiers (< 1.0) by definition have no
     # smiles in common with LogD.
-    band_mask = (logp_scored["tanimoto_sim"] >= tanimoto_min) & (logp_scored["tanimoto_sim"] <= tanimoto_max)
+    band_mask = (logp_scored["tanimoto_to_primary"] >= tanimoto_min) & (
+        logp_scored["tanimoto_to_primary"] <= tanimoto_max
+    )
     band = logp_scored.loc[band_mask].copy()
 
     n_band_total = len(band)
@@ -84,7 +87,7 @@ def build_one(
     log.info(
         f"[{suffix}]  Band [{tanimoto_min:.2f}, {tanimoto_max:.2f}]  "
         f"candidates={n_band_total:,}  selected={len(band):,}  "
-        f"sim_range={band['tanimoto_sim'].min():.3f}-{band['tanimoto_sim'].max():.3f}"
+        f"sim_range={band['tanimoto_to_primary'].min():.3f}-{band['tanimoto_to_primary'].max():.3f}"
     )
 
     # Take selected LogP rows and outer-merge with LogD primary.
@@ -122,14 +125,21 @@ def main(max_aux: int, seed: int, only: list[str] | None) -> None:
     logp_ids = logp.assign(id=logp["id"].astype(str).radd("logp_"))
     logd_ids = logd.assign(id=logd["id"].astype(str).radd("logd_"))
 
-    # Build the comparison once (logd=ref, logp=query) — gives each LogP a Tanimoto-to-LogD score.
-    # All three tiers are derived from the same scored set, so we only build the FP/UMAP model once.
-    log.info("Building DatasetComparison (logd=reference, logp=query) ...")
-    dc = DatasetComparison(
-        df_reference=logd_ids[["id", "smiles", "logd"]],
-        df_query=logp_ids[["id", "smiles", "logp"]],
-        reference_target="logd",
-        query_target="logp",
+    # Wide multi-task DataFrame: LogD-only rows + LogP-only rows. The shared FP/UMAP model
+    # is built once, and every LogP row gets a tanimoto_to_primary score against LogD.
+    mt_df = pd.concat(
+        [
+            logd_ids[["id", "smiles"]].assign(logd=logd_ids["logd"]),
+            logp_ids[["id", "smiles"]].assign(logp=logp_ids["logp"]),
+        ],
+        ignore_index=True,
+    )
+
+    log.info("Building MultiTaskAlignment (primary=logd, aux=logp) ...")
+    mta = MultiTaskAlignment(
+        mt_df,
+        primary="logd",
+        auxiliaries=["logp"],
         id_column="id",
     )
 
@@ -138,7 +148,7 @@ def main(max_aux: int, seed: int, only: list[str] | None) -> None:
         raise ValueError(f"--only filtered out all tiers; valid suffixes: {[t[0] for t in TIERS]}")
 
     for suffix, lo, hi in selected_tiers:
-        build_one(dc, logp, logd, suffix, lo, hi, max_aux, seed)
+        build_one(mta, logp, logd, suffix, lo, hi, max_aux, seed)
 
 
 if __name__ == "__main__":

workbench-0.8.332/data/public_data/overlap_summary.py

@@ -0,0 +1,128 @@
+"""Report LogP/LogD chemical-space overlap to inform multi-task modeling decisions.
+
+Goal: decide whether a chemprop multi-task LogP+LogD model is likely to outperform
+a LogD-only model. That depends on two things:
+
+  (1) Coverage — do LogD compounds have nearby LogP neighbors in chemical space?
+                 If LogP doesn't cover the LogD distribution, the auxiliary task adds
+                 little signal where it matters.
+  (2) Concordance — for similar compounds, do LogP and LogD agree (correlated)?
+                 Strong agreement suggests easy auxiliary signal; mild disagreement
+                 is exactly what multi-task learning can exploit.
+
+Thin CLI wrapper around `workbench.algorithms.dataframe.MultiTaskAlignment`:
+
+    MultiTaskAlignment(df, primary="logd", auxiliaries=["logp"], id_column="id")
+
+We then print the per-aux summary and the multi-task verdict. Per-compound details
+are written to `output/overlap_summary.csv` for downstream slicing/plotting.
+"""
+
+import logging
+from pathlib import Path
+
+import pandas as pd
+
+from workbench.algorithms.dataframe import MultiTaskAlignment
+
+log = logging.getLogger("workbench")
+
+DATA_DIR = Path(__file__).parent
+OUTPUT_DIR = DATA_DIR / "output"
+LOGP_PATH = OUTPUT_DIR / "logp" / "logp_all.csv"
+LOGD_PATH = OUTPUT_DIR / "logd" / "logd_all.csv"
+SUMMARY_CSV = OUTPUT_DIR / "overlap_summary.csv"
+
+
+def _print_summary(summary_row: pd.Series, n_logp: int, n_logd: int) -> None:
+    print("\n" + "=" * 72)
+    print("MultiTaskAlignment Summary  (primary=LogD, auxiliary=LogP)")
+    print("=" * 72)
+    print(f"  Primary (LogD):    {n_logd:>7,} compounds")
+    print(f"  Auxiliary (LogP):  {n_logp:>7,} compounds")
+    print(f"  Shared compounds:  {int(summary_row['n_shared']):>7,}")
+    print(f"  LogP-only:         {int(summary_row['n_aux_only']):>7,}")
+
+    print("\nLabel correlation on shared compounds")
+    r = summary_row["pearson_r"]
+    if pd.notna(r):
+        print(f"  Pearson r = {r:.3f}  ({summary_row['r_confidence']} confidence)")
+    else:
+        print(f"  Pearson r = N/A  ({summary_row['r_confidence']})")
+
+    print("\nChemical-space coverage (LogP -> nearest LogD)")
+    print(f"  mean Tanimoto:        {summary_row['tanimoto_coverage_mean']:.3f}")
+    print(f"  fraction sim >= 0.50: {100 * summary_row['frac_coverage_ge_05']:5.1f}%")
+    print(f"  fraction sim >= 0.30: {100 * summary_row['frac_coverage_ge_03']:5.1f}%")
+
+    print("\nZ-scored residual on aux-having rows")
+    print(f"  |residual| mean: {summary_row['residual_abs_mean']:.3f}")
+    print(f"  |residual| p95:  {summary_row['residual_abs_p95']:.3f}")
+
+
+def _print_verdict(summary_row: pd.Series) -> None:
+    print("\n" + "=" * 72)
+    print("Multi-Task Recommendation")
+    print("=" * 72)
+    print(f"  Overlap verdict:    {summary_row['overlap']}")
+    print(f"  Extension verdict:  {summary_row['extension']}")
+    print(f"  Recommendation:     {summary_row['recommendation']}")
+
+
+def main() -> None:
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s  %(message)s")
+
+    if not LOGP_PATH.exists():
+        raise FileNotFoundError(f"{LOGP_PATH} not found — run pull_logp_data.py first.")
+    if not LOGD_PATH.exists():
+        raise FileNotFoundError(f"{LOGD_PATH} not found — run pull_logd_data.py first.")
+
+    logp = pd.read_csv(LOGP_PATH)
+    logd = pd.read_csv(LOGD_PATH)
+
+    # ID disambiguation so the combined model can identify rows uniquely
+    logp = logp.assign(id=logp["id"].astype(str).radd("logp_"))
+    logd = logd.assign(id=logd["id"].astype(str).radd("logd_"))
+
+    # Build a wide multi-task DataFrame: one row per source compound with NaN where missing
+    mt_df = pd.concat(
+        [
+            logd[["id", "smiles"]].assign(logd=logd["logd"]),
+            logp[["id", "smiles"]].assign(logp=logp["logp"]),
+        ],
+        ignore_index=True,
+    )
+
+    mta = MultiTaskAlignment(
+        mt_df,
+        primary="logd",
+        auxiliaries=["logp"],
+        id_column="id",
+    )
+
+    summary = mta.summary()
+    summary_row = summary.iloc[0]
+    _print_summary(summary_row, n_logp=len(logp), n_logd=len(logd))
+    _print_verdict(summary_row)
+
+    results = mta.results()
+    keep = [
+        c
+        for c in [
+            "id",
+            "smiles",
+            "logd",
+            "logp",
+            "x",
+            "y",
+            "tanimoto_to_primary",
+            "residual_logp",
+        ]
+        if c in results.columns
+    ]
+    results[keep].to_csv(SUMMARY_CSV, index=False)
+    print(f"\nPer-compound results saved -> {SUMMARY_CSV}")
+
+
+if __name__ == "__main__":
+    main()

{workbench-0.8.330 → workbench-0.8.332}/docs/blogs/3d_descriptors.md

@@ -1,21 +1,33 @@
 # 3D Molecular Descriptors: Shape, Surface, and Pharmacophore Features
-!!! tip inline end "Combine 2D + 3D"
-    Run both the [2D descriptor endpoint](molecular_standardization.md) and the 3D endpoint, then concatenate the results for a ~388-feature set covering topological, electronic, and geometric properties.
+!!! tip inline end "When to reach for 3D"
+    On most TDC ADMET endpoints, careful 2D fingerprints + learned graph representations are highly competitive on their own. Use the 3D endpoint when you need shape-aware features for geometry-sensitive endpoints (passive permeability, P-gp / BCRP interactions, conformer-dependent solubility), typically alongside the 2D feature set, not as a replacement for it.
 
-2D molecular descriptors capture a lot about a molecule's chemistry from its connectivity graph alone -- molecular weight, hydrogen bond donors, topological polar surface area, and hundreds of other properties. But some of the most important ADMET properties depend on the molecule's *shape* in three dimensions: how it fits into a transporter binding site, whether it can fold to mask polar groups for membrane permeation, or how its charge distribution maps onto its surface.
+2D molecular descriptors capture a lot about a molecule's chemistry from its connectivity graph alone -- molecular weight, hydrogen bond donors, topological polar surface area, and hundreds of other properties. Some ADMET properties have geometric components that 2D descriptors capture only indirectly through their correlations with the molecular graph: how a molecule fits into a transporter binding site, whether it can fold to mask polar groups for membrane permeation, or how its charge distribution maps onto its surface. The 3D endpoint exposes these directly as engineered features -- whether they help on a given task is an empirical question, not a foregone conclusion.
 
 Workbench's 3D descriptor endpoints compute **74 conformer-based features** from SMILES strings, covering molecular shape, charged partial surface area, pharmacophore spatial distribution, and conformational flexibility. Like all Workbench endpoints, the contract is simple: **send a DataFrame, get a DataFrame back** -- the input DataFrame comes back with 74 descriptor columns appended. Two pipeline modes are available -- a **fast** endpoint for realtime inference and a **Boltzmann** endpoint for high-quality batch processing. Both produce the same 74 features so downstream models can consume either interchangeably.
 
-## Why 3D Descriptors?
+## When 3D Descriptors Help (and When They Don't)
 
-2D descriptors treat molecules as graphs -- atoms are nodes, bonds are edges. This misses geometry-dependent properties that matter for ADMET:
+2D descriptors treat molecules as graphs -- atoms are nodes, bonds are edges. There are ADMET-relevant properties whose geometric components 2D captures only indirectly:
 
 - **Membrane permeability** depends on molecular shape and the spatial separation of polar and nonpolar regions (amphiphilic moment)
 - **Transporter interactions** (P-gp, BCRP) correlate with molecular elongation, nitrogen spatial distribution, and overall size
 - **Protein-ligand binding** depends on 3D shape complementarity, not just functional group counts
-- **Intramolecular hydrogen bonds** enable "chameleonic" behavior where molecules mask polar groups in nonpolar environments -- a purely 3D phenomenon
+- **Intramolecular hydrogen bonds** enable "chameleonic" behavior where molecules mask polar groups in nonpolar environments -- a 3D phenomenon
 
-These properties can't be captured from the molecular graph. You need 3D coordinates.
+3D descriptors expose these directly as engineered features. Whether they actually move the needle on a given downstream model is a separate, empirical question -- and the answer is more nuanced than the bullets above suggest.
+
+### When 2D Alone Is Enough
+
+There is a real and well-supported skeptical position in the cheminformatics community: for most ADMET endpoints, well-engineered 2D fingerprints + learned graph representations are competitive with -- or better than -- 2D + 3D combined.
+
+The evidence:
+
+- On the [TDC ADMET leaderboards](https://tdcommons.ai/benchmark/overview/) through 2024-2026, top reproducible models (MapLight, MapLight+GNN, CaliciBoost, NovoExpert-2) use **ECFP/Avalon/ErG + 200 RDKit 2D physchem + GIN embeddings**, with **no explicit 3D features**. The Konopka et al. 2026 critical assessment of TDC leaderboard reproducibility makes this concrete.
+- **PharmaBench** (Niu et al., *Sci. Data* 2024) finds no statistically significant 2D-vs-3D difference on most ADMET endpoints across thousands of compounds.
+- **Bahia et al.** (*Mol. Inform.* 2023) report a 2D + 3D advantage over 2D alone -- but the delta is low single-digit AUC / R², not transformative.
+
+So treat the 3D feature stream honestly: a *complementary* set that may give modest, endpoint-dependent gains (most plausibly on passive permeability, P-gp / BCRP recognition, conformer-dependent solubility) on top of a strong 2D + learned-representation baseline. It is not a foregone conclusion that 3D features improve any specific ADMET model -- run an ablation. See the [Limitations & Future Work](#limitations-future-work) section for forward-looking upgrades and a fuller honest accounting.
 
 ## Two Pipeline Modes: Fast and Full
 
@@ -55,7 +67,7 @@ The full endpoint uses a datamol-style adaptive tiering, with the upper tiers bu
 | 8-12 | 300 | Moderate flexibility, needs broader sampling |
 | ≥ 13 | 500 | High flexibility, large conformational space |
 
-This ensures adequate sampling of the conformational landscape without wasting compute on rigid molecules. On 300-conformer runs at 13+ rotatable bonds we measured ~20% NPR1 variance across random seeds; bumping to 500 conformers uses the documented "more samples" path to reduce that stochastic spread by roughly √(500/300) ≈ 1.29×.
+This ensures adequate sampling of the conformational landscape without wasting compute on rigid molecules. On 300-conformer runs at 13+ rotatable bonds we measured ~20% NPR1 variance across random seeds; bumping to 500 conformers reduces *within-seed* variance by roughly √(500/300) ≈ 1.29× via the documented "more samples" path. It does not eliminate *cross-seed* variance — different random seeds will still produce slightly different Boltzmann averages on highly flexible molecules — but for most ADMET endpoints this residual is below downstream model noise.
 
 ## The Computation Pipeline
 
@@ -253,6 +265,22 @@ cached_endpoint = InferenceCache(end_fast, cache_key_column="smiles")
 df_cached = cached_endpoint.inference(big_df)  # Only computes uncached rows
 ```
 
+## Limitations & Future Work
+
+The pipeline is conservative by design — production ADMET targets stable, deterministic features over the latest research methods. A few areas worth flagging for downstream users and future iterations:
+
+**3D vs 2D in ADMET reality.** As noted in the introduction, top reproducible TDC ADMET models lean on 2D fingerprints + learned graph representations. The published evidence (PharmaBench *Sci. Data* 2024; Bahia *Mol. Inform.* 2023) is that 3D descriptors give marginal-but-real gains on geometry-sensitive endpoints and roughly neutral effects on most others. The 3D feature stream complements rather than replaces a strong 2D + learned-representation baseline.
+
+**Cross-seed variance on highly flexible molecules.** The 500-conformer top tier reduces within-seed stochastic variance, but different random seeds will still produce slightly different Boltzmann averages on 13+ rotatable-bond molecules. For most ADMET endpoints this residual is below downstream model noise; for tasks that genuinely depend on a single conformer geometry it is not.
+
+**Forward-looking upgrades** (evidence-backed; not yet implemented):
+
+1. **Single-point xTB / tblite re-ranking** of MMFF-optimized conformers before Boltzmann weighting. Kong et al. (*ChemPhysChem* 2025) show GFN2-xTB is currently the most suitable energy filter for drug-like conformer ranking. Pip-installable via `tblite-python`, deterministic.
+2. **CONFORGE as alternative embedder** for macrocycles and very-flexible scaffolds. Seidel et al. (*JCIM* 2023, CDPKit) — open source, slightly outperforms RDKit on small molecules and matches it on macrocycles where ETKDGv3 sampling plateaus.
+3. **Replace Gasteiger partial charges in CPSA** with AM1-BCC or an ML charge model (DASH; Mahmoud et al. 2023). Gasteiger is documented as the least accurate common partial-charge method, and CPSA accounts for 43 of our 52 Mordred 3D features — the highest-leverage upgrade for the existing feature set.
+
+Deliberately *not* on this list: ML conformer generators (ETFlow, GeoMol, Lyrebird) — research-stage with no proven ADMET benefit; MACE-OFF / ANI-2x routine optimization — too heavy for production throughput; tautomer/protomer ensemble enumeration — mature in research, niche in production. We may revisit any of these as the surrounding tooling matures.
+
 ## References
 
 **Conformer Ensemble Methods**
@@ -271,16 +299,27 @@ df_cached = cached_endpoint.inference(big_df)  # Only computes uncached rows
 - Landrum, G. *"Understanding conformer generation failures."* RDKit Blog (2023). [Blog post](https://greglandrum.github.io/rdkit-blog/posts/2023-05-17-understanding-confgen-errors.html)
 - Landrum, G. *"Scaling conformer generation."* RDKit Blog (2025). [Blog post](https://greglandrum.github.io/rdkit-blog/posts/2025-08-30-confgen-scaling.html)
 - Datamol conformer generation with adaptive tiering. [Documentation](https://docs.datamol.io/stable/tutorials/Conformers.html)
+- Seidel, T., Permann, C., Wieder, O., Kohlbacher, S. & Langer, T. *"High-Quality Conformer Generation with CONFORGE: Algorithm and Performance Assessment."* J. Chem. Inf. Model. 63, 5549-5570 (2023). [DOI: 10.1021/acs.jcim.3c00563](https://doi.org/10.1021/acs.jcim.3c00563)
 
 **Force Fields**
 
 - Tosco, P., Stiefl, N. & Landrum, G. *"Bringing the MMFF force field to the RDKit: implementation and validation."* J. Cheminform. 6, 37 (2014). [DOI: 10.1186/s13321-014-0037-3](https://doi.org/10.1186/s13321-014-0037-3)
+- Kong, Z., et al. *"Discriminating High from Low Energy Conformers of Druglike Molecules."* ChemPhysChem (2025). [DOI: 10.1002/cphc.202400992](https://doi.org/10.1002/cphc.202400992)
 
 **Descriptors**
 
 - RDKit 3D Descriptors: [Documentation](https://www.rdkit.org/docs/source/rdkit.Chem.Descriptors3D.html)
 - Mordred Community: [GitHub](https://github.com/JacksonBurns/mordred-community)
 - Stanton, D.T. & Jurs, P.C. *"Development and Use of Charged Partial Surface Area Structural Descriptors in Computer-Assisted Quantitative Structure-Property Relationship Studies."* Anal. Chem. 62, 2323-2329 (1990). [DOI: 10.1021/ac00220a013](https://doi.org/10.1021/ac00220a013)
+- Bleiziffer, P., Schaller, K. & Riniker, S. *"Machine Learning of Partial Charges Derived from High-Quality Quantum-Mechanical Calculations."* J. Chem. Inf. Model. 58, 579-589 (2018). [DOI: 10.1021/acs.jcim.7b00663](https://doi.org/10.1021/acs.jcim.7b00663)
+- Mahmoud, A.H., Masters, M.R., Lill, M.A., et al. *"DASH: Dynamic Attention-Based Substructure Hierarchy for Partial Charge Assignment."* J. Chem. Inf. Model. (2023). [DOI: 10.1021/acs.jcim.3c00800](https://doi.org/10.1021/acs.jcim.3c00800)
+
+**ADMET Benchmarks and 2D vs 3D Evidence**
+
+- Huang, K., et al. *"Therapeutics Data Commons: Machine Learning Datasets and Tasks for Drug Discovery and Development."* [TDC ADMET Leaderboards](https://tdcommons.ai/benchmark/overview/)
+- Konopka, M., et al. *"Critical Assessment of ML Models for ADMET Prediction in TDC Leaderboards."* bioRxiv (2026). [DOI: 10.64898/2026.02.26.708193](https://www.biorxiv.org/content/10.64898/2026.02.26.708193v1)
+- Niu, Z., et al. *"PharmaBench: Enhancing ADMET Benchmarks with Large Language Models."* Sci. Data 11, 985 (2024). [DOI: 10.1038/s41597-024-03793-0](https://doi.org/10.1038/s41597-024-03793-0)
+- Bahia, M.S., et al. *"Comparison Between 2D and 3D Descriptors in QSAR Modeling Based on Bio-Activities."* Mol. Inform. 42, 2200186 (2023). [DOI: 10.1002/minf.202200186](https://doi.org/10.1002/minf.202200186)
 
 **ADMET and Chameleonic Molecules**
 

workbench-0.8.332/docs/blogs/chemprop_multi_task.md

@@ -0,0 +1,129 @@
+# Multi-Task ChemProp: Two Mechanisms, One Model
+!!! tip inline end "ChemProp in Workbench"
+    Workbench supports four ChemProp variants — single-task, multi-task, hybrid, and foundation. See the [ChemProp Models](../models/chemprop_models.md) page for the full API.
+
+A multi-task ChemProp model trains **one** message-passing encoder to predict **several** target properties at once. Each target gets its own feedforward head, but they all share the molecular representation that the MPNN learns. That shared encoder is where the wins come from — and where the failures come from when the setup is wrong.
+
+In this blog we'll separate the two distinct mechanisms that drive multi-task lift:
+
+1. **Transfer learning** — gradients from auxiliary tasks regularize and enrich the shared encoder
+2. **Chemical-space expansion** — auxiliary datasets bring compounds the primary task never saw
+
+These mechanisms are independent. A multi-task setup can lift from one, the other, or both — and recognizing which is at play guides how you design the data.
+
+## Mechanism 1: Transfer Learning Through the Shared Encoder
+
+In a single-task ChemProp model, the encoder is supervised by exactly one signal — the gradient flowing back from one head, derived from one target column. In a multi-task model, the same encoder receives gradients from **every** head. Each task pulls the embedding space in its own direction, and the encoder converges to a representation that satisfies all of them at once.
+
+<figure style="margin: 20px 0; width: 100%; display: block;">
+<img src="../../images/chemprop_mt_transfer.svg" alt="Multi-task transfer learning: shared MPNN encoder with N task-specific FFN heads, gradients flowing back to the shared encoder" style="width: 100%; height: auto; display: block;">
+<figcaption><em>One encoder, N heads, N gradients flowing back into the same weights — auxiliaries co-supervise the primary's representation.</em></figcaption>
+</figure>
+
+This sounds like a tug-of-war, and it is — but a productive one when the tasks are related. The encoder ends up encoding the molecular features useful **across** the family of tasks, rather than overfitting cues specific to any single one. That is exactly the transfer-learning setup, except the "transfer" happens *during* training rather than after.
+
+Whether you get lift depends on the label relationship between primary and auxiliary on **shared** compounds (the rows where both have ground truth):
+
+<table style="width: 100%;">
+  <thead>
+    <tr>
+      <th style="background-color: rgba(58, 134, 255, 0.5); color: white; padding: 10px 16px;">Pearson r on shared rows</th>
+      <th style="background-color: rgba(58, 134, 255, 0.5); color: white; padding: 10px 16px;">Verdict</th>
+      <th style="background-color: rgba(58, 134, 255, 0.5); color: white; padding: 10px 16px;">What happens to the encoder</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr><td class="text-orange" style="padding: 8px 16px; font-weight: bold;">0.4 – 0.95</td><td style="padding: 8px 16px;">Beneficial</td><td style="padding: 8px 16px;">Sweet spot — heads predict related but distinct targets, encoder learns richer features</td></tr>
+    <tr><td class="text-orange" style="padding: 8px 16px; font-weight: bold;">&gt; 0.95</td><td style="padding: 8px 16px;">Neutral</td><td style="padding: 8px 16px;">Labels redundant; aux head just re-weights primary supervision — no harm, no new info</td></tr>
+    <tr><td class="text-orange" style="padding: 8px 16px; font-weight: bold;">&lt; 0.4</td><td style="padding: 8px 16px;">Harmful</td><td style="padding: 8px 16px;">Heads disagree; encoder gradients conflict — negative-transfer risk</td></tr>
+  </tbody>
+</table>
+
+The "sweet spot" is the regime where transfer learning earns its name. The auxiliary task is related enough that its gradient pushes the encoder toward useful chemistry, but distinct enough that it brings new information. Outside that band, multi-task either contributes nothing or actively hurts — picking auxiliaries blindly is the most common cause of a multi-task model that underperforms its single-task baseline.
+
+## Mechanism 2: Chemical-Space Expansion
+
+The transfer-learning story above plays out on the **overlap region** — compounds where both primary and auxiliary tasks have ground truth. But auxiliary datasets often contain compounds the primary task never saw at all. When the auxiliary is, say, public ChEMBL data while the primary is in-house screening data, the auxiliary brings in **chemistry** the primary doesn't have.
+
+<figure style="margin: 20px 0; width: 100%; display: block;">
+<img src="../../images/chemprop_mt_expansion.svg" alt="Chemical-space expansion: single-task encoder sees only primary chemistry; multi-task encoder also sees auxiliary chemistry that surrounds and extends beyond the primary" style="width: 100%; height: auto; display: block;">
+<figcaption><em>A test compound outside the primary's chemistry is an extrapolation for the single-task encoder; the multi-task encoder has already seen that region via auxiliary data.</em></figcaption>
+</figure>
+
+For aux-only compounds, the primary head's loss is masked — there's no primary label to supervise against — so the primary head sees no direct gradient on those rows. But the **encoder** does. Every aux-only molecule pushes the encoder to handle a broader range of chemistry, and the encoder's representation generalizes accordingly. The primary head, downstream, inherits a more capable encoder.
+
+Two important consequences follow:
+
+- **Extension can lift even with zero label correlation.** It doesn't require the auxiliary task to *predict* anything related to the primary — it just requires the auxiliary to bring relevant chemistry. A noisy public dataset on a marginally related endpoint can still expand the encoder's coverage.
+- **Extension cannot actively hurt the primary.** The worst case for an aux-only compound is "encoder learns features that don't help" — not "encoder learns features that mislead." There's no primary gradient on those rows to be corrupted.
+
+This is why extension and overlap are scored as independent axes in our pre-flight check.
+
+## Putting Both Axes Together
+
+Workbench's `assess_multi_task_data` utility scores each auxiliary along these two axes from labels alone — no training required — then combines them into an actionable recommendation:
+
+```python
+from workbench.utils.multi_task import assess_multi_task_data
+
+assessment = assess_multi_task_data(
+    df,
+    target_columns=["logd", "ksol", "hlm_clint", "mlm_clint"],
+)
+print(assessment[["auxiliary", "overlap", "extension", "recommendation"]])
+```
+
+The recommendation logic is straightforward once both axes are scored:
+
+<table style="width: 100%;">
+  <thead>
+    <tr>
+      <th style="background-color: rgba(58, 134, 255, 0.5); color: white; padding: 10px 16px;">Recommendation</th>
+      <th style="background-color: rgba(58, 134, 255, 0.5); color: white; padding: 10px 16px;">When</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr><td class="text-orange" style="padding: 8px 16px; font-weight: bold;">Use</td><td style="padding: 8px 16px;">Either axis is genuinely lifting and there are no harm signals</td></tr>
+    <tr><td class="text-orange" style="padding: 8px 16px; font-weight: bold;">Marginal</td><td style="padding: 8px 16px;">Limited contribution from either axis — proceed only if domain knowledge supports it</td></tr>
+    <tr><td class="text-orange" style="padding: 8px 16px; font-weight: bold;">Risky</td><td style="padding: 8px 16px;">Harmful overlap, but extension might rescue it — validate against single-task baseline</td></tr>
+    <tr><td class="text-orange" style="padding: 8px 16px; font-weight: bold;">Skip</td><td style="padding: 8px 16px;">No useful contribution from either axis</td></tr>
+  </tbody>
+</table>
+
+## When Multi-Task Helps in Practice
+
+The mechanism breakdown suggests a few practical heuristics:
+
+1. **Pick auxiliaries that are mechanistically related to the primary.** ADMET endpoints (LogD, solubility, clearance, permeability) all share underlying physicochemistry — a textbook MT setup. Mixing unrelated bioactivity with ADMET tends to fall in the negative-transfer zone.
+
+2. **Public data is often valuable purely for chemical-space expansion.** Even when public labels are noisy or measured under different protocols, the chemistry they bring can sharpen the encoder. Use SMILES-based merging via `combine_multi_task_data(..., merge_on_smiles=True)` to fold them in without needing a shared id namespace.
+
+3. **Use task weights deliberately.** When all tasks are end products, `compute_inverse_count_task_weights` balances each task's gradient contribution. When you have a primary + auxiliaries that are only there to help the primary, use manual primary-favored weights (e.g. `[1.0, 0.3, 0.3]`) so the auxiliaries don't swamp the primary signal.
+
+4. **Validate against single-task baselines.** Multi-task is an architectural choice, not a free lunch. Always train the matched single-task model and compare on a held-out test set — especially for "Risky" and "Marginal" auxiliaries.
+
+## Summary
+
+Multi-task ChemProp lift comes from two distinct mechanisms — transfer learning through the shared encoder, and chemical-space expansion from non-overlapping auxiliary compounds. Recognizing which mechanism is operating (or whether neither is) lets you design the multi-task setup deliberately rather than hoping it works. Workbench's pre-flight assessment quantifies both axes from the labels alone, before you spend the training cycles.
+
+## References
+
+**ChemProp v2** — The message-passing neural network framework that Workbench uses for molecular property prediction:
+
+- Chemprop v2 Paper: Graff, D.E., Morgan, N.K., Burns, J.W., et al. *"Chemprop v2: An Efficient, Modular Machine Learning Package for Chemical Property Prediction."* Journal of Chemical Information and Modeling 66(1), 28–33 (2026). [DOI: 10.1021/acs.jcim.5c02332](https://doi.org/10.1021/acs.jcim.5c02332)
+- GitHub: [https://github.com/chemprop/chemprop](https://github.com/chemprop/chemprop)
+
+**Multi-Task Learning** — Foundational and survey work on shared-representation learning:
+
+- Caruana, R. *"Multitask Learning."* Machine Learning 28, 41–75 (1997). [DOI: 10.1023/A:1007379606734](https://doi.org/10.1023/A:1007379606734)
+- Ruder, S. *"An Overview of Multi-Task Learning in Deep Neural Networks."* arXiv:1706.05098 (2017). [https://arxiv.org/abs/1706.05098](https://arxiv.org/abs/1706.05098)
+
+**Multi-Task Learning for Molecular Property Prediction**:
+
+- Ramsundar, B., Kearnes, S., Riley, P., et al. *"Massively Multitask Networks for Drug Discovery."* arXiv:1502.02072 (2015). [https://arxiv.org/abs/1502.02072](https://arxiv.org/abs/1502.02072)
+- Wenzel, J., Matter, H., Schmidt, F. *"Predictive Multitask Deep Neural Network Models for ADME-Tox Properties."* Journal of Chemical Information and Modeling 59(3), 1253–1268 (2019). [DOI: 10.1021/acs.jcim.8b00785](https://doi.org/10.1021/acs.jcim.8b00785)
+
+## Questions?
+<img align="right" src="../../images/scp.png" width="180">
+
+The SuperCowPowers team is happy to answer any questions you may have about AWS and Workbench. Please contact us at [workbench@supercowpowers.com](mailto:workbench@supercowpowers.com) or on chat us up on [Discord](https://discord.gg/WHAJuz8sw8)

{workbench-0.8.330 → workbench-0.8.332}/docs/blogs/index.md

@@ -12,6 +12,8 @@ Workbench blogs highlight interesting functionality and approaches that might be
 
 - **[SHAP Values for ChemProp Models](chemprop_shap.md):** How do you explain a graph neural network? In this blog we explore our per-bit ablation approach for computing [SHAP values](https://shap.readthedocs.io/) on [ChemProp](https://github.com/chemprop/chemprop) MPNN models. We walk through the technical approach, show key code snippets, and analyze real SHAP output from a LogD regression model — validating that the model independently learns known structure-lipophilicity relationships.
 
+- **[Multi-Task ChemProp: Two Mechanisms, One Model](chemprop_multi_task.md):** Multi-task ChemProp lifts come from two distinct mechanisms — transfer learning through a shared MPNN encoder, and chemical-space expansion from non-overlapping auxiliary compounds. We separate the two with conceptual diagrams, show when each one helps (vs. hurts), and tie the discussion back to Workbench's pre-flight `assess_multi_task_data` utility that scores both axes from labels alone.
+
 - **[Inside a Workbench AWS Endpoint](aws_endpoint_architecture.md):** A deep dive into endpoint architecture — comparing the default SageMaker stack (Nginx, Gunicorn, Flask) with Workbench's modern ASGI stack (Uvicorn, FastAPI). We cover the custom image with pre-loaded chemistry packages, and Workbench's binary-search error handling that isolates bad rows instead of failing entire batches.
 
 - **[Molecular Standardization](molecular_standardization.md):** Why molecular standardization matters for ML pipelines. We walk through Workbench's four-step pipeline — cleanup, salt handling, charge neutralization, and tautomer canonicalization. We also describe **2D** and **3D** molecular descriptors computed by our feature endpoints.

workbench-0.8.332/docs/images/chemprop_mt_expansion.svg

@@ -0,0 +1,111 @@
+<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 540" font-family="'Segoe UI', Roboto, 'Helvetica Neue', Arial, sans-serif">
+  <defs>
+    <radialGradient id="exp-primary" cx="50%" cy="50%" r="50%">
+      <stop offset="0%" stop-color="#1E88E5" stop-opacity="0.55"/>
+      <stop offset="100%" stop-color="#1565C0" stop-opacity="0.15"/>
+    </radialGradient>
+    <radialGradient id="exp-aux" cx="50%" cy="50%" r="50%">
+      <stop offset="0%" stop-color="#FB8C00" stop-opacity="0.40"/>
+      <stop offset="100%" stop-color="#EF6C00" stop-opacity="0.08"/>
+    </radialGradient>
+    <filter id="exp-shadow">
+      <feDropShadow dx="1" dy="2" stdDeviation="2" flood-opacity="0.18"/>
+    </filter>
+    <marker id="exp-arrow" markerWidth="9" markerHeight="7" refX="8" refY="3.5" orient="auto">
+      <path d="M0,0 L9,3.5 L0,7 Z" fill="#888"/>
+    </marker>
+  </defs>
+
+  <!-- Title -->
+  <text x="500" y="38" text-anchor="middle" font-size="26" font-weight="700" fill="#E0E0E0">Chemical-Space Expansion</text>
+  <text x="500" y="64" text-anchor="middle" font-size="17" fill="#9E9E9E">Auxiliary tasks bring novel chemistry the primary never saw</text>
+
+  <!-- ============================================================== -->
+  <!-- LEFT PANEL: Single-task encoder                                  -->
+  <!-- ============================================================== -->
+  <rect x="30" y="90" width="430" height="380" rx="14" fill="#1F2A33" stroke="#37474F" stroke-width="1.5"/>
+  <text x="245" y="121" text-anchor="middle" font-size="19" font-weight="700" fill="#E0E0E0">Single-Task Encoder</text>
+  <text x="245" y="143" text-anchor="middle" font-size="16" fill="#90A4AE">trained on primary task only</text>
+
+  <!-- Primary cluster -->
+  <ellipse cx="220" cy="295" rx="105" ry="90" fill="url(#exp-primary)" stroke="#1E88E5" stroke-width="2"/>
+  <text x="220" y="298" text-anchor="middle" font-size="16" font-weight="700" fill="#BBDEFB">Primary task</text>
+  <text x="220" y="319" text-anchor="middle" font-size="15" fill="#90CAF9">in-house chemotypes</text>
+
+  <!-- Primary dots -->
+  <g fill="#1E88E5">
+    <circle cx="180" cy="255" r="3.5"/><circle cx="210" cy="245" r="3.5"/>
+    <circle cx="245" cy="260" r="3.5"/><circle cx="275" cy="280" r="3.5"/>
+    <circle cx="170" cy="295" r="3.5"/><circle cx="200" cy="340" r="3.5"/>
+    <circle cx="240" cy="350" r="3.5"/><circle cx="265" cy="325" r="3.5"/>
+    <circle cx="155" cy="325" r="3.5"/><circle cx="225" cy="275" r="3.5"/>
+  </g>
+
+  <!-- Test compound (out-of-distribution for primary) -->
+  <g>
+    <circle cx="380" cy="200" r="7" fill="none" stroke="#FF5252" stroke-width="2.5"/>
+    <circle cx="380" cy="200" r="3" fill="#FF5252"/>
+    <text x="395" y="195" font-size="16" font-weight="700" fill="#FF5252">test compound</text>
+    <text x="395" y="215" font-size="14" fill="#EF9A9A">(out of distribution)</text>
+  </g>
+  <line x1="370" y1="207" x2="320" y2="240" stroke="#FF5252" stroke-width="1.5" stroke-dasharray="4,3"/>
+
+  <!-- Verdict -->
+  <rect x="55" y="408" width="380" height="50" rx="8" fill="#2D1B1B" stroke="#FF5252" stroke-width="1.5"/>
+  <text x="245" y="431" text-anchor="middle" font-size="16" font-weight="700" fill="#FF8A80">Encoder never saw this chemistry</text>
+  <text x="245" y="450" text-anchor="middle" font-size="14" fill="#EF9A9A">unreliable extrapolation, brittle predictions</text>
+
+  <!-- ============================================================== -->
+  <!-- RIGHT PANEL: Multi-task encoder                                  -->
+  <!-- ============================================================== -->
+  <rect x="540" y="90" width="430" height="380" rx="14" fill="#1F2A33" stroke="#37474F" stroke-width="1.5"/>
+  <text x="755" y="121" text-anchor="middle" font-size="19" font-weight="700" fill="#E0E0E0">Multi-Task Encoder</text>
+  <text x="755" y="143" text-anchor="middle" font-size="16" fill="#90A4AE">trained on primary + auxiliary tasks</text>
+
+  <!-- Aux extension cloud (large, behind primary) -->
+  <ellipse cx="755" cy="290" rx="180" ry="140" fill="url(#exp-aux)" stroke="#FB8C00" stroke-width="2" stroke-dasharray="6,3"/>
+  <text x="845" y="180" text-anchor="middle" font-size="17" font-weight="700" fill="#FFCC80">Aux chemistry</text>
+  <text x="845" y="200" text-anchor="middle" font-size="15" fill="#FFB74D">public / related endpoints</text>
+
+  <!-- Primary cluster (same as left panel) -->
+  <ellipse cx="700" cy="300" rx="80" ry="68" fill="url(#exp-primary)" stroke="#1E88E5" stroke-width="2"/>
+  <text x="700" y="305" text-anchor="middle" font-size="15" font-weight="700" fill="#BBDEFB">Primary</text>
+
+  <!-- Primary dots -->
+  <g fill="#1E88E5">
+    <circle cx="680" cy="275" r="3.5"/><circle cx="705" cy="267" r="3.5"/>
+    <circle cx="730" cy="283" r="3.5"/><circle cx="660" cy="300" r="3.5"/>
+    <circle cx="695" cy="335" r="3.5"/><circle cx="725" cy="325" r="3.5"/>
+    <circle cx="710" cy="300" r="3.5"/><circle cx="675" cy="325" r="3.5"/>
+  </g>
+
+  <!-- Aux-only dots (scattered around the larger cloud, NOT in primary) -->
+  <g fill="#FB8C00">
+    <circle cx="610" cy="225" r="3"/><circle cx="640" cy="205" r="3"/>
+    <circle cx="600" cy="275" r="3"/><circle cx="585" cy="325" r="3"/>
+    <circle cx="610" cy="370" r="3"/><circle cx="650" cy="390" r="3"/>
+    <circle cx="720" cy="400" r="3"/><circle cx="790" cy="390" r="3"/>
+    <circle cx="845" cy="365" r="3"/><circle cx="880" cy="325" r="3"/>
+    <circle cx="890" cy="265" r="3"/><circle cx="800" cy="205" r="3"/>
+    <circle cx="755" cy="200" r="3"/><circle cx="700" cy="205" r="3"/>
+    <circle cx="870" cy="290" r="3"/><circle cx="595" cy="300" r="3"/>
+    <circle cx="620" cy="340" r="3"/>
+  </g>
+
+  <!-- Test compound (now INSIDE the union, sitting on an aux dot) -->
+  <g>
+    <circle cx="820" cy="225" r="7" fill="none" stroke="#69F0AE" stroke-width="2.5"/>
+    <circle cx="820" cy="225" r="3" fill="#69F0AE"/>
+    <text x="836" y="222" font-size="16" font-weight="700" fill="#69F0AE">test compound</text>
+    <text x="836" y="241" font-size="14" fill="#A5D6A7">(covered by aux)</text>
+  </g>
+
+  <!-- Verdict -->
+  <rect x="565" y="408" width="380" height="50" rx="8" fill="#152418" stroke="#69F0AE" stroke-width="1.5"/>
+  <text x="755" y="431" text-anchor="middle" font-size="16" font-weight="700" fill="#A5D6A7">Encoder learned features across both</text>
+  <text x="755" y="450" text-anchor="middle" font-size="14" fill="#C8E6C9">extends to chemistry the primary never labeled</text>
+
+  <!-- Bottom annotation -->
+  <rect x="30" y="488" width="940" height="42" rx="10" fill="none" stroke="#FF9800" stroke-width="1.5" stroke-dasharray="6,3"/>
+  <text x="500" y="515" text-anchor="middle" font-size="16" fill="#BDBDBD">Aux-only compounds give the <tspan font-weight="700" fill="#FF9800">encoder</tspan> more chemistry to learn from — even without a primary gradient</text>
+</svg>