Gradient of draws from random variables

Question

In the context of neural networks there has recently been some interest in differentiation incarnations of random variables.

Example. Given a random variable $y \sim \mathcal{N}(\mu, \sigma^2)$. Now let $\hat{y}$ be a sample from $y$, how can we calculate ${\partial \hat{y} / \partial \mu}$ and ${\partial \hat{y} / \partial \sigma}$?

One way to do this is to use a helper distribution $\epsilon \sim \mathcal{N}(0, 1)$ and reparametrize $y' = \mu + \sigma \epsilon$. We can then use the simple rules of calculus to find ${\partial \hat{y} / \partial \mu}$ and ${\partial \hat{y} / \partial \sigma}$ because all stochasticity is now within $\epsilon$.

This can also be done with many more distributions.

Does anyone know how this trick is called or who used it first? Or was it unknown just until recently?

Answer 1

To ease the exposition, let us stick to families $\{P_\theta\,;\,\theta\in\Theta\}$ indexed by a single parameter $\theta$.

To compute $\frac{\mathrm d}{\mathrm d\theta}\hat y$ at a given $\theta$, it seems one should be able to build a sample $(y^t_k)$ of distribution $P_t$ from a sample $(y^\theta_k)$ of distribution $P_{\theta}$, at least for every $t$ in a neighborhood of $\theta$.

The gaussian case $\mathcal N(\mu,\sigma^2)$ with $\theta=\mu$ fits in this framework thanks to the identity $$y^m_k=y^\mu_k+m-\mu$$ for every $m$, which yields $$\hat y^m=\hat y^\mu+m-\mu,$$ and finally $$\frac{\mathrm d}{\mathrm d\mu}\hat y=1,$$ for every $\mu$. Likewise, for $\theta=\sigma$, one can make use of the realizations $$y^\sigma_k=\mu+\sigma\epsilon_k$$ for every $\sigma$, with $\epsilon_k$ i.i.d. standard normal. This yields $$y^s_k=\mu+s\sigma^{-1}(y^\sigma_k-\mu)$$ for every $s$, hence $$\hat y^s=\mu+s\sigma^{-1}(\hat y^\sigma-\mu)$$ and finally, $$\frac{\mathrm d}{\mathrm d\sigma}\hat y=\sigma^{-1}(\hat y-\mu).$$ Each time, the crucial tool seems to be the simultaneous realization on some common probability space of a sample $(y^\theta_k)$ for every value of the parameter $\theta$, using the same random input.

What comes to mind in this context is the notion of coupling of random variables or stochastic processes.

To be more specific, if such a coupling yields that, for every $\theta$, $$ \hat y=G(\theta,\epsilon), $$ where $\epsilon$ is some random variable or some stochastic process which does not depend on $\theta$, then $$ \frac{\mathrm d}{\mathrm d\theta}\hat y=\frac{\partial G}{\partial\theta}(\theta,\epsilon). $$