Categorical Reparameterization with Denoising Diffusion models

Learning models with categorical variables requires optimizing expectations over discrete distributions, a setting in which stochastic gradient-based optimization is challenging due to the non-differentiability of categorical sampling. A common workaround is to replace the discrete distribution with a continuous relaxation, yielding a smooth surrogate that admits reparameterized gradient estimates via the reparameterization trick. Building on this idea, we introduce ReDGE, a novel and efficient diffusion-based soft reparameterization method for categorical distributions. Our approach defines a flexible class of gradient estimators that includes the Straight-Through estimator as a special case. Experiments spanning latent variable models and inference-time reward guidance in discrete diffusion models demonstrate that ReDGE consistently matches or outperforms existing gradient-based methods. The code will be made available at https://github.com/samsongourevitch/redge.

💡 Research Summary

The paper tackles the long‑standing difficulty of differentiating expectations over categorical distributions, a problem that arises in reinforcement learning, variational inference with discrete latents, token‑level decisions in language models, and combinatorial optimization. Traditional solutions fall into two families: (i) score‑function estimators such as REINFORCE, which are unbiased but suffer from prohibitively high variance; and (ii) continuous relaxations like the Gumbel‑Softmax (Concrete) distribution, which enable the reparameterization trick but introduce bias that depends on a temperature parameter τ. As τ → 0 the bias vanishes but gradients become ill‑conditioned or vanish; as τ grows the gradients are well‑behaved but the objective deviates from the true discrete loss.

The authors propose a fundamentally different approach: they use denoising diffusion models (DDMs) to construct a soft reparameterization of a categorical distribution. A diffusion model defines a family of marginals (π_t)_{t∈