Support alternative ways of choosing normal approximations

[sethaxen]

Given an optimization trace, Pathfinder proposes the multivariate normal approximation constructed from the trace that maximizes the ELBO. It does this by approximating the ELBO at each point.

The discussion notes that instead of exhaustively approximating the ELBO at each point, Bayesian optimization could be used to optimize over (or even between) the points. More generally, we could allow alternative objective functions than ELBO and allow any discrete optimizer to be provided. While between points, we could interpolate means, we'd need to think a bit about how to interpolate covariances between points.

[mschauer]

Would this also cover say diagonal approximations to the covariance?

[sethaxen]

In principle, if one could specify a different way of choosing the best distribution, then yes, once could maximize ELBO (or some other objective) over some transformation of that distribution instead.

It might be cleaner though to allow the user to provide such a transformation, which would be applied to all constructed distributions before computing the objective. This would allow the user to control the distributions used for the components of the mixture model returned by multipathfinder.

[sethaxen]

In https://discourse.mc-stan.org/t/comparing-stans-adaptation-phase-to-that-of-nuts-rs/32259/5, @aseyboldt wrote:

I’d be really curious how a fisher divergence based pathfinder would work out by the way

If $q_i$ is the $i$th multivariate normal approximation, the two divergences to consider would be either $D(q_i \Vert p)$ or $D(p \Vert q_i)$, where for the score function $s_p(x) = \nabla_x \log p(x)$,

$$D(q \Vert p) = \mathbb{E}_{x \sim q}[\lVert s_p(x) - s_q(x) \rVert^2].$$

$D(p \Vert q)$

This is the form that nuts-rs's metric adaptation minimizes for multivariate normal $q$. The global optimum over all multivariate normal distributions $q$ will be the one that has the same mean and covariance as $p$ (assuming the first two moments exist).

Estimating this requires draws from $p$, which we don't have. I suppose we could importance weight the draws from $q_i$, but my guess is that we will then need many more draws from each $q_i$ than we currently use for the ELBO, which in and of itself is not expensive, but evaluating $s_p(x)$ for each draw can be. And of course, for a bad approximation we then also badly approximate the objective, which seems like a recipe for poorly selecting the optimal $q_i$. Yet another alternative might be to draw a small number of points from each $q_i$ and then use all of these draws to estimate each Fisher divergence using importance weights from a uniformly weighted mixture model, which at least would mean approximation error due to bad proposals would be shared by all of the estimates.

$D(q \Vert p)$

This just requires sampling from $q=\mathcal{N}(\mu, \Sigma)$. In fact, given draws $z_j \sim \mathcal{N}(0, I_d), j \in 1 \ldots J$,

$$\begin{aligned} D(q \Vert p) &= \mathbb{E}_{q}[\lVert s_p(x) \rVert^2 + 2 (x-\mu)^\top \Sigma^{-1} s_p(x)] + \mathrm{tr}(\Sigma^{-1})\\\ &\approx \frac{1}{J} \sum_{j=1}^J \lVert s_p(Lz_j + \mu) + L^{-\top} z_j \rVert^2, \end{aligned}$$

where $LL^\top = \Sigma$

Which is better?

I don't have a great intuition for which of these two divergences are most reasonable, but some simple experiments show that maybe only the first is sensible, e.g. we can fit a Normal distribution with fixed scale to a t-distribution:

using Plots, Distributions, ForwardDiff, QuadGK

function fisher_div(a::ContinuousUnivariateDistribution, b::ContinuousUnivariateDistribution)
    log_ratio(x) = logpdf(a, x) - logpdf(b, x)
    expectand(x) = pdf(a, x) * ForwardDiff.derivative(log_ratio, x)^2
    return quadgk(expectand, extrema(a)...)[1]
end

plot(
    plot(m -> fisher_div(TDist(2), Normal(m,1)), -50, 50; ylabel="D(p‖q)"),
    plot(m -> fisher_div(Normal(m,1), TDist(2)), -50, 50; ylabel="D(q‖p)");
    ylims=(0, NaN), yticks=[0], xlabel="m", suptitle="p=t(2), q=N(m,1)", label="",
    lw=2, dpi=300,
)

Here, even though the 2nd moment of the target doesn't exist, the global optimum of the single parameter is where we'd expect it to be. However, while the loss is concave for the first divergence, it's not for the 2nd, and it would be easy to select an arbitrarily bad approximation using the 2nd divergence.

[sethaxen]

Hm, then again here's an example where the target is N(0, 1) and the distribution we're fitting is N(0, s). In this case the metrics are all computable in closed-form (here using Mathematica):

fnorm[x_, m_, s_]:=PDF[NormalDistribution[m,s],x]
ft[x_,n_]:=PDF[StudentTDistribution[n],x]
score[p_,x_]:= D[Log[PDF[p,x]],x]
fisherdiv[p_,q_]:=Expectation[(score[p,x]-score[q,x])^2, x \[Distributed] p]
negelbo[p_,q_]:=Expectation[Log[PDF[p,x]] - Log[PDF[q,x]], x \[Distributed] p]

$Assumptions={s>0}
p = NormalDistribution[0,1]
q = NormalDistribution[0,s]
fisherdiv[p,q] // FullSimplify (* (s^2-1)^2/s^4*)
fisherdiv[q,p] // FullSimplify(*(s^2-1)^2/s^2*)
negelbo[q,p]//FullSimplify(*(s^2-1-2Log[s])/2 *)
LogLogPlot[{%%%,%%,%}, {s, 0, 100}, PlotLegends->{"DFisher[p,q]","DFisher[q,p]","-ELBO[q,p]"}]

So while $D(p \Vert q)$ still has the correct global maximum, at large $s$ the loss is essentially flat. I'd be concerned then that when selecting among the different $q_i$s, if all of them are fairly poor approximations to the target, then even if we had exact draws from $p$, approximation error would be the main determinant of which approximation is selected, and an arbitrarily bad approximation would look as good as the actually best approximation. Curiously, $D(q\Vert p)$ here looks reasonable.

[aseyboldt]

Hey, cool that you are looking into this!

Just two quick things I noticed while reading:

• During mass matrix adaptation we minimize D(p||q) where p is the posterior and q an approximation. In a few small experiments with VI (i.e., minimizing D(q || p)) I often ran into trouble where the approximation can shoot off to infinity, because then the gradients at a fixed location go to zero. I think that might be related to your flat cost? This can't happen if we minimize with respect to the actual posterior as in the mass matrix adaptation case though. It helped a bit to minimize D(q || p) + gamma KL(q || p). I think D(N(0, I), f^*p) + gamma KL(N(0, 1) || f^*p) still has a closed form solution for affine functions f.
• The norm that you choose to measure the distance of the score function really matters. So to define the fisher divergence you really need to specify three things: The two distributions and the norm. Each normal distribution defines a norm on the scores, as $\lVert s\rVert^2 = s^T \Sigma^{-1} s$. But since the posterior is typically not a normal distribution, I think you probably only have the option to use the norm induced by q.

[sethaxen]

Just two quick things I noticed while reading:

Thanks for the response and pointers! I've also been running NUTS experiments on posteriordb with swapping in single components of nuts-rs's adaptation and comparing with some other ideas so I may ping you on that when I have something to show/discuss.

• In a few small experiments with VI (i.e., minimizing D(q || p)) I often ran into trouble where the approximation can shoot off to infinity, because then the gradients at a fixed location go to zero. I think that might be related to your flat cost?

In the above example that seems to happen with D(p||q), whereas it seems like it wouldn't happen (for this example) with D(q || p), which surprised me.

This can't happen if we minimize with respect to the actual posterior as in the mass matrix adaptation case though.

Why can't it happen? Because you can analytically estimate the global minimizer, or another reason?

I think D(N(0, I), f^*p) + gamma KL(N(0, 1) || f^*p) still has a closed form solution for affine functions f.

Ah, that's worth experimenting with. Have you tried a symmetrized divergence D(N(0, I) || f^*p) + D(f^*p || N(0, 1)) (assuming computability)?

But since the posterior is typically not a normal distribution, I think you probably only have the option to use the norm induced by q.

Yeah, unfortunately I don't see any other reasonable option here. Plus I could be wrong, but I think this is the one I've seen most often in score matching papers.