# LEAPS: A discrete neural sampler via locally equivariant networks

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

## Key facts

- Cited by 31 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 55 references that are in the map.
- First posted on arXiv 15 Feb 2025.

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

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [A Continuous Time Framework for Discrete Denoising Models](https://mlsanity.com/paper/2205.14987/) (2022)
- [General E(2)-Equivariant Steerable CNNs](https://mlsanity.com/paper/1911.08251/) (2019)
- [Discrete Flow Matching](https://mlsanity.com/paper/2407.15595/) (2024)
- [Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design](https://mlsanity.com/paper/2402.04997/) (2024)
- [Path Integral Sampler: a stochastic control approach for sampling](https://mlsanity.com/paper/2111.15141/) (2021)
- [Protein Design with Guided Discrete Diffusion](https://mlsanity.com/paper/2305.20009/) (2023)
- [Denoising Diffusion Samplers](https://mlsanity.com/paper/2302.13834/) (2023)
- [Flow-based generative models for Markov chain Monte Carlo in lattice field theory](https://mlsanity.com/paper/1904.12072/) (2019)
- [Iterated Denoising Energy Matching for Sampling from Boltzmann Densities](https://mlsanity.com/paper/2402.06121/) (2024)
- [Improved sampling via learned diffusions](https://mlsanity.com/paper/2307.01198/) (2023)
- [Solving Statistical Mechanics Using Variational Autoregressive Networks](https://mlsanity.com/paper/1809.10606/) (2018)
- [Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2102.07501/) (2021)
- [Generator Matching: Generative modeling with arbitrary Markov processes](https://mlsanity.com/paper/2410.20587/) (2024)
- [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)
- [NETS: A Non-Equilibrium Transport Sampler](https://mlsanity.com/paper/2410.02711/) (2024)
- [Discrete Flows: Invertible Generative Models of Discrete Data](https://mlsanity.com/paper/1905.10347/) (2019)
- [Continual Repeated Annealed Flow Transport Monte Carlo](https://mlsanity.com/paper/2201.13117/) (2022)
- [A Diffusion Model Framework for Unsupervised Neural Combinatorial Optimization](https://mlsanity.com/paper/2406.01661/) (2024)
- [Flow Matching with General Discrete Paths: A Kinetic-Optimal Perspective](https://mlsanity.com/paper/2412.03487/) (2024)
- [Learning Interpolations between Boltzmann Densities](https://mlsanity.com/paper/2301.07388/) (2023)

## Built on by (most cited first, 30 of 31)

- [Adjoint Sampling: Highly Scalable Diffusion Samplers via Adjoint Matching](https://mlsanity.com/paper/2504.11713/) (2025)
- [Any-Order Flexible Length Masked Diffusion](https://mlsanity.com/paper/2509.01025/) (2025)
- [Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference](https://mlsanity.com/paper/2508.12511/) (2025)
- [Tilt Matching for Scalable Sampling and Fine-Tuning](https://mlsanity.com/paper/2512.21829/) (2025)
- [MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control](https://mlsanity.com/paper/2508.10684/) (2025)
- [Solving Inverse Problems via Diffusion-Based Priors: An Approximation-Free Ensemble Sampling Approach](https://mlsanity.com/paper/2506.03979/) (2025)
- [RNE: plug-and-play diffusion inference-time control and energy-based training](https://mlsanity.com/paper/2506.05668/) (2025)
- [Proximal Diffusion Neural Sampler](https://mlsanity.com/paper/2510.03824/) (2025)
- [Inference-Time Scaling of Discrete Diffusion Models via Importance Weighting and Optimal Proposal Design](https://mlsanity.com/paper/2505.22524/) (2025)
- [Discrete Neural Flow Samplers with Locally Equivariant Transformer](https://mlsanity.com/paper/2505.17741/) (2025)
- [CREPE: Controlling Diffusion with Replica Exchange](https://mlsanity.com/paper/2509.23265/) (2025)
- [Discrete Feynman-Kac Correctors](https://mlsanity.com/paper/2601.10403/) (2026)
- [Bridge Matching Sampler: Scalable Sampling via Generalized Fixed-Point Diffusion Matching](https://mlsanity.com/paper/2603.00530/) (2026)
- [RETRO SYNFLOW: Discrete Flow Matching for Accurate and Diverse Single-Step Retrosynthesis](https://mlsanity.com/paper/2506.04439/) (2025)
- [Scaling Autoregressive Models for Lattice Thermodynamics](https://mlsanity.com/paper/2603.14695/) (2026)
- [Flow Sampling: Learning to Sample from Unnormalized Densities via Denoising Conditional Processes](https://mlsanity.com/paper/2605.03984/) (2026)
- [Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion](https://mlsanity.com/paper/2605.23346/) (2026)
- [Discrete Adjoint Schr"odinger Bridge Sampler](https://mlsanity.com/paper/2602.08243/) (2026)
- [Energy-based generator matching: A neural sampler for general state space](https://mlsanity.com/paper/2505.19646/) (2025)
- [Sampling from Energy distributions with Target Concrete Score Identity](https://mlsanity.com/paper/2510.23106/) (2025)
- [Discrete diffusion samplers and bridges: Off-policy algorithms and applications in latent spaces](https://mlsanity.com/paper/2602.05961/) (2026)
- [A2D2: Fine-Tuning Any-Length Discrete Diffusion for Adaptive Decoding](https://mlsanity.com/paper/2606.13565/) (2026)
- [Free energy Estimation on Any State Space](https://mlsanity.com/paper/2605.31063/) (2026)
- [FrOGS: Discrete Neural Sampler for Independent Alloy Configurations Across Chemical Conditions](https://mlsanity.com/paper/2609.02948/) (2026)
- [FluxLite: Inference-Time Proposal Control for Discrete Diffusion Models](https://mlsanity.com/paper/2609.35947/) (2026)
- [Efficient Inference for Coupled Hidden Markov Models in Continuous Time and Discrete Space](https://mlsanity.com/paper/2510.12916/) (2025)
- [Particle Monte Carlo methods for Lattice Field Theory](https://mlsanity.com/paper/2511.15196/) (2025)
- [Boltzmann Reinforcement Learning for Noise resilience in Analog Ising Machines](https://mlsanity.com/paper/2602.09162/) (2026)
- [Stochastic Counterdiabatic Driving via Biorthogonal Liouvillian Eigenmodes](https://mlsanity.com/paper/2607.24393/) (2026)
- [Iterative Exact Discrete Guidance for Energy-Based Sampling](https://mlsanity.com/paper/2609.37043/) (2026)

## Abstract (arXiv)

We propose "LEAPS", an algorithm to sample from discrete distributions known up to normalization by learning a rate matrix of a continuous-time Markov chain (CTMC). LEAPS can be seen as a continuous-time formulation of annealed importance sampling and sequential Monte Carlo methods, extended so that the variance of the importance weights is offset by the inclusion of the CTMC. To derive these importance weights, we introduce a set of Radon-Nikodym derivatives of CTMCs over their path measures. Because the computation of these weights is intractable with standard neural network parameterizations of rate matrices, we devise a new compact representation for rate matrices via what we call "locally equivariant" functions. To parameterize them, we introduce a family of locally equivariant multilayer perceptrons, attention layers, and convolutional networks, and provide an approach to make deep networks that preserve the local equivariance. This property allows us to propose a scalable training algorithm for the rate matrix such that the variance of the importance weights associated to the CTMC are minimal. We demonstrate the efficacy of LEAPS on problems in statistical physics.