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

AI summary

Overview

Research area: Machine learning — specifically amortised sampling from unnormalised discrete distributions using discrete diffusion samplers, discrete Schrödinger bridges, and off-policy reinforcement learning techniques, with an application to posterior sampling in the discrete latent spaces of pretrained image generative models.

Technical level: Advanced. The paper assumes familiarity with diffusion models, Markov chain Monte Carlo, stochastic optimal control, variational inference, GFlowNets, and Schrödinger bridge / iterative proportional fitting theory.

One-sentence scope: The paper introduces off-policy training techniques for discrete diffusion samplers, extends discrete diffusion samplers to data-to-energy Schrödinger bridges for the first time, and demonstrates their use for data-free posterior sampling inside the discrete latent space of a pretrained VQ-VAE.

What This Paper Is About

The paper addresses the problem of drawing samples from a discrete distribution $p_{\text{target}}(x) = \frac{1}{Z}e^{-\mathcal{E}(x)}$ over a space $\mathcal{S} = {1,\dots,C}^d$, where the energy function can be queried pointwise but the normalising constant $Z$ is unknown and exact sampling is intractable. Continuous-space diffusion samplers are well developed, and off-policy reinforcement learning has already been shown to improve them, but discrete-space counterparts have not exploited these ideas.

The goal is threefold: to bring off-policy training into discrete diffusion samplers, to generalise discrete diffusion samplers from sampling to bridging between two arbitrary distributions (where one side is given only by an unnormalised density), and to show that discrete diffusion samplers can perform posterior inference in the discrete latent spaces of pretrained image generative models.

Key Contributions

1. Off-policy reinforcement learning for discrete diffusion samplers. The authors introduce replay buffers, importance-weighted (prioritised) buffers, and MCMC exploration from the buffer into discrete diffusion sampler training (Section 2.2) and show consistent benefits on established and new synthetic benchmarks (Sections 5.1 and 5.2).

2. Discrete data-to-energy Schrödinger bridges. The paper generalises discrete diffusion samplers to the bridge problem and, building on the continuous-space data-to-energy iterative proportional fitting procedure of Tamogashev and Malkin (2026), presents data-to-energy Schrödinger bridge training in the discrete domain for the first time (Section 3, Section 5.2, Figure 1).

3. Posterior sampling in discrete latent spaces. For the first time, the paper shows that discrete diffusion samplers can enable posterior inference in the discrete latent spaces of pretrained image generative models (Section 4 and Section 5.3), generalising the continuous-latent "noise outsourcing" approach of Venkatraman et al. (2025) to a VQ-VAE with an autoregressive prior.

4. Recognition of a previously unnoticed connection. The authors point out that algorithms equivalent to discrete diffusion samplers were already considered as early as Zhang et al. (2022a) in the language of generative flow networks (using trajectory balance and off-policy training), but that this connection had not been recognised until now.

Main Findings

• Off-policy training improves discrete samplers consistently: On the 16×16 Ising and Potts benchmarks (Table 1, mean ± std over 5 runs), off-policy methods "consistently outperform or match on-policy methods," and trajectory balance (TB) performs slightly better than the log-variance (LV) objective on most temperatures.

• Gains are largest at low temperature: The greatest difference between on-policy and off-policy methods occurs at the highest inverse temperatures. At the lowest-temperature Ising setting ($\beta = 1.2$), many methods experience severe mode collapse, but both MCMC-assisted methods successfully model the distribution well (e.g. TB + Buffer + MCMC reaches ELBO 615.03 ± 0.06, EUBO 615.14 ± 0.02, Sinkhorn 0.02 ± 0.00).

• TB off-policy beats MDNS and LV on Potts: On the Potts model with $q = 3$, the TB off-policy methods outperform both on-policy methods as well as LV and MDNS. At the lowest temperature ($\beta = 1.2$) many methods collapse to a single mode while both MCMC-assisted methods do not.

• TB without MCMC can be unstable: For both Ising and Potts models, TB without MCMC is sometimes unstable at low temperatures.

• Mode coverage on multimodal synthetic densities: For the discretised 40GMM with $d = 4 \times 8 = 32$, only the off-policy TB + Buffer + MCMC method discovers all 40 modes (Figure 3).

• Bridges at fixed discretisation: Figure 4 compares discrete-space Schrödinger bridges on 16-dimensional binary Gray-coded spatial data learnt by data-to-energy IPF with on-policy versus off-policy LV training, across three pairs: 3GMM ↔ 4GMM ($d = 2 \times 8 = 16$), S-curve ↔ 10GMM ($d = 2 \times 8 = 16$), and 10GMM ↔ 40GMM ($d = 2 \times 8 = 16$), each run for 20 steps from $X_0$ to $X_{20}$.

• Error correction to the existing literature: Equation (3)'s minimiser is unique for any full-support trajectory distribution $\mathbb{P}$, and when $\mathbb{P}$ is on-policy the gradient coincides up to a scalar with twice the gradient of the reverse KL divergence, for any choice of $c$.

• Practical discretisation insight: Extending results of Berner et al. (2026) to the discrete case, the authors find that unmasking diffusion samplers trained with few unmasking steps perform well when sampled with one unmasked position per transition (Appendix A.1).

• Latent-space results: Section 5.3 is described as evaluating outsourced posterior sampling, but the numeric results are not included in the truncated content provided here; the paper states only that this is the setting where discrete diffusion samplers are applied to pretrained image generative model latent spaces.

Methodology in Plain English

The researchers treat sampling as a sequential decision problem. A forward "denoising" process starts from a simple distribution and applies learned transition kernels that gradually reveal the entries of a discrete sequence; a backward "noising" process starts from the unknown target and applies known corruption kernels (for example, masking entries with a mask token $(\star,\dots,\star)$, or replacing them uniformly). If the distribution over the whole forward trajectory can be made to match the distribution over the whole backward trajectory, then the final samples of the forward process follow the target.

To measure how close the two trajectory distributions are, the authors use a second-moment (squared-log-ratio) divergence whose scalar offset $c$ can either be learned — recovering trajectory balance — or set to the value that minimises the empirical mean over a batch — recovering log-variance. Crucially, when the training trajectories come from the model itself (on-policy), this divergence has the same gradient as a KL divergence; but the minimiser is unique for any full-support distribution over trajectories, which licenses training on trajectories generated by something other than the current model.

That is the opening for off-policy training. The authors reuse past samples from a replay buffer, optionally weighting buffered states by the importance weights of the trajectories that produced them so that training concentrates on high-target-density states the sampler found improbable, and optionally refining buffered samples with a few steps of an MCMC kernel whose stationary distribution is the target (the kernel only needs energy evaluations, so it is cheap compared with sampler rollouts).

For the bridge problem, the initial distribution is no longer fixed to a simple $p_0$ and the backward kernel becomes a parametric object to be learned. The authors adapt iterative proportional fitting — alternately projecting onto the constraint that the forward process starts at $p_0$ and the backward process ends at $p_{\text{target}}$ — to the discrete case, using maximum-likelihood training for the first projection and a variant of the log-variance loss for the second, since no samples from the target are available. Masked diffusion is not applicable here, so a uniform diffusion process serves as the reference kernel.

Experiments cover 16×16 Potts and Ising energies (with the Swendsen-Wang algorithm used as the MCMC proposal), synthetic Gaussians and ManyWell densities converted to binary Gray codes by subsampling each of $D$ dimensions at $2^8$ values to give sequences of length $d = 8D$, and posterior sampling in the discrete latent space of a pretrained VQ-VAE with an autoregressive prior.

Why This Matters

Impact on research. The paper connects two previously separate lines of work — recent masked discrete diffusion samplers and the earlier off-policy, generative-flow-network-style algorithms of Zhang et al. (2022a) — and provides a unifying view linking discrete diffusion, any-order autoregressive modelling, stochastic optimal control, and hierarchical variational inference. It also extends Schrödinger bridge methods into the discrete data-to-energy regime, and moves discrete diffusion sampler evaluation beyond synthetic densities to realistic pretrained generative model latent spaces.

Real-world applications (as implied by the paper's settings):

• Statistical mechanics. Sampling from Ising and Potts energy models, which the paper uses directly as benchmarks and which it notes are the classical setting for energy-based models over discrete spaces.
• Bayesian (variational) inference. The paper describes the unnormalised-density sampling problem as "the typical setting in variational (Bayesian) inference."
• Conditional image generation and guidance. Posterior inference $p(x \mid y) \propto p(x)p(y \mid x)$ in the latent space of a pretrained VQ-VAE, where $y$ could be a class label supplied by a classifier, without any unbiased conditional samples and without differentiating through the decoder.
• Combinatorial sampling and approximate counting. The discrete state space setting connects to foundational work on sampling in combinatorial spaces (Jerrum and Sinclair, 1996).

Industry relevance. The methods are amortised — once trained, sampling is fast, since the sampler makes one pass rather than running a long MCMC chain. The MCMC-assisted exploration step is explicitly described as computationally inexpensive because it needs only energy evaluations, not model rollouts, which matters for scaling. The latent-space application is directly relevant to practitioners who already have a pretrained discrete-latent generative model and want to condition or guide it without retraining the decoder.

Future Directions

• Scaling to higher dimensions and longer sequences. The paper evaluates synthetic densities at $d = 16$, $32$, and $80$, and bridge tasks at $d = 16$, alongside 16×16 lattices; whether the off-policy advantages persist at much larger discrete state spaces is left open.

• Stability of off-policy TB at low temperatures. The authors report that TB without MCMC is sometimes unstable at low temperatures and that many methods collapse to a single mode at $\beta = 1.2$, suggesting room for better variance control or more robust objectives.

• Masking-granularity trade-offs. Because the paper allows multiple maskings per step, and notes that this introduces correlation error while increasing computational speed and easing memory requirements, the design space of how many positions to mask per step remains an open question (partially explored in Appendix D.1.4).

• Bridging posterior distributions in latent space. The paper explicitly states that one could similarly model Schrödinger bridges between posterior distributions in latent space defined via the VQ-VAE setup, which the paper introduces but does not fully develop.

• Fuller evaluation of the latent-space results. Only the setup for Section 5.3 is described in the content available; quantitative outcomes for the latent-space posterior sampling experiments are not reported here.

Target Audience

This paper is most useful to machine learning researchers working on diffusion models, diffusion samplers, generative flow networks, Schrödinger bridges, and sampling from unnormalised densities; to probabilistic machine learning and statistics researchers interested in Bayesian posterior inference and MCMC alternatives; to physicists and computational scientists working with discrete energy-based models such as Ising and Potts systems; and to practitioners of generative modelling who want to condition or guide pretrained discrete-latent models (such as VQ-VAEs) without retraining them. Readers will need a solid background in diffusion modelling, variational inference, and reinforcement learning to follow the derivations.

Authors’ abstract

Sampling from a distribution $p(x) \propto e^{-\mathcal{E}(x)}$ known up to a normalising constant is an important and challenging problem in statistics. Recent years have seen the rise of a new family of amortised sampling algorithms, commonly referred to as diffusion samplers, that enable fast and efficient sampling from an unnormalised density. Such algorithms have been widely studied for continuous-space sampling tasks; however, their application to problems in discrete space remains largely unexplored. Although some progress has been made in this area, discrete diffusion samplers do not take full advantage of ideas commonly used for continuous-space sampling. In this paper, we propose to bridge this gap by introducing off-policy training techniques for discrete diffusion samplers. We show that these techniques improve the performance of discrete samplers on both established and new synthetic benchmarks. Next, we generalise discrete diffusion samplers to the task of bridging between two arbitrary distributions, introducing data-to-energy Schrödinger bridge training for the discrete domain for the first time. Lastly, we showcase the application of the proposed diffusion samplers to data-free posterior sampling in the discrete latent spaces of image generative models.