Mechanistic Reaction Prediction via Discrete Flow Matching on Graph-Structured Electron Occupation

Nguyen Xuan-Vu; Octavian Susanu; Daniel Armstrong; Philippe Schwaller

cs.AIAug 27, 2026

Mechanistic Reaction Prediction via Discrete Flow Matching on Graph-Structured Electron Occupation

Authors: Nguyen Xuan-Vu, Octavian Susanu, Daniel Armstrong, Philippe Schwaller

Organizations: École Polytechnique Fédérale de Lausanne (EPFL) · National Centre of Competence in Research (NCCR) Catalysis

Abstract

Chemical reactions are fundamentally transformations in electron space, yet most machine learning approaches model them either through de novo generation of product molecules or through heuristic graph edits that operate directly on molecular topology. We introduce MAELLE (MechAnistic Edit fLow-matching on eLectron rEarrangements), which instead models reactions as discrete flow matching over electron occupation vectors. Concretely, we formulate the reactant-to-product mapping as a Continuous-time Markov Chain (CTMC) over the graph-structured integer-valued electron occupation space defined on all bonding, non-bonding, and hydrogen sites. To construct the interpolants between the reactants and products, we generalize the discrete flow matching mixture path to an edit-based formulation, where the electron moves are interpolated using Optimal Transport, yielding a mechanism-like set of moves without elementary step annotations. MAELLE achieves competitive performance on the USPTO-480K benchmark compared with leading reaction prediction models. Beyond in-distribution learning, we evaluate robustness across two out-of-distribution settings - structural complexity and reaction type - and find that MAELLE maintains strong performance where existing methods degrade. Finally, because the learned flow operates over the full electron redistribution, MAELLE naturally recovers mechanistic trajectories that align with known chemistry and can predict side products of a reaction.

Figures & tables

Figure 1: Conceptual overview of MAELLE: a. MAELLE’s sampling process with mechanism-like trajectories that model the transition from reactants to products as CTMC. b. The three elementary actions on electrons, FLOW , DEL , and ADD , demonstrated on an oxidation and a reduction reactions. At training time: c. Electron actions are interpolated using OT; d. interpolant graphs are sampled by sampling a subset of electron moves and applying them to the reactant graphs. e. MAELLE model architecture and learning objective.

Method	Input	Output	Mech. step	Directional	Broad	Atom
			label-free	elec. re-dist.	coverage	conservative
Molecular Transformer [ 9 ]	SMILES	SMILES	✓	✗	✓	✗
Graph2SMILES [ 11 ]	SMILES	SMILES	✓	✗	✓	✗
Chemformer [ 26 ]	SMILES	SMILES	✓	✗	✓	✗
MEGAN [ 12 ]	Graph	Graph edits	✓	✗	✓	✗
GTPN [ 13 ]	Graph	Graph edits	✓	✗	✓	✓

Table 1: Comparison of reaction prediction approaches.

Method	Prediction	Top-1	Top-3	Top-5	Top-10
	modality
Molecular Transformer † [ 9 ]	SMILES	0.886	0.935	0.942	—
Graph2SMILES † [ 11 ]	SMILES	0.903	0.940	0.948	0.953
MEGAN † [ 12 ]	Graph edits	0.863	0.924	0.940	0.954
GTPN † [ 13 ]	Graph edits	0.832	0.860	0.865	—
NERF † [ 27 ]	Electron edits	0.907	0.933	0.937	—

Table 2: Top-

k

accuracy comparison on USPTO-480K. Entries taken from the original papers are marked with " † "

Figure 2: Comparative performance of MAELLE, Graph2SMILES, Molecular Transformers, and MEGAN on our OOD benchmark. a. From left to right, the first three plots show top-k accuracy on in-distribution test set (IID), esterification-split test set (OOD Ester), and molecular-weight-split test set (OOD Mass). The last plot describes the top-1 accuracy divided into weight bins. b. An illustrative example from the OOD Mass split. c. An illustrative example from the OOD Ester split.

Figure 3: a. Top-k accuracy and plausibility rate, with the latter having mean and standard deviation reported for 3 runs of LLM evaluation. The LLM-judge has a ground truth coverage of $81.1\pm 0.1\%$ . b. Test case 39591 from the test set of USPTO-480K serves as an illustrative example where MAELLE successfully predicts the main product as top-1 and the side product as top-2, while Graph2SMILES and MolecularTransformers propose hallucinated outcomes. c. Two example trajectories that MAELLE samples leading to the two different outcomes.

Bond formation type	# ground-truth elementary steps	Recovered by MAELLE’s trajectories
Bond formation type	# ground-truth elementary steps	Correct	Top 1	All
Persistent	17,925	13,137 (73.3%)	14,236 (79.4%)	17,847 (99.6%)
Transient	19,658	3,042 (15.5%)	747 (3.8%)	5,814 (29.6%)
Total	37,583	16,179 (43.0%)	14,983 (39.9%)	23,661 (63.0%)

Table 3: MAELLE’s coverage of bond-forming elementary steps categorized into persistent and transient bonds.

Sampling mode	Elementary steps coverage	Directional agreement of elementary steps
Correct	16,179/37,583 (43.0%)	13,694/16,179 (84.6%)
Top-1	14,983/37,583 (39.9%)	13,342/14,983 (89.0%)
All	23,661/37,583 (63.0%)	18,505/23,661 (78.2%)

Table 4: MAELLE’s coverage of bond-forming elementary steps and their directional agreement with FlowER’s reference. For a MAELLE-predicted step and a ground-truth step that form the same bond, their directions agree if their sources and sinks are identical.

Figure 4: a. An example from FlowER’s test set with expert-labeled elementary steps. b. Pseudo mechanistic steps in MAELLE’s trajectory. Note that during training, MAELLE does not consume the labeled elementary steps.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

\displaystyle\mathcal{L}\;=\;\sum_{e\in\mathcal{E}}\bigl(\lambda_{e}^{\mathrm{flow}}+\lambda_{e}^{\mathrm{del}}\bigr)+\sum_{a\in\mathcal{V}}\lambda_{a}^{\mathrm{add}}\;-\;\frac{\kappa^{\prime}(t)}{1-\kappa(t)}\!\!\!\!\sum_{(e^{s},\,e^{d})\,\in\,\mathbf{M}_{\mathrm{target}}}\!\!\!\!\begin{cases}\log\lambda_{a}^{\mathrm{add}}+\log p^{\text{sink}}_{\theta}(e^{d}\mid a,\;\mathcal{E})&\text{if }e^{s}=\varnothing\\[4.0pt] \log\lambda_{e^{s}}^{\mathrm{del}}&\text{if }e^{d}=\varnothing\\[4.0pt] \log\lambda_{e^{s}}^{\mathrm{flow}}+\log p^{\text{sink}}_{\theta}(e^{d}\mid e^{s},\;\mathcal{E})&\text{otherwise}\end{cases}

Appendix

Algorithm 2 MAELLE ’s Training Loop

Split	$N$	MW (mean $\pm$ std) [Da]	OOD criterion
train	$304\,228$	$311.9\pm 90.6$	–
val	$37\,582$	$312.7\pm 90.6$	–
test_iid	$37\,578$	$311.3\pm 90.5$	none (IID)
test_ood_mass	$38\,909$	$585.2\pm 83.0$	product MW $>529$ Da
test_ood_ester	$1779$	$321.1\pm 112.1$	non-methyl ester chemistry

Appendix

Table 5: Split sizes, sampling MW (mean

\pm

std of the Gaussian over product molecular weight), and OOD criteria for uspto_unified_v3 .

Chemical Reaction Prediction

See all 18 papers

Learning Transferable Reaction Mechanisms from Visual Chemical Knowledge
Sep 27, 2026Yujian Yuan, Jiaxin Xu, Xin Cai +4Visual Representation Learning Transfer Learning
The Hong Kong University of Science and Technology · The Chinese University of Hong Kong · University of Edinburgh
Learning Mechanistic Reasoning for Chemical Reactions with Large Language Models
Jul 14, 2026Xingyu Dang, Haocheng Tang, Junmei Wang +1LLM Reasoning
Department of Computer Science, Princeton University, Princeton, NJ, 08540, United States · School of Pharmacy, University of Pittsburgh, Pittsburgh, PA, 15261, United States · Khoury College of Computer Science, Northeastern University, Boston, MA, 02115, United States +2
Neural Petri flows for chemical reactions
Oct 6, 2026Jose Eduardo Escrig Molina, Daniel ProbstAI-Assisted Scientific Research Chemical Reaction Prediction
Bioinformatics Group Bioprocess Engineering Group Wageningen Unversity Wageningen, the Netherlands

Sep 27, 2026J/K move · Enter open · S save

Learning Transferable Reaction Mechanisms from Visual Chemical Knowledge

Yujian Yuan, Jiaxin Xu, Xin Cai, Yufan Chen+3

The Hong Kong University of Science and Technology · The Chinese University of Hong Kong · University of Edinburgh

Visual Representation Learning Transfer Learning Molecular Representation Learning Chemical Reaction Prediction

Open paper AskPDF arXiv

Reaction mechanisms describe the step-by-step transformations underlying chemical reactions and are central to reaction analysis and synthesis. Learning-based models have achieved strong performance on established mechanism-prediction benchmarks, but transferring them to unseen chemistry remains challenging. Such transfer is difficult because familiar mechanisms must be applied to unfamiliar molecular structures, and some target mechanisms may be poorly covered by the training data. To address these challenges, we introduce MechaVLM, a visual framework that combines transferable chemical representations with external mechanistic knowledge. It learns reusable visual features through multiscale chemical grounding and cross-rendering contrastive learning. For open-book prediction, MechaVLM retrieves a fixed set of precedents from 70,384 literature mechanism figures and re-reads relevant visual evidence as the molecular state evolves, directly using the figures without symbolic mechanism parsing. An atom-indexed language decoder then recursively generates executable electron edits to construct the complete mechanism. We further introduce MechBench, a challenging literature-derived benchmark with 2,184 mechanisms and 9,146 elementary steps. Across cross-dataset and literature-derived benchmarks, MechaVLM establishes strong zero-shot mechanism prediction. Its closed-book model alone improves Step/Pathway Top-1 by 12.50/13.93 percentage points on FlowER-to-ReactMech transfer, while external visual precedents unlock further gains on challenging OOD reactions. The learned representation also generalizes beyond mechanism prediction to atom mapping and reaction center prediction.

Abstract

Chemical Reaction Prediction

Explore similar work

Learning Transferable Reaction Mechanisms from Visual Chemical Knowledge

Learning Mechanistic Reasoning for Chemical Reactions with Large Language Models

Neural Petri flows for chemical reactions