Categorical Reparameterization with Denoising Diffusion models

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Gourevitch, Samson, Durmus, Alain, Moulines, Eric, Olsson, Jimmy, Janati, Yazid
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866908822037069824
author Gourevitch, Samson
Durmus, Alain
Moulines, Eric
Olsson, Jimmy
Janati, Yazid
author_facet Gourevitch, Samson
Durmus, Alain
Moulines, Eric
Olsson, Jimmy
Janati, Yazid
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2601_00781
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Categorical Reparameterization with Denoising Diffusion models
Gourevitch, Samson
Durmus, Alain
Moulines, Eric
Olsson, Jimmy
Janati, Yazid
Machine Learning
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.
title Categorical Reparameterization with Denoising Diffusion models
topic Machine Learning
url https://arxiv.org/abs/2601.00781