Categorical Reparameterization with Denoising Diffusion models

2026-01-02

arXiv.org

AI Summary

Gradient-based optimization of discrete categorical variables has long been hindered by either the high variance of score function estimators or the bias introduced by continuous relaxations. This work proposes a novel, training-free soft reparameterization approach by introducing denoising diffusion mechanisms to categorical variable modeling. Specifically, it constructs a differentiable sampler via a closed-form denoiser derived from a Gaussian noising process, enabling efficient gradient estimation with low bias. The method achieves optimization performance on par with or superior to existing techniques across multiple benchmark tasks, establishing a new paradigm for discrete optimization that is training-free, differentiable, and low-bias.

Technology Category

• Search and Optimization: Mixed Discrete/Continuous Search
• Machine Learning: Optimization
• Constraint Satisfaction and Optimization: Mixed Discrete/Continuous Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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.

Problem

• categorical variables
• gradient-based optimization
• reparameterization
• score-function estimators
• continuous relaxations

Innovation

• categorical reparameterization
• denoising diffusion
• gradient estimation
• discrete optimization
• diffusion sampler