Discrete Diffusion Samplers and Bridges: Off-Policy Algorithms and Applications in Latent Spaces

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.

To understand how random processes behave, we often need examples generated by these processes (e.g., to compute statistics). In some cases, we do not have access to such examples, but we have a way to score how likely they are -- this is known as the sampling problem in probabilistic inference. In this work, we focus on processes that produce discrete outputs. We train models to learn how to sample by starting from simple random noise and gradually transforming it into outputs we are interested in by iteratively applying a transformation represented by a neural network. Learning the full diversity of some random processes is challenging. Previous work on continuous outputs showed that training the neural network with so-called off-policy reinforcement learning methods can greatly improve the diversity of generated samples. We extend these results and show that off-policy methods give a large improvement in the discrete setting as well. Building on previous work for continuous outputs, we also show that, in this discrete setting, algorithms can be trained to transform samples from one random process into samples from another when only a scoring function (but no samples) for one or both processes is available. Finally, we also show that these samplers can be used to perform conditional generation, where outputs are generated according to specified requirements. This allows us to draw samples from pretrained image models that satisfy desired constraints, even when no samples that satisfy these constraints are available.