###### tags: `unnormalised models` `one-offs`

# Bregman Divergences for Generalised Score Matching

**Overview**: The following exposition expands on a [tweet of mine](https://twitter.com/sam_power_825/status/1364276169709940744) which proposes some generalisations of the score matching method for parameter estimation. There are connections to the work of [Sugiyama-Suzuki-Kanamori](https://link.springer.com/article/10.1007/s10463-011-0343-8) and also that of [Gutmann-Hirayama](https://arxiv.org/abs/1202.3727).

## Introduction

For parameter estimation, a challenging class of models are so-called *unnormalised likelihood models*. These are models for which the likelihood function is only known up to an intractable normalising constant, i.e. we can write

\begin{align}
P(x|\theta) = f(x,\theta)/Z(\theta)
\end{align}

where $f$ can be evaluated easily, but $Z$ cannot. As such, traditional estimation methods may not easily apply.

# Towards Score Matching

One way of circumventing this difficulty is to replace maximisation of the likelihood by solving an alternative optimisation problem (or root-finding problem), such that in the large-data limit, the optimal value of $\theta$ is the same as the true parameter.

One approach is to consider the following objective

\begin{align}
J(\theta) = \mathbf{E}_q \left[ \frac{1}{2} \left| \nabla_x \log q(x) - \nabla_x \log P(x | \theta) \right|^2 \right],
\end{align}

where $q$ is the law of the data. The presence of the $\nabla_x$ operators makes it clear that the objective in question does not involve $Z(\theta)$. However, it is not yet clear how to estimate $J$ from samples, as $\nabla_x \log q (x)$ is not available in closed form.

The solution is to carry out integration-by-parts on the integral which defines $J$, such that the dependence on $\nabla_x \log q (x)$ is circumvented. In particular, first write

\begin{align}
J(\theta) &= \mathbf{E}_q \left[ \frac{1}{2} \left| \nabla_x \log q(x) - \nabla_x \log P(x | \theta) \right|^2 \right] \\
&= \mathbf{E}_q \left[ \frac{1}{2} \left| \nabla_x \log q(x) \right|^2 \right] - \mathbf{E}_q \left[ \langle \nabla_x \log q(x), \nabla_x \log P(x | \theta) \rangle \right] + \mathbf{E}_q \left[ \frac{1}{2} \left| \nabla_x \log P(x | \theta) \right|^2 \right].
\end{align}

The first term is independent of $\theta$, and so can be neglected. The third term can be estimated directly, as we have samples from $q$, and can evaluate $\left| \nabla_x \log P(x | \theta) \right|^2$ in closed form. The middle term requires more work:

\begin{align}
\mathbf{E}_q \left[ \langle \nabla_x \log q(x), \nabla_x \log P(x | \theta) \rangle \right] &= \int q(x) \langle \nabla_x \log q(x), \nabla_x \log P(x | \theta) \rangle \, dx \\
&= \int \langle \nabla_x q(x), \nabla_x \log P(x | \theta) \rangle \, dx \\
&= - \int q(x) \cdot \text{div}_x \left( \nabla_x \log P(x | \theta) \right) \, dx\\
&= -\mathbf{E}_q \left[ \Delta_x \log P(x | \theta) \right].
\end{align}

As such, given samples $x^1, \ldots, x^N$ drawn iid from $q$, we can estimate $J(\theta)$ up to a constant as

\begin{align}
J(\theta) \approx j(q) + \frac{1}{N} \sum_{i=1}^N \left\{ \Delta_x \log P(x^i | \theta) + \frac{1}{2} \left| \nabla_x \log P(x^i | \theta) \right|^2 \right\}
\end{align}

and then seek to minimise this numerically with respect to $\theta$.

## Beyond Quadratic Divergences

Looking at this derivation, we see that the key ingredients were that:

* Our ideal objective $J$ would be minimised at the true value of $\theta$, and
* We can estimate $J$ (up to a $\theta$-independent constant) from our samples, without knowing the normalising constant.

As such, we might seek generalisations. One direction might be to replace the quadratic penalty $\frac{1}{2} \left| \nabla_x \log q(x) - \nabla_x \log P(x | \theta) \right|^2$ by another penalty which is also amenable to integration by parts.

To construct such penalties, it is useful to introduce the notion of a *Bregman divergence*. Given a convex 'potential' $\Phi$, the Bregman divergence from $x$ to $y$ is defined as

\begin{align}
B_\Phi (x \to y) = \Phi (y) - \Phi(x) - \langle \nabla \Phi (x), y - x \rangle,
\end{align}

with appropriate modifications for non-smooth $\Phi$. Convexity ensures that $B_\Phi$ is nonnegative, and under strong convexity, one can bound $B_\Phi$ from below by an appropriate multiple of $\frac{1}{2} |x - y|^2$.

## Bregman Score Matching

To this end, consider the functional

\begin{align}
J(\theta) = \mathbf{E}_q \left[ B_{\Phi} ( \nabla_x \log P (x | \theta ) \to \nabla_x \log q(x) )\right].
\end{align}

Similar considerations to earlier can reassure us that this will also be minimised at the true value of $\theta$. Moreover, we can write it as

\begin{align}
J(\theta) &= \mathbf{E}_q \left[ B_{\Phi} ( \nabla_x \log P (x | \theta ) \to \nabla_x \log q(x))\right] \\
&= \mathbf{E}_q \left[ \Phi(\nabla_x \log q (x)) - \Phi( \nabla_x \log P (x | \theta ) ) \right] - \mathbf{E}_q \left[ \langle \nabla_s\Phi (\nabla_x \log P (x | \theta )), \nabla_x \log q(x) - \nabla_x \log P (x | \theta ) \rangle \right],
\end{align}

and again we treat the troublesome term by integration by parts

\begin{align}
-\mathbf{E}_q \left[ \langle \nabla_s\Phi (\nabla_x \log P (x | \theta )), \nabla_x \log q(x) \rangle \right] &= \mathbf{E}_q \left[ \text{div}_x (\nabla_s\Phi (\nabla_x \log P (x | \theta )) \right] \\
&= \mathbf{E}_q \left[ \langle\langle \nabla_s^2 \Phi ( \nabla_x \log P (x | \theta ) ), \nabla_x^2 \log P (x | \theta) \rangle\rangle \right],
\end{align}

which furnishes us with a usable objective.

## Conclusion

The above derivations justify that such an objective function may be reasonable to work with (e.g. tractable, potentially consistent), but do not argue for the benefits of such an approach, neither statistical nor computational, over the quadratic approach. There are also other axes along which one could generalise score matching, e.g. replacing $\nabla_x \log P (x | \theta)$ by some other normalising constant-free quantity. The eager reader may enjoy this work of [Lyu](https://arxiv.org/abs/1205.2629) for some steps in this direction.