-
-
Generalisation Sheet for Validation
-
MainPage
-
Docs
Inspiration
This project was inspired by "How Powerful are Graph Neural Networks?" by Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka (ICLR 2019). The paper introduces the Graph Isomorphism Network (GIN), which matches the expressive power of the Weisfeiler-Lehman graph isomorphism test by summing neighbor features instead of averaging them. Since molecules are inherently graphs, we applied that idea to mutagenicity prediction — a critical DNA-damage safety check for new drugs. Because the prediction depends so heavily on graph structure, the question of whether the model could be trusted felt just as important as whether it could predict.
What it does
MutagenAI takes a chemical input (a SMILES string or a common name) and returns a mutagenicity verdict built from four layers, not a single score:
- 3-model GIN ensemble — trained on the AMES dataset, giving both a probability and a confidence spread based on how much the three models disagree
- Expert structural alerts — deterministic SMARTS rules for known DNA-reactive chemical groups, combined with the model per ICH M7 guidelines
- Metabolite screening — catches "pro-mutagens," compounds that only become dangerous after the body metabolizes them
- Applicability-domain score — measures how similar a query is to the training data, so predictions on unfamiliar chemistry get flagged rather than silently trusted
Every prediction comes with atom-level explanations, an interactive test sheet, and OECD QMRF-style documentation, so the tool's limitations stay visible to whoever's using it.
How we built it
Development split into two main phases: getting the model to predict well, and then making sure that prediction could actually be trusted.
Phase 1 — Architecture. Trained GIN on the MUTAG dataset as a sanity check, reaching about 92% accuracy. But MUTAG skews heavily toward aromatic compounds, so that number didn't say much about performance on real, diverse chemistry.
Phase 2 — Generalization. Moved to the larger AMES dataset (~6,500 compounds). After 150 epochs, accuracy settled at 77–79% — a real drop from MUTAG's 92%, and the expected trade-off for a harder, much more diverse dataset.
Phase 3 — Reliability. ~78% accuracy isn't something you hand someone as a bare safety label, so three things got layered on top of the raw model.
At each layer \(k\), GIN updates a node's representation by combining it with the sum of its neighbors' representations — summing rather than averaging or max-pooling is exactly what gives GIN its extra expressive power:
$$ h_v^{(k)} = \text{MLP}^{(k)} \left( (1 + \epsilon^{(k)}) \cdot h_v^{(k-1)} + \sum_{u \in \mathcal{N}(v)} h_u^{(k-1)} \right) $$
where \(\epsilon^{(k)}\) is a learnable parameter controlling how much a node weighs its own previous representation against its neighbors. The whole molecule's representation is then built by summing all node features at every layer and concatenating across layers, so the model keeps information from each stage of message-passing rather than only the last one:
$$ h_G = \text{CONCAT}\left(\sum_{v \in G} h_v^{(k)} \;\middle|\; k = 0, 1, \dots, K\right) $$
To turn that into an honest confidence signal, the 3-model ensemble measures how much its members disagree — each model's predicted probability \(p_i\) is compared to the ensemble mean \(\bar{p}\), and the spread is reported as a variance:
$$ \sigma^2 = \frac{1}{N} \sum_{i=1}^{N} (p_i - \bar{p})^2 $$
A wide spread is the ensemble's way of saying "I'm not sure," which a single model can't give you on its own. On top of that, deterministic SMARTS alerts and the applicability-domain check flag chemistry the model is likely out of its depth on. The whole thing ships as a Dash app, deployed behind Gunicorn.
Challenges we ran into
- Generalization vs. accuracy. Choosing the diverse AMES dataset (~78% accuracy) over the narrower MUTAG dataset (92%) was a real product decision — it meant giving up a better-looking number for one that actually reflects real-world performance.
- Contextualizing a safety prediction. A single accuracy number isn't enough for something like genotoxicity. Ensembling, SMARTS alerts, and domain checks all exist because that one number needed real context.
- Deployment hurdles. Getting the model into production surfaced a Git LFS issue (the "checkpoint" that shipped was actually a 130-byte pointer file, not the real weights) and a memory bottleneck, since each Gunicorn worker loading its own copy of the ensemble was enough to OOM-kill a small container.
Accomplishments that we're proud of
- Turned a graph-isomorphism research paper into a tool that addresses a real chemical safety problem, not just a benchmark score
- Built a multi-layered trust system — alerts, metabolite screening, domain scores — that reports how much to trust a prediction, not just the prediction itself
- Every verdict is explainable at the atom level; nothing here is a black box
- Shipped a fully working, deployed, interactive app with a live test sheet, so anyone can check the pipeline against their own known compounds on the spot
What we learned
- Headline accuracy means little without knowing what data it was measured on — a narrow benchmark can flatter a model in ways real-world chemistry immediately exposes
- Confidence and reliability signals aren't optional for a safety-relevant tool — they're what separates "honest about its limits" from "just looks confident"
- Cheap ensembling for variance, combined with deterministic expert rules, covers a lot of the failure modes a purely learned model can't reliably catch on its own
- Getting a model into production is its own discipline, separate from training it — memory budgeting, file integrity, and deployment logistics mattered as much as the architecture did
What's next for MutagenAI
- Expand from one-step metabolite screening to multi-step bioactivation pathways
- Add stereochemistry- and salt-form-aware predictions
- Benchmark against external, held-out toxicology datasets to more rigorously validate generalization
- Explore uncertainty-calibration techniques (e.g. conformal prediction) on top of raw ensemble disagreement
- Extend the same explainable, multi-layer approach to other toxicity endpoints beyond mutagenicity
Log in or sign up for Devpost to join the conversation.