📊 statistics

Discrete diffusion samplers and bridges: Off-policy algorithms and applications in latent spaces

This paper introduces off-policy training techniques and a novel data-to-energy Schrödinger bridge framework for discrete diffusion samplers, demonstrating their effectiveness in improving sampling performance on synthetic benchmarks and enabling data-free posterior sampling within discrete latent spaces of image generative models.

Original authors: Arran Carter, Sanghyeok Choi, Kirill Tamogashev, Víctor Elvira, Nikolay Malkin

Published 2026-05-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Arran Carter, Sanghyeok Choi, Kirill Tamogashev, Víctor Elvira, Nikolay Malkin

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find the best seats in a massive, dark theater (the "target distribution"). You know the layout of the theater and where the good seats are (the "energy function"), but you don't know the total number of seats, and you can't just walk in and pick one because the lights are off. You need a guide to lead you to the good seats.

For years, scientists have had great guides for continuous spaces (like a smooth floor), but for discrete spaces (like a grid of specific, separate seats), the guides were often clumsy. They would get stuck in one section of the theater or miss entire rows of good seats.

This paper introduces a new, smarter way to train these guides using Discrete Diffusion Samplers. Here is the breakdown of their three main innovations, explained simply:

1. The "Replay Buffer" and the "Scout" (Off-Policy Training)

The Problem: Imagine a tour guide who only learns by walking the path they are currently on. If they get stuck in a dead end, they never learn about the great seats in the next room. They are "on-policy," meaning they only learn from their own immediate mistakes.

The Solution: The authors teach the guide to use Off-Policy techniques.

The Replay Buffer: Think of this as a memory bank. The guide saves every interesting path it has ever taken, even if it was a few weeks ago. When training, instead of just walking the current path, the guide reviews these old paths to learn from them.
The Scout (MCMC): Sometimes, the guide needs a little push to break out of a rut. The authors add a "Scout" (a Markov Chain Monte Carlo algorithm). This Scout is a local explorer that can hop around nearby seats to find better spots and feed that information back to the main guide.

The Result: By using this memory bank and the Scout, the guide learns much faster and, crucially, finds all the good seats (modes) in the theater, not just the ones it stumbled upon first. In the paper's tests, this method prevented the guide from getting stuck in a single corner of the room.

2. The "Bridge Builder" (Data-to-Energy Schrödinger Bridges)

The Problem: Usually, you want to get from Point A (a simple, known distribution) to Point B (your complex target). But what if Point B isn't a list of seats you can see? What if Point B is just a set of rules describing how good a seat is (an energy function), without showing you the seats themselves?

The Solution: The authors built a Bridge between these two worlds.

Imagine you have a map of a city (Point A) and a list of "best neighborhoods" defined only by their reputation scores (Point B).
The paper creates a "Schrödinger Bridge" that connects the known map to the reputation list. It learns the path to walk from the known city to the reputation-based neighborhood, even though you can't see the destination until you arrive.
They did this for the first time in a "discrete" world (where seats are distinct blocks, not a smooth street).

The Result: They successfully built a path from a simple starting point to a complex, rule-based destination, visualized in the paper as moving from a mix of three Gaussian shapes to a mix of two, represented as binary codes.

3. The "Translator" for Image Generators (Outsourced Sampling)

The Problem: Modern AI image generators (like those that make pictures of cats or digits) often work in a "latent space." Think of this as a secret code language the AI uses to understand images. Sometimes, you want to force the AI to generate a specific type of image (e.g., "only odd numbers"), but you can't easily tell the AI how to do that directly.

The Solution: The authors used their new guide to sample directly in this secret code language.

Instead of trying to fix the image pixel-by-pixel, they trained their guide to navigate the discrete latent space (the secret code) of a pre-trained image model (a VQ-VAE).
They told the guide: "Find the codes that, when decoded, look like the number '5' or '7'."

The Result: The guide successfully learned to navigate the secret code to produce images of specific digits (like 1, 5, 7) and categories (odd vs. even numbers) without needing to see the final images during the training process. It effectively "outsourced" the hard work of finding the right image to the guide working in the code space.

Summary

In short, this paper takes a powerful sampling technique used for smooth, continuous problems and adapts it for discrete, blocky problems (like grids or codes).

It makes the sampler smarter by letting it remember past paths and use local explorers to avoid getting stuck.
It builds a bridge to reach destinations defined only by rules, not by examples.
It proves this works for image generation, allowing the AI to find specific images by navigating their internal "secret code" language.

The paper claims these methods consistently outperform previous techniques, especially in difficult scenarios where the sampler tends to get stuck in just one solution instead of exploring all the good ones.

Technical Summary: Discrete Diffusion Samplers and Bridges

Problem Statement

The paper addresses the challenge of sampling from a distribution $p(x) \propto e^{-E(x)}$ over a discrete state space $\mathcal{S} = \{1, \dots, C\}^d$ , where the energy function $E$ is known but the normalizing constant $Z$ is intractable. While diffusion samplers have been successfully adapted for continuous spaces using stochastic optimal control and path-space matching, their application to discrete spaces remains underdeveloped. Existing discrete diffusion samplers often fail to leverage off-policy reinforcement learning (RL) techniques that are standard in continuous settings, leading to suboptimal mode coverage and sampling quality, particularly in highly multimodal distributions. Furthermore, the problem of bridging two arbitrary distributions (the Schrödinger bridge) in discrete spaces has largely been restricted to settings where samples from both distributions are available (data-to-data), leaving the "data-to-energy" setting (where one distribution is defined only by an unnormalized density) unexplored.

Methodology

The authors propose a framework that integrates off-policy training strategies into discrete diffusion samplers and generalizes these methods to discrete Schrödinger bridges.

1. Off-Policy Training for Discrete Diffusion Samplers

The core methodology involves training a forward kernel $\overrightarrow{p}_\theta$ to match a target distribution by minimizing a divergence between the forward process trajectory distribution and a backward (noising) process trajectory distribution. The objective functions studied are second-moment divergences, specifically:

Trajectory Balance (TB): Treats the normalizing constant as a learnable parameter.
Log-Variance (LV): Minimizes the variance of the log-ratio of forward and backward probabilities.

To improve exploration and mode coverage, the authors introduce three off-policy strategies adapted from RL and Generative Flow Networks (GFlowNets):

Replay Buffer: Stores sampled states $X_N$ and reuses them by sampling backward trajectories $\overleftarrow{p}^{\otimes N}$ for training.
Importance-Weighted Buffer: Assigns sampling probabilities to buffer elements based on the importance weights of the trajectories that generated them, prioritizing high-probability regions that the current sampler might miss.
MCMC Exploration: Refines samples from the replay buffer using Markov Chain Monte Carlo (MCMC) steps (e.g., Swendsen-Wang or Metropolis-Hastings) before they are used for training. This leverages the energy function directly to guide the sampler toward high-density regions without requiring model rollouts.

2. Discrete Data-to-Energy Schrödinger Bridges

The paper generalizes the sampling problem to finding a Schrödinger bridge (SB) between two arbitrary distributions $p_0$ and $p_N$ . Unlike previous discrete SB works that assume samples for both endpoints, this work addresses the data-to-energy setting where $p_N$ is defined only by an energy function.

The method employs an Iterative Proportional Fitting (IPF) procedure.
It alternates between fitting the backward kernel (via maximum likelihood on samples from $p_0$ ) and the forward kernel (via log-variance loss on samples from the energy-defined $p_N$ ).
Off-policy techniques (buffers and MCMC) are integrated into the IPF loop to facilitate exploration of the energy-defined distribution.

3. Application to Latent Spaces

The authors apply these samplers to "outsourced" posterior sampling. They train diffusion samplers in the discrete latent space of a pre-trained Vector Quantized Variational Autoencoder (VQ-VAE). The goal is to sample from the posterior $p(z|y) \propto p_{latent}(z)p(y|f(z))$ without observing unbiased conditional samples, effectively conditioning the generative model by learning to sample in its latent space.

Key Contributions

Off-Policy Algorithms for Discrete Samplers: The introduction of off-policy RL techniques (replay buffers, importance weighting, and MCMC exploration) into the training of discrete diffusion samplers. The paper demonstrates that these techniques significantly improve performance over on-policy baselines.
Data-to-Energy Schrödinger Bridges: The first generalization of discrete diffusion samplers to the data-to-energy Schrödinger bridge problem, enabling the construction of transport maps between a sample-based distribution and an energy-defined distribution.
Latent Space Posterior Sampling: The demonstration of discrete diffusion samplers for posterior inference in the discrete latent spaces of image generative models, showing the ability to condition on class labels (e.g., odd/even digits) without retraining the generative model.

Experimental Results

The authors evaluate their methods on synthetic benchmarks and real-world applications:

Ising and Potts Models: On $16 \times 16$ lattices, off-policy methods (specifically TB with Buffer + MCMC) consistently outperform on-policy methods and the baseline Masked Diffusion Neural Sampler (MDNS). They achieve lower Sinkhorn distances and magnetization/correlation errors, particularly at low temperatures (high $\beta$ ) where mode collapse is prevalent.
Discretized Synthetic Densities: On challenging multimodal distributions (40GMM and ManyWell) discretized via Gray codes, off-policy methods successfully discover all modes, whereas on-policy methods and MDNS often fail to cover the full distribution.
Schrödinger Bridges: In bridge tasks between simple and multimodal distributions, off-policy methods scale effectively to complex targets where on-policy samplers collapse.
Outsourced Sampling: Visualizations of MNIST digits generated from the latent space show that both on-policy and off-policy samplers can successfully generate digits of desired classes (e.g., odd vs. even), validating the utility of the approach for conditional generation.

Significance and Claims

The paper claims that integrating off-policy training techniques is a critical step forward for discrete diffusion samplers, addressing the issue of mode collapse and improving sampling efficiency in high-dimensional, multimodal discrete spaces. By bridging the gap between discrete diffusion and off-policy RL, the authors show that these samplers can handle complex statistical physics problems (Ising/Potts) and combinatorial optimization tasks more effectively than previous methods.

Furthermore, the work extends the theoretical scope of Schrödinger bridges to the data-to-energy setting in discrete domains, offering a new tool for transport between distributions where one is defined by an energy function. Finally, the application to latent spaces demonstrates a practical pathway for conditioning pre-trained generative models without requiring access to the full data distribution or retraining the generative model itself.

The authors acknowledge limitations, including the need to fix a noising kernel and the computational cost of trajectory storage, but suggest that training at coarser time scales (as supported by continuous-time theory extensions) may mitigate some costs. They identify future directions such as learning energy functions from data and scaling these methods to larger image generative models.