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

Noosaga Placements

• Machine Learningsubfield90%
  The paper introduces a new machine learning method (discrete flow matching) for reaction prediction, with a focus on learning a generative model over electron occupations. The core contributions are algorithmic and empirical, fitting within machine learning.
  We introduce MAELLE ... which instead models reactions as discrete flow matching over electron occupation vectors.MAELLE achieves competitive performance on the USPTO-480K benchmark compared with leading reaction prediction models.
• Deep Generative Modelsframework90%
  The paper uses discrete flow matching, a deep generative modeling technique, to model the distribution of electron occupations. The method is a form of generative modeling over discrete state spaces.
  We introduce MAELLE ... which instead models reactions as discrete flow matching over electron occupation vectors.Discrete flow matching (DFM) learns a Continuous-time Markov Chain (CTMC) that transports a source distribution p0 at t = 0 to a target p1 at t = 1.
• Deep Learningsubfield80%
  The method uses a deep learning architecture with graph neural networks and attention mechanisms, and the training involves deep generative modeling. The paper discusses architectural details and training procedures.
  We adopt an encoder–decoder architecture ... with graph neural networks and attention mechanisms.The full MAELLE model has approximately 16M trainable parameters ...
• Generative Deep Learningframework80%
  The method is a generative deep learning approach, as it learns to generate product molecules by sampling from a learned distribution over electron occupations. The paper discusses generative modeling and sampling trajectories.
  By modeling reaction as a generative problem, simulating multiple trajectories, and taking the most popular outputs, we are somewhat mimicking how a real chemistry process happensAt inference, the products are generated by the marginal CTMC starting from the reactants.
• Probabilistic Aisubfield70%
  The method is based on probabilistic generative modeling, specifically discrete flow matching and continuous-time Markov chains, which are probabilistic frameworks for modeling distributions over discrete states.
  we formulate the reactant-to-product mapping as a Continuous-time Markov Chain (CTMC) over the graph-structured integer-valued electron occupation spaceDiscrete flow matching (DFM) learns a Continuous-time Markov Chain (CTMC) that transports a source distribution p0 at t = 0 to a target p1 at t = 1.
• Attention Mechanisms and Transformersframework70%
  The model architecture uses attention mechanisms, including self-attention and cross-attention, at both atom and electron-site levels. The paper describes the use of Transformer blocks and attention layers.
  After the GNN backbone, both encoder and decoder refine representations through two stages of attention operating at different granularities.In the decoder, this is a TransformerDecoder: each layer applies self-attention over the decoder’s atom embeddings, followed by cross-attention to the encoder’s atom embeddings.

Abstract

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.

Paper Context