Categorical Reparameterization with Denoising Diffusion models

TL;DR

The paper tackles gradient estimation for expectations over discrete categorical distributions by introducing ReDGE, a diffusion-based soft reparameterization that yields a training-free differentiable sampling map from Gaussian noise to the target pmf. A key insight is that for simplex-vertex support the denoiser admits a closed-form expression, enabling efficient reverse transitions and a DDIM-inspired sampler that interpolates between continuous relaxation and exact discrete sampling. The authors derive practical gradient estimators, including hard variants and ReinMax-like second-order corrections, and propose parameter-dependent initializations to improve performance with minimal diffusion overhead. Empirically, ReDGE variants match or surpass strong baselines (ST, Gumbel–Softmax, ReinMax) across polynomial programming, Gaussian mixture variational inference, Sudoku, and Categorical VAE tasks, demonstrating competitive optimization performance and favorable convergence properties. This diffusion-based approach offers a flexible, low-bias alternative to traditional relaxations and opens avenues for bias-correction via REBAR/RELAX-style control variates and more sophisticated base distributions."

Abstract

Gradient-based optimization with categorical variables typically relies on score-function estimators, which are unbiased but noisy, or on continuous relaxations that replace the discrete distribution with a smooth surrogate admitting a pathwise (reparameterized) gradient, at the cost of optimizing a biased, temperature-dependent objective. In this paper, we extend this family of relaxations by introducing a diffusion-based soft reparameterization for categorical distributions. For these distributions, the denoiser under a Gaussian noising process admits a closed form and can be computed efficiently, yielding a training-free diffusion sampler through which we can backpropagate. Our experiments show that the proposed reparameterization trick yields competitive or improved optimization performance on various benchmarks.

Figures (6)

• Figure 1: Visualization of the DDIM transport for $\targ{\theta}{}{} = \theta\!\cdot\!\updelta_{-2} + (1 - \theta)\!\cdot\!\updelta_{2}$ with the linear schedule $(\alpha_t, \sigma_t) = (1-t, t)$. First two rows: DDIM trajectories with varying ${t_{1}}$ for two different values of $\theta \in [0,1]$. Third row: The DDIM map $\theta \mapsto T^\theta _0(x_1; t_{1:n-1})$ for fixed input quantiles $z$ and three different values of ${t_{1}}$. $\Phi$ stands for the standard Gaussian cdf.
• Figure 2: Polynomial programming benchmark for different values of the exponent $p$.
• Figure 3: Categorical VAE training curves for the configuration with $L=48,\, K=2$.
• Figure 4: Polynomial programming benchmark for different values of the exponent $p$ with the linear relaxation.
• Figure 5: Negative ELBO curves during optimization for the GMM experiment.

Theorems & Definitions (20)

• Proposition 3.1: Informal
• Proposition 1.1
• Corollary 1.2
• Lemma 1.3
• proof
• Lemma 1.4: Continuity of the margin function outside of $\mathsf{H}$ \ref{['eq:hyperplane-union']}
• proof
• Lemma 1.5: Softmax bound
• proof
• Lemma 1.6: Covariance control