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 ReDGE consistently matches or outperforms existing gradient-based methods.

Many AI systems need to make choices among separate options, such as choosing a word, a label, or a category inside a model. These choices are difficult to improve during training because modern learning methods work best when small changes produce smooth, measurable effects, while a discrete choice is more like flipping a switch. Existing methods often replace the hard choice with a softer approximation, but these approximations can be inaccurate, inefficient, or hard to use. We introduce ReDGE, a new method that uses ideas from diffusion models to create smoother versions of categorical choices. This allows learning algorithms to estimate how changing these choices would affect the final result. ReDGE also connects to a widely used existing method, the Straight-Through estimator, showing that it belongs to a broader family of approaches. In experiments on several types of models with discrete choices, ReDGE performs as well as or better than existing gradient-based methods, making it easier to train and guide machine-learning systems that rely on categorical decisions.