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

• Discrete variables generally lack a pathwise estimator unbiased for all differentiable functions, despite the method's appeal for continuous-variable Monte Carlo gradients.The continuous pathwise estimator is valued for unbiasedness, typically low variance, and implementation simplicity.Source: §1 · p. 1
• Fini extends pathwise gradients to discrete variables by requiring unbiasedness for polynomials up to degree D and selecting the least-norm solution.The weaker polynomial requirement admits multiple solutions; the least-norm principle selects among them.Source: §1 · p. 2
• Unlike categorical approximations for non-categorical variables, Fini is solved once at design time and evaluates independently of support size.Categorical approximations can add computational overhead; the Fini evaluation is fast.Source: §1 · p. 2
• Within the exactness class, Fini uniquely minimizes weight variance, derives a non-asymptotic bias bound beyond that class, and characterizes excess variance from categorical and compound representations.Source: §1 · p. 2
• Fini matches or outperforms tuned specialized and generalist baselines across neural sequence models, VAEs, and hierarchical topic models.This is the paper's reported cross-model empirical comparison.Source: §1 · p. 2

2.1 MONTE CARLO GRADIENT ESTIMATION AND SCORE FUNCTION GRADIENTS

• Score-function gradients estimate derivatives of expectations from Monte Carlo samples and apply to continuous and discrete variables, with unbiasedness under suitable regularity conditions.The estimator can nevertheless have high variance.Source: §2.1 · p. 3

2.2 PATHWISE GRADIENTS

• For continuous families, pathwise gradients use derivative information through a weight multiplying f′(Z), but common discrete variables admit no such weight unbiased for every differentiable function.An explicit reparameterization differentiates the sampling path with respect to distribution parameters.Source: §2.2 · p. 3

3 FINI: FINITE-ORDER PATHWISE GRADIENTS

• Fini requires the pathwise estimator w_D(Z)f′(Z) to be unbiased for polynomials of degree at most D, rather than every differentiable function.Linearity reduces the requirement to the monomials x, x^2, …, x^D.Source: §3 · p. 3
• The least-norm rule resolves the underdetermined constraints in L2(pθ), yielding a weight that is a linear combination of monomials up to degree D−1.The monomial basis supplies the Riesz representers for the constraints.Source: §3 · p. 3
• Fini-D is obtained by solving a D × D moment system, requiring finite moments through order 2D−2, differentiable first D moments, and a nonsingular Gram matrix.The resulting polynomial weight is evaluable in O(D) time, independent of support size.Source: §3 · pp. 3–4
• Fini-4 illustrates a variance cost of higher exactness: its variance under the log-rate parameterization diverges as λ ↓ 0.The construction extends to any prescribed order.Source: §3 · p. 4
• For Poisson variables, order-2 quadrature approximates the difference identity, and the Poisson shift identity replaces the shifted derivative in expectation to recover the weight.The order-3 rule uses interpolation nodes t = 0, 1, 2.Source: §3 · p. 4
• Fini-2 recovers fully exact pathwise gradients for scalar location–scale families and the minimum-energy affine pathwise field for Gaussian random vectors.For categoricals, a symmetry rule instead of the least-norm rule recovers ReinMax at temperature 1; Fini-2 corresponds to hard one-hot MVE.Source: §3 · p. 4

4 THEORY

• Under the stated assumptions, Fini’s degree-D exact solution uniquely minimizes both norm and weight variance among admissible weights.The result assumes the Gram matrix G_D is nonsingular.Source: §4 · p. 4
• For Poisson, categorical ReinMax at τ = 1 has a degree-2-exact large-truncation limit, whose categorical representation adds excess weight variance through a null component.The excess variance is characterized for ReinMax applied to Poisson and scalar exponential families.Source: §4 · pp. 4–5
• Rao–Blackwellization removes auxiliary noise in compound representations, but closes the weight-variance gap at order 2 and not at order 3.Conditioning the compound order-D Poisson weight gives the Negative-Binomial weight; Rao–Blackwellization alone suffices only when the conditioned weight is the least-norm solution.Source: §4 · p. 5
• Beyond degree-D polynomials, exactness cancels the Taylor polynomial and the bias is controlled by the first unmatched derivative.Least norm minimizes the bound’s weight-dependent term; the bound predicts small bias when the downstream computation is locally well approximated by low-degree polynomials.Source: §4 · p. 5

5 RELATED WORK

• Score-function estimators are unbiased but often high variance, while MVD and GO reduce variance at the cost of additional downstream function evaluations.Score-function variance reduction can also require more samples or carefully designed control variates.Source: §5 · p. 5
• MVD and GO retain arbitrary-function unbiasedness, whereas Fini is described as low-bias and uses one downstream evaluation.For Poisson, quadrature of the MVD identity yields Fini’s weights.Source: §5 · p. 5
• Categorical methods offer a common interface for discrete variables, but computation grows with the number of states; Fini instead applies directly to the target law.Gumbel–Softmax uses temperature-controlled continuous relaxations, whereas straight-through estimators preserve the hard sample and substitute a backward derivative.Source: §5 · pp. 5–6
• ReinMax replaces the linear straight-through approximation with a quadratic one and serves as Fini’s principal generalist comparator.The paper notes that MVE has theoretically lower weight variance than ReinMax but does not consistently outperform it empirically.Source: §5 · p. 6
• For Poisson, EAT-cubic is a distribution-specific specialist comparator, while Gamma–Poisson compounding transfers a Poisson estimator to Negative-Binomial variables.Fini instead applies its general construction directly to the target variable; the paper quantifies variance overhead from compound representations.Source: §5 · p. 6

6.1 LINEAR POISSON VAE: TESTING ORDER-2 EXACTNESS

• Fini-2 has about one-fifth the gradient-error standard deviation of tuned EAT-cubic and follows the analytic-gradient validation-ELBO trajectory more closely.The linear-decoder benchmark has a quadratic objective in latent counts, making Fini-2 exact.Source: §6.1 · p. 6, Figure 2 · p. 7

6.2 PERFORMANCE ON REALISTIC BENCHMARKS

• Temperature-free Fini-2 and Fini-3 achieve mean validation ELBOs close to tuned EAT-cubic on the synthetic POGLM benchmark, while other baselines perform worse.The comparison uses 30 matched seeds; the three highest table means differ by at most 0.02.Source: §6.2 · p. 7, Table 1 · p. 7
• Fini-3 leads in four of six nonlinear Poisson-VAE settings, including all Bernoulli-likelihood cases, with Fini-2 close behind.Both orders remain in the leading tier across all six settings, but higher exactness does not uniformly improve optimization under Gaussian likelihoods.Source: §6.2 · p. 7, Table 2 · p. 8
• On both topic-model datasets, Fini-2 and Fini-3 have the best train and test ELBOs; ReinMax leads heldout perplexity on 20 Newsgroups, while Fini-3 leads on RCV1.Fini closely tracks the low-variance MVD on both ELBO components; ReinMax at τ = 1.3 tilts toward worse KL and better log likelihood, a pattern that may explain its 20 Newsgroups perplexity result through bias.Source: §6.2 · p. 7, Table 3 · p. 8
• Fini-2 and Fini-3 match tuned EAT-cubic on POGLM, improve most of its ELBOs, and outperform ReinMax on every displayed ELBO.The preferred Fini order remains task dependent.Source: Table 1 · p. 7, §6.2 · p. 8, Table 2 · p. 8

6.3 GENERALITY ACROSS COMMON DISCRETE VARIABLES

• Across Poisson, Binomial, and Negative-Binomial VAEs, Fini-3 leads at higher latent dimensions, while Fini-2 leads in two of three 2-dimensional settings; the preferred order depends on the problem.For Negative-Binomial latents, native Fini-2 and Fini-3 have lower gradient variance than matched Gamma–Poisson counterparts; conditioning recovers Fini-2 exactly, but an order-3 Z-only component can remain.Source: §6.3 · pp. 8–9, Table 4 · p. 8, Figure 3 · p. 8
• For Binomial latents, categorical embedding adds computation, while Fini’s fixed-order weights have costs independent of support size; at 1024, 1D MVE takes 38× as long as Fini-2.The paper distinguishes this computational overhead from the statistical overhead of auxiliary-variable representations.Source: §6.3 · p. 9, Figure 4 · p. 9

6.4 SPEED COMPARISONS

• Across three benchmarks on L4 and A100 GPUs, Fini-2 and Fini-3 rank first and second in every workload–hardware pair, with Fini-2 consistently faster because its weight is affine rather than quadratic.Both methods require one forward–backward pass and evaluate affine or quadratic polynomials.Source: §6.4 · p. 9, Table 5 · p. 9

7 DISCUSSION

• Fini’s order controls a bias–variance trade-off, and higher order is not a uniform improvement; its fourth-order Poisson example can have worse variance at low rates.The method also requires relevant moments and their parameter derivatives, at least D support points for order-D exactness, and smoothness and finite moments for its bound beyond the prescribed class.Source: §7 · p. 9

A DISCRETE ALL-FUNCTION PATHWISE OBSTRUCTION

• On isolated discrete support, when the probability derivative at a support point is nonzero, no sample-measurable weight can give an exact pathwise identity for every smooth function.A smooth bump can have zero derivative on the support while its expectation still varies with that point’s probability.Source: A DISCRETE ALL-FUNCTION PATHWISE… · p. 13

B SOLVING THE OPTIMIZATION PROBLEM

• When the Gram matrix is nonsingular, the degree-D exact weight has a unique least-norm solution in L2(pθ).The construction selects among weights satisfying the degree-D exactness constraints.Source: B SOLVING THE OPTIMIZATION PROBLEM · p. 13

B.1 AUGMENTED REPRESENTATIONS AND RAO–BLACKWELLIZATION

• Conditioning a degree-D exact compound weight on the observed variable preserves exactness, while the native minimum-L2 weight has no greater variance than either compound version.If the Rao–Blackwellized weight lies in the exactness subspace, it equals the native minimum-L2 weight.Source: §B.1 · pp. 13–14

B.2 GAMMA–POISSON SPECIALIZATION

• For the Gamma–Poisson construction, the native order-2 negative-binomial weight is the Rao–Blackwellization of its compound counterpart; at order 3, the variance gap includes a native null-component contribution.The compound weight is formed by applying an order-D Poisson log-rate weight conditionally on the mixing variable.Source: §B.2 · p. 14

C.1 RECOVERING PATHWISE DERIVATIVES

• For Gaussian location-scale families, quadratic exactness recovers the pathwise derivative for sufficiently regular functions, while a pure location family needs only degree-1 exactness.The appendix also gives a Gaussian-vector recovery result for sufficiently regular functions.Source: §C.1 · pp. 14–15

C.2 CONNECTION TO MEASURE-VALUED DERIVATIVES AND QUADRATURE

• Polynomially exact quadrature of the Poisson finite-difference identity produces the same weights as the moment construction; the degree-3 rule uses nodes 0, 1, and 2.The analogous two-node construction also yields the Binomial quadratic weight.Source: §C.2 · p. 15

C.3 CONNECTION TO REINMAX AND GENERALIZED STRAIGHT-THROUGH METHODS

• Among quadratic-exact categorical weights, sample-wise symmetry uniquely selects ReinMax at temperature 1, while the least-norm rule gives a separate categorical estimator.The scalar degree-2 least-norm weight also coincides with the hard quadratically-unbiased one-dimensional minimum-variance estimator; the cited correspondence is specific to order 2.Source: §C.3 · pp. 16–18

D.1 PROOF OF PROPOSITION 4.1

• Any degree-D exact weight differs from the minimum-norm weight by a component orthogonal to the exactness subspace, and equality in the norm or variance objective requires that component to vanish.The proof identifies the minimum-norm element as unique.Source: §D.1 · p. 18

D.2 PROOF OF PROPOSITION 4.5

• The bias analysis assumes a bounded (D+1)th derivative, finite moments through order 2D+2, square-integrable weights and score, and a score identity for the function and degree-D polynomials.Its proof uses a degree-D Taylor polynomial, whose exactness reduces the bias to the remainder terms.Source: §D.2 · p. 18

D.3 ORDER-2 REINMAX IN A SCALAR NATURAL EXPONENTIAL FAMILY

• For a scalar natural exponential family, the order-2 τ=1 ReinMax variance gap is strictly positive when at least three support points have positive probability.The proof shows the gap can vanish only if a nonzero quadratic vanishes almost surely.Source: §D.3 · p. 19

D.4 POISSON REINMAX UNDER INCREASING CATEGORICAL TRUNCATION

• The τ=1 ReinMax weight for a Poisson law truncated to {0, …, K} converges in L2(pλ) as K increases.The proof uses convergence of the truncated mean and variance and finiteness of the Poisson fourth moment.Source: §D.4 · p. 19

E ADDITIONAL CLOSED-FORM ESTIMATORS

• The additional closed-form estimators are minimum-L2 solutions for the stated parameter coordinates, with order 2 affine and order 3 quadratic weights.The appendix introduces formulas for Poisson, Binomial, Negative-Binomial, and Beta-Binomial settings.Source: E ADDITIONAL CLOSED-FORM ESTIMATORS · p. 20, §E.1 · p. 20, §E.2 · p. 20

E.1 POISSON: ORDER 4

• The Poisson appendix gives an order-4 weight derived from moments through order 6.The displayed material also gives the resulting weight and variance, without a readable variance comparison in these passages.Source: §E.1 · p. 20

E.2 BINOMIAL: η = logit p

• For X ∼ Binomial(n, p), the appendix specifies η = logit p and gives an order-2 weight formula.Source: §E.2 · p. 20

E.3 NEGATIVE BINOMIAL: η = logit p

• For the Negative-Binomial convention with η = logit p, the appendix gives an order-2 estimator and its corresponding raw-p weight.Source: §E.3 · pp. 20–21

E.4 NEGATIVE BINOMIAL: ρ = log r

• For the Negative-Binomial parameterization ρ = log r, the appendix gives a corresponding raw-r weight.Source: §E.4 · p. 21

F TWO-PARAMETER STRAIGHT-THROUGH IMPLEMENTATION

• The two-parameter surrogate adds detached weights u and v to the hard sample’s straight-through corrections for optimized coordinates a and b.For the Negative-Binomial parameterization, a = η = logit p and b = ρ = log r.Source: F TWO-PARAMETER STRAIGHT-THROUGH… · pp. 22–23

G ADDITIONAL EXPERIMENTAL DETAILS

• ReinMax truncation uses the batch’s largest rate and probability p = 0.999; means and standard deviations use independently seeded, estimator-matched runs.The passage specifies this ReinMax truncation detail and the default reporting procedure.Source: G ADDITIONAL EXPERIMENTAL DETAILS · p. 23

G.1 LINEAR POISSON VAE

• −167.02 ± 0.05 for order 2, −167.11 ± 0.06 for EAT-cubic at temperature 0.05, and −166.99 ± 0.04 for the exact reference; the saturated benchmark primarily validates exactness and gradient variance.The EAT-cubic comparison has p = 0.034 by a two-sided Welch t-test.Source: §G.1 · p. 23

G.2 NEURAL SEQUENCE MODEL

• The neural sequence benchmark uses repository-default Conv1d initialization after the stated N(0, 0.52) initialization failed to reproduce the reported ELBO range; statistics use 30 matched seeds.With either one or five Monte Carlo samples, the stated initialization converged around −427 to −429, while the default initialization reproduced results near −330.5.Source: §G.2 · p. 23

G.3 NONLINEAR POISSON AND COUNT VAES

• The Poisson VAE reuses an earlier design with a Pois(1.0) prior, while count-family experiments use dynamically binarized MNIST and specified Binomial, Negative-Binomial, and sparse-Poisson settings.The count-family settings are p = 0.5 for Binomial and Negative-Binomial latents and λ = 0.1 for sparse Poisson; both experiment families use five matched seeds.Source: §G.3 · p. 23

G.4 BINOMIAL SUPPORT-SIZE TIMING

• The Binomial support-size timing benchmark compares native order-2 and order-3 weights with Hooper & Shekhovtsov’s released 1D MVE implementation at τ = 106.It evaluates batches of 200 × 128 latents on an NVIDIA L4 across n ∈ {4, 8, 16, 32, 64, 128, 256, 512, 1024}, with Figure 4 reporting milliseconds per forward–backward evaluation.Source: §G.4 · p. 23

G.5 DEF EXPERIMENT DETAILS

• The DEF experiments use reduced temperature grids and learning rates, AdamW with batch size 64, and five matched seeds; training lasts 500 epochs on 20 Newsgroups and 100 on RCV1.The selected-configuration grids are ReinMax {1, 1.3, 1.5}, EAT-cubic {0.02, 0.05, 0.2}, and learning rates {3 × 10−3, 10−3, 3 × 10−4}.Source: §G.5 · p. 23

G.6 REINMAX AND THE RATE–DISTORTION TRADEOFF

• Figure 5 shows our estimators and EAT-cubic tracking near-exact directions, while ReinMax favors lower distortion and higher rate; on 20 Newsgroups it has better held-out perplexity despite worse train and test ELBO.The displayed direction diagnostic points away from interpreting ReinMax’s perplexity advantage as more accurate ELBO-gradient estimation.Source: Figure 5 · p. 24, §G.6 · p. 24

G.5 DEF EXPERIMENT DETAILS

• Amortized inference avoids storing document-specific RCV1 latents; the topic model uses point-estimated decoder weights, 30 supertopics, 100 topics, and a Pois(0.1) sparsity prior.The encoder is a two-layer neural network mirroring the Gaussian-weight, softplus-link Poisson decoder.Source: §G.5 · p. 24
• The 20 Newsgroups corpus has 1,995 terms and 11,258 training/7,487 test documents; RCV1 retains 10,000 terms and uses 794,414 training/10,000 test documents.The RCV1 documents are pooled and randomly split.Source: §G.5 · p. 24

G.7 TIMING PROTOCOL

• Table 5 timing values average five post-warm-up repeats, with each workload’s model, batch size, hardware, and training configuration held fixed across estimators.Timing repeats cover each estimator–workload–hardware combination.Source: §G.7 · pp. 24–25