Why VAEs use reparameterisation instead of REINFORCE

A VAE does not need the reparameterisation trick in order to have an encoder gradient. We could train its encoder with REINFORCE instead. The estimate would still be unbiased; it would simply be much noisier.

The first structural difference appears before latent dimension enters the story. With respect to the objective $f$, REINFORCE is zeroth-order: it observes the value $f(z)$ at a random sample and combines that value with the sample’s log-probability gradient. Reparameterisation is first-order: it observes the local slope $\nabla_z f(z)$ and follows that slope through the sampling map. Even with one latent variable, a function value and a derivative contain different information.

High dimension adds a second problem. REINFORCE broadcasts one scalar value to every latent coordinate, whereas reparameterisation supplies a separate partial derivative to each coordinate. In the factorised $d$-dimensional setting analysed below, this cross-coordinate interference can make matching one pathwise sample require roughly $d$ score-function samples. The factor of $d$ is an amplification of the information gap, not its origin.

Policy gradients and VAEs share this mathematical shape: their parameters change the distribution that produces the samples. Imagine trying to improve the average exam score in a room when your only knob changes who enters, not any person’s score. The population moves as you turn the knob. Score-function and pathwise estimators are two ways to differentiate that moving average [1].

The reparameterisation trick is therefore not part of the VAE model. It is one way to estimate the model’s gradient. This post removes that estimator, replaces it with REINFORCE, and separates the two things we lose: local slope information already at $d=1$, and coordinate-local credit when $d>1$.

One expectation, two gradient estimators

Start with the common objective

Here $q_\phi$ is the distribution that produces a random sample $z$, $\phi$ controls that distribution, and $f(z)$ assigns a value to the sample. For the moment, assume that $f$ itself does not contain $\phi$. We want $\nabla_\phi F$: the direction in which changing the sampler improves the expected objective.

There are two standard ways to obtain this gradient. The first differentiates the probability of each possible sample. Using $\nabla_\phi q_\phi = q_\phi \nabla_\phi \log q_\phi$ gives

This is the score-function estimator; in reinforcement learning it is REINFORCE [2]. Here “score function” names the density derivative $\nabla_\phi \log q_\phi(z)$, not the objective value $f(z)$. The factor $f(z)$ says how good the sampled outcome was; the density score says how to change $\phi$ to make that outcome more likely. Multiplying them turns “this sample was good” into a parameter update.

The important advantage is that we never differentiate $f$. It may be a simulator, a discrete test, or any other black box. We only need to evaluate $f(z)$ and differentiate the log-probability of drawing $z$. In optimisation language, the estimator is zeroth-order with respect to $f$—even though it still differentiates the density $q_\phi$.

The second method rewrites the random draw. Suppose we can first sample parameter-free noise $\varepsilon \sim q_0$ and then produce $z$ through a differentiable map $z = T_\phi(\varepsilon)$. The noise distribution stays fixed when $\phi$ changes, so ordinary chain-rule differentiation applies:

This is the pathwise estimator. The right-hand side multiplies two quantities: $\nabla_z f$ tells us how the objective would change if the sample moved, and $\partial T_\phi / \partial \phi$ tells us how changing the parameters would move that sample. For a Gaussian, $T_\phi(\varepsilon) = \mu_\phi + \sigma_\phi \odot \varepsilon$; this special case is the reparameterisation trick used by VAEs.

The pathwise estimator asks for more structure. The map $T_\phi$ and the objective $f$ must both be differentiable. In return, it is first-order with respect to $f$: it gets the local slope $\nabla_z f$ instead of only the value $f(z)$.

Under the usual regularity conditions, (1) and (2) are exact expressions for the same gradient. Neither is an approximation to the other. In practice we approximate either expectation with Monte Carlo samples, and those sample estimates can have very different variance.#Other gradient estimators exist, including measure-valued derivatives and finite-difference methods, but score-function and pathwise estimators dominate this setting.

That distinction matters because an optimiser never sees the exact expectation. It sees a noisy estimate from one or a few samples and takes a step. So the useful question is not “which estimator has a gradient?” Both do. It is “which estimator gives the lower-variance gradient, and why?”

The gap already exists at d = 1

Consider the smallest possible example. Let $z=\mu+\varepsilon$ with $\varepsilon\sim\mathcal{N}(0,1)$, and let the objective be $f(z)=z$. There is one latent variable and one parameter. The exact gradient is as simple as it can be:

For REINFORCE, the Gaussian mean score is $\partial_\mu\log q_\mu(z)=z-\mu=\varepsilon$. At any fixed $\mu$, the variance-minimising baseline that is constant with respect to $z$ is $b=\mathbb{E}[f(z)]=\mu$. Treat $b$ as a detached control variate when forming the estimator; one sample then gives

The estimate is unbiased, but it fluctuates because it must infer the direction from two random quantities: whether the sampled value was above its baseline, and which side of the Gaussian mean produced it.

The pathwise estimator differentiates the same sample directly:

It receives the exact answer from every sample because this particular objective has a constant local slope. No cross-coordinate credit assignment exists—the entire difference comes from objective-value information versus objective-derivative information.

This example does not prove that every pathwise estimator has lower variance than every score-function estimator. Variance depends on $f$, $q_\phi$, the parameterisation and the available control variates. It proves the narrower point we need: setting $d=1$ removes cross-coordinate interference, but it does not turn a zeroth-order observation of $f$ into a first-order one.

Applying both estimators to a VAE

A VAE introduces an unobserved latent variable $z$ to explain an observed example $x$. Its decoder defines the joint model $p_\theta(x,z) = p(z)p_\theta(x \mid z)$, while its encoder $q_\phi(z \mid x)$ approximates the posterior over $z$. Training maximises the evidence lower bound, or ELBO:

The expectation is over latent samples drawn from the encoder. Inside the brackets, $\log p_\theta(x,z)$ rewards a latent value that explains the data under the model, while $-\log q_\phi(z \mid x)$ is the encoder-entropy contribution. Their difference is the sampled learning signal; write it as $A_{\theta,\phi}(x,z)$.

The decoder parameters $\theta$ appear only inside this signal, so their derivative passes through the expectation in the usual way. The encoder parameters $\phi$ are harder: they control both the distribution that draws $z$ and the $-\log q_\phi$ term inside the signal. This is the moving-distribution problem from the previous section, with one extra detail.

The direct derivative of $-\log q_\phi$ averages to zero. Indeed, $\mathbb{E}_q[\nabla_\phi(-\log q_\phi)] = -\int \nabla_\phi q_\phi \, dz = -\nabla_\phi 1 = 0$. Therefore a valid REINFORCE estimator for the encoder is

The baseline $b(x)$ is any value that does not depend on the sampled $z$. Subtracting it changes the estimator’s variance but not its expectation, because the expected score $\mathbb{E}_q[\nabla_\phi \log q_\phi]$ is zero. Equation (4) is the precise sense in which a VAE can be trained with REINFORCE.

One encoder serves many inputs. This is amortised inference. Because the typical ELBO can differ greatly from one input $x$ to another, one global baseline for the whole dataset is usually too crude. The natural baseline is a function $b(x)$.

A standard VAE also satisfies the pathwise estimator’s requirements. Its Gaussian latent can be written as $z = \mu_\phi(x) + \sigma_\phi(x) \odot \varepsilon$ with fixed noise $\varepsilon \sim \mathcal{N}(0,I)$, and its decoder is differentiable. This makes the VAE a particularly favourable case for reparameterisation.

These conditions are useful, not automatic. Score-function estimators for variational inference came first, and black-box variational inference was built on them [3]. Reparameterised autoencoders did not make (3) differentiable for the first time. They supplied a much lower-variance way to estimate a gradient that already existed [4], [5].

High dimension adds a second penalty

The one-dimensional example isolates the information gap. Now suppose the encoder factorises across $d$ latent coordinates, as a diagonal Gaussian encoder does:

Think of $\phi_i$ as the encoder parameters that control coordinate $z_i$. Its density score $s_i = \partial_{\phi_i}\log q_i(z_i;\phi_i)$ depends only on $z_i$. But in REINFORCE it is multiplied by the full objective value $f(z_1,\ldots,z_d)$. Every coordinate therefore receives the same scalar assessment of the entire latent vector.

The pathwise estimator uses $\partial f/\partial z_i$ instead. Coordinate $i$ receives the derivative specific to moving $z_i$. This creates the additional high-dimensional difference shown in figure 2: REINFORCE broadcasts one global value, while reparameterisation assigns coordinate-local derivatives.

We can make the cost of this additional broadcast exact. Define $\bar f(z_i)=\mathbb{E}[f(z)\mid z_i]$: fix coordinate $z_i$, average over all the other coordinates, and keep only the part of the objective value predictable from $z_i$. The following claim separates the one-coordinate score-function variance from the extra variance contributed by other coordinates. The proof is included for completeness; the paragraph after it gives the plain-language reading.

The first term in claim 1 is the variance that can remain with one coordinate. When $d=1$, $\bar f=f$ and the second term is zero, but REINFORCE must still convert the value $f(z)$ into a direction by multiplying it by the density score $s_i$. A baseline can reduce this variance, but it cannot supply the missing derivative $\partial f/\partial z_i$.

The second term is the genuinely high-dimensional penalty. It is the variation left after $z_i$ is fixed: in a factorised encoder, that variation comes from the other latent coordinates. It is noise charged to coordinate $i$ even though coordinate $i$ did not cause it. A baseline that cannot see $z$ cannot remove this term.

Now suppose the learning signal is a sum of $d$ similarly scaled coordinate contributions. This is exact in the appendix’s orthogonal linear model and is a useful approximation in some larger models. The conditional variance $\operatorname{Var}(f\mid z_i)$ then contains roughly $d-1$ unrelated contributions, so it grows like $d$. The gradient signal itself stays $O(1)$ for each coordinate, which makes the signal-to-noise ratio fall like $d^{-1/2}$.

Averaging $N$ independent Monte Carlo samples improves signal-to-noise by $\sqrt{N}$. Recovering a factor of $\sqrt d$ therefore takes about $N=d$ samples. This is where the “one pathwise sample costs roughly $d$ REINFORCE samples” rule comes from; it is a scaling statement for this factorised setting, not a universal constant.

The pathwise estimator does not broadcast the scalar $f(z)$. For a coordinatewise sampling map, $\partial z_k/\partial\phi_i=0$ when $k\ne i$, so (2) reduces to $\partial_{z_i}f(z)\,\partial_{\phi_i}T_i$. Other coordinates may still affect the value of $\partial_{z_i}f$ when the decoder couples them, but they do not enter as an additive global score. If $f$ is separable, they disappear completely.

Reparameterisation uses locality twice: a local slope already at $d=1$, and coordinate-specific credit when $d>1$.

What a baseline can and cannot buy

A baseline is a reference value for the objective. Instead of multiplying the density score $s_i$ by the raw learning signal $f$, REINFORCE multiplies it by the centred signal $f-b$. Since $\mathbb{E}[s_i]=0$, subtracting $b$ leaves the expected gradient unchanged. A good baseline says, in effect, “do not reward a sample merely for being typically good; reward it for being better than expected.”

But centring a function value does not turn it into a derivative. Even the best constant baseline leaves the first term of claim 1 unless the particular objective and parameterisation make that term degenerate. This is the limit already visible at $d=1$: a baseline can remove an irrelevant offset from $f(z)$, but it cannot reveal the local slope $\nabla_z f(z)$ that REINFORCE never evaluated.

This centring removes a large but avoidable source of variance. If the global objective has a mean of size $O(d)$, failing to subtract it creates $O(d^2)$ variance. An optimal scalar baseline removes that offset, but the other coordinates still fluctuate, leaving the cross-coordinate contribution at $O(d)$ in the quadratic Gaussian example.¹

A coordinate-specific baseline can go further. Let $z_{-i}$ denote every latent coordinate except $z_i$. Under the factorised encoder, a baseline $b_i(x,z_{-i})$ is independent of $s_i$, so $\mathbb{E}[b_i s_i]=0$ and the estimator remains unbiased. Choosing $b_i\approx\mathbb{E}[f\mid z_{-i}]$ removes the other coordinates’ additive contribution when $f$ is separable and can reduce it more generally. Rao–Blackwellised and local-expectation estimators use this idea [3], [6], but they typically require work for each coordinate. The trade is therefore about $d$ times as many samples or about $d$ times as much per-coordinate computation. Reparameterisation obtains the local derivative in one backward pass.

Amortisation makes the target move. One encoder serves inputs whose typical ELBO values may be very different, and the baseline accuracy required to expose a small gradient can tighten as the encoder approaches the posterior. A single global baseline can therefore look adequate early and become useless later. Learned input-dependent baselines are a requirement for a practical score-function VAE rather than a minor refinement [7]; the appendix gives the corresponding numbers in the solvable model.

Two VAE-specific control variates

Beyond a generic baseline, the VAE objective offers two useful pieces of estimator structure.

Keep the full learning signal. VAEs often compute (3) as a reconstruction term minus an analytic $\mathrm{KL}(q_\phi\Vert p)$. For a score-function estimator, applying REINFORCE only to reconstruction may increase variance. At the exact posterior, the full learning signal $A_{\theta,\phi}(x,z)$ equals the constant $\log p_\theta(x)$; the $-\log q_\phi$ term is acting as a control variate. Splitting it away keeps the noisy part and discards its cancellation.

Drop the zero-mean pathwise term. The standard pathwise loss differentiates $-\log q_\phi$ both through the sampled $z$ and directly through $\phi$. The direct component has zero expectation, so dropping it remains unbiased. This is the “sticking the landing” (STL) estimator [8]. Its variance becomes exactly zero at the exact posterior, whereas the standard pathwise estimator retains a variance floor. Near convergence, the meaningful comparison is therefore against this corrected pathwise estimator, not only the textbook version.

What to remember

Yes, a VAE can be trained with REINFORCE. The estimator in (4) is unbiased; it simply gives the optimiser a noisier update.

The reason begins at $d=1$. REINFORCE infers a direction from objective values, while reparameterisation reads the local derivative. High dimension then adds a second cost: one global value is broadcast across coordinates, so each coordinate inherits fluctuations caused by the others. In the solvable factorised model, matching one pathwise sample eventually costs roughly $d$ score-function samples.

Score-function estimators remain essential when the latent is discrete, the objective is a black box, or the sampling process cannot be differentiated. In those settings, baselines and stronger control variates are the price of working with less local information [7], [9], [10], [11].

The companion question is what happens when reinforcement learning does satisfy the pathwise requirements. First-order objective information should help even before dimension enters, and coordinate-local credit should add another advantage in high-dimensional action spaces. In practice neither gain is as automatic as it sounds, which is where the other post begins.

Appendix: the model behind the numbers

To put constants behind the scaling argument, we need a VAE whose gradient variance can be computed exactly. A linear Gaussian decoder provides one: probabilistic principal component analysis with a Gaussian encoder [12], [13]. The model is

Here $d$ is the latent dimension, $D$ is the data dimension, and the encoder parameters are $\phi=(\mu,\log s)$. Choose decoder columns satisfying $W^\top W=\alpha I_d$. This orthogonality makes the exact posterior diagonal and isotropic: $p(z\mid x)=\mathcal{N}(m^\star,\beta^{-1}I_d)$ with precision $\beta=1+\alpha/\sigma^2$. The encoder family therefore contains the exact posterior; the calculation can follow training all the way to a known optimum.

Write the encoder’s mean error as $m=\mu-m^\star$ and reparameterise a sample as $z=\mu+s\odot\varepsilon$, where $\varepsilon\sim\mathcal{N}(0,I_d)$. The sampled learning signal separates into the constant log evidence $\log p_\theta(x)$ plus independent coordinate terms:

In (5), each summand depends only on one standard-normal variable $\varepsilon_i$. That is the separable case of claim 1, and it reduces every variance calculation to Gaussian moments of order at most eight. The figure and table use $D=784$, $\alpha=1$, $\sigma=0.3$, and 64 inputs drawn from the model; $d=32$ except when latent dimension is swept.

The three curves in figure 3 show the predicted scalings. Pathwise signal-to-noise stays approximately constant as $d$ grows. A baselined score estimator loses a factor of $d^{1/2}$ because of the conditional-variance term. Without a baseline it loses roughly another factor of $d^{1/2}$ because the unremoved mean offset also grows. The final column of table 1 squares the pathwise-to-score signal-to-noise ratio to convert it into the number of score samples needed for equal precision.

What the numbers add

The $d=1$ row is the control that separates the two stories. With no other coordinate available to create assignment noise, the optimally baselined score estimator still needs $3.9$ samples per pathwise sample at this parameter point. That residual gap is model-dependent, not a universal constant, but it confirms that the estimator difference begins before dimension does.

The remaining rows measure the high-dimensional amplification. For $d=8,32,128,512$, the required score-function sample counts are $11.1,35.2,133.8,527.3$. Once $d$ is moderate, the exchange rate tracks latent dimension to within a few percent, as predicted by the conditional-variance term in claim 1.

The same calculation separates offset removal from credit assignment. At $d=32$, the unbaselined score estimator has about 1400 times the pathwise variance. A perfect scalar baseline cuts this to 33 times, but 88% of what remains is conditional variance from the other coordinates. Changing the data dimension from 64 to 3136 leaves the baselined signal-to-noise ratio unchanged in this model: data dimension changes an offset the baseline can remove, while latent dimension changes how many coordinates share one value.

A separate sweep from encoder initialisation toward the exact posterior exposes the amortisation problem. The baseline may initially miss the per-input ELBO by about 75 nats without dominating variance, but the tolerated error falls below one nat near the posterior, while ELBO values vary by about 50 nats across inputs. These constants belong to this model; the qualitative lesson is that the baseline target can become stricter as inference improves.

The script below implements the closed-form variance calculation for one sampled input. It then checks one coordinate against two million direct Monte Carlo draws; the two estimates agree within the relative errors reported below.²

import numpy as np

D, d, alpha, sigma = 784, 32, 1.0, 0.3
beta = 1.0 + alpha / sigma**2                        # posterior precision, per coordinate
rw, rx = np.random.default_rng(0), np.random.default_rng(1)
W = np.linalg.qr(rw.standard_normal((D, d)))[0] * np.sqrt(alpha)    # so W.T @ W = alpha I
x = rx.standard_normal(d) @ W.T + sigma * rx.standard_normal(D)

mu, s = np.zeros(d), np.ones(d)                      # the encoder at initialisation
m = mu - (x @ W) / (beta * sigma**2)                 # mismatch against the exact posterior
p1, p2 = -beta * m * s, 0.5 * (1.0 - beta * s**2)    # A = const + sum_i (p1 e_i + p2 e_i^2)
grad = np.sqrt((beta * m) @ (beta * m) + ((1 - beta * s**2) ** 2).sum())

def score_var(delta):
    """Per-coordinate exact variance of the score estimator; delta = b - ELBO(x), in nats.

    Returns the mean and scale blocks, each as (signal-carrying term, conditional variance) —
    the two halves of the claim. Only the first depends on the baseline.
    """
    q0 = -delta - p2
    cond = (p1**2 + 2 * p2**2).sum() - (p1**2 + 2 * p2**2)         # Var(f | z_i)
    return (((q0 + 3 * p2) ** 2 + 2 * p1**2 + 6 * p2**2) / s**2, cond / s**2,
            2 * (q0 + 5 * p2) ** 2 + 10 * p1**2 + 24 * p2**2, 2 * cond)

def path_var(stl):
    """Exact variance of the pathwise estimator, summed over every parameter."""
    g = (beta * s - 1 / s) ** 2 if stl else beta**2 * s**2
    h = 2 * (1 - beta * s**2) ** 2 if stl else 2 * beta**2 * s**4
    return float(g.sum() + (beta**2 * m**2 * s**2 + h).sum())

for name, v in [("score, b = 0", sum(t.sum() for t in score_var(480.63))),
                ("score, perfect b", sum(t.sum() for t in score_var(0.0))),
                ("pathwise", path_var(False)), ("pathwise + STL", path_var(True))]:
    print(f"{name:>18}  variance {v:10.4g}   SNR {grad / np.sqrt(v):.4f}")

signal, cond, _, _ = score_var(50.0)                 # the closed form, put in front of samples
eps = np.random.default_rng(7).standard_normal((2_000_000, d))
adv = eps @ p1 + (eps**2) @ p2 - (50.0 + p2.sum())
print(f"\nmean coordinate 0: sampled {(adv * eps[:, 0] / s[0]).var(ddof=1):.6g}, "
      f"exact {signal[0] + cond[0]:.6g}")