Split Gibbs Discrete Diffusion Posterior Sampling

TL;DR

This work tackles posterior sampling in discrete-state spaces by enabling plug-and-play conditioning with discrete diffusion priors. It introduces Split Gibbs Discrete Diffusion Posterior Sampling (SGDD), an alternating likelihood-prior scheme augmented by an auxiliary variable ${\mathbf{z}}$ and a regularization $D({\mathbf{x}},{\mathbf{z}};\eta)$, with convergence guarantees as $\eta\to0$. A theoretical bound based on KL divergence and relative Fisher information establishes convergence to the true posterior under realistic assumptions, while empirical results across synthetic data, DNA design, discrete image inpainting, and monophonic music infilling show consistent performance gains over baselines. The method enables reward-guided generation and solving inverse problems in discrete spaces, offering a principled, scalable, plug-and-play framework with available code for community use and extension.

Abstract

We study the problem of posterior sampling in discrete-state spaces using discrete diffusion models. While posterior sampling methods for continuous diffusion models have achieved remarkable progress, analogous methods for discrete diffusion models remain challenging. In this work, we introduce a principled plug-and-play discrete diffusion posterior sampling algorithm based on split Gibbs sampling, which we call SGDD. Our algorithm enables reward-guided generation and solving inverse problems in discrete-state spaces. We demonstrate the convergence of SGDD to the target posterior distribution and verify this through controlled experiments on synthetic benchmarks. Our method enjoys state-of-the-art posterior sampling performance on a range of benchmarks for discrete data, including DNA sequence design, discrete image inverse problems, and music infilling, achieving more than 30% improved performance compared to existing baselines. Our code is available at https://github.com/chuwd19/Split-Gibbs-Discrete-Diffusion-Posterior-Sampling.

Figures (8)

• Figure 1: An illustration of our method on inpainting discretized MNIST. Our method implements the Split Gibbs sampler, which alternates between likelihood sampling steps and prior sampling steps. Likelihood sampling steps enforce data consistency with regard to measurement ${\mathbf{y}}$, while prior sampling steps involve denoising uniform noise from $z^{(k)}$ by a pretrained discrete diffusion model starting from noise level $\eta_k$. Two variables converge to a sample from the posterior distribution as the noise level $\eta_k$ reduces to zero.
• Figure 2: Generating DNA sequences involves trade-offs between targeted objectives and consistency with the prior distribution. SGDD is plotted in blue and darker colors mean higher $\beta$.
• Figure 3: Diversified samples when the measurement ${\mathbf{y}}$ is sparse. Samples are generated by SGDD when solving an MNIST inpainting task.
• Figure 4: Convergence speed of different noise schedules. We compare our noise annealing scheduler $\{\eta_k\}_{k=1}^K$ (orange) to fixed noise schedules of various noise levels on the MNIST AND task. $y$-axis is the PSNR of variables ${\mathbf{x}}^{(k)}$ with respect to the ground truth. Our noise schedule converges faster than schedulers with fixed noises.
• Figure 5: Sampling quality vs. computing budget. The $x$-axis indicates the number of function evaluations of the discrete diffusion, while the $y$-axis shows the PSNR metric. Experiments are done with 10 discretized MNIST samples on XOR and AND tasks.

Theorems & Definitions (9)

• Theorem 1
• Lemma 1: Data processing inequality of Metropolis Hasting
• Lemma 2: Free-energy-rate-functional-relative-Fisher-information (FIR) inequality (from Theorem 6.2.3. in hilder2017fir)
• proof : Proof Sketch.
• Lemma 3
• proof
• Lemma 4
• proof
• proof