Extending Pathwise Gradients to Discrete Random Variables via Finite-Order Relaxation

Donghan He; Luhuan Wu

cs.LGOct 6, 2026

Extending Pathwise Gradients to Discrete Random Variables via Finite-Order Relaxation

Authors: Donghan He, Luhuan Wu

Organizations: Department of Applied Mathematics and Statistics, Johns Hopkins University

Abstract

Pathwise gradients are preferred for continuous random variables because they are unbiased, low variance, and work with a single sample. For discrete variables, however, the pathwise identity cannot generally be exact for every differentiable function. We propose a general framework to construct finite-order exact pathwise gradient estimators for a range of common discrete variables such as Poisson. The estimator is the least-norm solution among all solutions that are unbiased for polynomials of degree at most. The resulting estimators preserve the hard forward sample, require no temperature tuning, and can be implemented in a few lines of codes. Against other admissible solutions, our estimator is unique and minimizes weight variance; in contrast, prior works use categorical variables or augmented representations to approximate non-categorical variables that induces excess variance and computations. To understand approximation bias for functions beyond the prescribed class, we also derive a non-asymptotic bias bound. In experiments our low order methods match or improve tuned baselines across linear, nonlinear and hierarchical latent-variable models, while out-speeding competitors in every runtime benchmark.

Figures & tables

Panel of Figure 1 — Figure 1: Fini te-order pathwise gradients. Left : pathwise gradients are unbiased for all differentiable functions - this is too rigid for discrete variables, so we relax the unbiasedness requirement to low-order polynomials; Center : generally, there are many solutions to the finite-order unbiasedness requirement; we use the least-norm principle to select the order-D unbiased solution, avoiding excess weight variance; Right : as an example, Fini-2 and 3 for Poisson reduce to a few lines of code.

Figure 2: Linear Poisson VAE with 512 latent dimensions. Left : mean single-sample gradient error relative to the analytic gradient early (epoch 100) and partway through training (epoch 1,000), shown as dots with standard-deviation bars. Both means are near zero; Fini-2 has about one-fifth the standard deviation of tuned EAT-cubic. Right : validation ELBO over training, mean

\pm

standard deviation across five matched seeds. Fini-2 follows the analytic-gradient trajectory more closely.

Dataset	Fini-3	Fini-2	EAT-cubic (0.1)	GSM (0.5)	EAT-sigmoid (0.5)	Score
Synthetic	$\mathbf{-330.61\pm 0.79}$	$\mathbf{-330.60\pm 0.75}$	$\mathbf{-330.62\pm 0.79}$	$-331.94\pm 0.95$	$-332.18\pm 0.97$	$-333.23\pm 1.51$

Table 1: Synthetic POGLM validation ELBO (higher is better), mean

\pm

std. over 30 seeds. Parentheses give the selected temperature

\tau

. Boldface marks the three highest means, which differ by at most

0.02

	Bernoulli Likelihoods			Gaussian Likelihoods
	MNIST	Fashion-MNIST	Omniglot	MNIST	Fashion-MNIST	Omniglot
Fini-3	$\bm{-118.60\pm 0.26}$	$\bm{-251.98\pm 0.40}$	$\bm{-129.38\pm 0.07}$	$\underline{595.48\pm 3.72}$	$\bm{133.74\pm 2.80}$	$\underline{403.14\pm 2.46}$
Fini-2	$\underline{-118.88\pm 0.19}$	$\underline{-252.26\pm 0.18}$	$\underline{-129.93\pm 0.12}$	$594.47\pm 3.92$	$132.09\pm 3.29$	$\bm{403.80\pm 2.00}$
EAT-cubic	$-120.06\pm 0.17$ (0.5)	$-253.04\pm 0.22$ (0.5)	$-130.76\pm 0.07$ (0.2)	$\bm{595.66\pm 3.60}$ (0.5)	$\underline{132.42\pm 1.70}$ (0.5)	$400.84\pm 2.05$ (0.5)
ReinMax	$-125.87\pm 0.38$ (1.1)	$-260.56\pm 0.57$ (1.1)	$-140.07\pm 0.40$ (1.0)	$589.21\pm 2.78$ (1.0)	$104.06\pm 2.51$ (1.1)	$396.63\pm 1.53$ (1.0)

Table 2: Nonlinear Poisson-VAE training ELBO

\uparrow

, mean

\pm

standard deviation over five matched seeds. Columns pair each dataset with a Bernoulli or Gaussian observation likelihood. First is bold, second is underlined.

		20 Newsgroups			RCV1
Model	Training method	Train ELBO $\uparrow$	Test ELBO $\uparrow$	PPL $\downarrow$	Train ELBO $\uparrow$	Test ELBO $\uparrow$	PPL $\downarrow$
Poisson DEF	Fini-3	$\bm{-179.30\pm 2.81}$	$\underline{-199.31\pm 2.04}$	$660.7\pm 7.6$	$\underline{-269.27\pm 3.16}$	$\underline{-267.64\pm 3.73}$	$\bm{379.3\pm 5.4}$
	Fini-2	$\bm{-179.30\pm 2.62}$	$\bm{-198.98\pm 1.92}$	$\underline{659.7\pm 6.7}$	$\bm{-269.20\pm 3.05}$	$\bm{-267.60\pm 3.73}$	$\underline{381.9\pm 4.6}$
	EAT-cubic (0.2)	$\underline{-180.94\pm 2.59}$	$-199.77\pm 2.17$	$664.6\pm 8.5$	$-271.92\pm 3.18$	$-270.16\pm 3.82$	$388.4\pm 4.1$
	ReinMax (1.3)	$-191.25\pm 2.88$	$-210.01\pm 2.37$	$\bm{638.1\pm 10.5}$	$-301.86\pm 2.83$	$-300.05\pm 3.27$	$384.4\pm 4.8$

Table 3: Poisson DEF results on train and test ELBO (

\uparrow

) and heldout perplexity (PPL,

\downarrow

), mean

\pm

standard deviation over five matched seeds. Best is bold; second is underlined.

Family	Latent dim.	Fini-2	Fini-3	ReinMax
binomial $n=8,p=0.5$	2	$\mathbf{-174.51\pm 0.45}$	$\underline{-176.91\pm 1.07}$	$-179.67\pm 0.67$ (1.0)
	16	$\underline{-123.38\pm 0.14}$	$\mathbf{-123.11\pm 0.10}$	$-125.55\pm 0.10$ (1.0)
	128	$\underline{-118.73\pm 0.22}$	$\mathbf{-118.43\pm 0.17}$	$-121.45\pm 0.15$ (1.1)
negbin $r=8,p=0.5$	2	$\underline{-177.34\pm 0.15}$	$\mathbf{-177.31\pm 0.24}$	$-182.58\pm 0.48$ (1.0)
	16	$\underline{-129.29\pm 0.16}$	$\mathbf{-129.25\pm 0.20}$	$-136.58\pm 0.71$ (1.0)
	128	$\underline{-122.58\pm 0.06}$	$\mathbf{-122.54\pm 0.36}$	$-124.41\pm 0.24$ (1.1)

Table 4: Count-VAE training ELBO, mean

\pm

standard deviation over five independent matched seeds. Best is bold; second is underlined.

Figure 3: Negative-Binomial representation diagnostic. Native Fini-2 and Fini-3 have lower gradient variance than matched Gamma–Poisson counterparts; ReinMax has higher variance.

	Linear VAE		Nonlinear VAE		DEF
Estimator	L4	A100	L4	A100	L4	A100
Fini-3	$\underline{0.326\pm 0.010}$	$\underline{0.361\pm 0.008}$	$\underline{1.868\pm 0.074}$	$\underline{2.016\pm 0.083}$	$\underline{0.928\pm 0.021}$	$\underline{0.907\pm 0.016}$
Fini-2	$\bm{0.309\pm 0.005}$	$\bm{0.338\pm 0.004}$	$\bm{1.768\pm 0.017}$	$\bm{1.848\pm 0.046}$	$\bm{0.863\pm 0.014}$	$\bm{0.835\pm 0.016}$
EAT-cubic	$1.427\pm 0.005$	$0.581\pm 0.010$	$2.680\pm 0.040$	$2.942\pm 0.073$	$1.546\pm 0.039$	$1.565\pm 0.015$
ReinMax	$1.144\pm 0.006$	$0.568\pm 0.009$	$2.692\pm 0.035$	$2.859\pm 0.104$	$1.451\pm 0.024$	$1.610\pm 0.134$

Table 5: Wall-clock(s) per epoch (

\downarrow

), mean

\pm

standard deviation over five repeats.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Local rate–distortion update directions in selected high-performing runs. Our estimators closely track the near-exact MVD reference; EAT-cubic is close on 20 Newsgroups but deviates more on RCV1, while ReinMax tilts toward lower distortion and higher rate.

Explore similar work

OrderGrad: Optimizing Beyond the Mean with Order-Statistic Policy Gradient Estimation
Jun 4, 2026Paavo Parmas, Yongmin Kim, Kohsei Matsutani +5Mean
The University of Tokyo
GradInf: Gradient Estimation as Probabilistic Inference
Jul 8, 2026Gaurav Arya, Mathieu Huot, Moritz Schauer +2Probabilistic Inference Gradient
Carnegie Mellon University, Pittsburgh, USA · Massachusetts Institute of Technology, Cambridge, USA · Chalmers University of Technology & University of Gothenburg, Gothenburg, Sweden +1
Does "Do Differentiable Simulators Give Better Policy Gradients?'' Give Better Policy Gradients?
Apr 20, 2026Ku Onoda, Paavo Parmas, Manato Yaguchi +1Policy Gradient Differentiable Physics
The University of Tokyo, Tokyo, Japan

Jun 4, 2026, cs.LGJ/K move · Enter open · S save

OrderGrad: Optimizing Beyond the Mean with Order-Statistic Policy Gradient Estimation

Paavo Parmas, Yongmin Kim, Kohsei Matsutani, Shota Takashiro+4

The University of Tokyo

Mean

Open paper AskPDF arXiv

Policy-gradient methods usually optimize expected return, but many real world applications care about distributional properties of returns: tail risk, outlier robustness, or best-of-K discovery. We introduce OrderGrad, a family of likelihood-ratio and reparameterization gradient estimators for order-statistic objectives. OrderGrad optimizes finite-sample L-statistics, i.e., weighted averages of sorted rewards or costs, recovering objectives such as VaR, CVaR, trimmed means, medians, and top-m/best-of-K criteria by changing only the rank weights. For any fixed sample size and rank-weight vector, OrderGrad provides an unbiased gradient estimator for the corresponding order-statistic objective. The method is implemented as a simple reward transformation that can then be used in an otherwise standard policy-gradient or reparameterized update. We study the resulting estimator's variance behavior and evaluate it on tasks where mean optimization is mismatched to the deployment objective, including LLM math post-training and other tasks. OrderGrad provides a unified, plug-and-play route to risk-averse, robust, and exploratory learning. Code: https://github.com/paavo5/ordergrad

Abstract

Explore similar work

OrderGrad: Optimizing Beyond the Mean with Order-Statistic Policy Gradient Estimation

GradInf: Gradient Estimation as Probabilistic Inference

Does "Do Differentiable Simulators Give Better Policy Gradients?'' Give Better Policy Gradients?