Title: Globally Convergent Variational Inference
Venue: NeurIPS 2024
Accepted: yes
Keywords: forward KL divergence; neural posterior estimation; neural tangent kernel; convex optimization
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.

Globally Convergent Variational Inference

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.

1 Introduction

In variational inference (VI), the parameters η of an approximation to the posterior Q (Θ; η ) are selected to optimize an objective function, typically the evidence lower bound (ELBO) (Blei et al., 2017). However, the ELBO is generally nonconvex in η , even for simple variational families such as the family of Gaussian distributions, and so only convergence to a local optimum of the ELBO can be guaranteed (Ghadimi and Lan, 2015; Ranganath et al., 2014; Hoffman et al., 2013). As the number of such optima and the degree of suboptimality of each are generally unknown, the lack of global convergence guarantees constitutes a significant complication for practitioners and a longstanding barrier to the broader adoption of VI.

In this work, we present the first global convergence result for variational inference. We accomplish this in the context of an increasingly popular alternative objective for variational inference, the expected forward KL divergence:

Here, P ( X ) denotes a marginal of the model and P (Θ | X ) denotes the posterior. For each x ∈ X , the approximation Q (Θ; η ) to P (Θ | X = x ) is indexed by the distributional parameters η ∈ Y ⊆ R q , which are themselves the output of a neural network f ( x ; ϕ ) with weights ϕ ∈ Φ . Approximating a posterior distribution by minimizing L P has a long history (Section 2.1), and is sometimes known as neural posterior estimation (NPE) (Papamakarios and Murray, 2016) and the 'sleep' objective of reweighted wake-sleep (Bornschein and Bengio, 2015; Le et al., 2019). Minimization of this objective is straightforward: computing unbiased gradients requires only sampling θ, x ∼ P (Θ , X ) from the

joint model (Section 2.1), which is readily accomplished by ancestral sampling. This approach is 'likelihood-free' in that the density of P ( X | Θ) need not be evaluated, and therefore expected forward KL minimization is more widely applicable than ELBO-based optimization, which requires a tractable likelihood function. Analysis of the amortized problem (i.e., optimizing an objective that averages over P ( X ) ) is beneficial when considering the forward KL; for the non-amortized problem in which a single observation x is considered, only biased estimates of the gradient of the forward KL can be obtained using self-normalized importance sampling, making convergence difficult to establish (Bornschein and Bengio, 2015; Le et al., 2019; Owen, 2013). Our analysis considers a functional form of variational objective L P , given by

where L F : H → R is defined over a general reproducing kernel Hilbert space of functions H . We refer to (1) as the 'parametric objective', as its argument is the parameters ϕ ∈ Φ , and we refer to (2) as the 'functional objective' as its argument is a function f ∈ H . These objectives are closely related: under a given network parameterization, provided f ( · ; ϕ ) ∈ H , we have L P ( ϕ ) = L F ( f ( · ; ϕ )) . The objective L P has been considered in several related works (see Section 2.1). The formulation of L F and the analysis of its minimizer relative to those of L P is our main contribution.

We first demonstrate strict convexity of the functional objective L F when Q is parameterized as an exponential family distribution (Section 3). This implies the existence of a unique global optimizer f ∗ of L F for a large class of variational families. Afterward, we analyze kernel gradient flow dynamics using the neural tangent kernel to show that minimization of L P results (asymptotically) in an empirical mapping f that is at most ϵ -suboptimal relative to f ∗ , provided a sufficiently flexible neural network is used to parameterize f (Section 4). Together, these results imply that in the infinite-width limit, optimization of L P by gradient descent recovers a unique global solution.

Our analysis relies on fairly mild conditions, the most important of which are the positive definiteness of the neural tangent kernel and the structure of the variational family (e.g., an exponential family) (Section 6). Our proofs further assume a two-layer ReLU network architecture, but we conjecture that this assumption can be relaxed, and our experiments (Section 5) demonstrate global convergence for a wide variety of architectures. We illustrate that the minimization of L P converges to a global solution for problems with both synthetic and semi-synthetic data, and that finite network widths exhibit the behavior of the asymptotic regime (i.e., that of an infinitely wide network).

We further show that optimizing L P can produce better posterior approximations than likelihoodbased ELBO methods, which suffer from convergence to shallow local optima. These results suggest, surprisingly, that a likelihood-free approach to inference can outperform likelihood-based approaches. Further, even for practitioners interested in inference for a single observation x 0 , for whom amortization is not needed for computational efficiency, our approach may still be preferable to traditional ELBO-based inference due to the convergence guarantees of the former.

Related work. There is a large body of literature that analyzes the convergence of variational inference methods that target the ELBO. These works typically prove rates of convergence to the posterior (Zhang and Gao, 2020) or to a local optimum of the objective, as the ELBO is not amenable to global minimization because it is nonconvex (Domke, 2020; Domke et al., 2023; Kim et al., 2023). Liu et al. (2023) demonstrated nonconvexity of the ELBO in the context of small-object detection, and showed empirically that the expected forward KL was more robust to the pitfalls of numerical optimization. The nonconvexity of the variational objective has previously been addressed through workarounds such as convex relaxations (Fazelnia and Paisley, 2018) or iterated runs of the optimization routine to improve the quality of the local optimum (Altosaar et al., 2018). Our work differs from previous work in that our convergence result is global. Additionally, our approach is novel compared to related analyses because we consider the (arguably) more complicated problem of amortized inference, where the variational parameters are the weights of a neural network and are shared among observations.

2 Background

2.1 The Expected Forward KL Divergence

The expected forward KL objective is equivalent to the sleep-phase objective of Reweighted WakeSleep (RWS) (Bornschein and Bengio, 2015), and to the objective optimized by forward amortized variational inference (FA VI) (Ambrogioni et al., 2019). It is also a special case of the thermodynamic variational objective (TVO) (Masrani et al., 2019). Similar objectives have been referred to as neural posterior estimation (NPE) (Papamakarios and Murray, 2016; Papamakarios et al., 2019), though in these works the prior distribution, and thus the marginal P ( X ) , mutates during training.

Objectives based on the forward KL divergence generally result in variational posteriors that are overdispersed, a desirable property compared to reverse KL-based optimization (Le et al., 2019; Domke and Sheldon, 2018).

Unbiased gradient estimation for the parametric objective L P is straightforward. The outer expectation over P ( X ) allows for gradients to be computed as

where q is the density of Q with respect to Lebesgue measure (see Appendix B for details). Rewriting the left-hand side as an expectation over P (Θ) and P ( X | Θ) by Bayes' rule illustrates that samples can be drawn from this model by ancestral sampling of Θ followed by X .

Other methods targeting the forward KL, such as the wake-phase of RWS, often optimize over a different expectation, typically E X ∼D KL [ P (Θ | X ) || Q (Θ; f ( X ; ϕ ))] . Here, the outer expectation is over an empirical dataset D rather than P ( X ) . In this case, approximation techniques such as importance sampling are required to estimate the gradient, as sampling from P (Θ | X = x ) is intractable for any x . Relying on importance sampling results in biased gradient estimates, with which stochastic gradient descent (SGD) may not converge (Bornschein and Bengio, 2015; Le et al., 2019).

2.2 The Neural Tangent Kernel

A neural network architecture and the parameter space Φ of its weights together define a family of functions { f ( · ; ϕ ) : ϕ ∈ Φ } . Let ℓ ( x, f ( x )) denote a general real-valued loss function and consider selecting the parameters ϕ to minimize E P ( X ) ℓ ( X,f ( X ; ϕ )) , where P ( X ) is a distribution on the data space X . The neural tangent kernel (NTK) (Jacot et al., 2018) analyzes the evolution of the function f ( · ; ϕ ) while ϕ is fitted by gradient descent to minimize the above objective. Continuoustime dynamics are used in the formulation; ϕ ( t ) and f ( · ; ϕ ( t )) are defined for continuous time t . The parameters ϕ thus follow the ODE

Here, ˙ ϕ denotes the derivative with respect to t , and by the chain rule, the function values f ( x ; ϕ ( t )) evolve via

We define ℓ ′ ( X,f ( X )) := ∇ f ℓ ( X,f ( X )) to simplify the notation. The product of Jacobians above is known as the neural tangent kernel (NTK):

The seminal work of Jacot et al. (2018) defined and studied this kernel and established the convergence of K ϕ to a limiting kernel for certain neural network architectures as the layer width grows large.

2.3 Vector-Valued Reproducing Kernel Hilbert Spaces

Most existing NTK-based analyses consider neural networks with scalar outputs and squared error loss. Instead, we consider neural networks with multivariate outputs to accommodate the multidimensional

distributional parameter η , which parameterizes our variational distribution Q (Θ; η ) . Furthermore, we consider the objective functions L P and L F . Consequently, we rely on results from the vector-valued reproducing kernel Hilbert space (RKHS) literature, as these spaces contain vector-valued functions such as the network function f ( · ; ϕ ) . Carmeli et al. (2008, 2006) provide a detailed review, and in the following, we summarize the key properties.

Recall that η = f ( · ; ϕ ) and η ∈ R q . An R q -valued kernel on X is a map K : X × X → R q × q . The neural tangent kernel (4) is precisely such a kernel. If the R q -valued kernel K is positive definite, then K defines a unique Hilbert space of functions H whose elements are maps from X to R q (Carmeli et al., 2006) . This kernel K is called the reproducing kernel, and the corresponding space of functions is the RKHS associated with the kernel K .

In an R q -valued RKHS, the reproducing property takes on a more general form. For any x ∈ X , f ∈ H , and v ∈ R q ,

where ⟨· , ·⟩ H is the inner product of the Hilbert space H . For any fixed x ∈ X and v ∈ R q , K ( · , x ) v is a function mapping from X ↦→ R q , as is required to be an element of H .

3 Convexity of the Functional Objective

We now turn to the analysis of the functional objective L F given in Equation (2). We fix an RKHS H over which to minimize L F for now, specializing to the particular choice of H based on the neural tangent kernel subsequently. Let ℓ ( x, f ( x )) = KL[ P (Θ | X = x ) || Q (Θ; f ( x ))] . The functional L F then has the form L F = E P ( X ) ℓ ( X,f ( X )) ; we will use this form subsequently for our neural tangent kernel analysis. Our first result shows that targeting L F is highly desirable theoretically: L F admits a unique global minimizer if the variational family Q is an exponential family, as is common practice in VI.

Lemma 1. Suppose that Q (Θ; η ) is an exponential family distribution in minimal representation with natural parameters η , sufficient statistics T ( θ ) , and density q ( θ ; η ) with respect to Lebesgue measure λ (Θ) . Then, for any observation x ∈ X ⊆ R d , the loss function

is strictly convex in η , provided that P (Θ | X = x ) ≪ Q (Θ; η ) ≪ λ (Θ) for all η ∈ Y ⊆ R q .

A proof of Lemma 1, which follows quickly from the convexity of the log partition function in the natural parameter (Wainwright and Jordan, 2008, Proposition 3.1), is provided in Appendix A. Lemma 1 shows the strict convexity of the function ℓ in η . This implies the strict convexity of the functional L F ( f ) = E P ( X ) ℓ ( X,f ( X )) in f by the linearity of expectation, which in turn implies the existence of at most one global minimizer.

Corollary 1. Suppose that Q (Θ; η ) is an exponential family distribution. Then, under the conditions of Lemma 1, the functional objective

is strictly convex in f . Consequently, the set of global minimizers of L F is either a singleton set or empty.

We will assume that the set of global minimizers is nonempty (so that the minimization of L F is well-posed) and let f ∗ denote the global minimizer. We also assume that || f ∗ || H < ∞ so that f ∗ ∈ H . Hereafter, we use the term 'unique' to mean unique almost everywhere with respect to P ( X ) . Furthermore, in a slight abuse of notation, f ∗ will denote the unique equivalence class of functions that minimize L F ( f ) .

Whereas Lemma 1 establishes the convexity of the (non-amortized) forward KL divergence, Corollary 1 establishes the convexity of L F , an amortized objective, in function space. The convexity of the functional objective L F holds regardless of the distribution chosen for the outer expectation by the same linearity argument. Choices other than P ( X ) , however, may not permit unbiased gradient estimation, as is the case for the wake-phase updates of RWS (Section 2.1).

4 Global Optima of the Parametric Objective

In practice, we must directly minimize L P rather than L F , as optimizing the latter over the infinitedimensional space H directly is not tractable. Thus, in the second phase of our analysis, we consider convergence to f ∗ by minimizing the parametric objective L P with gradient descent. We define ϕ across continuous time as in Equation (3). Continuous-time dynamics simplify theoretical analysis; SGD with unbiased gradients follows a (noisy) Euler discretization of the continuous ODE (Santambrogio, 2017; Yang et al., 2021). Considering X ∼ P ( X ) for the outer expectation in both L P and L F is key in this context: this choice enables unbiased stochastic gradient estimation for L P (see Appendix B), whereas other choices require approximations that result in biased gradient estimates (see Section 2.1) and thus follow different gradient dynamics.

Analysis of the trajectories of the parametric objective L P throughout its minimization initially seems infeasible: the argument of this objective is the neural network parameters ϕ , and even well-behaved loss functions such as the mean squared error (MSE) are nonconvex in these parameters. Nevertheless, neural tangent kernel (NTK)-based results enable analysis of L P . We bridge the divide between the minimizers of the convex functional L F and the nonconvex objective L P using the limiting kernel, and show that in the large-width limit, the optimization path of L P converges arbitrarily close to f ∗ , the unique minimizer of the functional objective L F .

Theorem 1. Consider the widthp scaled 2-layer ReLU network, evolving via the flow

where f t denotes f ( · , ϕ ( t )) . Let f ∗ denote the unique minimizer of L = L F from Lemma 1, and fix ϵ > 0 . Then, under conditions (C1)-(C4), (D1)-(D4), and (E1)-(E5), there exists T > 0 such that almost surely

Regularity conditions (C1)-(C4), (D1)-(D4), and (E1)-(E5) are provided in Appendices C, D, and E, respectively. We consider a scaled two-layer ReLU network architecture (further detailed in Appendix C) and use this simple architecture to prove the results as the network width p tends to infinity. Our results may also be extended to multilayer perceptrons with other activation functions. Below, we briefly sketch the key ingredients needed to prove Theorem 1.

Recall the NTK K p ϕ from Equation (4), where we now let p denote the network width. For certain neural network architectures, Jacot et al. (2018) show that as the network width p tends to infinity, the neural tangent kernel becomes stable and tends (pointwise) towards a fixed, positive-definite limiting neural tangent kernel K ∞ .

Under suitable positivity conditions on the limiting kernel, we take the domain H of L F to be the RKHS with kernel K ∞ (Section 2.3). Because L F has a unique minimizer f ∗ , under mild conditions on K ∞ , f ∗ may be characterized as the solution obtained by following kernel gradient flow dynamics in H , that is, the ODE given by

In other words, starting from some function f 0 , following the limiting NTK gradient flow dynamics above minimizes the functional objective L F for sufficiently large T .

Lemma 2. Let f ∗ denote the minimizer of L F from Lemma 1, and ϵ > 0 . Fix f 0 , and let K ∞ denote the limiting neural tangent kernel. Let f 0 evolve according to the dynamics

Suppose that the conditions of Lemma 1 and (E1)-(E3) hold. Then, there exists T > 0 such that L ( f T ) ≤ L ( f ∗ ) + ϵ , where L is the loss functional of L F .

Appendix E enumerates regularity conditions (E1)-(E3) and provides a proof of Lemma 2. The characterization of f ∗ in Lemma 2 clarifies how the analysis of the parametric objective L P will proceed. The gradient flow L P causes the network function to similarly evolve according to a kernel

gradient flow via the empirical neural tangent kernel, that is,

as derived in Section 2.2. Comparison of the minimizers of L P and L F can be accomplished by comparing the two gradient flows above, i.e. kernel gradient flow dynamics that follow K p ϕ ( t ) and K ∞ , respectively. As K p ϕ ( t ) → K ∞ (Appendix D), these trajectories should not differ greatly: for any fixed T , the functions obtained by following the kernel gradient dynamics with K p ϕ ( t ) and K ∞ can be made arbitrarily close to one another, provided p is sufficiently large. The proof of Theorem 1 first selects a T using Lemma 2, and then bounds the difference in the trajectories on [0 , T ] for sufficiently large width p by convergence of the kernels K p ϕ ( t ) → K ∞ . Our proof differs from previous results in that it relies on uniform convergence of kernels (cf. Appendices C and D), enabling the analysis of population quantities such as E P ( X ) ℓ ( X,f ( X )) .

Theorem 1 proves convergence to an ϵ -neighborhood of the global solution when optimizing L P despite the highly nonconvex nature of this optimization problem in the network parameters ϕ . For sufficiently flexible network architectures, optimization of L P thus behaves similarly to that of L F , which we have shown is a convex problem in the function space H in Section 3.

5 Experiments

Having established conditions under which global convergence is guaranteed, the main aim of our experiments is to demonstrate approximate global convergence in practice, even for scenarios where the conditions assumed in our proofs are not satisfied exactly. Section 5.1 demonstrates that finite-neuron layer widths used in practice approximate the limiting behavior well, while Section 5.2 and Section 5.3 utilize problem-specific network architectures for amortized inference. Our results suggest that there may exist weaker assumptions under which global convergence is still guaranteed.

5.1 Toy Example

We first assess whether the asymptotic regime of Theorem 1 is relevant to practice with layers of finite width. We use a diagnostic motivated by the lazy training framework of Chizat et al. (2019), which provides the intuition that in the limiting NTK regime, the function f behaves much like its linearization around the initial weights ϕ 0 :

Liu et al. (2020) prove that equality holds exactly in the equation above if and only if f ( x ; ϕ ) has a constant tangent kernel (i.e., K ∞ ). Therefore, similarity between f and its linearization indicates that the asymptotic regime of the limiting NTK is approximately achieved. Note that even if f is linear in ϕ , as in the above expression, it may still be highly nonlinear in x .

We consider a toy example for which || x || 2 = 1 . The generative model first draws a rotation angle Θ uniformly between 0 and 2 π , and then a rotation perturbation Z ∼ N (0 , σ 2 ) , where we take σ = 0 . 5 . Conditional on Θ and Z , the data x is deterministic: x = [ cos ( θ + z ) , sin ( θ + z )] ⊤ . This construction ensures that the data lie on the sphere S 1 ⊂ R 2 , which guarantees the positivity of the limiting NTK for certain architectures (Jacot et al., 2018). We aim to infer Θ given a realization x , marginalizing over the nuisance latent variable Z . Our variational family Q (Θ; f ( x )) is a von Mises distribution, whose support is the interval [0 , 2 π ] . This family is an exponential family distribution, allowing for the application of Lemma 1. The encoder network f ( · ; ϕ ) is given by a dense two-layer network (that is, one hidden layer) with rectified linear unit (ReLU) activation, which we study as the network width p grows. The network outputs f ( x ; ϕ ) parameterize the natural parameter η .

We fit the neural network f ( x ; ϕ ) in two ways. First, we use SGD to fit the network parameters ϕ to minimize L P . Second, we fit the linearization f lin ( x ; ϕ ) = f ( x ; ϕ 0 ) + J ϕ f ( x ; ϕ 0 )( ϕ -ϕ 0 ) in ϕ . We perform both of these fitting procedures for various widths p . For both settings, stochastic gradient estimation was performed by following the procedure in Appendix B. For evaluation, we fix N = 1000 independent realizations x ∗ 1 , . . . , x ∗ N from the generative model with underlying ground-truth latent parameter values θ ∗ 1 , . . . , θ ∗ N , and evaluate the held-out negative log-likelihood (NLL), -1 N ∑ N i =1 log q ( θ ∗ i | f ( x ∗ i ; ϕ )) , for both functions: f ( x ; ϕ ) and f lin ( x ; ϕ ) . Figure 1 shows

the evolution of the held-out NLL across the fitting procedure for three different network widths p : 64 , 256 and 1024 . The difference in quality between the linearizations and the true functions at convergence diminishes as the width p grows; for p = 1024 , the two are nearly identical, providing evidence that the asymptotic regime is achieved.

0.735

0.735

0.735

0.730

64

0.730

256

0.730

1024

64-lin

256-lin

1024-lin

0.725

0.725

0.725

0.720

0.720

Z

Z

0.715

0.715

0.715

0.710

0.710

0.710

0.705

0.0

0.5

1.0

1.5

2.0

0.705

0.0

0.5

1.0

1.5

2.0

0.705

0.0

0.5

1.0

1.5

2.0

Gradient Step

×106

Gradient Step

×106

Gradient Step

×106

5.2 Label Switching in Amortized Clustering

In this experiment and the subsequent one, we consider a more diverse set of network architectures. Although Theorem 1 assumes a shallow, dense network architecture, these experimental results show that empirically, global convergence can be obtained for many other architectures as well. We consider the difficult problem of amortizing clustering. In this problem, we are given an unordered collection (set) of data x = { x 1 , . . . , x n } , x i ∈ R and our objective is to infer the locations and labels of the cluster centers from which the data were generated. The generative model first draws a shift parameter S ∼ N (0 , 100 2 ) followed by

In the above, σ 2 , τ 2 , µ ∈ R d and p ∈ R d are known variances, locations, and proportions, respectively. The global parameter S = s ∈ R shifts the d locations µ 1 , . . . , µ d to µ 1 + s, . . . , µ d + s . The cluster centers Z are then obtained by adding noise to these locations. An implicit labeling is imposed on the centers by assigning each a dimension in R d via the prior. Finally, the data are drawn independently from a mixture of d univariate Gaussians with centers Z i , i = 1 , . . . , d . We consider the tasks of inferring the scalar shift S and the vector of cluster centers Z . We fix µ = [ -20 , -10 , 0 , 10 , 20] ⊤ with d = 5 . We fix the hyperparameters σ = 0 . 5 and τ = 0 . 1 , and artificially fix the shift as S = 100 to generate n = 1000 independent realizations X = x from the generative model.

̸

Inferring S should be straightforward because the vector µ is known and fixed, ensuring that the joint likelihood p ( x, s ) (marginalizing over z ) is unimodal in s . However, inference on Z may be difficult for likelihood-based methods because the order of the entries of Z matters: the joint density p ( x, z, s ) = p ( x, π ( z ) , s ) for a permutation π , even though p ( x | z ) = p ( x | π ( z )) . 'Label switching' can thus pose a significant obstacle for likelihood-based methods, as any permutation of the cluster centers Z will still explain the data well. This problem formulation thus results in a likelihood, and hence a posterior density, with many local optima but a single global optimum (i.e., where the cluster centers have the correct labeling).

Now we show that likelihood-based approaches to variational inference, such as maximizing the evidence lower-bound (ELBO), result in suboptimal solutions compared to the minimization of L P . We take the variational distributions q ( S ; f 1 ( x ; ϕ 1 )) and q ( Z ; f 2 ( x ; ϕ 2 )) to be Gaussian. We fit the networks f 1 and f 2 . Due to the exchangeability of the observations x = { x 1 , . . . , x n } , we parameterize each as permutation-invariant neural networks, fitting both ϕ 1 and ϕ 2 to either minimize L P or to maximize the ELBO (see Appendix F). We perform 100 replications of this experiment across different random seeds, and consider two different parameterizations of the Gaussian variational

distribution: a mean-only parameterization with fixed unit variance and a natural parameterization with an unknown mean and unknown variance. Both of these variational families are exponential families, and so convexity (in the sense of Corollary 1) holds.

Figure 2 plots kernel-smoothed frequencies of point estimates of S , where each point estimate is the mode of the variational posterior for an experimental replicate. Both ELBO- and L P -based training estimate S = 100 well. However, the limitations of ELBO-based training are evident in the fidelity of the posterior approximation of Z . Table 2 plots the average ℓ 1 distance || ˆ Z -Z || 1 between the variational mode ˆ Z and the true latent draw Z . Optimization of the ELBO converges to a local optimum that is a permutation of the entries of Z , resulting in a large ℓ 1 distance on average. Minimization of L P , on the other hand, converges to the global optimum without any label switching. Table 1 indicates the degree of label switching, showing the proportion of trials in which the entries ˆ Z were correctly ordered.

5.3 Rotated MNIST Digits

We consider the task of inferring a shared rotation angle Θ for a set of N MNIST digits. The generative model is as follows. The rotation angle is drawn as Θ ∼ Uniform (0 , 2 π ) , and for all i ∈ [ N ] a noise term is drawn as Z i ∼ N (0 p , σ 2 I p ) . Finally, each image is drawn as

Here, RotateMNIST is fitted ahead of time and fixed throughout this experiment. Given a latent representation z and angle θ , RotateMNIST returns a 28 × 28 MNIST digit image rotated counterclockwise by θ degrees. (See Appendix F for additional experimental details.) We aim to fit the variational posterior q (Θ; f ( x 1 , . . . , x N ; ϕ )) , implicitly marginalizing over the nuisance latent variables Z 1 , . . . , Z N . The true data { x i } N i =1 are generated from the above model, with N = 1000 digits and an underlying counterclockwise rotation angle of θ = 260 degrees. We provide visualizations of some of these digits in Figure 3. The variational distribution q (Θ; f ( x 1 , . . . , x N ; ϕ )) is taken to be a von Mises distribution with natural parameterization, as in Section 5.1, although for this example the architecture of the encoder network makes f invariant to permutations of the inputs { x i } N i =1 , to reflect the exchangeability of the data. We fit f to minimize L P using 100,000 iterations of SGD.

We compare to fitting θ to directly maximize the likelihood of the data p ( θ, { x i } N i =1 ) . As this quantity is intractable, we maximize the importance-weighted bound (IWBO), which is a generalization of the ELBO (Burda et al., 2016), by maximizing the joint likelihood p ( θ, { z i } N i =1 , { x i } N i =1 ) in θ while fitting a Gaussian variational distribution on the variables Z .

The likelihood function is multimodal in θ due to the approximate rotational symmetry of several handwritten digits: zero, one, and eight are approximately invariant under rotations of 180 degrees. Additionally, the digits six and nine are similar following 180-degree rotations. These symmetries yield a multimodal posterior distribution on Θ . Likelihood-based fitting procedures, such as maximizing the IWBO, often get stuck in these shallow local optima, while fitting f to minimize the parametric objective L P finds a unique global solution. Figure 4 shows estimates of the angle θ conditional on the data { x i } N i =1 during training with the IWBO objective, with θ initialized to a

ELBO

NaturalParameterization

Mean Parameterization

80

90

100

110

120 80

90

100

110

120

350

300

()

250

200

Ang

100

50

0

200

400

Gradient Step

600

800

1000

1200

variety of values. For some initializations, the IWBO optimization converges quickly to near the correct value of 260 degrees, but in many others, it converges to a shallow local optimum.

We perform the same routine, fitting q (Θ; f ( x 1 , . . . , x N ; ϕ )) to minimize the expected forward KL divergence L P . Figure 5 shows that across a variety of initializations of the angles, this approach always converges to a unique solution and fits the posterior mode to the correct value of 260 degrees. L P minimization converges rapidly, and so Figure 6 zooms in on the initial few thousand gradient steps to show the various trajectories among initializations.

280

275

270

Degrees

265

260

0255

250

245

240

0

20000

40000

GradientStep

60000

80000

100000

5.4 Local Optima vs. Global Optima

In this experiment, we directly compare the quality of the variational approximation minimizing the expected forward KL objective to that of the local optima found by optimizing the ELBO. We consider an adapted version of the rotated MNIST digit problem outlined above. Each digit x i is drawn from the model x i ∼ N ( Rotate ( x 0 i , θ ) , τ 2 ) for i = 1 , . . . , 50 with angle set to θ true = 260 degrees. The unrotated digits x 0 i are fixed a priori, eliminating the nuisance latent variables Z in this setting. This allows us to directly perform ELBO-based variational inference on Θ , instead of merely maximizing the likelihood (or a bound thereof) as in Section 5.3. We fit an amortized Gaussian variational distribution with fixed variance σ 2 = 0 . 5 2 for inference on Θ , and use the same prior as above. This is an exponential family with only one location parameter to be learned.

Fitting q (Θ; f ( x 1 , . . . , x 50 ; ϕ )) is performed by minimizing either the negative ELBO or the expected forward KL divergence. The only differences between the two approaches are their objective functions: the data, network architecture, learning rates, and all other hyperparameters are the same. We optimize each objective function over 10,000 gradient steps, and measure the quality of the obtained variational approximations by a variety of metrics, including: the (non-expected) forward KL divergence; the reverse KL divergence; negative log-likelihood; and the angle point estimate. Intuition suggests that a global optimum should outperform local ones, and we find this to be the case. By any of

350

300

Degrees)

250

200

150

3

100

50

0

0

250

500

750

1000

1250

1500

1750

2000

GradientStep

the performance metrics we consider, the global solution of the expected forward KL minimization outperforms the variational approximations found by optimizing the ELBO.

Forward KL: KL(p(0| Ctrue) II q(0| true))

Negative Log Likelihood (— log q(Otrue | Ctrue)

40

30

20

20

10

0

0

0

2000

4000

6000

8000

10000

0

2000

4000

6000

8000

10000

Gradient Step

Gradient Step

Reverse KL: KL(q(0|Ctrue) I| p(0| true)

Angle Estimate (Variational Mode)

x107

4.0

M

300

3.5

200

3.0

100

2.5

0

0

2000

4000

6000

8000

10000

0

2000

4000

6000

8000

10000

Gradient Step

Gradient Step

Truth

Expected Forward KL

-ELBO

6 Discussion

In this work, we showed that in the asymptotic limit of an infinitely wide neural network, gradient descent dynamics on the expected forward KL objective L P converge to an ϵ -neighborhood of a unique function f ∗ , a global minimizer. Our results depend on several regularity conditions, the most important of which is the positive definiteness of the limiting neural tangent kernel, the compactness of the data space X , and the specific architecture considered, a single hidden-layer ReLU network. We conjecture and show experimentally that global convergence holds more generally. First, we illustrate that the asymptotic regime describes practice with a finite number of neurons (see Section 5.1). Second, our conditions allow for a wide variety of activation functions beyond ReLU (see Appendix D). Third, empirically, global convergence can be achieved with many network architectures. Beyond multilayer perceptrons, we illustrate global convergence for convolutional neural networks (CNNs) and permutation-invariant architectures in Section 5.2 and Section 5.3.

Expected forward KL minimization is a likelihood-free inference (LFI) method. For Bayesian inference, likelihood-based and likelihood-free methods are not typically viewed as competitors, but as different tools for different settings. In a setting where the likelihood is available, the prevailing wisdom suggests utilizing it. However, our results suggest that likelihood-free approaches to inference may be preferable even when the likelihood function is readily available. We find likelihood-based methods are prone to suffer the shortcomings of numerical optimization, often converging to shallow local optima of ELBO-based variational objectives. Expected forward KL minimization instead converges to a unique global optimum of its objective function.

There are situations in which ELBO optimization may nevertheless be preferable. First, if the likelihood function of the model is approximately convex and well-conditioned in the region of interest, ELBO optimization should recover a nearly global optimizer. Second, in certain situations, the ELBO can be optimized using deterministic optimization methods, which can be much faster than SGD. Third, if the generative model has free model parameters, with the ELBO they can be fitted while simultaneously fitting the variational approximation, with a single objective function for both tasks. Fourth, the ELBO can be applied with non-amortized variational distributions, which can have computational benefits in settings with few observations. Many important inference problems do not fall into any of these four categories. Even for those that do, the benefits of expected forward KL minimization, including global convergence, may outweigh the benefits of ELBO optimization.

Acknowledgments

We thank the reviewers for their helpful comments and suggestions. This material is based on work supported by the National Science Foundation under Grant Nos. 2209720 (OAC) and 2241144 (DGE), and the U.S. Department of Energy, Office of Science, Office of High Energy Physics under Award Number DE-SC0023714.

References

Altosaar, J., Ranganath, R., and Blei, D. (2018). Proximity variational inference. In International Conference on Artificial Intelligence and Statistics .

Ambrogioni, L., Güçlü, U., Berezutskaya, J., van den Borne, E., Güçlütürk, Y., Hinne, M., Maris, E., and van Gerven, M. (2019). Forward amortized inference for likelihood-free variational marginalization. In International Conference on Artificial Intelligence and Statistics .

Ba, J., Erdogdu, M., Suzuki, T., Wu, D., and Zhang, T. (2020). Generalization of two-layer neural networks: An asymptotic viewpoint. In International Conference on Learning Representations .

Blei, D. M., Kucukelbir, A., and McAuliffe, J. D. (2017). Variational inference: A review for statisticians. Journal of the American Statistical Association , 112(518):859-877.

Bornschein, J. and Bengio, Y. (2015). Reweighted wake-sleep. In International Conference on Learning Representations .

Burda, Y., Grosse, R. B., and Salakhutdinov, R. (2016). Importance weighted autoencoders. In International Conference on Learning Representations .

Carmeli, C., De Vito, E., and Toigo, A. (2006). Vector valued reproducing kernel Hilbert spaces of integrable functions and Mercer theorem. Analysis and Applications , 4(4):377-408.

Carmeli, C., Vito, E. D., Toigo, A., and Umanità, V. (2008). Vector valued reproducing kernel Hilbert spaces and universality.

Chizat, L., Oyallon, E., and Bach, F. (2019). On lazy training in differentiable programming. In Neural Information Processing Systems .

Domke, J. (2020). Provable smoothness guarantees for black-box variational inference. In International Conference on Machine Learning .

Domke, J., Gower, R. M., and Garrigos, G. (2023). Provable convergence guarantees for black-box variational inference. In Neural Information Processing Systems .

Domke, J. and Sheldon, D. R. (2018). Importance weighting and variational inference. In Neural Information Processing Systems .

Dragomir, S. S. (2003). Some Gronwall Type Inequalities and Applications . Nova Science Publishers.

Fazelnia, G. and Paisley, J. (2018). CRVI: Convex relaxation for variational inference. In International Conference on Machine Learning .

Ghadimi, S. and Lan, G. (2015). Stochastic approximation methods and their finite-time convergence properties. In Handbook of Simulation Optimization , pages 149-178. Springer.

Hoffman, M. D., Blei, D. M., Wang, C., and Paisley, J. (2013). Stochastic variational inference. Journal of Machine Learning Research , 14(40):1303-1347.

Jacot, A., Gabriel, F., and Hongler, C. (2018). Neural tangent kernel: Convergence and generalization in neural networks. In Neural Information Processing Systems .

Kim, K., Oh, J., Wu, K., Ma, Y., and Gardner, J. R. (2023). On the convergence of black-box variational inference. In Neural Information Processing Systems .

Le, T. A., Kosiorek, A. R., Siddharth, N., Teh, Y. W., and Wood, F. (2019). Revisiting reweighted wake-sleep for models with stochastic control flow. In International Conference on Uncertainty in Artificial Intelligence .

Lee, J., Lee, Y ., Kim, J., Kosiorek, A., Choi, S., and Teh, Y . W. (2019). Set transformer: A framework for attention-based permutation-invariant neural networks. In International Conference on Machine Learning .

Liu, C., Zhu, L., and Belkin, M. (2020). On the linearity of large non-linear models: when and why the tangent kernel is constant. In Neural Information Processing Systems .

Liu, R., McAuliffe, J. D., Regier, J., and the LSST Dark Energy Science Collaboration (2023). Variational inference for deblending crowded starfields. Journal of Machine Learning Research , 24(179):1-36.

Masrani, V., Le, T. A., and Wood, F. (2019). The thermodynamic variational objective. In Neural Information Processing Systems .

Owen, A. B. (2013). Monte Carlo Theory, Methods and Examples .

Papamakarios, G. and Murray, I. (2016). Fast ϵ -free inference of simulation models with bayesian conditional density estimation. In Neural Information Processing Systems .

Papamakarios, G., Sterratt, D., and Murray, I. (2019). Sequential neural likelihood: Fast likelihoodfree inference with autoregressive flows. In International Conference on Artificial Intelligence and Statistics .

Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., et al. (2019). PyTorch: An imperative style, high-performance deep learning library.

Ranganath, R., Gerrish, S., and Blei, D. (2014). Black box variational inference. In International Conference on Artificial Intelligence and Statistics .

Santambrogio, F. (2017). Euclidean, metric, and Wasserstein gradient flows: an overview. Bulletin of Mathematical Sciences , 7(1):87-154.

Shapiro, A. (2003). Monte Carlo sampling methods. In Stochastic Programming , volume 10 of Handbooks in Operations Research and Management Science , pages 353-425. Elsevier.

Srivastava, M. K., Khan, A. H., and Srivastava, N. (2014). Statistical Inference: Theory of Estimation . PHI Learning.

Wainwright, M. J. and Jordan, M. I. (2008). Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning , 1(1-2):1-305.

Yang, J., Hu, W., and Li, C. J. (2021). On the fast convergence of random perturbations of the gradient flow. Asymptotic Analysis , 122:371-393.

Zhang, F. and Gao, C. (2020). Convergence rates of variational posterior distributions. The Annals of Statistics , 48(4):2180 - 2207.

A Convexity of the Functional Objective

We prove Lemma 1 from the manuscript below.

Proof. Let the log-density be given by log q ( θ ; η ) = log h ( θ ) + η ⊤ T ( θ ) -A ( η ) . First, observe that under the conditions given, the function ℓ is equivalent (up to additive constants) to a much simpler expression, the expected log-density of q , via

where C is a constant that does not depend on η . Now, the mapping η → -log q ( θ ; η ) is convex in η because its Hessian is ∂A ∂η∂η ⊤ = Var ( T ( θ )) ≻ 0 (cf. Chapter 6.6.3 of Srivastava et al. (2014)). Strict convexity follows from minimality of the representation (cf. Wainwright and Jordan (2008), Proposition 3.1). We can show ℓ is (strictly) convex in η by applying linearity of expectation. We have for any λ ∈ [0 , 1]

where the second line follows from the convexity of the map η ↦→-log q ( θ ; η ) above for any value of θ . In fact, the inequality holds strictly as well by strict convexity above and the continuity of log q ( θ ; η ) in θ . So the function ℓ ( x, η ) is strictly convex in η .

B Unbiased Stochastic Gradients for the Parametric Objective

Computation of unbiased estimates of the gradient of the loss function L P ( ϕ ) with respect to the parameters ϕ is all that is needed to implement SGD for L P . Under mild conditions (see Proposition 1), the loss function L ( ϕ ) may be equivalently written as

for density functions p, q , where f ( · ; ϕ ) denotes a function parameterized by ϕ . Under the conditions of Proposition 1, differentiation and integration may be interchanged, so that

and unbiased estimates of the quantity can be easily attained by samples drawn ( θ, x ) ∼ P (Θ , X ) .

Proposition 1. Let (Ω 1 , B 1 ) , (Ω 2 , B 2 ) be measurable spaces on which the random variables X : Ω 1 → X and Θ : Ω 2 → O are defined, respectively. Suppose that for all x ∈ X and all ϕ ∈ Φ we have P (Θ | X = x ) ≪ Q (Θ; f ( x ; ϕ )) ≪ λ (Θ) , with λ (Θ) denoting Lebesgue measure and ≪ denoting absolute continuity. Further, suppose that log ( p (Θ | X ) q (Θ; f ( X ; ϕ )) ) is measurable with respect to the product space (Ω 1 × Ω 2 , B 1 ×B 2 ) for each ϕ ∈ Φ , and ∇ ϕ log q ( θ ; f ( x ; ϕ )) exists for almost all ( θ, x ) ∈ O × X . Finally, assume there exists an integrable Y dominating ∇ ϕ log q ( θ ; f ( x ; ϕ )) for all ϕ ∈ Φ and almost all ( θ, x ) ∈ O × X . Then for any B ∈ N and any ϕ ∈ Φ the quantity

is an unbiased estimator of the gradient of the objective L P , evaluated at ϕ ∈ Φ .

Proof. By the absolute continuity assumptions, for any x ∈ X the distributions P (Θ | X = x ) and Q (Θ; f ( x ; ϕ )) admit densities with respect to Lebesgue measure denoted p ( θ | x ) and q ( θ ; f ( x ; ϕ )) ,

respectively. We may then rewrite the KL divergence from Equation (1) as

because the Radon-Nikodym derivative dP/dQ is given by the ratio of these densities. Equation (1) is thus equivalent to

This expectation is well-defined by the measurability assumption on log ( p (Θ | X ) q (Θ; f ( X ; ϕ )) ) . To interchange differentiation and integration, it suffices by Leibniz's rule that the gradient of this quantity with respect to ϕ is dominated by a measurable r.v. Y . More precisely, there exists an integrable Y ( θ, x ) defined on the product space O × X such that ∣ ∣ ∣ ∣ ∇ ϕ log ( p ( θ | x ) q ( θ ; f ( x ; ϕ )) ) ∣ ∣ ∣ ∣ ≤ Y ( θ, x ) for all ϕ ∈ Φ and almost everywhereP (Θ , X ) . This is assumed in the statement of the proposition, and so we have

and the result follows by sampling.

The variance of the gradient estimator can be reduced at the standard Monte Carlo rate, and for any B Equation (12) can be used for SGD.

C The Limiting NTK

Before proceeding, we introduce the architecture specific to our analysis, a scaled two-layer network, and several theorems that we will use throughout the analysis.

The first result from Shapiro (2003) concerns optimization of the objective f ( x ) = E ξ ∼ P F ( x, ξ ) in x via its empirical approximation ˆ f n ( x ) = 1 n ∑ n i =1 F ( x, ξ i ) , ξ i iid ∼ P . We reproduce this result below. Theorem 2 (Proposition 7 of Shapiro (2003)) . Let C be a nonempty compact subset of R n and suppose that (i) for almost every ξ ∈ Ξ the function F ( · , ξ ) is continuous on C , (ii) F ( x, ξ ) , x ∈ C , is dominated by an integrable function, (iii) the sample ξ 1 , . . . , ξ n is iid. Then the expected value function f ( x ) is finite-valued and continuous on C , and ˆ f n ( x ) converges to f ( x ) with probability 1 uniformly on C .

The next two results are integral forms of Gronwall's inequality that we use in subsequent analysis. We refer to Dragomir (2003) for a detailed review and summarize the results therein below.

Theorem 3 (Gronwall's Inequality, Corollary 3 of Dragomir (2003)) . Let u ( t ) ∈ R be such that u ( t ) ≤ c 1 + c 2 ∫ t 0 u ( s ) ds for t > 0 and nonnegative c 1 , c 2 . Then

Theorem 4 (Theorem 57 of Dragomir (2003)) . Let u ( t ) ∈ R be such that u ( t ) ≤ c 1 + c 2 ∫ t 0 ∫ s 0 u ( v ) dvds for t > 0 and nonnegative c 1 , c 2 . Then

Now we turn to specifics of the architecture we consider. Assume the function f has the architecture of a (scaled) two-layer (single hidden layer) network mapping f : X → Y with X ⊆ R d and Y ⊆ R q . We consider this network architecture for a given width p , and study each of the i = 1 , . . . , q

coordinate functions of f . For a scaled two-layer network, the i th such function is

for i = 1 , . . . , q , where σ denotes an activation function. The scaling depends on the width of the network p . The parameters ϕ are thus ϕ = { w j , a ( · ) ,j } p j =1 where a ( · ) ,j denotes the vector [ a 1 j , . . . , a qj ] ⊤ (i.e. the j th coefficient for each component function i ). The individual parameters have dimensions as follows: w j ∈ R d and a ( · ) ,j ∈ R q , for all j = 1 , . . . , p , where again p denotes the network width and d the data dimension dim X . For ease hereafter, we write a j = a ( · ) ,j to refer to the entire j th vector of second layer network coefficients when the context is clear. As is standard, the first layer parameters are initialized as independent standard Gaussian random variables, i.e. w j iid ∼ N (0 , I d ) for all j = 1 , . . . , p . In other related works, the weight a ij is sometimes also drawn as a ij iid ∼ N (0 , 1) for all i = 1 , . . . , q, j = 1 , . . . , p , but in this work, we initialize these second-layer weights to zero for simplicity to ensure that at initialization, f ( · ; ϕ ) = 0 . A zero-initialized network function is used for analysis in several related works, e.g. Chizat et al. (2019) and Ba et al. (2020). For now, notationally we denote weights to be initialized as draws from an arbitrary distribution D , and we introduce specificity to D as required.

The neural tangent kernel (Equation (4)) can be computed explicitly for this architecture and is given in the lemma below, which proves pointwise convergence to the limiting NTK at initialization as the width p tends to infinity.

Lemma 3 (Pointwise Convergence At Initialization) . For the architecture above, consider any p . Let a ∈ R q , w ∈ R d be distributed according to a, w ∼ D for some distribution D such that a, w are integrable ( L 1 ) random variables. Assume X is compact, and σ ′ is bounded. Then, provided condition (C4) holds (see below), we have for any x, ˜ x ∈ X that

as p →∞ where K p ϕ (0) denotes the NTK at initialization constructed from draws a j , w j iid ∼ D and K ( x, ˜ x ; w,a ) ∈ R q × q is the q × q matrix whose k, l th entry is given by

for k, l = 1 , . . . , q

Proof. Consider the k th coordinate function of f . For any choice of p , the gradient is given by

̸

where we have imposed an arbitrary ordering on the parameters. In the above, we omitted partial derivatives ∂f k ∂a lj for l = k, j = 1 , . . . , p because these are all identically zero. From this, it follows that for any fixed x, ˜ x ∈ X , the k, l -th entry of K p ϕ (0) ( x, ˜ x ) is given by

The existence of the limiting NTK follows immediately: for each of the two terms above, each term is integrable by the compactness of X and domination (see (C4)). It follows that K ∞ ( x, ˜ x ) is the q × q matrix whose k, l th entry is given by

with w,a ∼ D . Convergence in probability pointwise follows from the weak law of large numbers, and almost sure convergence holds by the strong law of large numbers. K ( x, x ; a, w ) is integrable by the assumption (C4) (see below), so the expectation is well-defined.

The proof of the existence and pointwise convergence to the limiting NTK K ∞ above is rather straightforward, and this result has been previously established in other works (Jacot et al., 2018). For our analysis of kernel gradient flows in Theorem 1 for the expected forward KL objectives L P and L F , however, we require uniform convergence to K ∞ over the entire data space X .

We establish conditions under which this uniform convergence holds in two results, Proposition 2 and Proposition 3. Proposition 2, given below, concerns convergence at initialization to the limiting neural tangent kernel K ∞ (i.e. before beginning gradient descent). Proposition 3, proven in Appendix D, demonstrates that across a finite training interval [0 , T ] , the NTK changes minimally from its initial value in a large width regime. Generally, we refer to the first result as 'deterministic initialization' and the second as 'lazy training' following related works (Jacot et al., 2018; Chizat et al., 2019).

Below, we give suitable regularity conditions and state and prove Proposition 2.

(C1) The data space is X is compact.

(C2) The distribution D is such that w ∼ N (0 , I d ) and a = 0 w.p. 1. For j = 1 , . . . , p iid draws from this distribution, we thus have w j iid ∼ N (0 , I d ) and a ij = 0 w.p 1 for all i, j .

(C3) The activation function σ is continuous. Under (C2), this implies that the function K ( · , · ; a, w ) from Lemma 3 with a, w ∼ D is almost surely continuous.

(C4) The function K ( x, ˜ x ; a, w ) is dominated by some integrable random variable G , i.e. for all x, ˜ x ∈ X × X we have || K ( x, ˜ x ; a, w ) || F ≤ G ( a, w ) almost surely for integrable G ( a, w ) .

Proposition 2. Fix a scaled two-layer network architecture of width p , and let Φ denote the corresponding parameter space. Initialize ϕ (0) as independent, identically distributed random variables drawn from the distribution D in (C2). Let K p ϕ (0) : X × X → R q × q be the mapping defined by

( x, ˜ x ) ↦→ K ϕ (0) ( x, ˜ x ) = J ϕ f ( x ; ϕ (0)) J ϕ f (˜ x ; ϕ (0)) ⊤ . Then, provided conditions (C1)-(C4) hold, we have as p →∞ that

where K ∞ ( x, ˜ x ) := plim p →∞ K p ϕ (0) ( x, ˜ x ) is a fixed, continuous kernel.

Proof. The proof follows by direct application of Proposition 7 of Shapiro (2003). Precisely, we satisfy i) almost-sure continuity of K ( · , · ; a, w ) by (C3), ii) domination by (C4), and iii) the draws comprising K p ϕ (0) are iid by assumption. By this proposition, then, we have uniform convergence of K p ϕ (0) to K ∞ and obtain continuity of K ∞ as well.

D Lazy Training

Below, we prove several results that will aid in proving the 'lazy training' result of Proposition 3 (see below). Given the same architecture as above in Appendix C and a fixed width p and time T > 0 , we will begin by bounding || w j ( T ) -w j (0) || and || a kj ( T ) -a kj (0) || , || a lj ( T ) -a lj (0) || for all k, l = 1 , . . . , q and all j = 1 , . . . , p . As in Appendix C, there are several conditions that we impose and use in the following results. (D1)-(D2) are identical to (C1)-(C2), repeated for clarity.

(D1) The data space is X is compact.

(D2) The distribution D is such that w ∼ N (0 , I d ) and a = 0 w.p. 1. For j = 1 , . . . , p iid draws from this distribution, we thus have w j iid ∼ N (0 , I d ) and a ij = 0 w.p 1 for all i, j .

(D3) The function ℓ ( x, η ) = KL[ P (Θ | X = x ) || Q (Θ; η )] is such that ℓ ′ ( x ; η ) is bounded uniformly for all x and for all η ∈ { f ( x ; ϕ ( t )) : t > 0 } by a constant ˜ M , uniformly over the width p . We recall that this notation is shorthand for ∇ η ℓ ( x, η ) .

(D4) The activation function σ ( · ) is non-polynomial and is Lipschitz with constant C . Note that the Lipschitz condition implies σ has bounded first derivative i.e. | σ ′ ( r ) | ≤ C for all r ∈ R .

With these conditions in hand, we now prove several lemmas for individual parameters.

Lemma 4 (Lazy Training of w ) . For the width p scaled two-layer architecture above, assume that conditions (D1)-(D4) hold. Let ϕ evolve according to the gradient flow of the objective L P , i.e.

Fix any T > 0 . Then for all j = 1 , . . . , p ,

almost surely, where D p,T , E p,T are constants depending on p, T , and lim p →∞ D p,T = 0 and lim p →∞ E p,T = 0 .

Proof. First note that for any fixed j , we have

as required, where a ij ∈ R for i = 1 , . . . , q and x ∈ X ⊆ R d from (D1). We can bound the operator 2-norm of this matrix by observing that for any y ∈ R d we have

by observing || x || 2 2 is bounded by (D1) (and we absorb this term into the constant C ). By similar computations, we also have

Using condition (D4), it follows that

by Cauchy-Schwarz and (D1),(D4), where throughout the following we let K := | σ (0) | and again || x || terms are absorbed into the constant C . Now we will use these computations to bound the

variation on w j across the interval (0 , T ] . Fix any t ∈ (0 , T ] . Then we have

with c 1 = C ˜ M 2 KT 2 2 p + C 2 ˜ M 2 T 2 2 p || w j (0) || 2 and c 2 = C 2 ˜ M 2 p . Note that even though c 1 depends on T , this is constant as T is fixed. We write these quantities in this way to recognize a Gronwall-type inequality that we can use to bound the left hand side. Indeed, by direct application of Theorem 57 of Dragomir (2003) (see Theorem 4) we have that

giving the result for t = T if we take D p,T = C 2 ˜ M 2 T 2 2 p exp [ C 2 ˜ M 2 T 2 2 p ] and E p,T = C ˜ M 2 KT 2 2 p exp [ C 2 ˜ M 2 T 2 2 p ] . Clearly these constants satisfy lim p →∞ D p,T = 0 and lim p →∞ E p,T = 0 for any fixed T .

Lemma 5 (Lazy Training of a ) . Under the same conditions as Lemma 4, let ϕ evolve according to the gradient flow of problem L P , i.e.

Fix any T > 0 . Then we have for any j that

almost surely, where E p,T and F p,T are constants depending on p, T satisfying lim p →∞ E p,T = 0 and lim p →∞ F p,T = 0 .

Proof. We will use much of the same work as in Lemma 4. Namely, || a j ( t ) || 2 = || a j ( t ) -a j (0) || 2 almost surely by (D2), and thereafter for any t ∈ (0 , T ] we have

Clearly, these constants satisfy lim p →∞ F p,t → 0 and lim p →∞ G p,t → 0 (to see this, simply plug in the forms of D p,s and E p,s from Lemma 4 above) and we have the result by taking t = T .

Now with these results in hand, we may state and prove Proposition 3, given below.

Proposition 3. Under the same conditions as Proposition 2, fix any T > 0 . For any t ∈ (0 , T ] let K p ϕ ( t ) : X × X → R q × q be the mapping defined by ( x, ˜ x ) ↦→ K ϕ ( t ) ( x, ˜ x ) = J ϕ f ( x ; ϕ ( t )) J ϕ f (˜ x ; ϕ ( t )) ⊤ . Then provided conditions (D1)-(D4) hold, we have as p →∞ that

Proof. Let us examine the k, l th term of the q × q matrix given by K p ϕ ( t ) ( x, ˜ x ) -K p ϕ (0) ( x, ˜ x ) for fixed x, ˜ x , and some t ∈ (0 , T ] . The k, l th term is given by (see the work in Appendix C):

Above, we have explicitly made clear the dependence of the parameters on time, e.g. w j ( t ) vs. w j (0) . We aim to show that the quantity above tends to zero as p →∞ . We first prove this holds pointwise, and will consider the red and blue terms one at a time for a fixed x, ˜ x .

̸

First consider the j th summand of the red term. We will bound its absolute value. If k = l , we're done, so assume k = l . We have for any j that

and by the Lipschitz assumption on σ ( · ) and Cauchy-Schwarz, we can bound the quantity above as follows

where again || x || 2 has been absorbed into the constant C by (D1). Using Lemma 4, we can bound all terms above almost surely using || w j (0) || 2 as follows.

Recalling that w j (0) iid ∼ N (0 , I d ) , we have that || w j (0) || 2 and || w j (0) || 2 2 are integrable with expectations denoted µ and ν , respectively. All our work has allowed us to show that

by conditions on D p,t and E p,t from Lemma 4, the strong law of large numbers, and the classic result from analysis that lim n →∞ a n b n = (lim n →∞ a n ) (lim n →∞ b n ) , provided both limits on the right hand side exist. Lastly, we can achieve the same result for the blue term quickly. Because a ij (0) = 0 w.p. 1 by (D2), we have almost surely that

because for all j , we have | a kj | , | a kj | are dominated by || a j || 2 . From here, we have by Lemma 5 that we can bound each term in the sum above by

as p →∞ by similar logic to the above. Together, these results combine to show that | K p ϕ ( t ) ( x, ˜ x ) kl -K p ϕ (0) ( x, ˜ x ) kl | a.s. → 0 as p → ∞ . As k, l were arbitrary k, l ∈ 1 , . . . , q , we have || K p ϕ ( t ) ( x, ˜ x ) -K p ϕ (0) ( x, ˜ x ) || F a.s. → 0 . This establishes pointwise convergence for some fixed t ∈ (0 , T ] . Uniform convergence over all of X × X and all t ∈ (0 , T ] follows easily in this case. Firstly, the numbers D p,t , E p,t , F p,t , and G p,t are monotonic in t , so we can bound uniformly for all t ∈ (0 , T ] by taking t = T in the expressions above. Secondly, in our work above, our bounds on the red and blue terms were independent of the choice of point ( x, ˜ x ) . More precisely, the supremum over x, ˜ x can be accounted for in the bounds easily by observing that sup x, ˜ x ∈X ,t ∈ (0 ,T ] || K p ϕ ( t ) ( x, ˜ x ) -K p ϕ (0) ( x, ˜ x ) || F

can be bounded above by

by the same work as above.

E Kernel Gradient Flow Analysis

We rely on additional regularity conditions outlined below. We will consider the following three flows in our proof of Theorem 1 (for some choice of p ):

where f t is shorthand for f ( · ; ϕ ( t )) . The three flows above can be thought of as corresponding to L P , L F , and a 'lazy' variant, respectively. The flow of h t is 'lazy' because it follows the dynamics of a fixed kernel, the kernel at initialization. The flow of g t also follows a fixed kernel, but the limiting NTK K ∞ instead. The flow of f t is that obtained in practice, where the kernel K p ϕ ( t ) changes continuously as the parameters ϕ ( t ) evolve in time. The flow in h t is used to bound the differences between f t and g t in the proof of Theorem 1. We now enumerate the regularity conditions.

(E1) The functional L F ( f ) satisfies inf f L F ( f ) > -∞ .

(E2) The limiting NTK K ∞ is positive definite (so that the RKHS H with kernel K ∞ is welldefined).

(E3) Under (E1) and (E2), the function f ∗ minimizing L F satisfies || f ∗ || H < ∞ , so that f ∗ ∈ H .

(E4) For any choice of p , we have for all t, x that ℓ ′ ( x ; f t ( x )) , ℓ ′ ( x ; g t ( x )) , and ℓ ′ ( x ; h t ( x )) are bounded by a constant ˜ M .

(E5) The function ℓ is ˜ L -smooth in its second argument, i.e. || ℓ ′ ( x, η 1 ) -ℓ ′ ( x, η 2 ) || ≤ ˜ L || η 1 -η 2 || .

We first prove Lemma 2 from the manuscript.

Proof. Let f ∗ ∈ argmin L F ( f ) , where L F ( f ) is the functional objective. Hereafter L denotes L F . Then L ( f ∗ ) > -∞ by (E1). Then if f t evolves according to the kernel gradient flow (21) above (i.e.,

the flow with kernel K ∞ ), we have (from the chain rule for Fréchet derivatives) that

By definition, ∂f t ∂t = ˙ f t . We also have ∂L ∂f t : h ↦→ E P ( X ) ℓ ′ ( X,f t ( X )) ⊤ h ( X ) . Plugging this in yields

by the positiveness of the kernel K ∞ (from (E2)). Now define ∆ t = 1 2 || f t -f ∗ || 2 H , where H is the vector-valued reproducing kernel Hilbert space corresponding to the kernel K ∞ (see Carmeli et al. (2006) for a detailed review). We let ⟨· , ·⟩ denote the inner product on the space H . It follows that ∂ ∆ t ∂f t : h ↦→⟨ f t -f ∗ , h ⟩ . Then by the chain rule we have

To go from the second to the third line, we used the reproducing property of the vector-valued kernel, the definition of inner product, and the linearity of integration. More precisely, the reproducing property (cf. Eq. (2.2) of Carmeli et al. (2006)) tells us for any functions g, h and fixed x ,

and so the third line results from the second by exchanging the integral and inner product. In the second-to-last line, we used the convexity of ℓ in its second argument (from Lemma 1 of the manuscript). Now consider the Lyapunov functional given by

by the above work because i) t ˙ L ( f t ) ≤ 0 and ii) L ( f t ) -L ( f ∗ ) + ˙ ∆ t ≤ 0 , implying that E ( t ) ≤ E (0) for all t . Evaluating, thus

and so we have that there exists a sufficiently large T such that | L ( f T ) -L ( f ∗ ) | ≤ ϵ as desired.

Using this result and our previous results, we are now able to prove Theorem 1 from the manuscript.

Proof. We will examine the three gradient flows

Differentiating, we have

and establish the result by the triangle inequality, i.e.

where L denotes the functional objective L F . Let ϵ > 0 . The flow in h t will be used to help bound the first term, but we begin with the second term. By Lemma 2, pick T > 0 sufficiently large such that | L ( g T ) -L ( f ∗ ) | ≤ ϵ/ 2 . Fix this T . This provides a suitable bound on the second term.

Turning to the first term, by continuity of L ( f ) in f , there exists δ > 0 such that y ∈ B ( g T , δ ) = ⇒ | L ( y ) -L ( g T ) | ≤ ϵ/ 2 . We will show that there exists P sufficiently large such that p > P implies || f T -g T || ≤ δ almost surely, yielding the desired bound on the first term of the decomposition above. Throughout, || · || denotes the L 2 norm of a function with respect to probability measure P (i.e. the marginal distribution of P ( X ) from our joint model P (Θ , X ) ).

To show that there exists sufficiently large P such that || f T -g T || ≤ δ , we use another application of the triangle inequality

and construct bounds on the two terms on the right hand side using Proposition 2 and Proposition 3. Observe first that by (C2)/(D2), at initialization we have almost surely that f 0 = g 0 = h 0 = 0 . Also note that by continuity of K ∞ (established in Lemma 3) on the compact domain X × X we have sup x, ˜ x || K ∞ ( x, ˜ x ) || 2 < M for some M . Finally, note that by (E5) the function ℓ ′ ( x, η ) is Lipschitz in its second argument with constant ˜ L . Below, we let || · || 2 denote the 2-norm of a vector or matrix, depending on the argument, and || · || F the Frobenius norm of a matrix. For functions, as stated || · || denotes the L 2 norm with respect to measure P ( X ) , i.e. || f || 2 = ∫ f ( X ) ⊤ f ( X ) dP ( X ) . From here, we have

Now, we note the following facts. Firstly, for any kernel K that is uniformly bounded (i.e. || K ( x, y ) || 2 ≤ M for any x, y ), the L 2 norm of the function || K ( · , X ) v || for fixed X,v can be bounded by

Secondly, we have again for any fixed v and X that

since the matrix (spectral) 2-norm is dominated by the Frobenius norm. Plugging these facts into Equation (28) above, we have

By Proposition 2, there thus exists P 1 such that for all p > P 1 we have || g T -h T || ≤ δ 2 almost surely. We proceed nearly identically for the term || h T -f T || . We need only note that for sufficiently large p , say p > P 2 , we can bound K p ϕ (0) uniformly (almost surely) by a constant A > M . To see this, observe that by Proposition 2 we have that there exists almost surely a sufficiently large P such that sup x,y || K ∞ ( x, y ) -K p ϕ (0) ( x, y ) || F < A -M and so by triangle inequality we have

Thereafter,

and we can similarly switch from E P ( X ) || h t ( X ) -f t ( X ) || 2 to the L 2 norm || h t -f t || as above using Jensen's inequality, yielding

Clearly, by the same logic as the above there exists P 3 such that p > P 3 implies ˜ MT sup x,y,t ∈ (0 ,T ] || K p ϕ (0) ( x, y ) -K p ϕ (0) ( x, y ) || exp( A ˜ LT ) ≤ δ/ 2 by Proposition 3. Then for all p > max( P 1 , P 2 , P 3 ) , we have almost surely that || h T -f T || ≤ δ/ 2 . This completes the proof, as in this case we have by the triangle inequality that || f T -g T || ≤ δ and so | L ( f T ) -L ( g T ) | ≤ ϵ/ 2 by construction.

F Experimental Details

Our code is publicly available at https://github.com/declanmcnamara/gcvi_neurips . We used PyTorch (Paszke et al., 2019) for our experiments in accordance with its license, and NVIDIA GeForce RTX 2080 Ti GPUs.

F.1 Toy Example

Recall the generative model for this problem, given by the following:

The variable σ is a hyperparameter of the model that we take to be σ = 0 . 5 . The model is constructed in such a way that x ∈ S 1 satisfies assumptions (C1) and (D1), respectively. One thousand pairs of data points { θ i , x i } 1000 i =1 were generated independently from the model above and fixed as the dataset for which the ground-truth latent parameter values are known.

We constructed scaled, dense single hidden-layer ReLU networks of varying widths, with 2 j neurons for j = 6 , . . . , 12 with the same architecture as in Appendix C and the initialization described in condition (C2). All networks were trained to minimize the expected forward KL objective; stochastic gradients were estimated using batches of 16 independent simulated ( θ, x ) pairs from the generative model, and SGD was performed using the Adam optimizer with a learning rate of ρ = 0 . 0001 . We employ a learning rate scheduler that scales the learning rate as O (1 /I ) , where I denotes the number of iterations. All models were fitted for 200,000 stochastic gradient steps, and the execution time is less than one hour. The natural parameter for the von Mises distribution is parameterized as η = f ( x ; ϕ ) + 0 . 0001 . This small perturbation must be added because f ( · ; ϕ ) = 0 at initialization and because the value of η = 0 lies outside the natural parameter space for this variational family.

For the linearized neural network models, all training settings were the same except for the architecture. For these models, we first constructed neural networks as above for each width to compute the Jacobian evaluated at the initial weights J ϕ ( x ; ϕ 0 ) . Thereafter, the model in ϕ is fixed as

where ϕ, ϕ 0 are flattened vectors of parameters from the neural network architectures. Using this linearized model above, the parameter ϕ is fitted by SGD as above.

The plots in Figure 1 of the manuscript are constructed by evaluating the average negative loglikelihood on the dataset at each iteration, i.e. for the fixed n = 1000 pairs of observations above, we evaluate the finite-sample loss for the expected forward KL divergence. Up to fixed constants, this quantity is given by

where ϕ is the current iterate of the parameters (either the neural network parameters or the flattened vector of parameters of the same size for the linearized model). The horizontal red line in Figure 1 is set at the value -1 n ∑ n i =1 log p ( θ i | x i ) , where p denotes the exact posterior distribution (computed using numerical integration over a fine grid of evaluation points).

F.2 Label Switching in Amortized Clustering

Recall the amortized clustering model from the manuscript, given by

We take σ = 1 , τ = 0 . 1 and artificially fix s = 100 as our realization from the prior on S . Thereafter, we generate N = 1000 independent, identically distributed samples { x i } 1000 i =1 from the model above, conditional on observing S = 100 . The other non-random hyperparameters are the cluster centers µ = [ -20 , -10 , 0 , 10 , 20] ⊤ , as well as the mixture weights for the mixture of Gaussians, which are uniform, i.e. p j = 1 d for all j . These draws constitute the 'data' for this problem for which inference on S is desired conditionally on the draws x 1 , . . . , x 1000 ; separately, we seek to infer the random variable Z conditional on the data as well. Although Z and S are closely related, we impose a mean-field variational family q ( S, Z | X 1 = x 1 , . . . , X N = x n ) = q ( S | X 1 = x 1 , . . . , X N = x N ) q ( Z | X 1 = x 1 , . . . , X N = x N ) for ease, with both components taken to be isotropic Gaussian distributions.

We use two different parameterizations for each of q ( S ) and q ( Z ) . The mean parameterization fixes the variance at unity and aims to learn only the mean - in this case, the natural parameter of this family of distributions is simply the mean, so the network parameterizes the variational means µ S , µ Z for each of these distributions. We also use a natural parameterization with an unknown mean and variance, and in this case, the natural parameterizations η S , η Z are output by the networks,

which we denote f ( · ; ϕ ) and f ( · ; ψ ) for S and Z , respectively. Both amortized encoder networks f ( · , ϕ ) , f ( · ; ψ ) take the entire set of points { x 1 , . . . , x 1000 } as input. The architecture should thus be permutation-invariant, as the order of these points does not matter for inference on the latent quantities of interest. We accomplish this by using a set transformer architecture for the networks, which achieves permutation invariance using self-attention blocks (Lee et al., 2019).

Parameters are fitted to minimize L P using SGD with stochastic gradients estimated using simulated draws of the same size as the observed data. We use a learning rate of 0.001 for fitting both q ( Z ) and q ( S ) , with schedules that decrease these as training progresses across 50,000 total gradient steps. We perform one hundred replicates of this experiment across different random seeds, each time generating a new dataset and refitting the networks. Accordingly, there is a high computational cost: using 10 parallel processes, the experiment takes about 8 hours to run. We compare L P -based minimization to fitting both networks to maximize an evidence lower bound (ELBO), using the IWBO with K = 10 importance samples as the objective, but otherwise keeping all other hyperparameters the same. Figure 2 in the manuscript plots the mode of q ( S ; f ( x 1 , . . . , x N ; ϕ )) across one hundred replicates of the experiment with different random seeds. The estimate should be approximately 100 to match the ground truth, and both ELBO-based and L P -based training perform well.

For inference on Z , we extract the mode ˆ Z ∈ R 5 as the point estimate and compute the ℓ 1 distance to the true draw Z = z for our dataset (which is known a priori because the data were generated synthetically). ELBO-based training illustrates label-switching behavior, converging to a vector which is a permutation of the entries of the true draw z , resulting in a large ℓ 1 distance. L P -based fitting does not suffer from this pathological behavior and instead converges rapidly to the global optimum.

F.3 Rotated MNIST

We use the MNIST dataset, freely available from torchvision under the BSD-3 License 1 to fit a GAN generative model of MNIST digits. The simplistic GAN uses dense networks for both the discriminator and generator and is fit to the data with binary cross-entropy loss. Training was stopped when the generator produced realistic images, sufficient for our problem (see Figure 3).

The variational distribution q (Θ; f ( x 1 , . . . , x N ; ϕ )) on the shared rotation angle is taken to be von Mises for minimization of L P . We parameterize q using its natural parameterization as in Appendix F.1. The encoder for Θ uses three Conv2D layers to increase the number of channels to 64 (with a kernel size of 3 and a stride length of 2), followed by a flattening layer and a three-layer dense network with LeakyReLU activations. Permutation invariance is imposed by a naive averaging step across all observations in the dataset { x i } N i =1 that are provided to the network. The neural network architecture for this example thus differs significantly from the simple two-layer ReLU network, yet still demonstrates the same global convergence behavior.

We compare to a semi-amortized approach for maximizing the IWBO for this example. To compute the IWBO, the practitioner must posit a variational distribution on z i for each image x i . This is because only the full joint likelihood p ( θ, { z i } N i =1 , { x i } N i =1 ) is readily available. Accordingly, for this method we construct a second network f ( · ; ψ ) used to construct a variational distribution q ( Z i ; f ( X i ; ψ )) on Z i for a given X i . This network is amortized over all images, yielding q ( Z i ; f ( x i ; ψ )) for any image x i in the dataset { x i } N i =1 . We parameterize these distributions on the latent representation as multivariate (isotropic) Gaussian with mean and log-scale parameters for each dimension the outputs of an encoder network. The architecture for this encoder consists of three Conv2D layers, each with kernel size 3 and stride of length 2. Across the three layers we increase the number of channels from a single channel (the input) up to 8 channels. After the convolutional layers, we perform a flattening layer followed a 2-layer dense network with ReLU activations. As the latent dimension for the data is known to be 64, the outputs of the encoder are 128-dimensional to parameterize the mean and log-scale across each dimension.

As there is only one rotation angle of interest we directly maximize the joint likelihood p ( θ, { z i } N i =1 , { x i } N i =1 ) in θ by targeting the IWBO. Conceptually, this approach is the same as placing a point mass variational family on Θ , i.e., we simply initialize θ 0 prior to training and update its value directly to maximize the IWBO. This approach thus fits the parameters { ψ, θ 0 } of q ( Z i ; f ( x i ; ψ )) and δ θ 0 , the distributions on the latent representations and the angle, respectively.

These parameters are optimized using the Adam optimizer with initial learning rate 0.01 and square summable learning rate schedule. The angle parameter value is initialized at a variety of angles in different trials to produce Figure 4.

The advantageous marginalization properties of L P -based fitting allow it to perform inference on θ without performing inference on the latent representations Z i ; thus, the distributions q ( Z i ; f ( x i ; ψ )) need not be fitted when using this approach. We use the Adam optimizer for all parameter updates with initial learning rate of 0.0001 (and a square summable learning rate schedule) for L P -based training, which only fits the parameters ϕ of q (Θ; f ( x 1 , . . . , x N ; ϕ )) .

Minimization of L P trains for 100,000 gradient steps, although convergence is rapid, and so trajectories are truncated to produce Figure 4. Runs were performed in parallel across multiple processes for both methods and completed in less than one hour.

F.4 Local Optima vs. Global Optima

For this experiment, the setting above is modified slightly to remove the nuisance latent variables { Z i } and put ELBO-based optimization on more equal footing with expected forward KL minimization.

The data are generated as follows: for i = 1 , . . . , 50 , we fix latents z i iid ∼ p ( z ) , and θ true = 260 degrees is set artificially. The digits x 0 i are computed via the generative adversarial network and thereafter fixed for the rest of the problem (the superscript denotes a rotation angle of zero degrees). With the unrotated digits fixed, the only random variables in the model are Θ and the rotated versions of the digits.

Conditioning on the observed data { x i } 50 i =1 , Θ is thus the only latent to be inferred. In this setting, the ELBO can compute the full likelihood function p ( θ, { x i } 50 i =1 ) for any value of Θ = θ (the same Uniform (0 , 2 π ) is placed on the angle). The simulation-based approach of expected forward KL minimization can proceed as usual: we draw Θ ∼ Uniform (0 , 2 π ) and simulated digits are obtained by rotating the digits according to the drawn angle and adding noise.

The variational distribution, as stated in the text, is Gaussian with fixed variance instead of the von Mises distribution used in the local optima section. This is necessary for a one-to-one comparison between the two methods, as the von Mises distribution is not reparameterizable for ELBO-based training. The mean of the Gaussian is parameterized via the same neural network architecture above for fitting by both the ELBO and the expected forward KL. We use the Adam optimizer with learning rate 0.001 for 10,000 gradient steps.

For computing some of the metrics in Figure 7, we rely on approximations as these quantities are difficult to compute exactly. Letting x true denote the observed data { x i } 50 i =1 , the forward KL divergence KL ( p ( θ | x true ) || q ( θ | x true )) is estimated using self-normalized importance sampling with K = 1 , 000 samples drawn from the prior as a proposal. The reverse KL divergence KL ( q ( θ | x true )) || p ( θ | x true ) is computed as the difference log p ( x ) -ELBO ( q ) ; the log-evidence is approximated via the importance weighted bound (IWBO) with K = 1 , 000 .

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract and introduction outline the main contributions of our work relative to preceding work, making clear that our results are primarily theoretical while still having implications for practitioners. We make clear that for our results to hold, regularity conditions must be satisfied.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: In Section 6, we articulate the main conditions necessary for results to hold. These are enumerated even more explicitly in Appendices C, D, and E. We acknowledge in Section 6 that in practice these assumptions may not be met, and in our experiments we anticipate several settings where the assumptions are not met exactly, but our results still approximately hold.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We include several lemmas and/or theorems in the main text. Due to space constraints, we do not enumerate all assumptions in the statements, but we do reference the assumptions numerically and list them explicitly in the relevant appendices. Our appendices contain detailed proofs of all lemmas and theorems provided. As the proof of Theorem 1 is rather involved, we devote a large portion of Section 4 to sketching the proof and providing intuition for the reader.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide detailed information in Appendix F to allow for reproduction of results. If accepted, we will publicly release code to reproduce our experiments exactly.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: In the appendix, we include a link to the GitHub repository with all code necessary to reproduce experiments.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide detailed information on all experiments in-text, with additional implementation details in the relevant appendices. We also plan to publicly release code for this paper on GitHub, where experiments can be replicated and various details (e.g., network architectures) can be examined to understand the results.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Where reasonable, we show statistical significance by aggregating runs across different random seeds. We do this in Section 5.2. In our other experiments, which illustrate optimization trajectories, we omit error bars as this type of data does not readily admit uncertainty estimates.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: In Appendix F, we detail the hardware used (GPUs), the memory requirements, and the approximate runtime for each of our experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have read the Code of Ethics and verify that the paper complies.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Our work is primarily theoretical in nature, although we hope practitioners will utilize it to improve inference in probabilistic modeling. While the actions of individual practitioners have the potential for positive or negative societal impacts, our work is several stages removed from direct paths to applications yielding undesirable outcomes for society.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We do not release any models with this paper that carry risk for misuse.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We credit and cite PyTorch in our appendices, used with license to implement experiments. Our other main asset used is the MNIST dataset, whose license we name in Appendix F.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: No new assets are released with this work.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: There are no human participants or crowdsourcing used in this work.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: There are no human participants or crowdsourcing used in this work.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
The paper studies an alternative objective for variational inference, the expected forward KL divergence. Under some technical assumptions, convexity is shown, which facilitates global optimization. Moreover, a tractable surrogate objective is presented and it is shown that the approximation error can be made arbitrarily small. Finally, an experimental evaluation suggests that global convergence may even occur when the technical assumptions are violated.

STRENGTHS:
The paper addresses a very interesting problem and the methodology is clever. The technical development seems very careful.

WEAKNESSES:
The paper is very technical and dense. It would be helpful to provide more intuition.

QUESTIONS:
1. How restrictive are the assumptions?
2. Can you explain the intuition more how your approach achieves convexity?
3. I understand you presented an example where the expected forward KL outperforms the standard ELBO. How does optimization of these two objectives compare more generally in practice? I imagine there are situations where optimizing the ELBO still yields better approximations to the actual posterior. Do you have any insights into these more practical aspects?

LIMITATIONS:
Technical limitations are stated but it would be helpful to discuss more prominently how restrictive the assumptions are.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-06-25

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work addresses a common problem of non-convexity while approximating posterior distributions using variational inference. Although using ELBO as a variational objective is popular, this paper considers a forward KL(FKL) divergence objective. Its first main contribution is to show that when the variational family belongs to the exponential family of distributions, the FKL objective is strictly convex in the variational parameters. In particular, the paper parameterizes the exponential variational family by a neural network. Their second contribution is to show that under certain conditions on the neural architecture, the solution of the FKL objective is only \\e-suboptimal to the global functional solution of the  FKL objective. The paper also demonstrates the efficacy of the proposed with the help of some experiments.\n\nSTRENGTHS:\n1. Using expected forward KL with the exponential family is simple and interesting when dealing with ELBO's non-convexity. \n2. The idea of using an exception with P(X) also seems interesting, as it simplifies the efforts of sampling from the posterior while using the FKL objective. Is this idea novel? If yes, the authors should signpost it.\n3. The connection between NTK and the gradient dynamics of the functional objective is also interesting and novel.\n4. The paper is well written.\n\nWEAKNESSES:\n1. I think the authors underestimate their findings in Lemma 1, i.e., the convexity of the variational objective when forward KL is used with the exponential family. Can they demonstrate its efficacy on a toy problem that doesn’t use any neural network? \n2. It is slightly misleading to say that the FKL objective with NTK finds global optima. Since the NN is highly non-convex, the result only shows that the local solution is close to the global one and that too in the limit of NN width. \n3. I don't understand how Lemma 1 shows that L_F is strictly convex. I believe it just proves convexity.\n\nQUESTIONS:\nMinor comments. \nL45: L_p to L_P\nL771: What does \\eta over equality mean?\nL843: a,\\in - remove\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-04"}

comments: SUMMARY:
The authors established the global convergence of a particular VI method, which is based on forward KL and variational family parameterized by neural network. The analysis techniques are extended from widely studied NTK for two layer neural network. The authors also conducted experiments to verify the theoretical results.

STRENGTHS:
The paper provided rigorous global convergence guarantee of amortized VI based on forward KL with exponential variational family and NTK. Although the results were asymptotic in the width of neural network, various numerical experiments were conducted to show that the theoretical results seem to hold with finite width.

WEAKNESSES:
The major concern is about novelty. What is the difference between the analysis in this paper and existing analysis of two layer neural network in supervised learning setting? Since the authors assumed that the variational family is exponential family, the optimization objective terms out to be a convex function of the output of two layer neural network, which has been widely studied. Although this is new in the literature of VI, the techniques used are not novel and can be easily extended from supervised learning setting to VI based on forward KL.

QUESTIONS:
The setting for variational inference in this paper assumes that data can be simulated during training (or with known latent variables), which is a bit different from standard VI where the data is fixed. How useful would it be for more general VI problems where the forward KL seems to be challenging to use?

LIMITATIONS:
The limitations are pointed out in the paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors study convergence of forward-KL variational inference in the neural posterior estimation (NPE) setting with an exponential family variational distribution, whose natural parameters are produced by a neural network. For this setting, it is known that the forward-KL is convex in the natural parameters of the variational distribution. The authors show how this extends to the NPE setting, i.e. expected forward-KL, by linearity of the outer expectation and establish the existence of a global minimizer for the functional NPE objective. The authors further show that the limiting neural tangent kernel (NTK) of the network can be used to construct a Reproducing Kernel Hilbert space in which optimization via kernel gradient flow converges to the unique minimizer. Finally the authors show how this finding can be extended to the parametric setting by noting that, for certain network architectures, the parametric NTK tends to the limiting NTK as the width of the network tends to infinity. Interestingly, the authors demonstrate empirically that in practice infinite width might not be required to converge to a region close to the global minimizer.\n\nSTRENGTHS:\nThe results are novel and theoretically interesting to the variational inference community. I found the paper, given its theoretical nature, well written and relatively easy to follow along.\n\nWEAKNESSES:\nThe result stated in Lemma 1 is well known (unless I am missing something), please consider citing a textbook or relevant review paper for the fact that the forward-KL divergence is convex in the natural parameters of the variational distribution in the exponential family case.\n\nThe experiments are not directly compared to ELBO/IWBO based optimization, hence, it is unclear how big the actual gap in performance is. In settings where the true posterior is a member of the variational family, the minimizer of the forward-KL divergence also minimizes the reverse-KL and hence maximizes the ELBO and vice-versa. A plot that reports the symmetric-KL over training for both the ELBO and FKL objective would be insightful. In the end, if FKL-optimization indeed finds the minimizer (or close to it), the learned variational distribution should outperform ELBO/IWAE-optimization no matter which divergence (or corresponding bound) is used for comparison.\n\n--- Edit ---\n\nThe authors addressed my concerns by adding appropriate references and providing additional results that demonstrate the utility of their approach. Consequently, I have raised my score to an Accept (7).\n\nQUESTIONS:\n-\n\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}