Retention Time Prediction for Small Molecules in Untargeted Metabolomics

TL;DR

Graph neural networks that predict LC retention times of small molecules from structure alone, to raise confidence in peak annotations for untargeted metabolomics LC-MS/MS. Ten graph architectures were benchmarked against fingerprint and descriptor baselines; AttentiveFP won on accuracy and is ~120× faster per molecule than the descriptor pipeline it replaces. It ships as a batch CLI and a REST API, both containerised.

Key results

• Trained on ~5,000 molecules with retention times from a 10-minute C18 HPLC method
• AttentiveFP reaches 0.334 min test MAE, best of ten graph architectures
• Duplicate compounds differ by 0.196 min on the same instrument, so the model lands within 1.7× of the experiment's own reproducibility floor
• 0.5 s per molecule at inference versus 60 s for Mordred descriptors — the difference between a usable API and a batch job
• Mean absolute error used as the loss, being robust to the large deviations that dominate this dataset

Results

Test-set MAE in minutes, all ten graph architectures, lower is better (results/all_model_results.json):

Model	Validation MAE	Test MAE
AttentiveFP	0.321	0.334
GAT	0.347	0.379
GIN (contextpred)	0.409	0.401
GIN (edgepred)	0.439	0.413
MPNN	0.398	0.442
GIN (masking)	0.450	0.455
GCN	0.492	0.458
GIN (infomax)	0.456	0.471
Weave	0.475	0.482
NF	1.202	1.222

Against the non-graph baselines (validation MAE, from the summary figure above):

Representation	Validation MAE	Inference time / molecule
Morgan fingerprints	0.62	fast
Mordred descriptors + XGBoost	0.41	60 s
AttentiveFP (graph + attention)	0.32	0.5 s

Two things matter here beyond the accuracy ranking. First, the four self-supervised pretrained GIN variants did not beat a plain supervised AttentiveFP on this task — pretraining bought nothing at this dataset size. Second, descriptor calculation dominates the cost of the descriptor approach, so the graph model is not just more accurate but the only one of the two that can answer a request interactively.

The 0.196 min reproducibility floor is the number to judge all of this against: it is measured in Step 1 from compounds appearing more than once, and no model can be expected to beat it.

Background

Unidentified peaks are the central problem in untargeted metabolomics by LC-MS/MS. Matching MS/MS spectra alone leaves many candidates; adding a retention-time expectation narrows them considerably. This project predicts that expectation directly from molecular structure.

Dataset. ~5,000 molecules with retention times measured on a 10-minute C18 HPLC method, split 8:1:1 into train/validation/test.

Representations compared. Three families, ten graph architectures — see Results for the numbers.

Approach	Representation
Fingerprints	Morgan bit vectors
Descriptors	Mordred molecular descriptors + XGBoost
Graph	GCN, GAT, MPNN, Weave, NF, AttentiveFP, and four pretrained GIN variants

Splits are chemically stratified via PCA over descriptor space rather than drawn at random, so the test set is not dominated by close analogues of training compounds. Hyperparameters for AttentiveFP were then tuned with Bayesian optimization (hyperopt) over layer count, time steps, graph feature size, and dropout.

Installation

Requires Python 3.10+.

# uv (recommended)
uv venv && source .venv/bin/activate    # .venv\Scripts\activate on Windows
uv pip install -e .

# or pip
pip install -e .

The DGL stack needs extra care. dgl-lifesci is not published on PyPI at all, and dgl's PyPI releases do not always include a wheel for every platform (2.2.1, for instance, is Windows-only). Install both from conda-forge, or take dgl from the project's own wheel index:

conda install -c conda-forge dgl-lifesci   # brings dgl with it

# alternatively, for dgl only:
pip install dgl -f https://data.dgl.ai/wheels/repo.html

For development (pytest + ruff):

uv pip install -e ".[dev]"

Quickstart

Predictions need trained weights, which are not committed to this repository — only the per-column configure.json files under models/ are tracked. Point RT_MODEL_DIR at a directory laid out like this:

models/
  c18/{configure.json, model.pth}
  pfp/{configure.json, model.pth}

Batch prediction over a CSV of SMILES:

RT_MODEL_DIR=./models python scripts/predict_batch.py \
    -f data/input.csv -sc SMILES -tp c18 -ip results/

Serve the API:

RT_MODEL_DIR=./models python scripts/api_server.py
curl "http://localhost:9002/gnn?smiles=CC(=O)OC1=CC=CC=C1C(=O)O&column=C18"

{
  "canonical_smiles": ["CC(=O)Oc1ccccc1C(=O)O"],
  "inchikey": ["BSYNRYMUTXBXSQ-UHFFFAOYSA-N"],
  "RT": 3.609
}

Column names resolve case-insensitively, so c18 and C18 are equivalent.

Docker

docker build -f docker/Dockerfile.api   -t rt-prediction-api   .
docker build -f docker/Dockerfile.batch -t rt-prediction-batch .

See docker/README.md for run commands, batch-mode annotation, and the full option reference.

Workflow

1. Data cleaning and QC — SMILES standardization (neutralize, desalt, canonical tautomer), duplicate removal, and the retention-time reproducibility analysis that establishes the 0.196 min error floor.
2. Baselines and feature engineering — Mordred descriptors and fingerprints with XGBoost; PCA for dimensionality reduction and to construct chemically-stratified train/test splits.
3. GNN architecture comparison — GCN, GAT, MPNN and AttentiveFP under early stopping and dropout, starting from published defaults.
4. AttentiveFP fine-tuning — Bayesian hyperparameter search on the winning architecture. (Exported HTML; the source notebook was not retained.)
5. Deployment — batch CLI and REST API.

Chemically-stratified split, for reference:

Project structure

├── src/rt_prediction/       # Installable package
│   ├── data/                # SMILES standardization, identifiers, QC
│   ├── gnn/                 # Featurization, model loading, prediction
│   ├── inference/           # Batch helpers, model-path resolution, HTTP client
│   ├── models/              # Hyperparameter search spaces
│   └── utils/               # Filesystem helpers
├── scripts/                 # Entry points: training, batch prediction, API
├── notebooks/               # Workflow steps 1-3 and the RT filter
├── configs/model_configures/# Per-architecture GNN configs
├── models/                  # Per-column configure.json (weights not tracked)
├── docker/                  # Dockerfile.api, Dockerfile.batch, deploy docs
├── docs/                    # Exported Step 4 notebook
├── results/                 # Figures and metric summaries
└── tests/                   # Test suite

Development

pytest                # run the suite
ruff check .          # lint
ruff format .         # format

Tests that require the DGL/PyTorch stack call pytest.importorskip and report as skipped when it is absent, so the suite runs without a GPU-scale install. src/rt_prediction/gnn/utils.py is vendored from DGL-LifeSci and excluded from formatting so upstream diffs stay readable.

Code attribution

This codebase includes modified portions of Amazon's DGL-LifeSci, adapted for retention-time regression.

Original code

• Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
• SPDX-License-Identifier: Apache-2.0

Modifications

• Adapted GNN model loading and prediction utilities for retention-time regression
• Added SMILES standardization and quality-control pipelines
• Implemented batch processing and REST API deployment
• Extended hyperparameter search spaces for retention-time prediction

License

Apache-2.0 — see LICENSE.

Citation

@software{wang_rt_prediction_gnn,
  title  = {Retention Time Prediction for Small Molecules using Graph Neural Networks},
  author = {Wang, Shunyang},
  year   = {2024},
  url    = {https://github.com/Shunyang2018/rt-prediction-gnn}
}