PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2608.27429v1cs.AI

TL;DR

Reaction prediction often uses de novo product generation or topological graph edits, which can obscure electron conservation and mechanistic pathways. MAELLE models reactant-to-product changes as discrete flow matching over graph-structured electron occupations, using CTMC trajectories and Optimal Transport without elementary-step labels. It is competitive on USPTO-480K, robust across molecular-complexity and reaction-type shifts, and supports plausible side-product prediction.

  • Problem

    De novo generation and topological graph edits can lack explicit reactant-product correspondence, atom conservation, or mechanistic pathways.

  • Method

    MAELLE models directional electron redistribution with discrete flow matching over integer-valued electron occupations, CTMC trajectories, and Optimal Transport-derived moves.

  • Results

    MAELLE is competitive on USPTO-480K and outperforms all baselines on molecular-weight and reaction-type out-of-distribution splits.

  • Takeaways & Limitations

    The learned electron flow provides mechanistically interpretable trajectories and supports chemically plausible side-product prediction through multiple CTMC trajectory samples.

  • Takeaways & Limitations

    OT-derived trajectories are pseudo-mechanisms rather than energy-optimal elementary steps, and the framework currently predicts forward reactions only.

Abstract

from arXiv · show

Chemical reactions are fundamentally transformations in electron space, yet most machine learning approaches model them either through \textit{de novo} generation of product molecules or through heuristic graph edits that operate directly on molecular topology. We introduce MAELLE (\textbf{M}ech\textbf{A}nistic \textbf{E}dit f\textbf{L}ow-matching on e\textbf{L}ectron r\textbf{E}arrangements), 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 intermediate edit trajectories, we generalize the discrete flow matching mixture path to discrete electron rearrangements using Optimal Transport, yielding a sequence of mechanistically interpretable edit moves without requiring elementary step annotations. MAELLE achieves competitive performance on the USPTO-480K benchmark compared with leading reaction prediction models. Beyond in-distribution accuracy, 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.

1 Introduction

Reaction prediction matters for automating synthesis, but prevailing representations treat reactions as opaque mappings and can exploit shortcuts rather than reactivity reasoning. MAELLE instead models electron rearrangements as interpretable flow trajectories and reports competitive benchmark performance with stronger out-of-distribution robustness.

  • Motivation: Reliable reaction prediction is central to drug discovery, materials design, and automating synthetic chemistry because viable routes may require many individually successful steps.The space of possible reactions exceeds the already intractable space of drug-like molecules, motivating calibrated uncertainty.
  • Motivation: Existing approaches progress from fingerprints and sequence translation to graph encoders and generative graph edits, but generally map reactants to products opaquely.These representations reflect successive machine-learning paradigms rather than an explicit account of chemical transformation.
  • Motivation: Standard top-k accuracy can conflate reactivity reasoning with patent-style and substrate-level shortcuts, making out-of-distribution evaluation more faithful.Examples include stylistic signatures in USPTO and a single free hydroxyl among protected ones indicating alcohol-selective transformation.
  • MAELLE: MAELLE models reaction prediction as discrete flow matching over graph-structured integer-valued electron occupations and CTMC trajectories, without elementary-step annotations.Its contribution includes chemistry-aware electron edits and optimal-transport-derived trajectories.
  • Results: MAELLE is competitive on USPTO-480K and outperforms all baselines on molecular-weight and reaction-type out-of-distribution splits.The model also produces calibrated, ranked predictions and supports side-product prediction through generative trajectory sampling.

2 Related Work

Reaction prediction methods range from de novo product generation and topological graph edits to mechanistic electron redistribution. MAELLE extends the latter by modeling directional, atom-conserving electron transfers that can represent cyclic mechanisms without elementary-step labels.

  • Translation-based reaction prediction: SMILES-based methods generate products de novo, lacking explicit reactant-product structure and risking non-atom-conservative or mechanistically invalid transformations.They also degrade on molecules whose complexity lies outside the training domain.
  • Reactions as graph edits: Graph-edit methods expose transformations and can preserve atom conservation, but their purely topological edits do not guarantee a mechanistic pathway.For bond formation, they insert a bond without specifying where the bonding electrons originate.
  • Reactions as electron redistribution: Mechanistic electron modeling tracks source-to-sink redistribution, but explicit elementary steps traditionally require expert labels or substantial physical calculation.Prior datasets and models therefore trade mechanistic detail against annotation and computational demands.
  • Reactions as electron redistribution: ELECTRO removes mechanism labels but is restricted to linear electron-flow topologies and cannot cover cyclic mechanisms.This limits its coverage relative to reactions with cyclic electron flow.
  • Reactions as electron redistribution: NERF predicts net bond-electron changes in one shot, leaving transfer direction implicit and excluding lone pairs and other valence electrons.Its representation records bond gains or losses without identifying electron sources and sinks.
  • Reactions as electron redistribution: MAELLE uses FLOW, ADD, and DEL operations over integer-valued electron occupations, explicitly representing directional transfers while guaranteeing heavy-atom conservation.Optimal Transport supplies net moves, and discrete flow matching trains on sampled interpolant graphs without mechanism labels.

3 Background

Discrete flow matching learns a CTMC that transports coupled source and target distributions through a conditional mixture path. Edit Flows extend this construction to variable-length sequences by factorizing in an auxiliary aligned space and mapping back to observed sequences.

  • 3.1 Discrete Flow Matching: Discrete flow matching learns a CTMC that transports a source distribution at t = 0 to a target distribution at t = 1.
  • 3.1 Discrete Flow Matching: The conditional mixture path is defined between coupled endpoints and has delta distributions at the source and target boundaries.
  • 3.1 Discrete Flow Matching: For common-length sequences, each token independently switches from its source value to its target value according to a monotone schedule κ(t).
  • 3.1 Discrete Flow Matching: The resulting CTMC rate uses the posterior over target tokens and the hazard coefficient κ′(t)/(1 −κ(t)) for positions not yet changed.
  • 3.2 Edit Flows: Edit Flows remove fixed-length and position-correspondence requirements by using padded aligned sequences with insertion, deletion, and substitution operations.
  • 3.2 Edit Flows: The factorized auxiliary-space path is pushed to data space by deleting padding, making it equivalent to sampling edit operations on the source sequence.

4 Pseudo mechanistic reaction prediction with MAELLE

MAELLE represents reactions as electron-occupation changes on an augmented molecular graph and learns a CTMC over sparse electron edits. Optimal Transport aligns reactant and product occupations, while stochastic trajectories provide ranked products and side-product predictions.

  • Electron-space representation: The augmented graph includes heavy atoms, adjacency, and electron sites for bonds, lone pairs, and pooled per-atom hydrogens.
  • Electron-space representation: Reaction prediction acts on an integer occupation vector because heavy-atom adjacency is recovered deterministically from occupied bonding sites.
  • Electron edits: MAELLE uses three sparse edits: FLOW transfers an electron pair, DEL removes one, and ADD introduces one at a site.
  • Optimal-Transport move construction: Balanced integer Optimal Transport aligns electron sites, with a virtual sink/source handling reactions whose total electron count changes.
  • Discrete flow matching: Each move is independently applied according to κ(t), yielding an interpolant that recovers the reactant with no moves and the product with all moves applied.
  • Training: Training predicts marginal move rates from the reactant, current interpolant graph, and time, using an encoder–decoder conditioned on the reactant throughout.
  • Inference and ranking: Inference runs stochastic CTMC trajectories from the reactant, canonicalizes products, and ranks unique products by sampling frequency.
  • Inference and ranking: Multiple trajectories let MAELLE represent competing pathways and recover both high-confidence main products and plausible side products.

5 Results

MAELLE is evaluated on USPTO-480K using standard and out-of-distribution splits, where it remains competitive in-distribution and is more robust to molecular complexity and reaction novelty. Its electron-space trajectories also recover directional, product-defining flows and chemically plausible alternative outcomes without elementary-step supervision.

  • 5.1 Experiments: USPTO-480K contains approximately 480K patent reactions and is evaluated using standard in-distribution and molecular-weight- and reaction-type-based out-of-distribution partitions.The OOD Ester split withholds non-methyl esterifications, while OOD Mass tests molecular complexity.
  • 5.2 USPTO-480K benchmark: 87.2% top-1 accuracy places MAELLE near graph-edit baselines, within 3.5 points of Graph2SMILES, and below NERF on the standard split.MAELLE reaches 93.0%, 93.9%, and 94.6% at top-3, top-5, and top-10, respectively.
  • 5.3 Out-of-distribution benchmark: MAELLE and MEGAN maintain stable accuracy across molecular-weight bins, while Molecular Transformers degrade rapidly and Graph2SMILES declines less severely on OOD Mass.On IID, all models perform similarly, with top-10 accuracy around 90%.
  • 5.3 Out-of-distribution benchmark: MAELLE is the most performant model on OOD Ester, outperforming MEGAN despite not seeing the exact held-out reaction type during training.The authors attribute this to the CTMC framework leveraging intermediate graph states from similar reactions such as amide coupling.
  • 5.4 Plausibility and multiple outcomes: MAELLE’s top-1 plausibility rate is 80.7% versus 60.3% top-1 accuracy, and its top-2 and top-3 plausibility rates exceed those of the SMILES-based models.Its sampled trajectories can distinguish major and minor products and recover the main product at top-1 and a side product at top-2 in an illustrative case.
  • 5.5 Recovery of net directional electron flows: MAELLE’s FLOW operation explicitly identifies electron sources and sinks, producing directional trajectories that approximate arrow-pushing mechanisms rather than only net bond changes.The reported trajectories recover persistent, product-defining electron flows without consuming elementary-step labels.
  • 5.5 Recovery of net directional electron flows: Across 64 sampled trajectories, MAELLE recovers 43.0%, 39.9%, and 63.0% of annotated bond-forming steps under CORRECT, TOP-1, and ALL selection, respectively.ALL sampling recovers 99.6% of persistent bonds but 29.6% of transient ones; directional agreement is 84.6%, 89.0%, and 78.2% for the same schemes.

6 Conclusion

MAELLE predicts reactions in electron occupation space using discrete flow matching and recovers plausible side products through trajectory sampling. Its main limitations are restricted electron-pair modeling, forward-only prediction, and pseudo-mechanistic rather than energy-optimal trajectories.

  • MAELLE models reactions as CTMCs over an optimal-transport-derived auxiliary move space, connecting machine learning with arrow-pushing representations.The model operates in electron occupation space via discrete flow matching.
  • MAELLE achieves competitive USPTO-480K in-distribution performance and outperforms all baselines on molecular-weight and reaction-type generalization.
  • Multiple CTMC trajectory samples allow MAELLE to predict chemically plausible side products, unlike deterministic or beam-search-based approaches.
  • MAELLE models electron pairs rather than individual electrons, excluding radical reactions involving unpaired electrons.Single-electron resolution is identified as a possible extension.
  • The framework predicts forward reactions only, while its optimal-transport trajectories represent pseudo-mechanisms rather than energy-optimal elementary steps.The learned path captures likely combinatorial transitions rather than the true potential energy surface.

A Equivalence of the edit-based and augmented-space mixture paths

The appendix proves that the augmented-space mixture path and direct edit-subset sampling define the same conditional interpolant. This equivalence applies to both sequence edits and MAELLE’s electron-occupation edits.

  • Proposition 1 shows that the augmented-space path is equivalent to independently sampling edit moves and applying the selected subset.
  • The result is setting-agnostic and applies to insertion, deletion, substitution, and electron-occupation FLOW, DEL, and ADD edits.
  • The construction uses one binary coordinate per move, with each move included independently according to the interpolation schedule κt.
  • The interpolant has the correct boundary conditions: it equals the source state at t = 0 and the target state at t = 1.
  • For electron spaces, sites represent bonds, lone pairs, and pooled hydrogen attachments, with virtual sites available for unrealized structures.
  • The training loop obtains electron moves from optimal transport, samples an intermediate graph, and trains rates for flow, deletion, and addition moves.

E Neural Network Architecture

MAELLE uses a conditional encoder-decoder that represents both atoms and electron sites. The decoder processes time-dependent intermediate graphs while cross-attending to fixed reactant representations.

  • The encoder embeds the fixed reactant graph, while the decoder embeds the intermediate graph and conditions on encoder outputs through cross-attention.
  • Both networks operate at atom and electron-site granularities using GNN backbones followed by attention modules.
  • The decoder incorporates flow time t through a sinusoidal embedding and MLP before the GNN layers.
  • Entry embeddings encode valence, hydrogen, and bond information, including electron counts and topological distance.
  • The decoder predicts non-negative flow, deletion, and addition rates from entry embeddings.
  • Sink distributions use local candidate sets, restricting FLOW destinations by shared atoms and limiting ADD candidates to at most |V| + 1.

E.4 Training details

The full MAELLE model has approximately 16M trainable parameters and is trained for about 50 epochs on a single NVIDIA H100 GPU.

  • MAELLE has approximately 16M trainable parameters with hidden dimension d = 256 throughout.
  • Training lasts approximately 50 epochs, requires about 72 hours, and uses a single NVIDIA H100 GPU with effective batch size 64.

E.5 LLM-as-a-judge for benchmarking side product prediction

The section evaluates predicted-product chemical plausibility with an LLM judge using anonymized, balanced candidate lists and repeated runs. It also reports ground-truth coverage to calibrate the judge’s conservativeness.

  • Judge protocol: Gemini-3-flash independently judges whether each candidate product is mechanistically plausible from the reactant SMILES.The prompt considers valence, connectivity, common reactivity, and likely mechanisms, while accepting plausible minor or side products.
  • Judge protocol: The judge returns a plausibility verdict and brief mechanistic reasoning for every candidate in a structured JSON format.Each candidate is evaluated independently rather than ranked only by comparison with the other products.
  • Aggregation: Top-k plausibility rate is the fraction of a model’s first k predictions marked plausible, averaged with standard deviation over three LLM runs.The metric extends exact-match evaluation by measuring chemical plausibility among ranked predictions.
  • Aggregation: 81.1 ± 0.1% ground-truth coverage indicates that the judge is moderately conservative but stable across runs.The ground-truth product is injected as a sanity-check item to detect systematic over- or underpermissiveness.
  • Candidate control: Candidate products are anonymized, shuffled, canonicalized with RDKit, and capped to the smallest unique-list size across models.This balances comparisons among MAELLE, Molecular Transformer, and Graph2SMILES while reducing surface-form and list-size biases.

F Dataset

The dataset pipeline standardizes USPTO-480K reactions before splitting, with atom remapping supporting MAELLE’s electron-site alignment and cached optimal-transport moves supporting training. Baselines receive equivalent map-free reactions despite using different input formats.

  • Dataset preprocessing: All experiments use the USPTO-480K reaction corpus with standard and re-split partitions for in-distribution and out-of-distribution evaluation.The preprocessing steps are applied once to every reaction before any split.
  • Dataset preprocessing: RXNMapper replaces the corpus’s noisy atom maps to provide consistent reactant–product correspondence for MAELLE’s occupation targets.Correct correspondence is required for MAELLE’s electron-site alignment.
  • Dataset preprocessing: Optimal-transport electron moves are solved once before training, cached, and reused across epochs.Training then samples the mixture path per step using Bernoulli sampling.
  • Dataset preprocessing: Baselines receive the same map-free reactions, while only MAELLE retains atom correspondence for its model inputs and targets.This difference reflects each model’s input format rather than different underlying reactions.
  • Dataset preprocessing: Exact duplicate reactions are removed using a canonical reactant-to-main-product key, and overlap across splits is verified to be zero.The main product is selected by largest heavy-atom count after map stripping and RDKit canonicalization.
  • Partitioning: The low-MW split samples product molecular weight with Gaussian mean 250 Da and variance 1643.782 Da2, giving standard deviation approximately 40.5 Da.This construction separates the training distribution from the high-molecular-weight out-of-distribution test axis.

F.3 Out-of-distribution test axes

The out-of-distribution evaluation separates structural complexity and reaction-type generalization. It uses a high-molecular-weight test set and a dedicated non-methyl-ester chemistry set with selected ambiguous classes excluded.

  • OOD Mass: The high-molecular-weight test_ood_mass set contains 38 909 reactions sampled with Gaussian mean 750 Da and variance 6575.13 Da2.Its median product molecular weight is 576 Da, and the sampling parameters separate it from the training-validation distribution.
  • OOD Ester: The test_ood_ester set contains 1779 reactions covering ester-forming transformations that produce linkages other than methyl esters.The axis targets chemistry underrepresented in training relative to its synthetic importance.
  • OOD Ester: The ester OOD axis includes transformations such as ethyl esterification, O-acetylation, Yamaguchi lactonization, Shiina macrolactonization, and Baeyer–Villiger oxidation.These classes define the dedicated non-methyl-ester test chemistry.
  • OOD Ester: Eight ambiguous ester-producing reaction classes are removed from every split to prevent soft leakage of ester chemistry.Excluded classes include Fischer–Speier, Steglich, Mitsunobu, transesterification, and diazoalkane esterification.
  • Split design: Table 5 summarizes split sizes, molecular-weight sampling statistics, and OOD criteria for uspto_unified_v3.The table provides the compact specification of the partitioning design.
Loading 2608.27429v1…