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

2026-02-05

AI Summary

This work addresses the challenge of efficiently sampling from unnormalized energy distributions in discrete spaces by proposing a novel discrete diffusion sampling method that integrates off-policy training with Schrödinger bridges. It introduces, for the first time, the data-to-energy Schrödinger bridge framework into discrete domains, enabling bridging between arbitrary distributions and facilitating data-free posterior sampling in the discrete latent spaces of image generative models. By optimizing the sampler through off-policy reinforcement learning, the method significantly enhances both sampling efficiency and quality. Extensive experiments on multiple synthetic benchmarks demonstrate its superiority, effectively bridging a critical gap in both the theoretical foundations and practical applications of discrete diffusion sampling.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search
• Planning, Routing, and Scheduling: Mixed Discrete/Continuous Planning

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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\"odinger 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.

Problem

• discrete diffusion samplers
• off-policy training
• Schrödinger bridge
• unnormalised density sampling
• discrete latent spaces

Innovation

• discrete diffusion samplers
• off-policy training
• Schrödinger bridge
• latent space sampling
• unnormalised density