The Reparameterization Trick for Gaussian Latents

Why Reparameterize a Gaussian Latent?

In a count VAE, an encoder assigns a distribution to the latent representation of a cell, rather than a single fixed latent value. An objective such as the evidence lower bound (ELBO) includes an expectation over that distribution. To optimize the objective, we need gradients with respect to the parameters that determine the distribution. A direct random draw can obscure how those parameters affect the sampled value.

The Gaussian reparameterization trick separates the randomness from the distribution parameters. It expresses the latent draw as a deterministic, differentiable function of those parameters and a random variable whose distribution does not depend on them. The resulting gradient is called a pathwise gradient: it follows the change in the sampled latent value as the parameters change.

Deriving the Pathwise Gradient

Let \(\theta\) denote parameters that determine \(\mu_\theta\) and \(\sigma_\theta>0\), and let \(f\) be a differentiable function of the latent value. We want to differentiate the expectation \(J(\theta)=\mathbb{E}_{Z\sim N(\mu_\theta,\sigma_\theta^2)}[f(Z)]\). Using the Gaussian representation, the same objective is

The distribution of \(\epsilon\) is fixed: it does not change with \(\theta\). Under conditions that allow differentiation inside the expectation, such as differentiability of the integrand and an integrable bound on its derivative in a neighborhood of the parameter value, the chain rule gives

The derivative inside the expectation has two contributions. Changing \(\theta\) can move the latent distribution's mean, through \(\nabla_\theta\mu_\theta\), and change its scale, through \(\nabla_\theta\sigma_\theta\). The same draw of \(\epsilon\) is held fixed when taking this derivative. Randomness is still present across draws, so a gradient calculated from a finite number of draws is generally noisy.

In practice, an encoder often outputs a mean and a log standard deviation. Write \(\rho_\theta=\log\sigma_\theta\), so \(\sigma_\theta=\exp(\rho_\theta)\). This parameterization ensures a positive scale, and the sampled latent is

For a scalar parameter \(\theta\), the derivative of one sampled value is

This expression makes the gradient route explicit: the objective depends on the sample \(Z\), and the sample depends differentiably on the encoder parameters.

A One-Dimensional Encoder Gradient

Consider a hypothetical cell with count vector \(x=2\), used here as a single numerical input to a simple encoder. Let its mean be \(\mu_w(x)=wx+b\), and let its log standard deviation be \(\rho\). For fixed \(x\), a latent draw is \(z=wx+b+\exp(\rho)\epsilon\). To show the gradient calculation transparently, use the scalar function \(f(z)=-\frac{1}{2}(z-y)^2\), with \(y=1\). This function is a simple differentiable objective, not a count likelihood.

For this example, the expected objective can also be calculated exactly, which checks the meaning of the sampled gradients. Since \(Z\) has mean \(\mu\) and variance \(\sigma^2\),

Here \(\mu=2w+b=0.5\), \(\sigma^2=\exp(2\rho)=\exp(-0.4)\approx0.6703\), and \(y=1\), so the expected objective is approximately \(-0.46\). Its exact derivatives are

The single-draw gradients \(0.3450\) and \(0.0565\) need not equal these exact derivatives. Their role is to provide random estimates; averaging gradients from independent draws targets the exact derivatives.

Extension to Several Gaussian Latent Dimensions

For a diagonal Gaussian with \(d\) latent dimensions, let \(\boldsymbol{\mu}\) be a \(d\)-vector of means and \(\boldsymbol{\sigma}\) a \(d\)-vector of positive standard deviations. Generate independent standard-normal components in \(\boldsymbol{\epsilon}\), then set each latent component to \(z_j=\mu_j+\sigma_j\epsilon_j\). In vector notation, this is \(\mathbf{z}=\boldsymbol{\mu}+\boldsymbol{\sigma}\odot\boldsymbol{\epsilon}\), where \(\odot\) denotes componentwise multiplication. The chain rule applies to each component and to any differentiable scalar objective \(f(\mathbf{z})\).

Using the Trick in an ELBO and Checking Its Assumptions

In the ELBO derived for a count VAE, one term is an expectation over the encoder's latent distribution of a function involving the count likelihood and the prior. When that function is differentiable in the latent variable, reparameterization lets a sampled latent contribute a pathwise gradient to the expectation. A KL term that has a closed-form expression can instead be differentiated directly; reparameterization is needed for an expectation estimated using latent samples, not automatically for every term in the objective.

The basic calculation relies on several conditions and implementation choices:

• The Gaussian scale must be positive. Using \(\sigma=\exp(\rho)\) enforces this mathematically; numerical implementations should also avoid scales that underflow or become excessively large.
• The objective must be differentiable with respect to the sampled latent, or at least differentiable almost everywhere in a way that supports the gradient calculation. Discrete sampling steps or nondifferentiable operations can interrupt the pathwise gradient.
• Differentiating under the expectation must be justified by suitable integrability and smoothness conditions. The formula should not be assumed valid for an arbitrary unbounded or irregular objective.
• In a Monte Carlo implementation, draw \(\epsilon\) from a parameter-independent standard normal distribution and preserve its computational connection through \(z=\mu+\sigma\epsilon\). If the sample is detached from the calculation graph, gradients cannot pass through the draw.

A useful diagnostic is to fix a noise value and compare an automatic-differentiation gradient with a finite-difference calculation that uses that same noise value. For a parameter \(\theta\), the finite difference is \([f(z(\theta+h))-f(z(\theta-h))]/(2h)\), with \(\epsilon\) held constant in both evaluations. Resampling noise for the two evaluations adds random variation to the comparison and can make a correct gradient appear wrong. This diagnostic checks the derivative of a sampled path; it does not by itself establish that a finite-sample average has low Monte Carlo error.

Alternatives, Limitations, and Common Mistakes

A principal alternative is the score-function estimator. For a differentiable density \(q_\theta(z)\), under conditions permitting differentiation under the integral, it uses \(\nabla_\theta\mathbb{E}_{q_\theta}[f(Z)]=\mathbb{E}_{q_\theta}[f(Z)\nabla_\theta\log q_\theta(Z)]\). Unlike the pathwise calculation, this expression does not require differentiating \(f\) with respect to \(z\), so it can apply when the sampled variable or objective is not differentiable in a useful way. Its estimates can have high variance, particularly if \(f(Z)\) varies substantially. The pathwise estimator exploits a differentiable transformation and is often more direct for Gaussian latents, but it also has Monte Carlo variability and depends on a valid differentiable path.

A common mistake is to think that a sampled latent is “random, so it cannot be differentiated.” The correction is that the random input \(\epsilon\) is held fixed for the derivative, while the transformed value \(z=\mu_\theta+\sigma_\theta\epsilon\) changes with \(\theta\). Another error is to confuse variance with standard deviation: if a Gaussian is written \(N(\mu,\sigma^2)\), its sampling scale is \(\sigma\), not \(\sigma^2\). Finally, a pathwise gradient from one noise draw is not the exact gradient of the expectation; it is a stochastic estimate.

The trick does not remove the need to choose a suitable latent distribution, objective, or number of Monte Carlo draws. Nor does a successful gradient calculation show that a latent representation captures biologically meaningful structure. It addresses how gradients pass through continuous random Gaussian latents, not what conclusions should be drawn from the learned latents.