# Transition Path Sampling with Improved Off-Policy Training of Diffusion Path Samplers

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 65 references that are in the map.
- First posted on arXiv 30 May 2024.

## Builds on (most cited first, 20 of 20)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation](https://mlsanity.com/paper/2106.04399/) (2021)
- [Trajectory balance: Improved credit assignment in GFlowNets](https://mlsanity.com/paper/2201.13259/) (2022)
- [Path Integral Sampler: a stochastic control approach for sampling](https://mlsanity.com/paper/2111.15141/) (2021)
- [Denoising Diffusion Samplers](https://mlsanity.com/paper/2302.13834/) (2023)
- [Improved sampling via learned diffusions](https://mlsanity.com/paper/2307.01198/) (2023)
- [Solving high-dimensional Hamilton-Jacobi-Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space](https://mlsanity.com/paper/2005.05409/) (2020)
- [VarGrad: A Low-Variance Gradient Estimator for Variational Inference](https://mlsanity.com/paper/2010.10436/) (2020)
- [Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Time-Coarsened Dynamics](https://mlsanity.com/paper/2302.01170/) (2023)
- [OpenMM 8: Molecular Dynamics Simulation with Machine Learning Potentials](https://mlsanity.com/paper/2310.03121/) (2023)
- [Stochastic Optimal Control for Collective Variable Free Sampling of Molecular Transition Paths](https://mlsanity.com/paper/2207.02149/) (2022)
- [Accurate transition state generation with an object-aware equivariant elementary reaction diffusion model](https://mlsanity.com/paper/2304.06174/) (2023)
- [Learning to Scale Logits for Temperature-Conditional GFlowNets](https://mlsanity.com/paper/2310.02823/) (2023)
- [Diffusion-based Generative AI for Exploring Transition States from 2D Molecular Graphs](https://mlsanity.com/paper/2304.12233/) (2023)
- [Reinforcement learning of rare diffusive dynamics](https://mlsanity.com/paper/2105.04321/) (2021)
- [Differentiable Simulations for Enhanced Sampling of Rare Events](https://mlsanity.com/paper/2301.03480/) (2023)
- [Diffusion Methods for Generating Transition Paths](https://mlsanity.com/paper/2309.10276/) (2023)
- [Generative methods for sampling transition paths in molecular dynamics](https://mlsanity.com/paper/2205.02818/) (2022)
- [Accelerated Sampling of Rare Events using a Neural Network Bias Potential](https://mlsanity.com/paper/2401.06936/) (2024)

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

- [Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference](https://mlsanity.com/paper/2508.12511/) (2025)
- [Action-Minimization Meets Generative Modeling: Efficient Transition Path Sampling with the Onsager-Machlup Functional](https://mlsanity.com/paper/2504.18506/) (2025)
- [Temperature-Annealed Boltzmann Generators](https://mlsanity.com/paper/2501.19077/) (2025)
- [Control Consistency Losses for Diffusion Bridges](https://mlsanity.com/paper/2512.05070/) (2025)
- [Functional Adjoint Sampler: Scalable Sampling on Infinite Dimensional Spaces](https://mlsanity.com/paper/2511.06239/) (2025)
- [Transferable Learning of Reaction Pathways from Geometric Priors](https://mlsanity.com/paper/2504.15370/) (2025)
- [LatentFlow: A General Framework for Conditioning Stochastic Processes](https://mlsanity.com/paper/2607.12922/) (2026)
- [Learning Collective Variables from BioEmu with Time-Lagged Generation](https://mlsanity.com/paper/2507.07390/) (2025)
- [Entangled Schr"odinger Bridge Matching](https://mlsanity.com/paper/2511.07406/) (2025)
- [Stochastic Control Policies for Robust Molecular Transition Path Sampling](https://mlsanity.com/paper/2608.13800/) (2026)

## Abstract (arXiv)

Understanding transition pathways between two meta-stable states of a molecular system is crucial to advance drug discovery and material design. However, unbiased molecular dynamics (MD) simulations are computationally infeasible because of the high energy barriers that separate these states. Although recent machine learning techniques are proposed to sample rare events, they are often limited to simple systems and rely on collective variables (CVs) derived from costly domain expertise. In this paper, we introduce a novel approach that trains diffusion path samplers (DPS) to address the transition path sampling (TPS) problem without requiring CVs. We reformulate the problem as an amortized sampling from the transition path distribution by minimizing the log-variance divergence between the path distribution induced by DPS and the transition path distribution. Based on the log-variance divergence, we propose learnable control variates to reduce the variance of gradient estimators and the off-policy training objective with replay buffers and simulated annealing techniques to improve sample efficiency and diversity. We also propose a scale-based equivariant parameterization of the bias forces to ensure scalability for large systems. We extensively evaluate our approach, termed TPS-DPS, on a synthetic system, small peptide, and challenging fast-folding proteins, demonstrating that it produces more realistic and diverse transition pathways than existing baselines.