KLMinScoreGradDescent

Description

This algorithms aim to minimize the exclusive (or reverse) Kullback-Leibler (KL) divergence via stochastic gradient descent in the space of parameters. Specifically, it uses the the score gradient estimator, which is similar to the algorithm that was originally referred to as black-box variational inference (BBVI; ^{[RGB2014][WW2013]}). (The term BBVI has also recently been used to refer to the more general setup of maximizing the ELBO in parameter space. We are using the more narrow definition, which restricts to the use of the score gradient.) However, instead of using the vanilla score gradient estimator, we differentiate the "VarGrad" objective^[RBNRA2020], which results in the score gradient variance-reduced by the leave-one-out control variate^{[SK2014][KvHW2019]}. KLMinScoreGradDescent is also an alias of BBVI.

KLMinScoreGradDescent(adtype; optimizer, n_samples, averager, operator)

callback(; rng, iteration, restructure, params, averaged_params, restructure, gradient)

The associated objective value can be estimated through the following:

estimate_objective([rng,] alg, q, prob; n_samples, entropy)

Methodology

This algorithm aims to solve the problem

\[ \mathrm{minimize}_{q \in \mathcal{Q}}\quad \mathrm{KL}\left(q, \pi\right)\]

where $\mathcal{Q}$ is some family of distributions, often called the variational family, by running stochastic gradient descent in the (Euclidean) space of parameters. That is, for all $q_{\lambda} \in \mathcal{Q}$, we assume $q_{\lambda}$ there is a corresponding vector of parameters $\lambda \in \Lambda$, where the space of parameters is Euclidean such that $\Lambda \subset \mathbb{R}^p$.

Since we usually only have access to the unnormalized densities of the target distribution $\pi$, we don't have direct access to the KL divergence. Instead, the ELBO maximization strategy maximizes a surrogate objective, the evidence lower bound (ELBO; ^[JGJS1999])

\[ \mathrm{ELBO}\left(q\right) \triangleq \mathbb{E}_{\theta \sim q} \log \pi\left(\theta\right) + \mathbb{H}\left(q\right),\]

which is equivalent to the KL up to an additive constant (the evidence).

Algorithmically, KLMinRepGradDescent iterates the step

\[ \lambda_{t+1} = \mathrm{operator}\big( \lambda_{t} + \gamma_t \widehat{\nabla_{\lambda} \mathrm{ELBO}} (q_{\lambda_t}) \big) , \]

where $\widehat{\nabla \mathrm{ELBO}}(q_{\lambda})$ is the score gradient estimate^{[G1990][KR1996][RSU1996][W1992]} of the ELBO gradient and $\mathrm{operator}$ is an optional operator (e.g. projections, identity mapping).

Let us describe the score gradient estimator^{[G1990][KR1996][RSU1996][W1992]} of the ELBO gradient, also known as the score function method and the REINFORCE gradient. For variational inference, the use of the score gradient was proposed in ^{[WW2013][RGB2014]}. Unlike the reparameterization gradient, the score gradient does not require the target log density to be differentiable, and does not differentiate through the sampling process of the variational approximation $q$. Instead, it only requires gradients of the log density $\log q$. For this reason, the score gradient is the standard method to deal with discrete variables and targets with discrete support. In more detail, the score gradient uses the Fisher log-derivative identity: For any regular $f$,

\[\nabla_{\lambda} \mathbb{E}_{z \sim q_{\lambda}} f = \mathbb{E}_{z \sim q_{\lambda}}\left[ f(z) \log q_{\lambda}(z) \right] \; .\]

The ELBO corresponds to the case where $f = \log \pi / \log q$, where $\log \pi$ is the target log density. Instead of implementing the canonical score gradient, KLMinScoreGradDescent internally uses the "VarGrad" objective^[RBNRA2020]:

\[\widehat{\mathrm{VarGrad}}(\lambda) = \mathrm{Var}\left( \log q_{\lambda}(z_i) - \log \pi\left(z_i\right) \right) \; ,\]

where the variance is computed over the samples $z_1, \ldots, z_m \sim q_{\lambda}$. Differentiating the VarGrad objective corresponds to the canonical score gradient combined with the "leave-one-out" control variate^{[SK2014][KvHW2019]}.