Categorical Reparameterization with Denoising Diffusion models

TL;DR

ReDGE introduces a diffusion-based soft reparameterization technique for categorical variables, enabling efficient gradient estimation and outperforming existing methods in latent variable modeling and discrete diffusion tasks.

Contribution

The paper presents ReDGE, a novel diffusion-based reparameterization method that improves gradient estimation for categorical variables, including a flexible class of estimators.

Findings

ReDGE matches or outperforms existing gradient estimators.

ReDGE is effective in latent variable models.

ReDGE improves inference in discrete diffusion models.

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 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.