Why Does Machine Learning Fail for Small-Molecule Mass Spectrometry?

Code, analysis and evaluation pipelines used to benchmark machine-learning models that predict molecular fingerprints from small-molecule MS/MS spectra, and to diagnose why they fail: retrieval baselines that expose how much of the task is memorisation, adversarial learning-to-split partitions that expose the hardest generalisation gap, and analyses of out-of-vocabulary chemical formulas and experimental conditions.

If you want to run the same analysis on your own model, jump to Evaluating your own model.

---

Repository layout

ML_MS_analysis/
├── benchmarked_models/
│   ├── nearest_neighbour/     # Retrieval baselines: same-formula NN, DreaMS NN, formula-aware fallbacks, MIST comparison
│   ├── learning_to_split/     # `learning_to_split` package: losses + sampler for adversarial train/test splitting
│   ├── other_baselines/       # Binned-MLP, MS-transformer and formula-transformer fingerprint predictors (+ LS drivers)
│   └── mist/                  # Re-implementation of MIST used for error analysis (+ LS driver)
├── diagnosis/                 # Notebooks that turn cached predictions into the tables and figures of the paper
├── figures/                   # Figures produced by the diagnosis notebooks
├── data/                      # (not tracked) processed datasets, splits and MGF caches -- see Data
├── results/                   # (not tracked) all cached predictions and metrics -- created by the scripts
└── download.sh                # Helper to fetch the original CANOPUS export

Every folder under benchmarked_models/ has its own README with the exact commands.

Data

All processed datasets, split definitions and intermediate files are available on Google Drive:

https://drive.google.com/drive/folders/1v11lTwFSdlSRJ6ETLHqkbT809Ji9w0OY?usp=drive_link

Download them into data/ so that the layout is

data/
├── processed_data/
│   ├── NPLIB1.pkl                 # list of records (see below)
│   └── massspecgym.pkl
├── splits/{NPLIB1,massspecgym}/{scaffold,random,LS}.json   # {"train": [...], "val": [...], "test": [...]} record ids
└── MGF_files/{dataset}/{split}/{train,test}.mgf              # same spectra, same order, for DreaMS

Each record in processed_data/*.pkl is a dict with (at least) id_, smiles, formula, formula_corrected (precursor ion formula), precursor_type (adduct), precursor_MZ_final, instrument_type, collision_energy, peaks (list of {"mz", "intensity", "intensity_norm", "comment": {"f_pred": <sub-formula>}}) and FPs (MACCS, morgan4_{256,1024,2048,4096}, morgan6_{...} as bit strings).

Fingerprint definition. morgan4_4096 is an RDKit Morgan fingerprint of radius 2 (diameter 4) with 4096 bits computed on the kekulized molecule with aromatic flags cleared (MS_processing/dataprocessing/utils/chem_utils.py). It is not bit-identical to the plain RDKit GetMorganGenerator(radius=2, fpSize=4096) fingerprint that MIST is trained on; see the errata below before comparing numbers across models.

The NIST 2023 library cannot be redistributed; the scripts that mention nist2023 require your own licensed copy.

Environment

The retrieval and analysis code needs only numpy scipy scikit-learn tqdm rdkit matplotlib pandas (plus dreams for DreaMS embeddings and torch transformers matchms for the retrieval-task script). Model training additionally needs torch pytorch_lightning wandb pyyaml and, for MIST, the original mist package. Install the learning-to-split package with

pip install -e benchmarked_models/learning_to_split

Weights & Biases logging is used by the training scripts: run wandb login or export WANDB_API_KEY first.

Quick start

# 1. Nearest-neighbour baselines (fast, CPU only)
cd benchmarked_models/nearest_neighbour
python 02_compute_nn.py                                   # same-formula NN on binned spectra
python 04_compute_nn_formula_fallback.py --dataset NPLIB1 --split scaffold --rep binned   # formula-aware fallbacks
python 06_build_full_test_nn.py                           # full-test-set predictions + comparison with MIST
python 07_rescore_nn_with_mist_fp.py                      # same comparison with MIST's fingerprint definition

# 2. Fingerprint predictors (GPU)
cd ../other_baselines && python train.py --config_file w_meta_config.yaml
cd ../mist            && python train.py --config_file NPLIB1_scaffold_original.yaml

# 3. Learning to split (GPU)
cd ../other_baselines && python train_LS.py --config_file LS_config.yaml
cd ../mist            && python train_LS.py --config_file LS_config.yaml

# 4. Tables and figures
jupyter lab ../../diagnosis

Evaluating your own model

The analyses only need your model's per-spectrum predictions on the provided splits, in the same format as the models benchmarked here.

1. Train on the provided splits. Use data/splits/{dataset}/{split}.json (train/val/test ids). Do not touch test during model selection. The LS.json split is the learned adversarial split.
2. Export predictions to results/FP_prediction/<your_model>/{dataset}/<run_name>_{split}/:
   • test_results.pkl: dict mapping test id (string) to {"pred": [4096 floats], "GT": [4096 floats], "jaccard": float},
   • test_performance.json: {"jaccard": <mean over the test split>}. jaccard is the Tanimoto similarity between pred > 0.5 and GT. The run name must end with _{split} (_scaffold, _random, _LS). Use the same fingerprint definition for GT as the baselines you compare against (see the errata below), or re-score them with 07_rescore_nn_with_mist_fp.py as a template.
3. Compare with the retrieval baselines. diagnosis/01b_get_FP_prediction_results.ipynb lists every model folder found under results/FP_prediction/. benchmarked_models/nearest_neighbour/06_build_full_test_nn.py reports the nearest-neighbour baseline on the full test split and on the formula-matched / no-formula-match subsets; report your model on the same subsets (the ids are the first element of results/nearest_neighbour/nn_sim/{dataset}_{split}.pkl for the formula-matched subset).
4. Learn the hardest split for your model. Implement a splitter head for your architecture and reuse learning_to_split (losses + sampler) with the outer loop in benchmarked_models/other_baselines/train_LS.py. benchmarked_models/learning_to_split/README.md describes the interface in detail; the output best_split.pkl can be turned into a splits/{dataset}/LS.json and analysed with diagnosis/02a_analyze_LS_split.ipynb.

Errata: earlier MIST vs nearest-neighbour comparison

Earlier versions of our comparison between MIST and the nearest-neighbour baselines were not apples-to-apples. Two things differed:

issue	what was done before	what is done now
Fingerprint	Nearest-neighbour Jaccard used the dataset morgan4_4096 bits (kekulized Morgan r=2); MIST was trained and scored on plain RDKit Morgan r=2 / 4096. The two ground truths agree with a Jaccard of only ~0.5.	All nearest-neighbour rows are re-scored with MIST's fingerprint (07_rescore_nn_with_mist_fp.py).
Test spectra	On NPLIB1, MIST was trained on a subset of the training split (12,726 vs 12,909 scaffold; 12,665 vs 12,800 random) and evaluated on 2,684 / 2,721 of the 2,689 / 2,744 test spectra. MassSpecGym splits were identical.	Numbers are reported on the full test split and on the MIST-common spectra; both are given. No train/test leakage was found.

Corrected full-test-set mean Tanimoto (MIST fingerprint definition); previous values in parentheses:

dataset / split	NN, best-match fallback	NN, sub-formula fallback	DreaMS, best-match fallback	DreaMS, sub-formula fallback	MIST
NPLIB1 / scaffold	0.197 (0.194)	0.297 (0.293)	0.259 (0.258)	0.294 (0.291)	0.241
NPLIB1 / random	0.696 (0.695)	0.747 (0.746)	0.745 (0.744)	0.758 (0.756)	0.547
MassSpecGym / scaffold	0.258 (0.256)	0.347 (0.344)	0.308 (0.306)	0.349 (0.346)	0.267
MassSpecGym / random	0.929 (0.929)	0.934 (0.934)	0.945 (0.944)	0.947 (0.946)	0.674

The conclusions are unchanged (all shifts are below 0.006), but only the corrected numbers should be quoted. Two remaining asymmetries are documented rather than removed: the MIST scaffold-split runs embed the instrument type (embed_instrument: true) while the retrieval baselines never use it, and the DreaMS checkpoint was pre-trained (self-supervised, no labels) on a public corpus that may contain some test spectra. Full details: benchmarked_models/nearest_neighbour/formula_fallback_report.md.

License

See LICENSE.