# Learning Flexible Forward Trajectories for Masked Molecular Diffusion

arXiv: https://arxiv.org/abs/2505.16790 · HTML: https://mlsanity.com/paper/2505.16790/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 71 of its 87 references that are in the map.
- First posted on arXiv 22 May 2025.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 71)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Wan: Open and Advanced Large-Scale Video Generative Models](https://mlsanity.com/paper/2503.20314/) (2025)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers](https://mlsanity.com/paper/2401.08740/) (2024)
- [Junction Tree Variational Autoencoder for Molecular Graph Generation](https://mlsanity.com/paper/1802.04364/) (2018)
- [Muse: Text-To-Image Generation via Masked Generative Transformers](https://mlsanity.com/paper/2301.00704/) (2023)
- [A Comprehensive Overview of Large Language Models](https://mlsanity.com/paper/2307.06435/) (2023)
- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions](https://mlsanity.com/paper/2102.05379/) (2021)
- [DiGress: Discrete Denoising diffusion for graph generation](https://mlsanity.com/paper/2209.14734/) (2022)
- [A Continuous Time Framework for Discrete Denoising Models](https://mlsanity.com/paper/2205.14987/) (2022)
- [DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models](https://mlsanity.com/paper/2210.08933/) (2022)
- [GraphVAE: Towards Generation of Small Graphs Using Variational Autoencoders](https://mlsanity.com/paper/1802.03480/) (2018)
- [Graph Convolutional Policy Network for Goal-Directed Molecular Graph Generation](https://mlsanity.com/paper/1806.02473/) (2018)

## Built on by (most cited first, 7 of 7)

- [TUBE: Tangent Upper Bound on Evidence for Discrete Diffusion Language Models](https://mlsanity.com/paper/2605.24292/) (2026)
- [VQ-SAD: Vector Quantized Structure Aware Diffusion For Molecule Generation](https://mlsanity.com/paper/2605.00354/) (2026)
- [From Scores to Gibbs Correctors: Accelerating Uniform-Rate Discrete Diffusion Models](https://mlsanity.com/paper/2605.27352/) (2026)
- [A Systematic Evaluation of Co-folding Model Representations for Small-Molecule Learning](https://mlsanity.com/paper/2602.13249/) (2026)
- [MolHIT: Advancing Molecular-Graph Generation with Hierarchical Discrete Diffusion Models](https://mlsanity.com/paper/2602.17602/) (2026)
- [Graph Forward Distribution Matching for Molecular Inverse Design](https://mlsanity.com/paper/2609.32056/) (2026)
- [MIRRA: Fast, High-Fidelity Molecular Generation with Chemistry-Informed Discrete Flow Matching](https://mlsanity.com/paper/2610.17244/) (2026)

## Abstract (arXiv)

Masked diffusion models (MDMs) have achieved notable progress in modeling discrete data, while their potential in molecular generation remains underexplored. In this work, we explore their potential and introduce the surprising result that naively applying standards MDMs severely degrades the performance. We identify the critical cause of this issue as a state-clashing problem-where the forward diffusion of distinct molecules collapse into a common state, resulting in a mixture of reconstruction targets that cannot be learned using typical reverse diffusion process with unimodal predictions. To mitigate this, we propose Masked Element-wise Learnable Diffusion (MELD) that orchestrates per-element corruption trajectories to avoid collision between distinct molecular graphs. This is achieved through a parameterized noise scheduling network that assigns distinct corruption rates to individual graph elements, i.e., atoms and bonds. Extensive experiments on diverse molecular benchmarks reveal that MELD markedly enhances overall generation quality compared to element-agnostic noise scheduling, increasing the chemical validity of vanilla MDMs on ZINC250K from 15% to 93%, Furthermore, it achieves state-of-the-art property alignment in conditional generation tasks.