Globally Convergent Variational Inference

TL;DR

The work addresses global convergence in variational inference by reframing VI as forward KL minimization with a neural-network parameterization of the variational family. It leverages the neural tangent kernel to analyze function-space dynamics, proving that in the infinite-width limit with a PD NTK the functional objective has a unique minimizer, and that gradient descent on the parametric objective converges to near this minimizer. Empirically, the method demonstrates approximate global convergence at finite widths across synthetic and semi-synthetic tasks, and outperforms ELBO-based optimization in situations prone to local optima and label switching. The approach highlights the potential of likelihood-free VI to yield robust, globally convergent posteriors and suggests broader applicability beyond the strict NTK regime.

Abstract

In variational inference (VI), an approximation of the posterior distribution is selected from a family of distributions through numerical optimization. With the most common variational objective function, known as the evidence lower bound (ELBO), only convergence to a local optimum can be guaranteed. In this work, we instead establish the global convergence of a particular VI method. This VI method, which may be considered an instance of neural posterior estimation (NPE), minimizes an expectation of the inclusive (forward) KL divergence to fit a variational distribution that is parameterized by a neural network. Our convergence result relies on the neural tangent kernel (NTK) to characterize the gradient dynamics that arise from considering the variational objective in function space. In the asymptotic regime of a fixed, positive-definite neural tangent kernel, we establish conditions under which the variational objective admits a unique solution in a reproducing kernel Hilbert space (RKHS). Then, we show that the gradient descent dynamics in function space converge to this unique function. In ablation studies and practical problems, we demonstrate that our results explain the behavior of NPE in non-asymptotic finite-neuron settings, and show that NPE outperforms ELBO-based optimization, which often converges to shallow local optima.

Figures (9)

• Figure 1: Negative log-likelihood across gradient steps, for network widths $64, 256$, and $1024$ neurons. NLL for the exact posterior is denoted by the red line.
• Figure 1: Proportion out of one hundred replicates where posterior mode of $q(Z; f_2(x; \phi_2))$ was a vector in increasing order.
• Figure 2: Average $\ell_1$ distance $||\hat{Z}-Z||_1$ (and std. deviations) across one hundred replicates.
• Figure 3: Mode of $q(S; f_1(x; \phi_1))$ across experimental replicates.
• Figure 4: 100 of the $N=1000$ data observations with counterclockwise rotation of 260 degrees.

Theorems & Definitions (22)

• Lemma 1
• Corollary 1
• Theorem 1
• Lemma 2
• proof
• Proposition 1
• proof
• Theorem 2: Proposition 7 of Shapiro2003MonteCarloSampling
• Theorem 3: Gronwall's Inequality, Corollary 3 of Dragomir2003GronwallInequalities
• Theorem 4: Theorem 57 of Dragomir2003GronwallInequalities