Technical Summary: Discrete Langevin-Inspired Posterior Sampling (∆LPS)

Problem Statement

The paper addresses the challenge of solving inverse problems in discrete state spaces. Inverse problems involve recovering an unknown signal $x$ from indirect, sparse, and noisy measurements $y$ (modeled as $y = A(x) + n$). While continuous diffusion models have become the standard for solving such problems via posterior sampling, their discrete counterparts remain underexplored.

Existing methods for discrete posterior sampling suffer from three primary limitations:

1. Continuous Relaxations: Methods like G2D2 relax categorical variables into continuous spaces to use gradient guidance, altering the inference geometry and risking that projected discrete tokens do not decode consistently.
2. Sequential Updates: Methods like SGDD rely on Gibbs-style updates, which are sequential and slow to mix in high-dimensional spaces.
3. Specialization: Methods like APS are tightly coupled to specific corruption processes (e.g., masked diffusion) and training paradigms, limiting their generality.

The goal is to develop a posterior sampler that remains fully discrete, updates all token dimensions in parallel, and is agnostic to the specific training paradigm of the discrete diffusion prior.

Methodology: ∆LPS

The authors propose ∆LPS (Discrete Langevin-Inspired Posterior Sampler), a training-free algorithm that utilizes gradient information to guide discrete transitions without leaving the discrete state space.

Core Formulation

The problem is formulated in the discrete latent space $z_0 \in \{1, \dots, K\}^L$ of an encoder-decoder model. The goal is to sample from the posterior $p(z_0|y) \propto p(y|z_0)p_\theta(z_0)$, where $p_\theta$ is a pretrained discrete diffusion prior. The sampler approximates the posterior transition $p(z_s|z_t, y)$ by first inferring a clean latent $z_0$ consistent with the measurement and the prior, then re-noising it to the next time step.

Discrete Langevin Proposal

The key innovation is adapting the Langevin dynamics to discrete spaces. In continuous settings, a Langevin update moves a sample $z_0$ in the direction of the gradient of the log-posterior potential $U(z_0) = \log p(y|z_0) + \log p_\theta(z_0; z_t)$. Since $z_0$ is discrete, the authors do not add a gradient vector directly. Instead, they use the Gaussian proposal distribution from continuous Langevin dynamics to define transition probabilities over valid discrete candidates $z'_0$:

$$q(z'_0|z_0; z_t, y) \propto \exp\left( -\frac{1}{4\eta} \|z'_0 - z_0 - \eta \nabla_{z_0} U(z_0; z_t, y)\|_2^2 \right)$$

By expanding the exponent and normalizing over the finite set of valid tokens, the proposal factorizes across token coordinates. This allows the sampler to compute a categorical distribution for each token position $\ell$ independently:
$$q_\ell(z'_0[\ell]) = \text{Cat}\left( \sigma\left( \frac{1}{2} \nabla_{z_0[\ell]} U \cdot (k - z_0[\ell]) - \frac{(k - z_0[\ell])^2}{4\eta} \right) \right)$$

This mechanism enables fully parallel updates across all $L$ token dimensions, avoiding the sequential bottleneck of Gibbs samplers.

Generalization and Preconditioning

• Prior Agnosticism: The method only requires predictive probabilities over categorical states, making it compatible with masked diffusion (MDLM) and uniform-state diffusion (Duo) priors.
• Blind Inverse Problems: For non-differentiable or unknown forward operators, the likelihood gradient is replaced by a learned contrastive surrogate (using encoders for data and measurements) to provide a smooth guidance signal.
• Adam Preconditioning: To handle poorly scaled gradients in high-dimensional latent spaces, the raw gradient is replaced with an Adam-preconditioned direction. This adaptively rescales updates, damping noisy large gradients while preserving consistently useful ones, improving mixing and convergence.

Key Contributions

1. Fully Discrete Posterior Sampling: ∆LPS is a training-free sampler that keeps every iterate in the discrete state space, avoiding soft-token optimization and projection errors.
2. Efficient Parallel Updates: By factorizing the proposal across coordinates, the method updates all token dimensions in parallel, significantly improving efficiency over sequential Gibbs-style samplers.
3. Prior-Agnostic Design: The sampler is not tied to a specific corruption process (masked vs. uniform) and works with any pretrained discrete diffusion prior that provides predictive probabilities.

Experimental Results

The authors evaluate ∆LPS on a diverse set of inverse problems across MNIST, CIFAR-10, and FFHQ, covering linear, nonlinear, and blind settings.

• Image Restoration (FFHQ): On tasks including inpainting, Gaussian deblurring, super-resolution, HDR, and motion deblurring, ∆LPS achieves performance comparable to or better than continuous baselines (e.g., DPS, PSLD) and outperforms existing discrete baselines (SGDD, G2D2, APS) in perceptual quality (LPIPS).
• Discrete Tasks (MNIST/CIFAR): The method demonstrates robustness across varying difficulty levels (Easy, Medium, Hard). It consistently matches or surpasses SGDD and G2D2 in token accuracy and PSNR, even with significantly lower computational budgets (fewer function evaluations) than the "full" SGDD variant.
• Blind Inverse Problems: In a spatial mapping task (inferring floorplans from walking trajectories where the forward operator is unknown), ∆LPS outperforms continuous baselines (CoGuide, DPS variants) in F1 score and Intersection over Union (IoU).

Significance and Claims

The paper claims that fully discrete, gradient-informed posterior samplers offer a scalable and general path for solving inverse problems over discrete representations.

• Scalability: The parallel update mechanism addresses the mixing speed limitations of discrete samplers in high-dimensional spaces.
• Generality: The method decouples the inference algorithm from the specific training paradigm of the prior, allowing it to leverage the growing ecosystem of discrete diffusion models (text, code, images, audio).
• Practicality: By avoiding continuous relaxations, ∆LPS ensures that the final samples are valid discrete tokens, preventing the "projection gap" where a continuous approximation fails to decode into a consistent discrete signal.

The authors conclude that as discrete generative models scale, principled inference methods like ∆LPS will become increasingly critical for applications in biomolecular design, symbolic music, code completion, and other domains where data is inherently discrete.