KLMinRepGradDescent

Description

This algorithm aims to minimize the exclusive (or reverse) Kullback-Leibler (KL) divergence via stochastic gradient descent in the space of parameters. Specifically, it uses the the reparameterization gradient estimator. As a result, this algorithm is best applicable when the target log-density is differentiable and the sampling process of the variational family is differentiable. (See the methodology section for more details.) This algorithm is also commonly referred to as automatic differentiation variational inference, black-box variational inference with the reparameterization gradient, and stochastic gradient variational inference. KLMinRepGradDescent is also an alias of ADVI .

KLMinRepGradDescent(adtype; entropy, 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 reparameterization gradient estimate^{[HC1983][G1991][R1992][P1996]} of the ELBO gradient and $\mathrm{operator}$ is an optional operator (e.g. projections, identity mapping).

The reparameterization gradient, also known as the push-in gradient or the pathwise gradient, was introduced to VI in ^{[TL2014][RMW2014][KW2014]}. For the variational family $\mathcal{Q}$, suppose the process of sampling from $q_{\lambda} \in \mathcal{Q}$ can be described by some differentiable reparameterization function $T_{\lambda}$ and a base distribution $\varphi$ independent of $\lambda$ such that

\[z \sim q_{\lambda} \qquad\Leftrightarrow\qquad z \stackrel{d}{=} T_{\lambda}\left(\epsilon\right);\quad \epsilon \sim \varphi \; .\]

In these cases, denoting the target log denstiy as $\log \pi$, we can effectively estimate the gradient of the ELBO by directly differentiating the stochastic estimate of the ELBO objective

\[ \widehat{\mathrm{ELBO}}\left(q_{\lambda}\right) = \frac{1}{M}\sum^M_{m=1} \log \pi\left(T_{\lambda}\left(\epsilon_m\right)\right) + \mathbb{H}\left(q_{\lambda}\right),\]

where $\epsilon_m \sim \varphi$ are Monte Carlo samples.

Handling Constraints with Bijectors

As mentioned in the docstring, KLMinRepGradDescent assumes that the variational approximation $q_{\lambda}$ and the target distribution $\pi$ have the same support for all $\lambda \in \Lambda$. However, in general, it is most convenient to use variational families that have the whole Euclidean space $\mathbb{R}^d$ as their support. This is the case for the location-scale distributions provided by AdvancedVI. For target distributions which the support is not the full $\mathbb{R}^d$, we can apply some transformation $b$ to $q_{\lambda}$ to match its support such that

\[z \sim q_{b,\lambda} \qquad\Leftrightarrow\qquad z \stackrel{d}{=} b^{-1}\left(\eta\right);\quad \eta \sim q_{\lambda},\]

where $b$ is often called a bijector, since it is often chosen among bijective transformations. This idea is known as automatic differentiation VI^[KTRGB2017] and has subsequently been improved by Tensorflow Probability^[DLTBV2017]. In Julia, Bijectors.jl^[FXTYG2020] provides a comprehensive collection of bijections.

\[\mathrm{ADVI}\left(\lambda\right) \triangleq \mathbb{E}_{\eta \sim q_{\lambda}}\left[ \log \pi\left( b^{-1}\left( \eta \right) \right) + \log \lvert J_{b^{-1}}\left(\eta\right) \rvert \right] + \mathbb{H}\left(q_{\lambda}\right)\]

This is automatically handled by AdvancedVI through TransformedDistribution provided by Bijectors.jl. See the following example:

using Bijectors
q = MeanFieldGaussian(μ, L)
b = Bijectors.bijector(dist)
binv = inverse(b)
q_transformed = Bijectors.TransformedDistribution(q, binv)

By passing q_transformed to optimize, the Jacobian adjustment for the bijector b is automatically applied. (See the Basic Example for a fully working example.)

Entropy Gradient Estimators

For the gradient of the entropy term, we provide three choices with varying requirements. The user can select the entropy estimator by passing it as a keyword argument when constructing the algorithm object.

The requirements mean that either Distributions.entropy or Distributions.logpdf need to be implemented for the choice of variational family. In general, the use of ClosedFormEntropy is recommended whenever possible. If entropy is not available, then StickingTheLandingEntropy is recommended. See the following section for more details.

The StickingTheLandingEntropy Estimator

The StickingTheLandingEntropy, or STL estimator, is a control variate approach ^[RWD2017].

StickingTheLandingEntropy()

It occasionally results in lower variance when $\pi \approx q_{\lambda}$, and higher variance when $\pi \not\approx q_{\lambda}$. The conditions for which the STL estimator results in lower variance is still an active subject for research.

The main downside of the STL estimator is that it needs to evaluate and differentiate the log density of $q_{\lambda}$, logpdf(q), in every iteration. Depending on the variational family, this might be computationally inefficient or even numerically unstable. For example, if $q_{\lambda}$ is a Gaussian with a full-rank covariance, a back-substitution must be performed at every step, making the per-iteration complexity $\mathcal{O}(d^3)$ and reducing numerical stability.

In this example, the true posterior is contained within the variational family. This setting is known as "perfect variational family specification." In this case, KLMinRepGradDescent with StickingTheLandingEntropy is the only estimator known to converge exponentially fast ("linear convergence") to the true solution.

Recall that the original ADVI objective with a closed-form entropy (CFE) is given as follows:

n_montecarlo = 16;
b = Bijectors.bijector(model);
binv = inverse(b)

q0_trans = Bijectors.TransformedDistribution(q0, binv)

cfe = KLMinRepGradDescent(
    AutoReverseDiff();
    entropy=ClosedFormEntropy(),
    optimizer=Adam(1e-2),
    operator=ClipScale(),
)
nothing

The repgradelbo estimator can instead be created as follows:

stl = KLMinRepGradDescent(
    AutoReverseDiff();
    entropy=StickingTheLandingEntropy(),
    optimizer=Adam(1e-2),
    operator=ClipScale(),
)
nothing

We can see that the noise of the repgradelbo estimator becomes smaller as VI converges. However, the speed of convergence may not always be significantly different. Also, due to noise, just looking at the ELBO may not be sufficient to judge which algorithm is better. This can be made apparent if we measure convergence through the distance to the optimum:

We can see that STL kicks-in at later stages of optimization. Therefore, when STL "works", it yields a higher accuracy solution even on large stepsizes. However, whether STL works or not highly depends on the problem^[KMG2024]. Furthermore, in a lot of cases, a low-accuracy solution may be sufficient.

Advanced Usage

There are two major ways to customize the behavior of KLMinRepGradDescent

• Customize the Distributions functions: rand(q), entropy(q), logpdf(q).
• Customize AdvancedVI.reparam_with_entropy.

It is generally recommended to customize rand(q), entropy(q), logpdf(q), since it will easily compose with other functionalities provided by AdvancedVI.

The most advanced way is to customize AdvancedVI.reparam_with_entropy. In particular, reparam_with_entropy is the function that invokes rand(q), entropy(q), logpdf(q). Thus, it is the most general way to override the behavior of RepGradELBO.

reparam_with_entropy(rng, q, q_stop, n_samples, ent_est)

To illustrate how we can customize the rand(q) function, we will implement quasi-Monte-Carlo variational inference^[BWM2018]. Consider the case where we use the MeanFieldGaussian variational family. In this case, it suffices to override its rand specialization as follows:

using QuasiMonteCarlo
using StatsFuns

qmcrng = SobolSample(; R=OwenScramble(; base=2, pad=32))

function Distributions.rand(
    rng::AbstractRNG, q::MvLocationScale{<:Diagonal,D,L}, num_samples::Int
) where {L,D}
    (; location, scale, dist) = q
    n_dims = length(location)
    scale_diag = diag(scale)
    unif_samples = QuasiMonteCarlo.sample(num_samples, length(q), qmcrng)
    std_samples = norminvcdf.(unif_samples)
    return scale_diag .* std_samples .+ location
end
nothing

(Note that this is a quick-and-dirty example, and there are more sophisticated ways to implement this.)

By plotting the ELBO, we can see the effect of quasi-Monte Carlo. We can see that quasi-Monte Carlo results in much lower variance than naive Monte Carlo. However, similarly to the STL example, just looking at the ELBO is often insufficient to really judge performance. Instead, let's look at the distance to the global optimum:

QMC yields an additional order of magnitude in accuracy. Also, unlike STL, it ever-so slightly accelerates convergence. This is because quasi-Monte Carlo uniformly reduces variance, unlike STL, which reduces variance only near the optimum.