score function

The score function is the gradient of a log-density with respect to its parameters:

It is the direction that we would move the parameters in order to make the current point $x$ more likely. Equivalently, it measures the sensitivity of the likelihood at $x$ to each dimension of the parameter $\theta$. It is used to define Fisher information, to estimate policy gradients, and generally comes up all over the place in statistical analysis.

Note that $\nabla_\theta \log p(x; \theta) = \frac{\nabla_\theta p(x; \theta)}{p(x; \theta)}$ by basic properties of the derivative. Inside an expectation, the $p(x; \theta)$ in the denominator cancels with the $p(x; \theta)$ term of the expectation; the is the fundamental 'trick' of many score function derivations. For example,

Theorem: the score function has expectation 0:

(applying the above 'trick'):

assuming that the dominated convergence theorem holds to allow us to swap the integral with the limit.

Gradients of expectations

The score function comes up as the general-purpose gradient estimator for expectations:

and applying the product rule and the score 'trick' from above

See Mohamed et al. (2019) for much much more on this.

Note that we can replace $f_\theta(x)$ in the above by any constant-shifted expression $f_\theta(x) - \beta$; this follows from the above result that $\beta \nabla_\theta \log p_\theta(x)$ has expectation zero. Here $\beta$ is known as a 'baseline' or control variate.

Surrogate objectives

If we naively sampled $x\sim p_\theta$ and then apply automatic differentiation to the expression inside the expectation $f_\theta(x)$, we would get an incorrect result that ignores the dependence through the sampling distribution $p_\theta$ (see limitations of autodiff). The reparameterization trick is one approach to fixing this. A more general approach is to write the objective as

where $\tilde{\theta} = \mathtt{stop\_gradient}(\theta)$ is a frozen copy of the parameters. This brings the dependence on $\theta$ explicitly into the objective, so that autodiff (even repeated autodiff for higher-order derivatives) does the right thing, while maintaining the objective's value.This trick is described in the DiCE paper, though it predates that paper and has been independently invented multiple times. If desired we can introduce a baseline $\beta$ as

so that the objective itself is again unchanged, but autodiff recovers the gradient estimator including $\beta$ as a control variate (note that the second $\beta$ term may be omitted if we only care about the surrogate objective's gradients rather than its value directly).

Optimal baseline

Let $g_\beta(x)$ denote the the gradient estimator (the quantity inside the expectation) with constant baseline $\beta$,

What choice for $\beta$ minimizes the variance of this estimate? For simplicity, let's assume $\theta$ is scalar, so that gradients become simple derivatives. Then we have

where we note that the second term is just the square of the derivative being estimated, which we can ignore because it does not depend on $\beta$. Expanding out the remaining expression, and sweeping aside similar constant terms,

Setting the derivative with respect to $\beta$ to zero gives the optimality condition

This is a bit unwieldy, but it simplifies dramatically if we allow ourselves to assume that $f_\theta(x)$ and $p_\theta(x)$ are independent. Then we can factor the expectations, so that the term

vanishes, using the property that the score function has expectation zero, and we get cancellation in the remaining terms

giving the intuitive result that the minimum-variance baseline $\beta^*$ is just the value of the expectation whose gradient we're trying to estimate. Note that this derivation relies on the assumption that $p_\theta(x)$ and $f_\theta(x)$ are independent, which may not hold in practice.For example, in reinforcement learning the probability of being in state $x$ is hopefully not independent of the reward of state $x$, since the task is explicitly to find policies that spend more time in high-reward states!

Unnormalized densities

What if we have $p_\theta(x)$ available only in unnormalized form, as

where $Z_\theta = \int_\Omega r_\theta(x) dx$? Then it seems like we require the gradient of the log normalizer $\log Z_\theta$. That's not a big deal, because this can itself be estimated from samples:

leaving us with

One way to understand this is that the score function itself must have mean zero because the distribution is normalized. So if we have an unnormalized density, we just need to subtract the mean of its proto-score-function in order to get the mean-zero score function.

I suppose that for an unbiased estimate, we'd need to estimate the inner expectation with a different independent sample $x\sim p$ than we used for the outer expectation. This is essentially contrastive divergence, except that we probably don't need MCMC because we can sample explicitly from $p$.