Chapter 12Differentiating through integration

In this chapter, we study how to differentiate through integrals, with a focus on expectations and solutions of ordinary differential equations.

12.1 Differentiation under the integral sign

Given two Euclidean spaces $\Theta$ and $\cY$, and a function $f: \Theta \times\cY \rightarrow \RR$, we often want to differentiate an integral of the form $$F(\thetav) \coloneqq \int_\cY f(\thetav, \y)d\y.$$ Provided that we can swap integration and differentiation, we have $$\nabla F(\thetav) = \int_\cY \nabla_\thetav f(\thetav, \y)d\y.$$ The conditions enabling us to do so are best examined in the context of measure theory. We refer the reader to e.g. (Cohn, 2013) for a course on measure theory and Flanders (1973) for an in-depth study of the differentiation under the integral sign. Briefly, if $\Theta = \cY = \RR$, the following conditions are sufficient.

1. $f$ is measurable in both its arguments, and $f(\theta, \cdot)$ is integrable for almost all $\theta \in \Theta$ fixed,

2. $f(\cdot, y)$ is absolutely continuous for almost all $y \in \cY$, that is, there exists an integrable function $g(\cdot, y)$ such that $f(\theta, y) = f(\theta_0, y) + \int_{\theta_0}^\theta g(\tau, y)d\tau$,

3. $\partial_1 f(\theta, y)$ (which exists almost everywhere if $f(\cdot, y)$ is absolutely continuous), is locally integrable, that is, for any closed interval $[\theta_0, \theta_1]$, the integral $\int_{\theta_0}^{\theta_1} \int |\partial_1 f(\theta, y)| dy d\theta$ is finite.

Any differentiable function $f:\Theta \times \cY \rightarrow \RR$ is absolutely continuous. However, the conditions also hold if $f$ is just absolutely continuous, that is, if $f(\cdot, y)$ is differentiable for almost all $y$. This weaker assumption can be used to smooth out differentiable almost-everywhere functions, such as the ReLU, as we study in Section 14.4.

12.2 Differentiating through expectations

A special case of differentiating through integrals is differentiating through expectations. We can distinguish between two cases, depending on whether the parameters $\thetav$ we wish to differentiate are involved in the distribution or in the function, whose expectation we compute.

12.2.1 Parameter-independent distributions

We first consider expectations of the form $$F(\thetav) \coloneqq \EE_{Y \sim p}[g(Y, \thetav)] = \int_{\cY} g(\y, \thetav) p(\y) d\y,$$ for a random variable $Y \in \cY \subseteq \RR^M$, distributed according to a distribution $p$, and a function $g \colon \cY \times \Theta \to \R$. Importantly, the distribution is independent of the parameters $\thetav$. Under mild conditions recalled in Section 12.1, we can swap differentiation and integration to obtain $$\begin{align*} \nabla F(\thetav) &= \nabla_\thetav \int_{\cY} g(\y, \thetav) p(\y) d\y \\ &= \int_{\cY} \nabla_\thetav g(\y, \thetav) p(\y) d\y \\ &= \EE_{Y \sim p}[\nabla_\thetav g(Y, \thetav)]. \end{align*}$$ Generally, the expectation cannot be computed in closed form. However, provided that we can sample from $p$, we can define a Monte-Carlo estimator of the value $$\widehat{F}_N(\thetav) \coloneqq \frac{1}{N}\sum_{i=1}^{N} g(Y_i, \thetav)$$ and of the gradient

$$\nabla \widehat{F}_N(\thetav) = \frac{1}{N}\sum_{i=1}^{N} \nabla_\thetav g(Y_i, \thetav),$$ for $N$ i.i.d. samples $Y_1, \ldots, Y_N$ from $p$. These estimators are unbiased, meaning that $\EE[\widehat{F}_N(\thetav)] = F(\thetav)$ and $\EE[\nabla \widehat{F}_N(\thetav)] = \nabla F(\thetav)$, and they converge to the true quantity as $N \rightarrow +\infty$. This suggests a simple implementation in an autodiff framework of the approximation of $\nabla F(\thetav)$:

1. Sample $y_1, \ldots, y_N$ from $p$.

2. Compute $\widehat{F}_N(\thetav) = \frac{1}{N}\sum_{i=1}^{N} g(y_i, \thetav)$.

3. Compute the gradient $\nabla \widehat{F}_N(\thetav)$ by automatic differentiation.

Computing higher-order derivatives follows the same principle: to get an approximation of $\nabla^2 F(\thetav)$, we can simply compute $\nabla^2 \widehat{F}_N(\thetav)$ by autodiff. As such, the implementation delineated above is akin to the “discretize-then-optimize” approach used to differentiate through the solution of an ODE (Section 12.6): we implement an approximation of the objective and simply call autodiff on it.

12.2.2 Parameter-dependent distributions

A more challenging case arises when the distribution depends on the parameters $\thetav$: $$E(\thetav) \coloneqq \EE_{Y \sim p_\thetav}[g(Y)] = \int_{\cY} g(\y) p_\thetav(\y) d\y,$$ where $Y \in \cY \subseteq \RR^M$ is a random variable, distributed according to a distribution $p_\thetav$ parameterized by $\thetav \in \Theta$ and where $g \colon \cY \to \R$ is, depending on the setting, potentially a blackbox function (i.e., we do not have access to its gradients). Typically, $\thetav \in \Theta$ could be parameters we wish to estimate, or it could indirectly be generated by $\thetav = f(\x, \w) \in \Theta$, where $f$ is a neural network with parameters $\w \in \cW$ we wish to estimate. The main difficulty in computing $\nabla E(\thetav)$ stems from the fact that $\thetav$ are the parameters of the distribution $p_\thetav$. Estimating an expectation $E(\thetav) = \EE_{Y \sim p_\thetav}[g(Y)]$ using Monte-Carlo estimation requires us to sample from $p_\thetav$. However, it is not clear how to differentiate $E$ w.r.t. $\thetav$ if $\thetav$ is involved in the sampling process.

Continuous case

When $\cY$ is a continuous set (that is, $p_\thetav(\y)$ is a probability density function), we can rewrite $E(\thetav)$ as $$E(\thetav) = \int_{\cY} p_\thetav(\y) g(\y) d\y.$$ Provided that we can swap integration and differentiation (see Section 12.1), we then have $$\begin{align*} \nabla E(\thetav) &= \nabla_\thetav \int_{\cY} p_\thetav(\y) g(\y) d\y \\ &= \int_{\cY} \nabla_\thetav p_\thetav(\y) g(\y) d\y. \end{align*}$$ Unfortunately, this integral is not an expectation and it could be intractable in general.

Discrete case

When $\cY$ is a discrete set (that is, $p_\thetav(\y)$ is a probability mass function), we can rewrite $E(\thetav)$ as $$E(\thetav) = \sum_{\y \in \cY} p_\thetav(\y) g(\y).$$ We then obtain $$\nabla E(\thetav) = \sum_{\y \in \cY} g(\y) \nabla_{\thetav} p_\thetav(\y).$$ Again, $\nabla E(\thetav)$ is not an expectation. We therefore cannot use Monte-Carlo estimation to estimate the gradient. Instead, we can compute it by brute force, i.e., by summing over all possible $\y \in \cY$. However, this is clearly only computationally tractable if $|\cY|$ is small or if $p_\thetav$ is designed to have sparse support, i.e., so that the set $\{\y \in \cY \colon p_\thetav(\y) \neq 0\}$ is small. Moreover, even if these conditions hold, summing over $\y$ could be problematic if $g(\y)$ is expensive to compute. Therefore, exact gradients are seldom used in practice.

In Sections 12.3 and 12.4, we review the score function and pathwise gradient estimators, to (approximately) compute $\nabla E(\thetav)$, allowing us to optimize $\thetav$ (or $\w$ using the chain rule) by gradient-based algorithms.

12.2.3 Application to expected loss functions

Differentiating through expectations is particularly useful when working with expected loss functions of the form $$L(\thetav; \y) \coloneqq \EE_{\hat Y \sim p_\thetav}[\ell(\hat Y, \y)],$$ where $\y$ is some ground truth. Equivalently, we can set $\ell = -r$, where $r$ is a reward function. As we shall see, the score function estimator will support a discrete loss function $\ell \colon \cY \times \cY \to \R$, while the pathwise gradient estimator will require a differentiable loss function $\ell \colon \RR^M \times \cY \to \R$. Intuitively, $L(\thetav; \y)$ will be low if $p_\thetav$ assigns high probability to predictions $\widehat \y$ with low loss value $\ell(\widehat \y, \y)$.

In the classification setting, where $\cY = [M]$, $p_\thetav$ is often chosen to be the Gibbs distribution, which is a categorical distribution induced by a softargmax $$p_\thetav(y) \coloneqq \frac{\exp(\theta_{y})}{\sum_{i \in [M]} \exp(\theta_i)} = \left[\mathrm{softmax}(\thetav)\right]_{y} \in (0, 1),$$ where $\theta_{y} \coloneqq f(\x, y, \w) \in \RR$ are logits produced by a neural network $f$. More generally, in the structured prediction setting, where $\cY \subseteq \R^M$ but $|\cY| \gg M$, we often use the distribution $$p_\thetav(\y) \coloneqq \frac{\exp(\langle \phi(\y), \thetav \rangle)}{\sum_{\y' \in \cY} \exp(\langle \phi(\y'), \thetav \rangle)},$$ where $\thetav = f(\x, \w) \in \RR^M$.

Given a distribution $\rho$ over $\cX \times \cY$, we then want to minimize the expected loss function, also known as risk, $$R(\w) \coloneqq \EE_{(X,Y) \sim \rho}[L(f(X, \w); Y)].$$ Typically, minimizing $R(\w)$ is done through some form of gradient descent, which requires us to be able to compute $$\begin{aligned} \nabla R(\w) &= \EE_{(X, Y) \sim \rho}[\nabla_\w L(f(X, \w); Y)] \\ &= \EE_{(X, Y) \sim \rho}[\partial_2 f(X, \w)^* \nabla L(f(X, \w); Y)]. \end{aligned}$$ Computing $\nabla R(\w)$ therefore boils down to computing the gradient of $L(\thetav; \y)$, which is the gradient of an expectation.

12.2.4 Application to experimental design

In experimental design, we wish to minimize a function $g(\lambdav)$, which we assume costly to evaluate. As an example, evaluating $g(\lambdav)$ could require us to run a scientific experiment with parameters $\lambdav \in \RR^Q$. As another example, in hyperparameter optimization, evaluating $g(\lambdav)$ would require us to run a learning algorithm with hyperparameters $\lambdav \in \RR^Q$. Instead of solving the problem $\argmin_{\lambdav \in \RR^Q} g(\lambdav)$, we can lift the problem to probability distributions and solve $\argmin_{\thetav \in \RR^M} E(\thetav)$, where $E(\thetav) = \EE_{\lambdav \sim p_\thetav}[g(\lambdav)]$. This requires the probability distribution $p_\thetav$ to assign high probability to $\lambdav$ values that achieve small $g(\lambdav)$ value. Solving this problem by stochastic gradient descent requires us to be able to compute estimates of $\nabla E(\thetav)$. This can be done for instance with SFE explained in Section 12.3, which does not require gradients of $g$, unlike implicit differentiation explained in Chapter 11. This approach also requires us to choose a distribution $p_\thetav$ over $\lambdav$. For continuous hyperparameters, a natural choice would be the normal distribution $\lambdav \sim \mathrm{Normal}(\muv, \Sigmav)$, setting $\thetav = (\muv, \Sigmav)$. Once we have obtained $\thetav$ by minimizing $E(\thetav)$, we need a way to recover $\lambdav$. This can be done for example by choosing the mode of the distribution, i.e., $\argmax_{\lambdav \in \RR^Q} p_\thetav(\lambdav)$, or the mean of the distribution $\EE_{\lambdav \sim p_\thetav(\lambdav)}[\lambdav]$. Of course, in the case of the normal distribution, they coincide.

12.3 Score function estimators, REINFORCE

12.3.1 Scalar-valued functions

The key idea of the score function estimator (SFE), also known as REINFORCE, is to rewrite $\nabla E(\thetav)$ as an expectation. The estimator is based on the logarithmic derivative identity $$\nabla_\thetav \log p_\thetav(\y) = \frac{\nabla_\thetav p_\thetav(\y)}{p_\thetav(\y)} \Longleftrightarrow \nabla_\thetav p_\thetav(\y) = p_\thetav(\y) \nabla_\thetav \log p_\thetav(\y).$$ Using this identity, we obtain the following gradient estimator.

The gradient of the log-PDF w.r.t. $\thetav$, $\nabla_\thetav \log p_\thetav(\y)$, is known as the score function, hence the estimator name. SFE is suitable when two requirements are met: it is easy to sample from $p_\thetav$ and the score function is available in closed form. Since the SFE gradient is an expectation, we can use Monte-Carlo estimation to compute an unbiased estimator of $\nabla E(\thetav)$: $$\nabla E(\thetav) \approx \widehat{\gammav}_N(\thetav) \coloneqq \frac{1}{N} \sum_{i=1}^N g(Y_i) \nabla_\thetav \log p_\thetav(Y_i), \tag{12.1}$$ where $Y_1, \dots, Y_N$ are sampled from $p_\thetav$.

Interestingly, the gradient of $g$ is not needed in this estimator. Therefore, there is no differentiability assumption about $g$. This is why SFE is useful when $g$ is a discrete loss function or more generally a blackbox function.

Another example when $\nabla_\thetav p_\thetav(\y)$ is available in closed form is in the context of reinforcement learning, where $p_\thetav(\y)$ is a Markov Decision Process (MDP) and is called the policy. Applying the SFE leads to the (vanilla) policy gradient method (Sutton et al., 1999) and can then be used to compute the gradient of an expected cumulative reward. However, SFE is more problematic when used with the Gibbs distribution, due to the explicit normalization constant.

As another example, we apply SFE in Section 14.4 to derive the gradient of perturbed functions.

Differentiating through both the distribution and the function

Suppose both the distribution and the function now depend on $\thetav$. When $g$ is scalar-valued and differentiable w.r.t. $\thetav$, we want to differentiate $$E(\thetav) \coloneqq \EE_{Y \sim p_\thetav} [g(Y, \thetav)].$$ Using the product rule, we obtain $$\nabla E(\thetav) = \EE_{Y \sim p_\thetav} [g(Y, \thetav) \nabla_\thetav \log p_\thetav(Y)] + \EE_{Y \sim p_\thetav} [\nabla_\thetav g(Y, \thetav)].$$

Differentiating through joint distributions

Suppose we now want to differentiate through $$E(\thetav) \coloneqq \EE_{Y_1 \sim p_\thetav, Y_2 \sim q_\thetav}[g(Y_1,Y_2)].$$ The gradient is then given by $$\nabla E(\thetav) = \EE_{Y_1 \sim p_\thetav, Y_2 \sim q_\thetav}[(\nabla_\thetav \log p_\thetav(Y_1) + \nabla_\thetav \log q_\thetav(Y_2)) g(Y_1,Y_2)],$$ which is easily seen by applying Proposition 12.1 on the joint distribution $\rho_\thetav \coloneqq p_\thetav \cdot q_\thetav$. The extension to more than two variables is straightforward.

12.3.2 Variance reduction

We discuss in this section various techniques for reducing the variance of estimators of $\nabla E(\thetav)$, where $E(\thetav) \coloneqq \EE_{Y \sim p_\thetav}[g(Y)]$.

Bias and variance

Recall the definition of $\widehat{\gammav}_N$ in Eq. (12.1). SFE is an unbiased estimator, meaning that $$\nabla E(\thetav) = \EE[\widehat{\gammav}_N(\thetav)],$$ where the expectation is taken with respect to the $N$ samples drawn. Since the gradient is vector-valued, we need to define a scalar-valued notion of variance. We do so by using the squared Euclidean distance in the usual variance definition to define $$\begin{align*} \VV[\widehat{\gammav}_N(\thetav)] &\coloneqq \EE [\|\widehat{\gammav}_N(\thetav) - \nabla E(\thetav)\|_2^2] \\ &= \EE[\|\widehat{\gammav}_N(\thetav)\|_2^2] - \|\nabla E(\thetav)\|^2_2. \end{align*}$$ The variance naturally goes to zero as $N \to \infty$.

Baseline

SFE is known to suffer from high variance (Mohamed et al., 2020). This means that this estimator may require us to draw many samples from the distribution $p_\thetav$ to work well in practice. One of the simplest variance reduction technique consists in shifting the function $g$ with a constant $\beta$, called a baseline, to obtain $$\nabla E(\thetav) = \EE_{Y \sim p_\thetav}[(g(Y) - \beta) \nabla_\thetav \log p_\thetav(Y)].$$ The reason this is still a valid estimator of $\nabla E(\thetav)$ stems from $$\begin{align*} \EE_{Y \sim p_\thetav}[ \nabla_\thetav \log p_\thetav(Y) ] &= \EE_{Y \sim p_\thetav}\left[ \frac{\nabla_\thetav p_\thetav(Y)}{p_\thetav(Y)} \right] \\ &= \nabla_\thetav \EE_{Y \sim p_\thetav}[\ones] \\ &= \nabla_\thetav \ones \\ &= \zeros, \end{align*}$$ for any valid distribution $p_\thetav$. The baseline $\beta$ is often set to the running average of past values of the function $g$, though it is neither optimal nor does it guarantee to lower the variance (Mohamed et al., 2020).

REINFORCE leave-one-out (RLOO)

A simple way to reduce the variance is REINFORCE leave-one-out (RLOO) (Ahmadian et al., 2024; Kool et al., 2019). The key idea is to draw multiple independent samples $Y_1, \dots, Y_N$ from $p_\thetav$ and for each sample $Y_i$, to use the remaining samples $(Y_j)_{j \neq i}$ as a baseline. The RLOO gradient estimator is defined as $$\hat{\Gamma}_N(\thetav) \coloneqq \frac{1}{N} \sum_{i=1}^N \big(g(Y_i) - \frac{1}{N-1} \sum_{j \neq i} g(Y_j)\big) \nabla_\thetav \log p_\thetav(Y_i).$$ This is an unbiased estimator, meaning that $$\nabla E(\thetav) = \EE[\hat{\Gamma}_N(\thetav)],$$ where the expectation is taken over $Y_1, \dots, Y_N$. This comes from the independence of $Y_i, Y_j$ for $i\neq j$, such that $\EE[g(Y_j) \nabla_\thetav\log p_\thetav(Y_i)] = \EE[g(Y_j)]\EE[\nabla_\thetav\log p_\thetav(Y_i)]= 0$ since $\EE[\nabla_\thetav\log p_\thetav(Y_i)]= 0$ by Bartlett’s first identity (12.2).

It can be checked that the estimator can equivalently be rewritten as $$\hat{\Gamma}_N(\thetav) = \frac{1}{N-1} \sum_{i=1}^N \big(g(Y_i) - \frac{1}{N} \sum_{j=1}^N g(Y_j)\big) \nabla_\thetav \log p_\thetav(Y_i).$$ This is slightly more convenient to implement, as we can use $\beta \coloneqq \frac{1}{N} \sum_{j=1}^N g(Y_j)$ as a shared baseline across samples.

Control variates

Another general technique are control variates. Let us denote the expectation of a function $h \colon \RR^M \to \RR$ under the distribution $p_\thetav$ as $$H(\thetav) \coloneqq \EE_{Y \sim p_\thetav}[h(Y)].$$ Suppose that $H(\thetav)$ and its gradient $\nabla H(\thetav)$ are known in closed form. Then, for any $\gamma \ge 0$, we clearly have $$\begin{align*} E(\thetav) &= \EE_{Y \sim p_\thetav}[g(Y)] \\ &= \EE_{Y \sim p_\thetav}[g(Y) - \gamma (h(Y) - H(\thetav))] \\ &= \EE_{Y \sim p_\thetav}[g(Y) - \gamma h(Y)] + \gamma H(\thetav) \end{align*}$$ and therefore $$\begin{align*} \nabla E(\thetav) &= \nabla_\thetav \EE_{Y \sim p_\thetav}[g(Y) - \gamma h(Y)] + \gamma \nabla H(\thetav). \end{align*}$$ Applying SFE, we then obtain $$\nabla E(\thetav) = \EE_{Y \sim p_\thetav}[(g(Y) - \gamma h(Y)) \nabla_\thetav \log p_\thetav(Y)] + \gamma \nabla H(\thetav).$$ Examples of function $h$ include a bound on $f$ or a second-order Taylor expansion of $f$, assuming that these approximations are easier to integrate than $f$ (Mohamed et al., 2020).

12.3.3 Vector-valued functions

It is straightforward to extend the SFE to vector-valued functions.

Differentiating through both the distribution and the function

If $\thetav$ now influences both the distribution and the function, $$E(\thetav) \coloneqq \EE_{Y \sim p_\thetav} [g(Y, \thetav)],$$ then, we obtain $$\partial E(\thetav) = \EE_{Y \sim p_\thetav} [g(Y, \thetav) \otimes \nabla_\thetav \log p_\thetav(Y)] + \EE_{Y \sim p_\thetav} [\partial_\thetav g(Y, \thetav)].$$

12.3.4 Second derivatives

Using the previous subsection with $g(\y, \thetav) = g(\y) \nabla_\thetav \log p_\thetav(\thetav)$, we easily obtain an estimator of the Hessian.

This can also be derived using the second-order log-derivative $$\nabla^2_\thetav \log p_\thetav(\y) = \frac{1}{p_\thetav(\y)} \nabla^2_\thetav p_\thetav(\y) - \frac{1}{p_\thetav(\y)^2} \nabla_\thetav p_\thetav(\y) \otimes \nabla_\thetav p_\thetav(\y)$$ so that $$\nabla^2_\thetav p_\thetav(\y) = p_\thetav(\y) \left[ \nabla^2_\thetav \log p_\thetav(\y) + \nabla_\thetav \log p_\thetav(\y) \otimes \nabla_\thetav \log p_\thetav(\y) \right].$$

Link with the Bartlett identities

The Bartlett identities are expressions relating the moments of the score function (gradient of the log-likelihood function). Using Proposition 12.1 with $g(\y) = 1$ and $\int_\cY p_\thetav(\y) d\y = 1$, we obtain $$\EE_{Y \sim p_\thetav}[\nabla_\thetav \log p_\thetav(Y)] = \zeros, \tag{12.2}$$ which is known as Bartlett’s first identity. Similarly, using Proposition 12.3, we obtain $$\begin{aligned} &\EE_{Y \sim p_\thetav} [\nabla^2_\thetav \log p_\thetav(Y)] + \EE_{Y \sim p_\thetav} [\nabla_\thetav \log p_\thetav(Y) \otimes \nabla_\thetav \log p_\thetav(Y)] \\ =& \EE_{Y \sim p_\thetav} [\nabla^2_\thetav \log p_\thetav(Y)] + \mathrm{cov}[\log p_\thetav(Y)] \\ =& \zeros, \end{aligned} \tag{12.3}$$ which is known as Bartlett’s second identity.

12.3.5 Autodiff-friendly implementation

So far, we derived expressions of score function estimators for first and second derivatives. From a user perspective, instead of manually implementing these expressions, it would be more convenient if we could just call autodiff on some function and recover the correct expressions of the score function estimators. This can be achieved using the so-called magicbox operator (Foerster et al., 2018). This operator is defined as $$\cM(\tau) \coloneqq \exp(\tau - \cS(\tau)),$$ where $\cS$ is the stop gradient operator.

During the evaluation of a program, the stop gradient operator acts as the identity. During automatic differentiation, that is, during forward or reverse mode autodiff, this operator stops any flow of derivatives. In other words, its Jacobian is zero on any inputs despite evaluating as the identity. Such an operator does not exist mathematically, but it can be defined in an autodiff framework using custom derivation rules. Evaluations and derivatives of programs making $\cS$ intervene are then to be understood first in terms of the autodiff framework and its custom rules.

Coming back to the magicbox operator, its evaluation and its derivatives in an autodiff framework are then for any $\tau \in \cR$ $$\begin{align*} \cM(\tau) & = \exp(\tau - \tau) = 1\\ \cM'(\tau) & = (\tau - \cS(\tau))' \cM(\tau) = 1 \cdot 1 = 1\\ \cM''(\tau) & = \dots = 1 \end{align*}$$ The “magic” happens when we compose the magicbox operator with another function. Consider a (twice) differentiable function $f \colon \RR^M \to \RR$. We then have $$\begin{align*} (\cM \circ f)(\thetav) &= 1 \\ \nabla (\cM \circ f)(\thetav) &= \cM'(f(\thetav))\nabla f(\thetav) = \nabla f(\thetav) \\ \nabla^2 (\cM \circ f)(\thetav) &= \cM''(f(\thetav)) \nabla f(\thetav) \otimes \nabla f(\thetav) + \cM'(f(\thetav)) \nabla^2 f(\thetav) \\ &= \nabla f(\thetav) \otimes \nabla f(\thetav) + \nabla^2 f(\thetav). \end{align*}$$ Let $Y$ be a sample from $p_\thetav$ and let us define $f(\thetav) \coloneqq \log p_\thetav(Y)$. Then, $$\begin{align*} (\cM \circ f)(\thetav) &= 1 \\ \nabla (\cM \circ f)(\thetav) &= \nabla_\thetav \log p_\thetav(Y) \\ \nabla^2 (\cM \circ f)(\thetav) &= \nabla_\thetav \log p_\thetav(Y) \otimes \nabla_\thetav \log p_\thetav(Y) + \nabla^2_\thetav \log p_\thetav(Y), \end{align*}$$ correctly recovering the expressions for the first derivative in Proposition 12.1 and for the second derivative in Proposition 12.3. Let us define the functions $$\begin{align*} E(\thetav) &\coloneqq \EE_{Y \sim p_\thetav}[g(Y)] \\ \bar{E}(\thetav) &\coloneqq \EE_{Y \sim \cS(p_\thetav)}[g(Y) \cM(\log p_\thetav(Y))]. \end{align*}$$ We then correctly obtain for all $\thetav \in \RR^M$ $$\begin{align*} E(\thetav) &= \bar{E}(\thetav) \\ \nabla E(\thetav) &= \nabla \bar{E}(\thetav) \\ \nabla^2 E(\thetav) &= \nabla^2 \bar{E}(\thetav). \end{align*}$$ The key benefit is that $\bar{E}$ is compatible with autodiff frameworks, while $E$ is not. In practice, we can draw $N$ samples $Y_1, \dots, Y_N$ (without differentiating through the sampling) and define $$\bar{E}_N(\thetav) \coloneqq \frac{1}{N} \sum_{i=1}^N g(Y_i) \cM(\log p_\thetav(Y_i)).$$ By Monte-Carlo estimation, we then have $$\begin{align*} E(\thetav) &= \EE[\bar{E}_N(\thetav)] \\ \nabla E(\thetav) &= \EE[\nabla \bar{E}_N(\thetav)] \\ \nabla^2 E(\thetav) &= \EE[\nabla^2 \bar{E}_N(\thetav)], \end{align*}$$ where the expectation is taken over $Y_1, \dots, Y_N$.

12.4 Path gradient estimators, reparametrization trick

As we saw previously, the main difficulty in computing gradients of expectations arises when the parameters $\thetav$ play a role in the distribution $p_\thetav$ being sampled. The key idea of path gradient estimators (PGE), also known as reparametrization trick, is to rewrite the expectation in such a way that the parameters are moved from the distribution to the function, using a change of variable.

12.4.1 Location-scale transforms

The canonical example of path gradient estimator is differentiating through the expectation $$E(\mu, \sigma) \coloneqq \EE_{U \sim \mathrm{Normal}(\mu, \sigma^2)}[g(U)],$$ where $g \colon \RR \to \RR$ is a differentiable function. If we let $Z \sim \mathrm{Normal}(0, 1)$, it is easy to check that $U = \mu + \sigma Z$. We can therefore write $$E(\mu, \sigma) = \EE_{Z \sim \mathrm{Normal}(0, 1)}[g(\mu + \sigma Z)].$$ The key advantage is that we can now easily compute the derivatives by mere application of the chain rule, since the parameters $\mu$ and $\sigma$ are moved from the distribution to the function: $$\begin{align*} \frac{\partial}{\partial \mu} E(\mu, \sigma) &= \EE_{Z \sim \mathrm{Normal}(0, 1)}[g'(\mu + \sigma Z)] \\ \frac{\partial}{\partial \sigma} E(\mu, \sigma) &= \EE_{Z \sim \mathrm{Normal}(0, 1)}[Z \cdot g'(\mu + \sigma Z)]. \end{align*}$$ The change of variable $$U \coloneqq \mu + \sigma Z \tag{12.4}$$ is called a location-scale transform. Such a transformation exists, not only for the normal distribution, but for location-scale family distributions, i.e., distributions parametrized by a location parameter $\mu$ and a scale parameter $\sigma > 0$, such that $U$ is distributed according to a distribution in the same family as $Z$ is distributed. Besides the normal distribution, examples of location-scale family distributions include the Cauchy distribution, the uniform distribution, the logistic distribution, the Laplace distribution, and Student’s $t$-distribution.

We can easily relate the cumulative distribution function (CDF) and the probability density function (PDF) of $Z$ to that of $U$, and vice-versa.

12.4.2 Differentiable transforms

We can generalize the idea of path gradient estimator (PGE) to any change of variable $$U \coloneqq T(Z, \thetav),$$ where $T \colon \RR^M \times \RR^Q \to \RR^M$ is a differentiable transformation. For example, if we gather $\mu$ and $\sigma$ as $\thetav \coloneqq (\mu, \sigma)$, we can write the location-scale transform as $$U = T(Z, \thetav) = \mu + \sigma Z.$$ We can derive the path gradient estimator for any such differentiable transformation $T$.

The path gradient estimator (a.k.a. reparametrization trick) gives an unbiased estimator of $\nabla E(\thetav)$. It has however two key disadvantages. First, it assumes that $g$ is differentiable (almost everywhere), which may not always be the case. Second, it assumes that $g$ is well-defined on $\R^M$, not on $\cU$, which could be problematic for some discrete loss functions, such as the zero-one loss function or ranking loss functions.

As an example of differentiable transform, in machine learning, we can sample Gaussian noise $Z$ and make it go through a neural network with parameters $\w$ to generate an image $X \coloneqq T(Z, \w)$. In statistics, many distributions are related to each other through differentiable transforms, as we recall below.

12.4.3 Inverse transforms

The inverse transform method can be used for sampling from a probability distribution, given access to its associated quantile function. Recall that the cumulative distribution function (CDF) associated with a random variable $Y$ is the function $F_Y \colon \RR \to [0,1]$ defined by $$F_Y(y) \coloneqq \PP(Y \le y).$$ The quantile function is then a function $Q_Y \colon [0,1] \to \RR$ such that $Q_Y(\pi) = y$ for $\pi = F_Y(y)$. Assuming $F_Y$ is continuous and strictly increasing, we have that $Q_Y$ is the inverse CDF, $$Q_Y(\pi) = F^{-1}_Y(\pi).$$ In the general case of CDF functions that are not strictly increasing, the quantile function is usually defined as $$Q_Y(\pi) \coloneqq \inf \{y \in \RR \colon \pi \le F_Y(y) \}.$$ Given access to the quantile function $Q_Y(\pi)$ associated with a distribution $p$, inverse transform sampling allows us to sample from $p$ by first drawing a sample from the uniform distribution and then making this sample go through the quantile function.