BOOM: Benchmarking Out-Of-distribution Molecular Property Predictions of Machine Learning Models

In plain words

Scientists use machine learning to guess a molecule's properties, such as how it stores energy or how it absorbs light. A good guess lets them screen many candidate molecules without a slow simulation or a lab test. The hard case is the one that matters most: finding a molecule whose property is more extreme than anything the model has seen. Most benchmarks hide this case, because their test molecules come from the same range as the training molecules.

BOOM is a benchmark that tests this case directly. For each property, the molecules with the most unusual values are kept out of training and used as a separate test set. A model is judged on those molecules and on ordinary test molecules, so the two results can be compared.

Across 15 models, from simple tree ensembles to large transformers and 3D graph networks, no model did well on every property. Models that were accurate on ordinary molecules were often much worse on the unusual ones. Pretrained chemical language models did not extrapolate better than models trained from scratch. The authors present the benchmark as a challenge for future models.

TL;DR

BOOM builds out-of-distribution (OOD) test sets for 10 molecular properties: the molecules in the lowest-density region of each property's kernel density estimate are held out §2 app. A.1. It reports RMSE on in-distribution (ID) and OOD test sets for 15 models, averaged over 3 training runs §2 tab. 7. The best model's average OOD error is about 3x its ID error, and no model leads on every task [abstract] §4.1. Pretraining ChemBERTa, MoLFormer and Regression Transformer improved ID error by 12% to 35% but gave no significant OOD gain §4.2. Code is Apache-2.0 LICENSE.

Problem

Molecular discovery needs predictions for molecules whose properties lie outside the range seen in training. The authors argue that current benchmarks mostly measure ID performance on datasets that are not hard enough, and that the OOD behaviour of chemistry models is largely unexplored §1. The MoleculeNet protocols it cites (random, scaffold, stratified and time splits) are the standard ones §3, and the paper argues that distribution-based splits are missing §1.

The paper defines OOD in property space. The OOD test set is a low-density region of the property distribution, so the test targets are rarer than the training targets §2. Depending on the property, this can also cause a shift in chemical structure §2. The definition follows a complement-distribution framing from prior work, which the paper cites §2.

Method

Benchmark construction

• Properties. Ten targets: eight QM9 properties (polarizability α, heat capacity C_v, HOMO, LUMO, HOMO-LUMO gap, dipole moment μ, electronic spatial extent ⟨R²⟩, zero-point vibrational energy ZPVE) and two from the 10k dataset (heat of formation HoF and density) §2 app. A.1.
• Source size. The 10k dataset has 10,206 molecules §2. The Chemprop pretraining experiment uses 133,885 QM9 training examples app. C.3. Pitfalls covers the QM9 split size.
• OOD split. Fit a Gaussian kernel density estimate to each property's values, using Scott's bandwidth rule. The molecules with the lowest estimated density form the OOD test set. The paper uses 1,000 molecules for the 10k dataset and the lowest 10% for QM9 §2 app. A.1. The code uses 10,000 molecules for QM9 prepare_splits_qm9.py.
• ID split. From the remaining molecules, 5% are drawn at random as the ID test set and the rest are used for training and fine-tuning app. A.1. The text says 10% for QM9 app. A.1; the code uses 5% for both datasets prepare_splits_qm9.py prepare_splits_10k.py.
• Extra dataset. Lipophilicity (MoleculeNet, 4,200 experimental octanol/water measurements), used on a subset of models as an experimentally measured OOD check §2 app. C.9.

Key equations as written in the paper:

P(\mathcal{M}, y \mid \mathcal{D}_{\text{test}}) \neq P(\mathcal{M}, y \mid \mathcal{D}_{\text{train}}) eq. 1

0 < P(y_{\text{test}}) \leq P(y_{\text{train}}) eq. 2

Eq. 2 states that the test-target density is non-zero and no higher than the training-target density §2. The released code does not evaluate eq. 2 directly. It ranks molecules by kernel density and cuts the lowest ones, as the pseudocode below shows prepare_splits_qm9.py.

Split procedure, as implemented in the released code. Per property, with the same logic for QM9 and the 10k dataset:

# BOOM split for one property  [§2] [app. A.1]
# Input:  values (N,)  one property value per molecule
#         n_ood        1000 (10k), 10000 (QM9, in the released code)
# Output: label per molecule: "ood", "id" or "train"

def split_property(values, n_ood, rng):
    # Gaussian kernel density, Scott's bandwidth  [app. A.1]
    # [code: boom/data/prepare_splits_qm9.py#L171-L172]
    density = fit_gaussian_kde(values, bandwidth="scott")
    # Density at each molecule's own property value
    # [code: boom/data/prepare_splits_qm9.py#L178]
    scores = density.score_samples(values)
    # Lowest-density n_ood molecules form the OOD test set  [§2]
    # [code: boom/data/prepare_splits_qm9.py#L180-L186]
    ood_index = indices_of_lowest(scores, n_ood)
    labels = ["train"] * len(values)
    for i in range(len(values)):
        if i in ood_index:
            labels[i] = "ood"
        # 5% of the rest go to ID test, the others to training
        # [code: boom/data/prepare_splits_qm9.py#L198-L204]
        elif rng.random() < 0.05:
            labels[i] = "id"
    return labels

The released code calls random.seed(42) before each property prepare_splits_qm9.py prepare_splits_10k.py. The paper does not state a seed.

Metrics

• RMSE on the ID and OOD test sets. Models whose OOD RMSE is close to their ID RMSE generalize well §2.
• Binned R² on OOD. The paper argues that OOD predictions span a wide range, so the sample mean is far from most samples and plain R² is inflated. It therefore averages R² over the OOD molecules that fall in each end of the property range §2. The tail fraction is not stated, and no implementation was found in the released repository.
• Each result is the mean of 3 training runs, with variance reported §2.

Models

Fifteen models in three groups §2 app. B.1 tab. 3:

• Traditional: Random Forest and an MLP, both on RDKit descriptors app. B.2.
• Transformers: ChemBERTa, MoLFormer, Regression Transformer, ModernBERT app. B.3. The paper calls these SMILES-based tab. 3.
• Graph and 3D: Chemprop (2D bond graph), TGNN, IGNN, EGNN, MACE, GotenNet app. B.4 tab. 3.
• Hybrid: Graphormer-3D, GeoFormer, Equivariant Transformer (ET) app. B.5 tab. 3.

Architectures were used off the shelf. Hyperparameter sweeps covered only non-architectural settings such as learning rate and training steps, with 10% of the training data held out for selection app. C.1.

Model summary tab. 3

Parameter counts are as the paper states them.

The GNN message-passing forms are given as equations in the appendix: a topological update with bond and node terms, an invariant update that adds a learned function of inter-atomic distance, and an equivariant update that also updates coordinates app. B.4. These are not reproduced as equation numbers here; the paper cites them in that appendix section.

Hyperparameters

Values as stated in the paper or taken from the released code. Where only the code gives a value, the card says so.

Runtimes per epoch on 10k and QM9 are given in tab. 4.

Results

Author-reported. Setup: 10 molecular property tasks, 15 models, RMSE on ID and OOD test sets, mean of 3 training runs §2 tab. 7.

• Headline gap. The top-performing model's average OOD error is about 3 times its ID error [abstract].
• Best OOD models. GotenNet has the lowest OOD RMSE on 7 of 10 tasks, and MACE on 2 (α and C_v) §4.1 tab. 7. The paper's text does not mention EGNN, which has the lowest OOD RMSE on density (0.0279) tab. 7.
• Best ID models. GeoFormer has the lowest ID RMSE on 3 tasks (HOMO, tied with ET; LUMO; C_v) §4.1 tab. 7.
• ID and OOD are weakly linked. Across tasks, the distance correlation between ID and OOD performance is moderate (d = 0.4678, p = 0.0010). High ID scores do not guarantee OOD scores §4.7.
• Electronic properties. No model achieves good OOD performance on HOMO, LUMO, gap and μ §4.1.
• Pretraining. Pretraining improved ID RMSE by 31% for ChemBERTa, 35% for MoLFormer and 12% for Regression Transformer, but OOD RMSE did not change significantly. Binned OOD R² fell by 53% for MoLFormer and 39% for ChemBERTa §4.2 fig. 4.
• Cross-property pretraining. Chemprop pretrained on one QM9 property and fine-tuned on another degraded OOD performance when the Pearson correlation between the two properties was below about 0.35, and improved it when they were strongly correlated §4.2 fig. 5.
• Representation. In the paper's averages, 3D models have lower OOD RMSE than graph models on all 10 columns. Graph models beat SMILES models on 9 of 10 §4.4 tab. 1. These are averages over the models in each group, so they mix architectures.
• Hyperparameter search. Searching 50 random Chemprop configurations per property lowered OOD RMSE for density (−60%), heat of formation (−23%) and ZPVE (−50%). The paper does not report a broad gain §4.3 fig. 6.
• Data augmentation. Adding a few thousand OOD molecules to training improved OOD performance on 7 of 8 QM9 tasks. The gain was sizeable with about 4% augmented data. μ did not improve §4.5 fig. 7.
• ModernBERT. Among transformers, ModernBERT reduced OOD RMSE for heat of formation by more than 58% and for C_v by more than 78% against the best other transformer §4.6 fig. 8.
• Lipophilicity. Chemprop had ID RMSE 0.463 and OOD RMSE 1.051. MoLFormer had the lowest OOD RMSE of the five models tested (0.956) tab. 9.
• LLMs. Llama 3.1 8B without fine-tuning gave very large OOD errors, for example 262.99 for heat of formation. The authors caution against comparing these numbers with the main table because of possible training-data overlap app. D.2 tab. 11.

Averaged RMSE by representation tab. 1

RMSE on ID and OOD test sets tab. 7

Mean of 3 runs. In the paper, the first row of each model is the ID split, the second the OOD split. The paper's table labels Random Forest and MLP as SMILES; see Pitfalls.

Mean batched R² (ID and OOD) tab. 5

The appendix calls this "batched R²" and the main text calls the OOD version "binned R²". The paper does not say whether these are the same metric (see Pitfalls). Rows match the RMSE table above.

Standard deviations of these values tab. 6 are in the paper and are not transcribed here.

Lipophilicity RMSE tab. 9

Mean of 3 runs ± standard deviation.

Statistical tests on OOD performance by model family tab. 8

Kruskal-Wallis across geometric, transformer and graph groups per property, then Mann-Whitney U per pair. N = 39 per property.

Runtimes tab. 4

Seconds per epoch, as reported.

Llama 3.1 8B, no fine-tuning tab. 11, tab. 12

The HOMO column repeats the HoF values exactly in both rows, which looks like a typesetting error (see Pitfalls).

Paper vs code

Repository: https://github.com/FLASK-LLNL/BOOM at 2bb1d835e62cd2c6a0d1d0321e2f283f7d802ea3, Apache-2.0 LICENSE. The owner is the FLASK-LLNL organisation, and the repository's CONTRIBUTORS file lists most of the paper's authors. The paper links the repository in its abstract. The repository is therefore official.

• Chemprop version. The paper says Chemprop v1.4.0 was used for pretraining and HPO app. C.3 app. C.4. The repository builds on Chemprop 1.7.1 with a modified training loop readme.md.
• L1 penalty. The Chemprop fork adds an L1 term to the training loss, weighted by l1_lambda train_FLASK.py. The paper does not describe it. The default is 0.0, so the penalty only matters if the paper runs set it args_FLASK.py.
• QM9 OOD size. The paper says the lowest 10% of QM9 molecules form the OOD test set app. A.1. The released code fixes it at 10,000 molecules prepare_splits_qm9.py. Figure 7 also refers to 10,000 OOD examples fig. 7.
• QM9 ID fraction. The paper says 10% of QM9 molecules form the ID test set app. A.1. The code uses 5% for QM9 prepare_splits_qm9.py, as it does for the 10k dataset.
• Random Forest and MLP features. The paper's results table labels them SMILES tab. 7. The paper's model summary and appendix say RDKit descriptors tab. 3 app. B.2. The code uses RDKit descriptors dataset.py MLP_train.py.
• Validation fraction. The paper says 10% holdout for hyperparameter selection app. C.1. The ChemBERTa script uses 20% run_experiment_density.py. The MACE config uses 10% density.yaml.
• Number of model-task combinations. The abstract says over 150 [abstract]. The repository README says over 140 README.md.
• Missing code. The binned R² metric is not implemented. The Chemprop OOD hyperparameter sweep and the Chemprop augmentation experiment are not in the repository. The Chemprop readme links OOD_Experiments directories that are absent readme.md.
• Data preparation. The README advises against rerunning data preparation, because the seed may not reproduce the same split README.md.
• Author name. The arXiv metadata lists "Busra Sahin"; the LaTeX source has "Busra Demirci". The card uses the arXiv metadata.

Pitfalls

• QM9 split size and ID fraction disagree. Use the released QM9 split files, or copy the code's choice (10,000 OOD, 5% ID), and say which one you used. The paper's text (10% OOD, 10% ID) would not reproduce the released data app. A.1 prepare_splits_qm9.py.
• "⟨R²⟩" in every results table is the electronic spatial extent, not the coefficient of determination. Its values are in RMSE units such as 52.3 tab. 7.
• Units look swapped. The dataset table pairs heat of formation with g/cc and density with kcal/mol tab. 2. The LLM prompt table pairs the same units tab. 10. Check units against the data before using them.
• Possible copy error in Table 7. The Graphormer OOD row is identical to the MoLFormer OOD row in every column except HoF tab. 7. Treat Graphormer's OOD numbers as unverified until checked against outputs.
• Possible copy error in Table 11. The HOMO column duplicates the HoF column in both LLM rows tab. 11.
• Random Forest and MLP labelled SMILES. Use the RDKit descriptor pipeline in the code when reproducing them tab. 7 dataset.py.
• Binned R² is under-specified. The paper does not give the tail fraction, the binning rule or the handling of the two tails §2. Implement it only with a stated convention, and report that convention.
• Table 8 caption disagrees with its values. The caption says only μ and GAP fail to reject the null. At 0.05, HOMO (0.05119) and ZPVE (0.05163) also do not reject tab. 8.
• Batched R² and binned R² may be different metrics. The appendix table uses "batched R²" and the main text uses "binned R²" tab. 5 §2. The paper does not define the link, so do not mix the two names.
• Chemprop L1 term is not in the paper. Confirm l1_lambda for any Chemprop run you reproduce train_FLASK.py.
• Two ET configurations. The code has lr 1e-4 and lr 1e-6 at batch 64 for 20 epochs, and the paper does not say which was used run_experiment.py.
• Tuning and testing use the same set. The paper picks the best Chemprop configuration per property by minimum OOD test RMSE fig. 6 app. C.4. The reported gains are therefore on the same test set used for selection, which inflates them.
• Augmentation selection is one sentence. Augmented molecules are drawn from the QM9 OOD set with properties beyond the 25th and 75th quantiles §4.5. The paper does not say which side is used for each property, or how the 5,000 held-out molecules are chosen fig. 7.
• LLM results may leak training data. The authors say the density values of some OOD molecules may appear in the LLM's training corpus app. D. Do not treat these as a fair comparison with the other models.
• Lipophilicity split is not described. Table 9 reports ID and OOD RMSE, but the split procedure for Lipophilicity is not stated app. C.9.
• Metadata author names. The arXiv author list and the LaTeX author list differ for one author (see Paper vs code).

Lineage

• Builds on the MoleculeNet benchmark and its split protocols, which the paper contrasts with its own property-space split §3.
• Uses or compares against ChemBERTa, MoLFormer, Regression Transformer, ModernBERT, Chemprop, EGNN, MACE, GotenNet, GeoFormer, Graphormer-3D and Equivariant Transformer §2 app. B.
• Related benchmarks named in the related work section: MatBench, MatFold, DrugOOD and the OOD-generalization work of Mole-OOD §3.
• Superseded by: not known from this wiki.