Categorical Reparameterization with Denoising Diffusion Models

Learning models with categorical variables requires optimizing expectations over discrete distributions, a setting in which stochastic gradientbased 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 diffusionbased 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 is available at https://github.com/ samsongourevitch/redge .