Categorical Reparameterization with Denoising Diffusion models

Samson Gourevitch, Eric Moulines, Alain Oliviero Durmus, Jimmy Olsson, Yazid Janati

Keywords: Diffusion models, discrete reparameterization trick, discrete diffusion guidance

TL;DR: We propose a diffusion-based reparameterization trick for categorical distributions

Abstract: 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 \algoname, a novel and efficient diffusion-based soft reparameterization method for categorical distributions. Our approach defines a flexible class of gradient estimators that includes the \sthrough~estimator as a special case. Experiments spanning latent variable models and inference-time reward guidance in discrete diffusion models demonstrate that \algoname\ consistently matches or outperforms existing gradient-based methods.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 91