CheMeleon-Rxn

A pretrained Chemprop Condensed-Graph-of-Reaction (CGR) D-MPNN encoder for reaction property prediction. It turns an atom-mapped reaction SMILES into a 2048-d reaction embedding, which you fine-tune on your own reaction-property data (barrier heights, reaction energies, yields, rate constants, …).

The model accompanies:

Zalte, A. S.; Fang, C.; Zheng, Y.; Metwally, E.; Cheng, A.; Glisman, A.; Pang, H.-W. Descriptor-Based Pre-Training Improves Reaction Property Prediction. ChemRxiv 2026. 10.26434/chemrxiv.15008692

Code: https://github.com/MSDLLCpapers/chemeleon-rxn

Also archived on Zenodo: 10.5281/zenodo.23923145

This checkpoint was trained and is distributed by MIT. It was trained by Akshat S. Zalte (Massachusetts Institute of Technology), an independent re-training of the paper's recipe done under MIT's own Pistachio license. Its downstream performance is statistically comparable to the results reported in the paper.

Model

• Encoder:
  • Chemprop 2.3.1 BondMessagePassing with mean aggregation.
  • Hidden size 2048, depth 6, LeakyReLU.
  • 8.9 M parameters.
• Input: an atom-mapped reaction SMILES.
  • Featurized with Chemprop's CondensedGraphOfReactionFeaturizer in REAC_DIFF mode.
  • Explicit hydrogens are kept (keep_h=True).
• Pretraining: self-supervised regression onto Mordred descriptors of the reaction, [Σ reactants ‖ (Σ products − Σ reactants)].
  • 70% of the targets are randomly masked at each step.
  • The corpus is about 2 M atom-mapped reactions.

Training data. The pretraining corpus is drawn mainly from Pistachio, a commercial reaction database licensed from NextMove Software. The training data therefore cannot be released. The rest of the corpus comes from public reaction datasets.

Usage

Requires chemprop==2.3.1. The checkpoint uses the same format as CheMeleon's chemeleon_mp.pt:

import torch
from huggingface_hub import hf_hub_download
from chemprop import data, nn, models
from chemprop.featurizers import CondensedGraphOfReactionFeaturizer

ckpt = torch.load(hf_hub_download("akshatzalte/chemeleon-rxn", "chemeleon_rxn_mp.pt"), weights_only=True)
mp = nn.BondMessagePassing(**ckpt["hyper_parameters"])
mp.load_state_dict(ckpt["state_dict"])

featurizer = CondensedGraphOfReactionFeaturizer(mode_="REAC_DIFF")

dps = [data.ReactionDatapoint.from_smi(smi, y=[y], keep_h=True) for smi, y in zip(smiles, targets)]
train_set = data.ReactionDataset(dps, featurizer=featurizer)
scaler = train_set.normalize_targets()

model = models.MPNN(
    message_passing=mp,
    agg=nn.MeanAggregation(),
    predictor=nn.RegressionFFN(input_dim=2048, hidden_dim=1024, activation="leakyrelu",
                               output_transform=nn.UnscaleTransform.from_standard_scaler(scaler)),
    init_lr=1e-4, max_lr=1e-4, final_lr=1e-4,   # flat LR, see below
)

Fine-tune with a flat learning rate of 1e-4. Set init_lr = max_lr = final_lr = 1e-4. Chemprop's default schedule (Noam, peaking at 1e-3) is meant for training from scratch. It can destroy the pretrained weights.

Limitations

• Reactions must be atom-mapped.
• The featurization must match pretraining: REAC_DIFF mode and keep_h=True.

License

MIT License.

Citation

@article{zalte2026chemeleonrxn,
  title   = {Descriptor-Based Pre-Training Improves Reaction Property Prediction},
  author  = {Zalte, A. S. and Fang, C. and Zheng, Y. and Metwally, E. and Cheng, A. and Glisman, A. and Pang, H.-W.},
  journal = {ChemRxiv},
  year    = {2026},
  doi     = {10.26434/chemrxiv.15008692}
}

The model builds on CheMeleon:

Burns, J. W.; Zalte, A. S.; Abreu, C. R. A.; Sieg, J.; Feldmann, C.; Mathea, M.; Green, W. H. Deep Learning Foundation Models for Low-Data Regimes from Classical Molecular Descriptors. J. Chem. Inf. Model. 2026. 10.1021/acs.jcim.6c01546