Globally Convergent Variational Inference
Authors: Declan McNamara
Jackson Loper
Jeffrey Regier
Venue: NeurIPS 2024 poster (2024)
Keywords: forward KL divergence; neural posterior estimation; neural tangent kernel; convex optimization
Primary area: probabilistic_methods
TL;DR: Utilizing convexity of the forward KL divergence, we establish a global convergence result for fitting an amortized variational posterior by analyzing neural tangent kernel (NTK) dynamics in the large-width setting.

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
Declan McNamara Jackson Loper Jeffrey Regier Department of Statisties University of Michigan
{declan, jaloper,
Abstract
In variational inference (VI), approximation of the posterior distribution is selected from a family of distributions through numerical optimization. With the most common variational objective function, known the evidence lower bound (ELBO), only convergence to 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 vari- ational distribution that is parameterized by a neural network. Our convergence result relies on the neural tangent kernel (NTK) to characterize the gradient dynam- ies 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 unique solution in a reproducing kernel Hilbert space (RKHS). Then, we show that gradient descent dynamics in function space converge this unique function. In ablation studies and practical problems, we demonstrate that our results explain the behavior NPE in non-asymptotic finite-neurod settings, and show that NPE outperforms ELBO-based optimization, which often converges to shallow local optima.
1 Introduction
In variational inference (VI), the parameters v an approximation to the posterior Q 17) are selected to optimize an objective function, typically the evidence lower bound (ELBO) (Blei et al. 2017). However, the ELBO generally nonconvex in 1). even for simple variational families such as the family of Gaussian distributions. and so only convergence to local optimum of the ELBO can be guaranteed (Ghadimi and Lan 2015 Ranganath et al., 2014 Hoofman et al., 2013). As the number of such optima and the degree of suboptimality of each 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 KI, divergence:
Lp(d) Er(x) [P(O) X) || f(x:)))
(I)
Here, P(X) denotes marginal of the model and P(e X) denotes the posterior. For each x E X,
the approximation Q(O: n)) to P(O X 2) is indexed by the distributional parameters 1j E ) C TR9
which are themselves the output of a neural network (=; Φ) with weights Φ E क. Approximating posterior distribution by minimizing Lp has 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 [015]] Le et al. 2019). Minimization of this objective straightforward: computing unbiased gradients requires only sampling 8, x P(O, X) from the
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 KI. minimization is more widely applicable than ELBOO-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 KI;; for the non-amortized problem in which a single observation z 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, 2013 Le et al.) 2019; Owen 2013). Our analysis considers a functional form of variational objective L Lp. given by
Lf(f) ==
[P(O) x) ||
(2)
where Lp: H + - is defined over a general reproducing kernel Hilbert space of functions H. We
refer to (I) the "parametric objective", its argument is the parameters ó E o and refer to (
as the "functional objective" as its argument is a function ʃ E H. These objectives are closely related:
under given network parameterization, provided ((: H. we have (Ø) = (p([(; $)). The
objective Lp has been considered in several related works (see Section 2. The formulation of L.
and the analysis of its minimizer relative to those of Lp is our main contribution.
We first demonstrate strict convexity of the functional objective LE when Q is parameterized as an exponential family distribution (Section 3. This implies the existence of unique global optimizer f* of Lp for a large class of variational families. Afterward, we analyze kernel gradient flow dynamics using the neural tangent kernel to show that minimization of Lp results (asymptotically) in an empirical mapping ff that is at most e-suboptimal relative to f", provided a sufficiently flexible neural network is used parameterize I (Section[4) Together, these results imply that in the infinite-width limit, optimization of Lp 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 (c.g., an exponential family) (Section 6). Our proof' further assume two-layer ReLU network architecture, but we conjecture that this assumption can be relaxed, and our experiments (Section demonstrate global convergence for wide variety of architectures. We illustrate that the minimization of Lp converges to 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 likelihood- based 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 single observation 2o, 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 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 not amenable to global minimization because it is nonconvex (Domke) 2020 Domke et al., 2023 Kim et a 2023). Liu et al. (023) demonstrated nonconvexity of the ELBO in the context of small-object detection, and showed empirically that the expected forward KI 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| 20018) 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 neural network and are shared among observations.
2

---

2 Background
2.1 The Expected Forward KL Divergence
The expected forward KL objective is equivalent the sleep-phase objective of Reweighted Wake- Sleep (RWS) (Bornschein and Bengio 2015), and to the objective optimized by forward amortized variational inference (FAVI) (Ambrogioni et 2019). It is also a special case of the thermodynamic variational objective (TVO) (Masrani et al. (019)) Similar objectives have been referred to as neural
posterior estimation (NPE) (Papamakarios and Murray 2016 Papamakarios al. 2019). though in these works the prior distribution, and thus the marginal P(X) ), mutates during training.
Objectives based the forward KL divergence generally result in variational posteriors that 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 Lp is straightforward. The outer expecta- tion over P(X) allows for gradients to be computed as
V VLL [P(e] X) || f(;; d)) = BM(x))) V Vologg(0:
where g is the density of Q with respect to Lebesgue measure (see Appendix for details). Rewriting the left-hand side expectation over P(6) and P(X ()) by Bayes' rule illustrates that samples can be drawn from this model by ancestral sampling of e followed by X.
Other methods targeting the forward KL, such as the wake-phase of RWS, often optimize over different expectation, typically Ex~D (P. X) Q(O; 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(O
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(:;)) : ó E }). Let f(x,f(x) fx)) denote a general real-valued loss function and consider selecting the parameters & to minimize Ep(x) [(X; σ)), where P(X) is 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 minimize the above objective. Continuous- time dynamics are used in the formulation; @(t) and f(= ((t)) are defined for continuous time The parameters & thus follow the ODE
d(t) -V.E (X
(3)
Here, ó denotes the derivative with respect to f, and by the chain rule, the function values f(z; ((t)) evolve via
f(:; EP(x) Jof(æ; (())) 6(0) ((x)(()))))
NTK
We define (X, ((x.f(x)) V) to simplify the notation. The product Jacobians above known as the neural tangent kernel (NTK):
Ka(x,d') (o
The seminal work of Jacot et al. (2018) defined and studied this kernel and established the convergence of Kg to a limiting kernel for certain neural network architectures as the layer width grows large.
2.3 Vector- Vector-Valued Reproducing Kernel Hilbert Spaces
Most existing -TK--ased analyses consider neural networks with scalar outputs and squared emror loss. Instead, we consider neural networks with multivariate outputs to accommodate the multidimensiooall
3

---

distributional parameter which parameterizes our variational distribution Q(@: m)). Furthermore, we consider the objective functions Lp and Lp. 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 the following, we summarize the key properties.
Recall that 7 = f(:: and 7) E IK*, An R8--valued kernel on X is map K x X + R2xy, The neural tangent kernel (4) is precisely such a kernel. If the IR9valued kernel K is positive definite, then K defines unique Hilbert space of functions H whose elements are maps from to R (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 R9-valued RKHS, the reproducing property takes on more general form. For any x E X, f E H. and v R2
f(x) v == u)) = (f()) K(,
where (x h is the inner product of the Hilbert space H. For any fixed æ E X' and e E RR,K(( K(,, 2)) is a function mapping from X' I88, 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 Ly given in Equation (2). We fix an RKHS H which to minimize L for now, specializing the particular choice of H based on the neural tangent kernel subsequently. Let f(x, [(x)) = KT. [P(@) X = x) || Q(;:
The functional
will this form subsequently for neural langent kernel analysis. Our first result shows that targeting LF is highly desirable theoretically: Le admits unique global minimizer if the variational family is an exponential family, as is common practice in VI.
LE then has the form IFF E Exxx (XX):
Lemma Suppose that Q(9:n) n) is an exponential family distribution minimal representation with natural parameters T, sufficient statistics T(0), and density g(0:n) with respect to Lebesgue measure A()). Then, for any observation E X C Rdd the loss function
((r,n) = KL P X = 2²) Q(9:n)
strictly convex in n. provided that P(O X x) < Q(::n) < ((0) for all n E CR.
A proof of Lemma 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 I shows the strict convexity of the function in This implies the strict convexity of the functional ( Er(x;e(X, f()) in f by the linearity of expectation, which in turn implies the existence of at most one global minimizer.
Corollary Suppose that Q(0; n) is an exponential family distribution. Then, under the conditions of Lemma 7 the functional objective
Lr(f) == KL [P(O) x) (())
is strictly convex in f. Consequently, the set of global minimizers of Le is either a singleton set or empty.
We will assume that the set of global minimizers is nonempty (so that the minimization of LF is well-posed) and let f denote the global minimizer. We also assume that || so that f E 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 Lx()).
Whereas Lemma] I establishes the convexity of the (non mon-aaorttzedd forward KL divergence, Corollary I esttablishs the convexity Lr, an amortized objective, in function space. The convexity of the Functional objective I. 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 Lp rather than Ly, optimizing the latter the infinite- dimensional space H directly is not tractable. Thus, in the second phase of our analysis, we consider convergence to f by minimizing the parametric objective Lp with gradient descent. We define & across continuous time as in Equation (3). Continuous-time dynamics simplify theoretical analysis; SGD with unbiased gradients follows (noisy) Euler discretization of the continuous ODE (Santambrogio 2017 Yang et al. 2021). Considering X ~ P(X) for the outer expectation in both Ip and 1p key this context: this choice enables unbiased stochastic gradient estimation for Lp (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 the trajectories of the parametric objective Lp throughout its minimization initially seems infeasible: the argument of this objective is the neural network parameters 6, 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 Lp. We bridge the divide between the minimizers of the convex functional LF and the nonconvex objective Lp using the limiting kernel, and show that in the large-width limit, the optimization path of Lp converges arbitrarily close f", the unique minimizer of the functional objective L.p.
Theorem 1. Consider the width-p scaled 2-layer ReLU network, evolving via the flow
K f(()) -Ep( K (X,
where fr denotes f(., k(t). Let f denote the unique minimizer of L = LF from Lemma 1 and fix e O. Then, under conditions (CI),C4), (D4), and (EI))E5) there exists T > 0 such that almost surely
lim T.(TT) s +(F)++. p-000
Regularity conditions (C4)-(C4), (D1)- and (EI)-(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 prove the results the network width p tends to infinity. Our results may also extended to multilayer perceptrons with other activation functions. Below, we briefly sketch the key ingredients needed to prove Theorem
Recall the NTK KP from Equation (4), where we now let p denote the network width. For certain neural network architectures, Jacot et al. (208)) show that as the network width p tends to infinity, the neural tangent kernel becomes stable and tends (pointwise) towards fixed, positive-defintte limiting neural tangent kernel Koo-
Under suitable positivity conditions on the limiting kernel, we take the domain H. of LF to be the RKHS with kernel Kox (Section 2.3). Because LF has a unique minimizer f*, under mild conditions on Koc. f may be characterized as the solution obtained by following kernel gradient flow dynamics in H. that is, the ODE given by
f((x) (x,
In other words, starting from some function fo. following the limiting NTK gradient flow dynamics above minimizes the functional objective Lp for sufficiently large T.
Lemma 2. Let I denote the minimizer of LF from Lemma ! and > O. Fix So. and let Koo denote the limiting neural langent kernel. Let Jo evolve according to the dynamics
j((x) = KoP( (X, (X,
Suppose that the conditions of Lemma 7 and (EI)-(E3) hold. Then, there exists T > such that L(fr) I L(f*) + where L is the loss functional of Lg.
Appendix F enumerates regularity conditions (E1)- and provides proof of Lemma 2 The characterization of f Lemma clarifies how the analysis of the parametric objective T.p will proceed. The gradient flow Lp causes the network function to similarly evolve according to a kerel
5

---

gradient flow via the empirical neural tangent kernel, that is.
f(z; - -P( K 200 ((X
as derived in Section 2.2 Comparison of the minimizers of Lp and Lp can be accomplished by comparing the two gradient flows above, i.e. kernel gradient flow dynamies that follow ( and Kac. respectively. As KP ((() - Kx (Appendix D). these trajectories should not differ greatly: for any fixed T, the functions obtained by following the kernel gradient dynamics with K and K can be made arbitrarily close to one another, provided p sufficiently large. The proof of Theorem first selects a I using Lemma 2 and then bounds the difference in the trajectories on [0, 7" for sufficiently large width p by convergence of the kernels KP - Ko Our proof differs from previous results in that it relies on uniform convergence of kernels (cf. Appendices C and DD. enabling the analysis of population quantities such as ( ((x, ((X)).
Theorem ! proves convergence to an e-neighborhood of the global solution when optimizing Lp despite the highly nonconvex nature of this optimization problem in the network parameters ø. For sufficiently flexible network architectures, optimization of Lp thus behaves similarly to that of Lp, which we have shown is a convex problem in the function space H. in Section
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 demonnsttss that ninite-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 D 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, behaves much like its lincarization around the initial weights Φ0:
f(x; 6) Rs f(a; 0o) + Jof(x; (o) (0)) Φ0)-
Liu et al. (2020) prove that equality holds exactly in the equation above if and only if f(x; Φ) has constant tangent kernel (i.e., J. Therefore, similarity between f and its linearization indicates that the asymptotic regime of the limiting NTK approximately achieved. Note that even if f is linear in Φ, as in the above expression, it may still be highly nonlinear in æ.
We consider a toy example for which ||œ||2 1. The generative model first draws rotation angle @ uniformly between 0 and 2n, and then a rotation perturbation Z ~ oO), where we take 0.5. Conditional on θ and Z, the data 2 is deterministic: == [cos + 2), sin(@ + =)]" This construction ensures that the data lic on the sphere St c IR? which guarantees the positivity of the limiting NTK for certain architectures (Jacot et al. 2018). D. We aim to infer O given a realization
marginalizing over the nuisance latent variable 2. Our variational family QO;; f(x) is a von Mises distribution, whose support is the interval [0, 2n). This family is exponential family distribution, allowing for the application of Lemmal] [ 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 the network width grows. The network outputs J(x; Φ) parameterize the natural parameter
We fit the neural network [(;; σ)) in two ways. First, we use SGD to fit the network parameters ø to minimize Lp. Second, we fit the lincarization fin σ) = f(r; do) + Jof(x: do) - do) 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 s; ....... 2% from the generative model with underlying ground-truth latent parameter values or...., , Bx and evaluate the held-out negative log- ligeliillooodd (NLL). v EN logg (0) f(x); Ø)), for both functions: f(x; Φ) and fin (z; Φ). Figure 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 grows; for p 1024, the two are nearly identical, providing evidence that the asymptotic regime is achieved.
0.780
0,735
0.756
wilth
Watth
64
0.730
1084
65-7
20--nn
1606-1n
0770
0.725
0.725
0,70
0.720
6.720
8716
0.73
6.75
0710
0.710
0.710
1.5
0,790
0.5
1.0
2.1
3.0
2.00
Goaliaal Step
X10
Gualion Stop
Ghulkai Stop
Figure 1: Negative ligeliiel across gradient steps, for network widths 64, 256, and 1024 neurons. NLL for the exact posterior is denoted by the red line.
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 [ assumes 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 = (x1,.... In,, I2 E 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 shift parameter S ~ (0, 1002) followed by
z | (S =) ~ (IR1 + s,.... ................................ Pd + ,
d
X₁ (z 2) itd Σ [e,.(o)
In the above, o², 72 IRd and p E Idd are known variances, locations, and proportions, respectively. The global parameter == s E IR shifts the d locations ,i...., Ma to μu + 8,.... M4 + 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 RRd via the prior, Finally, the data are drawn independently from mixture of d univariate Gaussians with centers Z,, i ....., 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 to generate n = 1000 independent realizations X = x from the generative model.
0.5 and 0.1, and artificially fix the shift as S 100
Inferring S should be straightforward because the vector is known and fixed, ensuring that the joint likelihood p(z,s) s) (marginalizing over 2) is unimodal in s. However, inference on may be difficult for likelihood-based methods because the order of the entries of Z matters: the joint density p(x, 2, 8) * p(æ, TT=), 8) for a permutation n, even though p(æ =)) = p(x x())). "Label switching" can thus pose a significant obstacle for likelihood-basee methods, as any permutation of the cluster centers Z will still cxplain the data well. This problem formulation thus results in a likelihood, and hence posterior density, with many local optima but 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 the minimization of Lp. We take the variational distributions q(S: fi (I: Φ1)) and gZ: f2(;; Φ2)) to be Gaussian. We fit the networks fi and f2. Due to the exchangeability of the observations x {xi....; wn). we parameterize each as permutation-invarin
neural networks, fitting both Φi and Φ2 to either minimize I.p or to maximize the ELBOO (see Appendix FF). We perform 100 replications of this experiment across different random seeds, and consider two different parameterizations of the Gaussian variational
7

---

distribution: a mean-only -anl 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 D holds.
Figure 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--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 4 distance || Z between the variational mode 2 and the true latent draw Z. Optimization of the ELBO converges to local optimum that is permutation of the entries of Z. resulting in large L distance average. Minimization of Lp, on the other hand, converges to the global optimum without any label switching. Table I indicaaess the degree of label switching, showing the proportion of trials in which the entries Z were correctly ordered.
ELBO
LP
ELBO
Mau Panuarc cietii Natrrll Passeeer satiin
Mean
0.03
1.00
Natural
0.02
1.00
Table 1:
Proportion out of one hundred repli- cates where posterior mode of (Z; [2(2; 02)) was a vector in increasing order.
ELBO
LP
Mean
77.0 (45.2)
1.8 (2.3)
Natural
86.0 (26.9)
2.9 (0.8)
Table 2: Average l, distance ||ž Z|I1 (and
Figure 2: Mode of g(S; fa d1)) across ex- perimental replicates.
std. deviations) across one hundred replicates.
5.3 Rotated MNIST Digits
We consider the task of inferring a shared rotation angle θ for set of N MNIST digits. The generative model is as follows. The rotation angle is drawn θ Uniform(0, 2T), and for all i E [N] a noise term is drawn as Z; ~ (N I²I)). Finally, each image is drawn as
(8) Here, RotateMNIST is fitted ahead of time and fixed throughout this experiment. Given latent representation z and angle 0, RotateMNIST returns a 28 x 28 MNIST digit image rotated counter- clockwise by 9 degrees. (See Appendix [ for additional experimental details.) We aim to fit the variational posterior q(;; f (I11.... Z2: o)), implicitly marginalizing over the nuisance latent vari- ables Z1, ....... Z The true data { generated from the above model, with N 1000 digits and an underlying counterclockwise rotation angle of f = 260 degrees. We provide visualizations of some of these digits in Figure 3 The variational distribution @G:; f ... NN: ; ②)) is taken to be von Mises distribution with natural parameterization, as in Section [.I] although for this example the architecture of the encoder network makes ff invariant to permutations of the inputs {z; f==1 to reflect the exchangeability of the data. We fit / to minimize Lp using 100,000 iterations of SGD.
X; (Z
θ 0) M (RotateMNIST(2, +),+;) i E [N].
We compare to fitting 0 to directly maximize the likelihood of the data p(0, (x. 1 As this quantity is intractable. we maximize the importance-weit weighted bound (IWBO), which is a generalization of the ELBO (Burda et al. 2016)). by maximizing the joint likelihood p(0, NN (2.) in 0 while fitting a Gaussian variational distribution on the variables Z.
The likelihood function is multimodal in 0 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 multimodal posterior distribution on O Likelihood-based fitting procedures, such as max- imizing the IWBO, often get stuck in these shallow local optima, while fitting ʃ to minimize the parametric objective I.P finds a unique global solution. Figure shows estimates of the angle conditional on the data {2, NN during training with the IWBO objective, with 0 initialized to
8

---

7777
2
O b
7
6 37777 1 22
3
350
3
4
3 8
7 4
Angle (Degrees)
250
6
0 4
47
150
1/1
28
I 6
4
" 5
S
50
Figure 3: 100 of the
Gradient Step
1000 data observations with
counterclockwise rotation
Figure 4: Estimate of angle 0 across gradient steps, with fitting performed to maximize the IWBO.
260 degrees.
variety of values. For some initializations, the IWBO optimization converges quickly near the correct value of 260 degrees, but in many others, it converges to shallow local optimum.
We perform the same routine, fitting g (fi...... NN: 6)) to minimize the expected forward KL divergence Lp. Figure P shows that across a variety of initializations of the angles, this approach always converges
unique solution and fits the posterior mode the correct value of 260 degrees. Lp minimization converges rapidly, and so steps to show the various trajectories among initializations.
in on the initial few thousand gradient
250
270
Augle (Degrees)
30
200
18
200
20
90300 Grarlient Step
SC300
300003
2200 2003
1104
1000
3750
2000
Grodicut Step
Figure 5: Mode of (0;; (x1...... NN; ②)) across training (starting at different initializa- tions) when minimizing objective Lp.
Figure Zoomed-in trajectories across the first 2000 gradient steps, showing similar estimates regardless of initialization.
5.4 Local Optima 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; is drawn from the model x; ~ N (Rotate (x.. 0), 72) for 1,..., 50 with angle set to 6 260 degrees. The unrotated digits xi are fixed a priori, eliminating setting. This allows to directly perform BI-BO-based variational inference on θ, instead of merely maximizing the likelihood (or a bound thereof) as in Section 5.3
nuisance latent variables in this
We fit an amortized Gaussian variational distribution with fixed variance o² 0.52 for inference on O, and use the same prior as above. This is an exponential family with only one location parameter to be learned.
Fitting g f( Tso; Φ)) 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 variety of metrics, including: the (non-expected) forward KI. divergence; the reverse KI. divergence; negative log-likelihen likeelooooo;;; and the angle point estimate. Intuition suggests that global optimum should outperform local ones, and find this to be the case. By any of

---

the performance metrics we consider, the global solution of the expected forward KL minimization outperforms the variational approximations found by optimizing the ELBO.
Forward KI. KLipi8 ]troe]
( Zeroe))
Negative Log Likelihood (- logg g(trre Frme))
30
10
0
2000
1000
6000
8000
1000
2000
1000
6000
8000
1000
Gruciient Strp
Grudioet Stop
X100
Reverse KL: KLig(O [(rue) p(0 Ftrie)
Angle Estimaace Varistional Mode)
4.0
300
3.5
200
3.0
mm
2.5
2000
4000
10000
2000
4000
6000
8000
10000
Gruliint Step
Grulicut Step
Truh
Expected Forward KL
ELBO
Figure 7: Forward and reverse KL divergences to the true posterior across fitting for minimization of the expected forward KL (blue) or the negative ELBO (green). We also plot the negative log likelihood of the true angle, as well as the variational mode (true angle Owe is plotted in red.)
Discussion
In this work, showed that in the asymptotic limit of an infinitely wide neural network, gradient descent dynamics on the expected forward KL objective Lp converge to an e-neighborhood of unique function f*, global minimizer. Our results depend 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, single hidden-layer ReLU network. We conjecture and show experimentally that global convergence holds more generally. First, illustrate that the asymptotic regime describes practice with finite number of neurons (see Section 5.1). Second, our conditions allow for 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 -nnvaarriinnnnniiant architectures 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-fred free methods as different tools for different settings. In a setting where the likelihood available, the prevailing wisdom suggests utilizing it. However, our results suggest that likelihood-fres 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 shallow local optima of ELBO-based variational objectives. Expected forward KL minimization instead converges unique global optimum of its objective function.
typically viewed competitors, but
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 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 fall into any of these four categories. Even for those that do, the benefits of expected forward KI. minimization, including global convergence, may outweigh the benefits of ELBO optimization.
10

---

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 2241 144 (DGE), and the U.S. Department Energy, Office 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 Statisties.
Ambrogioni, L., Güçlü, U., Berezutskaya, J., van den Borne E., Güçlütürk, Y., Hinne, M., Maris, E., and 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, 11(5188)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 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, F. D., Toigo, A., and Umanità, (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 Interna- tional Conference on Machine Learning.
Domke, J., Gower, R. M., and Garrigos, G. (2023). Provable convergence guarantees for black-box variational inference. 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, 130-13:13
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 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- permutation-invariann 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., McAuliffee 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-30
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 c -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 likelihood- free 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 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):2 :1:0: 2207.
12

---

A Convexity of the Functional Objective
We prove Lemma] | from the manuscript below.
Proof. Let the log-density be given by log q((; n) = log h()) + + TT() - A(n). 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
E(x, 1) = KL =))
= - log ((; + CC
where C/ is constant that does not depend on Now. the mapping + - ( is convex
in because its Hessian is ōngny Var(T(0)) > (cf. Chapter 6.6.3 of Srivastava et al. (0144).
AA
Strict convexity follows from minimality of the representation (cf. Wainwright and Jordan (2008),
Proposition 3.1 1) We can show is (strictly) convex by applying linearity of expectation. We
have for any À E [0, 1]
E(x, Am + (1 - A)12) - ((pe)x==) ( Am + - (1-X)m)
(9)
< -Bp(e (Alog <B (Alog + (1 A) log (()
(10)
>l(x, (1-
(11)
where the second line follows from the convexity of the map 1 + - lo of 0. In fact, the incquality holds strictly well by strict convexity above and the continuity of log q(0; n) in 0. So the function {(x, 11) is strictly convex in
q(;; 17) above for any value
B Unbiased Stochastic Gradients for the Parametric Objective
Computation of unbiased estimates of the gradient of the loss function Lp(Ø) with respect to the parameters ø is all that needed implement SGD for Lp Under mild conditions (see Proposition
1 the loss function L()) may be equivalently written as
p(O X)
p(O X)
Lp(Ø) = LE(EEEXX log O /X; x))
Ep(e.X) log g(O: [(X;)))
for density functions where f(; denotes a function parameterized by Φ. Under the conditions of Proposition differentiation and integration may be interchanged, so that
p(O) X)
V (o E V log
q(O; (X: 6)) -E(e.. V logo(0;((X:d)))
and unbiased estimates of the quantity be easily attained by samples drawn (0,x) P(O, X). Proposition 1. Let (9,, B,), (Ωg, B2) be measurable spaces which the random variables X O + X and f2 > O are defined, respectively. Suppose that for all x E X and all ó E & have P X =x) x)) Q(); [x:; d)) < X()), with X(Q) denoting Lebesgue measure and < denoling absolute continuity. Further, suppose that log
is measurable with respect to [[::[[(;)]]
p(O]X)
the product space (Ωl x 92 Bi x B2) for each & E ② and V a log q(: ((z: exists for almost all (0, x) O X. Finally, assume there exists an integrable
dominating V log q(0; f(r; Φ)) for all 4 and almost all (0, x) E O x X. Then for any B E N and any o E &b the quantity
B
( B
1
V
log q(0;: ([00;:((m;) f(x;))) (0,, itd P(O,))
(12)
is an unbiased estimator of the gradient of the objective Lp evaluated 4
Proof. By the absolute continuity assumptions, for any X' the distributions P(O = +))) and Q(O; (æ; 6)) admit densities with respect to Lebesgue measure denoted p(0 x)) and g(0; f(:; o)).
13

---

respectively, We may then rewrite the KL divergence from Equation as
dP(@) X = 5) dQ(e; 6)) p( 1z)
KL p p X - =
110( = (r(e)x==) log
= = log
((0:f(zz
because the Radon-Nikodyy derivative dP/dQ is given by the ratio of these densities. Equation thus equivalent to
( IX) (0: f(X: x)(,,)) log
p I
E Ep(e)))(eex)) log
(0;
This expectation is well-defined by the measurability assumption on log
p()
To inter-
((o,f(X;)))
change differentiation and integration, it suffices by Leibniz's rule that the gradient of this quantity with respect to ő is dominated by measurable n.v. Y. More precisely, there exists an integrable
(0, æ) defined on the product space O X' such that | log ((;=f(=;)))
p((]2)
((,, x) for all
c and almost everywhere-P(e) X). This is assumed in the statement of the proposition, and so we have
p(0 X) ((::)( x)( Ve log
p X)
V a p() log
f0;
= -Pp V log ((;;(X)))
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 (2) can be used for SGD.
C The Limiting NTK
Before proceeding, introduce the architecture specific to our analysis, scaled two-layer network, and several theorems that we will use throughout the analysis.
The first result from (Shaprro 2003 concerns optimization of the objective f(x) = Eg~pF E) in x via its empirical approximation fa(x) Σ. F(xx,)),& sid P. We reproduce this result below.
Theorem 2 (Proposition 7 of Shapiro (2003), ). Let C be nonempty compact subset of R" and suppose that (i) for almost every & the function F(-.E) f) continuous on C. (ii) F(z, E). x C. is dominated by an integrable function, (iii) the sample En is iid. Then the expected value function (x) is falute and continuous on C, and fa(x:) converges to f(x) with probability I uniformly on C.
The next two results are integral forms of Gronwall's inequality that in subsequent analysis. We refer Dragomir (2003) for detailed review and summarize the results therein below.
Theorem (Gronwall's Inequality, Corollary of Dragomir (2003). Let u(t) E IR. be such that v(() < e + 02 Jo u(s)ds for & > 0 and nonnegative Cı, Then
u(t)= < C1 exp[c2t].
Theorem (Theorem 57 of Dragomir (003)). Let u(t) Then
R be such that u(t) I c1 +
Ie lo u(u)duds for t >> 0 and nonnegative
u(t) S ci explc2t//2]
Now we turn to specifics of the architecture we consider. Assume the function f has the architecture of (scaled) two-layer (single hidden layer) network mapping f X -+ with X' C Rd and C R9. We consider this network architecture for given width p. and study each of the i L.....g
14

---

coordinate functions of f. For scaled two-layer network, the ith such function is
f(z;)), = VP 1=1
Σ agj (x w.)
for i = 1....,.. 9, where a denotes an activation function. The scaling depends on the width of network p. The parameters & are thus & {wi, "(),)))= where "()), denotes the vector [aij...., "aj]" (i.e. the jth coefficient for each component function i). The individual parameters have dimensions as follows: wy E Rd and ((). E IR" for all j = ......P. P. where again denotes the network width and d the data dimension dim X. For ease hereafter, we write aj a().) to refer to the entire jth 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. wy sid (O, Ia) for all 1......2 p. In other related works, the weight dy is sometimes also drawn az (0, 1) for all i ......................... 9. 1..... p. but in this work, initialize these second 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. CChizat et al. (2019) and [a al. (2020). For now, notationally we denote weights to be initialized as draws from an arbitrary distribution 7), and we introduce specificity to D required.
The neural tangent kernel (Equation (4) can be computed explicitly for this architecture and 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 Let E R2 E R be distributed according to are integrable (In) random variables. Assume X is compact, and o' is bounded. Then, provided condition (C4) holds (see below), we have for any x. 5t X that
D for some distribution D such that a,
K *(0) (x, x) " EpK(,,
(13)
as p + xo where K2t0) denotes the NTK at initialization constructed from draws az wj sid and K Kxx, ā; w. a) E 12x4 is the q X matrix whose k. lth entry is given by
1= (x TT)) a (TTw) w) + agaio (x T () o' (5 w) (x 8)]
for k, = 1.....g.
Proof. Consider the keth coordinate function of f. For any choice of p, the gradient is given by
(z(())
Boss
σ (x T w1)
0f (anc)
1
σ (I T wp)
vof((=;))=
Dekp afa (aie)
vp
axlo' (xr (wTwa) x
dee
akpo' ' (xc (TTppp) x
where we have imposed an arbitrary ordering on the parameters. In the above, we omitted partial 8fg for l * k,j = L,..., p because these are all identically zero. From this, it follows derivatives Bnzy that for any fixed x, X, the k. 1-th entry of K2/0) (xr, .) is given by
vof((1:(())) Vo.f o(0)
Σ I (x (T)y) (Tww,)
p
akzayoo' (x (w)) of (i (T,w (xTa).
j=1
15

---

The existence of the limiting NTK follows immediately: for each of the two terms above, each term integrable by the compactness of X and domination (see (C4)). It follows that K0 (z, ă) is the q x q matrix whose k, Ith entry is given by
BuganD [ (x w)
(a ("ww)==- agajo' (x w) o' o' w) (x a)]
with 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. [(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 rather straightforward, and this result has been previously established in other works (Jacot er al. 2018)) For our analysis of kernel gradient flows in Theorem for the expected forward KI. objectives 7.p and LE 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 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 finite training interval [0, T), the NTK changes minimally from its initial value in a large width regime. Generally, we refer 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
(CI) The data space is X' is compact.
(C2) The distribution D is such that w N0, Ia) and a w.p. For j .,...,p from this distribution, thus have wj rid (O, Ia) and oij = 0 w.p for all i,
p iid draws
(C3) The activation function a is continuous. Under (C2). this implies that the function K(-,;a, a, w) from Lemma 3 with a, w
almost surely continuous.
(C4) The function (xx, ā;; a, w) is dominated by some integrable random variable G, i.e. for all x. 8 E X x X' we have || (z, f;; a, w)lle I G(a, w) almost surely for integrable G(a, w). Proposition Fix a scaled two-layer network architecture of width p. and let & denote the corre- sponding parameter space. Initialize (()) as independent, identically distributed random variables drawn from the distribution in (C2). Let KP 0(0) X x X RR999 be the mapping defined by (z. Φ) Kotof (I, T) = J&f(x: (()) (()))) we have as p > 00 that
6(0)) T. Then, provided conditions (CI)-(C4) hold,
sup || (x, .) (I, " O. 6(0)
(DI)
x,
where (x, plo KP (o (a, σ) 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(,,aa, a, w) by (C3), ii) domination by (C4), and iii) the draws comprising K'P are iid by assumption. By this proposition, then, we have uniform convergence of ((0)
KP to Ko and obtain continuity of Koc as well.
D Lazy Training
Below, we prove several results that will aid in proving the "lazy training" result of Proposition (see below). Given the same architecture above in Appendix
we will begin by bounding (T)(T) wy (0)l and [ax;(T) C and fixed width p and time
(D)ll) [[or k 1, 1.....4 q and all j 1,...,7 pp. As in Appendix C. there are several conditions that impose and use in the following results. (DI)-(D) are identical to (CI)-(C2), repeated for clarity.
at (0)l for all
(D1) The data space is X' compact.
(D2) The distribution D is such that NO Ta) and 0 w.p. 1. For 1,..... p iid draws from this distribution, we thus have wy ~ (/ Ia) and ot = 0 w.p for all i.
16

---

(D3) The function ((.,n) - [P = x) Q(@:n)) is such that (((2;) is bounded uniformly for all and for all 7 E {f(z; ((t) t 0} by constant M, uniformly over the width p. We recall that this notation is shorthand for v,l(a,
(D4) The activation function o(-) is non- non-polynomial and is Lipschitz with constant C. Note that the Lipschitz condition implies a has bounded first derivative i.e. (r(r) & C' for all r e R.
With these conditions in hand, we now prove several lemmas for individual parameters.
Lemma (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 Lp i.e.
é(t) -VoL((().
Fix any T > 0. Then for all j = =1,.....
1[w,(T) wy (0)ll2 = 10,1(0)122 Dp.T + E.TT
(14)
almost surely, where Dp.r, Ep.r are constants depending on p, T, and limp-+oo Dp.r = and 1 Ep.T O.
Proof. First note that for any fixed j, we have
V (x;
Ta (x (xTw;) x T
1
Juy (x; ④)
goxd
vp
V w. 5 5o
(2; T)
O (x w,) xT
as required, where Oaj E IR for i = 1...., 2 and x E X C Rd from (DI). We can bound the operator
2-norm of this matrix by observing that for any y E Rd we have
1
[Jws 0)u1
o' NeTwg) (x y)
p
C²
||,
11
||0:13
11vll2 by (D4) and Cauchy-Schwarz
p
C
|| f(;; 1
||ار|
vp
by observing || is bounded by (D1) (and we absorb this term into the constant C). By similar computations, also have
xx
fx;
diag
:
R
V [a
xx
Using condition (D4), it follows that
| T w.y) VP
|| fl; 2)12
|o(0)] + Clx TanTT w.
vp
K + Clx T Ww
VP
K + Cl[wsl12
VP
by Cauchy-Schvee and (D4), where throughout the following we := (a(0) and again |||||| terms are absorbed into the constant C. Now will use these computations to bound the
17

---

variation on W across the interval (0, TJ. Fix any E (0, []. Then we have
||w (t) wy(0)ll -
[s)r(s)]lds
$ M ! I(X: by (D3)
CM
vD / (1a(()]12ds by above work
CM
VD ! [1a)(s) a,(0)ll by (D2)
CM
VD / I" S' (1)((v)ddd
CM ( S. [Fpx)/ll,/(Kxx))lII(XX///// 12deds
/P
CM2 { !-" Ja, by (D3)
VP
C2M² 1." 1-- (1))(())llddds by above work
2p
CM
C²
S.¹ S/ [1w)(0) w,(0)]l2
(0)(0)) + (1))(0)]ldddds
2p
P
CM
C²
C²M² / [ [1w,(w) w, (2)
[u,(0)]12
2p
2p
CT
CRM2T2
pt [00](0]11= + C2M2 S l [lw](w) (0)(llldddds
2p
2p
c1 1/¹ 1-"
02 c2llug(() (v) w/ (0)lladeds
with C1 = CM 2p + C2 MrT22 [|w (0) and C²M Note that even though ci depends on T. this is constant as T is fixed. We write these quantities in this way to recognize Gronwall-type inequality that we can use bound the left hand side. Indeed, by direct application of Theorem 57 of Dragomir (2003) (see Theorem 4
we have that
S. codvds
||0y (4) wy(0))12 ei exp
C2t
ci exp
CM2
C2M21
C M2222
||,,(0]]|2
exp
2p
2p
2p
giving the result for 1 T if we take Dp.r
C² gr
exp
and
Ep.
2p
2p
CT2 2p
[ 2p
exp
Clearly these constants satisfy limp Dp.r = and limp Ep,r =
for any fixed T.
18

---

Lemma 5 (Lazy Training of a). Under the same conditions as Lemma let ó evolve according to the gradient flow of problem
é(t) = VVLp(())
any T > Then we have for any j that
I|a, & ||x(00]l22 F.. Gp.TT
(15)
almost surely, where Ep. and Fp. are constants depending p, T satisfying p Pp. 0 and limp- Pp. 0.
Proof. We will use much of the same work as in Lemma Namely, ||a,()]22 ||a,(1) ag (0)122 almost surely by (D2), and thereafter for any E (0,7) we have
[ (v)llads
[1a)(t) - a0000) -
vp !
1
VP ( E ((X;
M
KMU MC
[ [|w5())]l2ds by work in Lemma
p
KMt +
MC
[( (1,()) [1oo,(0)]l2d)
P
KML MCI
||u,(0)]l2+ J Dr.s (k,]| Es.. by Lemma
MC P
P P
KMt MCt
[|w,(0)]lz S: Ep. + MC
MC
11w, 110,(00012 S Dp.ads
p
MCC MC S Dr.ads)
KMT
+ MC So" Eppods
-101(00)12
P
def || (0)l22 Fp.t + G...c
Clearly, these constants satisfy limp Fg,t + 0 and p Gp.t. - 0 (to see this, simply plug in the forms of Dp.s and Ep., from Lemma 4emma4jabowe above) and we have the result by taking = T.
Now with these results in hand, we may state and prove Proposition given below.
Proposition 3. Under the same conditions Proposition 2 fix any T > O. For any
(0,T) let Ko : X x X + Raxa be the mapping defined by (I, + Kon (I, = (,04 Then provided conditions (D1)((4) hold, we have as p that
sup G. (E
KKP b0y (I, 2)
0(0) (z. KP
K)|l a 0.
(LT)
19

---

Proof. Let us examine the k. Ith term of the q x q matrix given by KP o01) (z, 2) - KP (2, for
000)
fixed 2, T. and some & [). The k, lth term is given by (see the work in Appendix C
1
(x wy)()) (0 m,y(t))
(16)
P
(xTTw,(0)) σ (3 w,(())
(17)
anz (1)(())11(10)) (1)o (x T wy ws(1)) of (5T w)(()) (xT )
(18)
(=:(())ar(0))' (xTwz (w)()) of (a my(()) (x (xTa) a)
(19)
Above, we have explicitly made clear the dependence of the parameters on time, c.g. wy(t) vs. wy( We aim to show that the quantity above tends to zero as p + co. We first prove this holds pointwise. and will consider the red and blue terms one at a time for fixed æ, Z.
First consider the jth summand of the red term. We will bound its absolute value. If k: * I, we're done, so assume l: = I. We have for any that
(x wy (t) a w() - (x wy (o) a (TTw)(0))
o (x ()((t)))))))))))))))) o (o m,(()) a (x (,((t) σ (a w0) +
o (xTw,(t)) (0)) (x wy (0)) σ (i wy (0)) w,
(e lo (i w(()) - a (mTw, (0) + lo (i wy (0)) - | (x (t)]((t)) - a (eTw,(0))
and by the Lipschitz assumption on (-) and follows
we can bound the quantity above as
(K +
C| wy)0l) + (K + C||=||aa| Cl|=]|al|er,)() - (0)(2)
K
(()
+ [[,((()l2 [lw,(0)
since | ||||| are bounded by (D1)
C
K 2 + ((w - wy (0)ll2 + [0) by triangle inequality
S <C9]|ws)() - w. (0)1l2
=2CK | (t) w (0)l2 + C² (ww w, + (0) (t) wy (0)1l2
where again ||r||2 has been absorbed into the constant C by (D1). Using Lemma we bound all terms above almost surely using (0)l2 as follows.
+ 2CK (0))12 Ep.)) + C² (Dp.llu, Ep,)) +22 (Dp, (0)1 + (p))l) (2C2Dpp. + C2D3,.) ||w, (0)112 + (2CK Dp. 2C Dp,, Fopt + 20 F.p)) ||m, + (2CKF + C2 F2F22
20

---

Recalling that (0 Rd (N, Ia), we have that [1)](0)]2 and /0)(0)] are integrable with expecta- tions denoted and v, respectively. All our work has allowed us to show that
(x wy(()) (xTu((t)) a (& (e)(t))))))))))))))))))))) a (a ToTw,y (y)) a (o (TTw((0)))
(20²D p.t + CD (0)1 + (2CKDp, + 20 Dp Dipt + C.P.).) (0)l2))l2
P ==1
+ (2CK t + C² E3pt))
=
lim 2C Dn + m-00
CDD +C²
u+
lim 2CKDp, + 2C Dr.ppt + 2C Epx p 1-000
lim 2CK Ep,t + C²EE.11 100
0
by conditions on Dp.t and Ep.: from Lemma the strong law of large numbers, and the classic result from analysis that l b, = (lim,-co an) (lim,-so b,), provided both limits on the right hand side exist. Lastly, we can achieve the same result for the blue term quickly. Because aiy (0) w.p. by (D2), we have almost surely that
ass (t)ar(()))))'' (x w)])(())) of (TTw((t) (x (xTx)) )
aky (0)a (0)o' (a Te (0)) o' (3 (""w,(0)) (x
laky (t)(t))la)(())1)) (a (wy( 11o (# T (,) ||||||||/||
1=1
C
(t)
p
C²
|ay (1)B
p
because for all j. we have lazzl |4kyl are dominated by ||a/|/2- From here, we have by Lemma 5 that we can bound each term the sum above by
(1w)((0))2Fp... Gp.e)
P
C
P
F2 F [uy (0)] + 2F (p + G
(y)0)
P
3 0
a + oo by similar logic to the above. Together, these results combine to show that K (00) (x, 8)
KP Ko (I, R) " as p + co. As k, were arbitrary k, E 1,..., q, we have KPP @(t) (2, K2 (I, 2) 4.5. 0. This establishes pointwise convergence for some fixed 1. E (0, 7]. Uniform convergence over all of X x X' and all & E (0, T] follows easily in this case. Firstly, the numbers Dp.1, Ep,t, Fp.., and Gp.r are monotonic in t. so we can bound uniformly for all E (0, T by taking t= in the expressions above. Secondly, in our work above, our bounds the red and blue terms were independent of the choice of point (z, ž). More precisely, the supremum over x; H can be accounted for in the bounds easily by observing that sups, || KP x) KP
(0) (2, Fr
21

---

can be bounded above by
Σ 1k=1 (o wy(t)) ( (sTw((tt)
P
S sup a,EE
1=1
(o (o(
(oy(0)) (xTTx(())
(x (8Tn,(()) wo,
Oaj (L)(())1)(() (5 (,((()) of (# (dT) (x*x)
+ sup x,
(x))0)a;;(0)) (x' (,,w(()) of ( ( (x (oTx))
1
(20 DP, C² D2 (0)1 (2CK 20 DO.DDppT Dp.T 20 Ep,T) [|w,(0)]12
sup
a, P
+ (2CKEp,2 + E2
C²
+ sup 2, p
Σ
F2 1 + 2F P....Gpprii Gp.rllu;(0)) 12 Go
j=1
0
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 (for some choice of p):
f(()) Ex(x 201) (x, ((x)),((x))
(20)
(c(x) = K (z, (X,
(21)
ha(a)) PP)( KP s(o) (x. X) (X, (X))
(22)
where fü is shorthand for f(= ((t)). The three flows above can be thought of as corresponding to Lp Ly and "lazy" variant, respectively. The flow of h; is "lazy" because follows the dynamics of a fixed kernel, the kernel at initialization. The flow of also follows fixed kernel, but the
limiting NTK Ko0 instead. The flow of ft is that obtained in practice, where the kernel KT d(e) changes
continuously as the parameters e(t) evolve in time. The flow in fú is used to bound the differences We now enumerate the regularity conditions.
between fo and gi the proof of Theorem I
(E1) The functional L (f) satisfies inf Ls > -00.
(E2) The limiting NTK Koo is positive definite (so that the RKHS H with kernel Koo is well- defined).
(E3) Under (EI) and (E2), the function f minimizing LF satisfies
lI < co. so that f E H.
(E4) For any choice of p. have for all f that f(x)), ((;)) and (x; ha(x:) bounded by constant M.
(E5) The function is f-smooth in its second argument, i.e. [1 - nz)ll < imm - ml1--
We first prove Lemma from the manuscript.
Proof. Let 1 E argmin Lr(,), where Lr(l) is the functional objective. Hereafter denotes L-- Then L(f*) -c0 by (El). Then if f. evolves according to the kernel gradient flow (2 above (i.c..
22

---

the flow with kernel Kx), we have (from the chain rule for Fréchet derivatives) that
aL 05,
L(f)=
Oft dt
By definition, af: dt
h + E (X, (())) 1(X). Plugging this in yields ft. We also have AL öf.
I(f() = Ex~P(x) (2 (((x))
--Ex,x-p (X, Jx()) ( (X)
00
by the positiveness of the kernel 3 (from (E2)). Now define A, = 2 llfe - 1'1 where H. the vector- valued reproducing kernel Hilbert space corresponding to the kernel Koc (see Carmeli et al. (2006) for a detailed review). We let ( denote the inner product on the space H. It follows that
82. + (fe ,², h). Then by the chain rule we have 01.
- - -(f f
= -
(X (x(x) -
2 (EX ((X,
= I(f))
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 (ef. Eq. (2.2) of Carmeli et al. (2006)) tells us for any functions g, and fixed x,
(g. Ko (,)))(x)) = ((x))((e
and so the third line results from the second by exchanging the integral and inner product. In the second-to-last line, used the convexity of e in its second argument (from Lemma T of the manuscript). Now consider the Lyapunov functional given by
{()) = (T(()) - [((*)) A..
(23)
Differentiating, we have
- 8(1) I(f) - -(J*++/++++--))) ( A. 5
by the above work because i) ti(())) < 0 and ii) for all t. Evaluating, thus
+ A < 0, implying that E(t) E())
+[L(f)) - L(f") + < A ([4(f) - L())) < A. - A (L(f)) - L(f") < Ao since A, 0 0
1 -
(L(A+) =(()') s
and have that there exists sufficiently large T such that [(fr) I(f*) S desired.
Using this result and our previous results, able to prove Theorem I from the manuscript.
Proof. We will examine the three gradient flows
f((x) E 801) (z. f(())
(24)
(r(ar) = -EP(X) Koc (x, X)E
(25)
he(z) -E 22 (z, X)) (Xx,))())
(26)
23

---

and establish the result by the triangle inequality, i.e.
[L(fr) - -(ff))((ff" & (L(fr) - L(gr)] + +I(gr) I(f*),
(27)
where denotes the functional objective LF. Let e > O. The flow in h4 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 [[(rr) [((*'| S e/2. Fix this T. This provides suitable bound on the second term.
Turning to the first term, by continuity of L(f) in f. there exists ò > O such that y E Bigr, 6) + (L(y) ((gr) S e/22. We will show that there exists sufficiently large such that p > P implies Ilfr- 9T < 8 almost surely, yielding the desired bound on the first term of the decomposition above. Throughout, ||-1| denotes the L norm of function with respect to probability measure P (i.e. the marginal distribution of P(X) from our joint model P)),
To show that there exists sufficiently large P such that I| 9r || S 6, we use another application of the triangle inequality
I| srll s |1ʃr hr + || gr
and construct bounds on the two terms the right hand side using Proposition 2 and Proposition 3 Observe first that by (C2)(D2), at initialization we have almost surely that fo = gu = hg = O. Also note that by continuity of K (established in Lemma 3) on the compact domain X' x X' we have sup,,2 a0 (z, l2 < M for some M. Finally, note that by (E5) the function (' 77 is Lipschitz in its second argument with constant L. Below, we let || |l2 denote the 2-norm of a vector matrix, depending on the argument, and || F the Frobenius norm of a matrix. For functions, as stated || Ill denotes the L norm with respect to measure P(X), i.c [1f11 = J f(x) f(X)dP(X). From here, we have
ll9r hrll
[[ go) - -h)1l
K2 (00 (x hx(()) dt (x x),((x - KT ((0) (, (xx))((X (x))
EP(X) [Koll
- K (x (XX,
(X)) K2 (1. (xx))(XX,,,)))
$ Koo X) (x - (x(x,,(X)))
X)L - Ko () (xx)))(Xxxx())) (x
(28)
Now, we note the following facts. Firstly, for any kernel K that is uniformly bounded (i.e. [[K( y)ll2 I M for any x, y), the L norm of the function [ (yyll for fixed X, v can be bounded by
[K( X)o11 (TKKYY (XY,, S /
S 11
= [|K(;X)o|1 S M/1el12
24

---

Secondly, we have again for any fixed u and X that
()) - x (1 0112 - / T [Y, -K K2 (()0 (Y, K (,, (Y, vdP(Y)
/ Ko (X, K (Y, x) (
2
sup
K (v.
KT 6(0) (u: x3 1
2.21
11 kool K2 ((o) (, x)
w|l S sup
| (0 (v.
kTh )(0) (y, x) ||
||ll2
a.y
since the matrix (spectral) 2-norm is dominated by the Frobenius norm. Plugging these facts into Equation (28) above, we have
/ P (1(((x,((x)) - (xx (x))))) + sup (x, y) K 10 (r, -1ee (x) (a1x))-(1((X,,a
2.31
$ /'" - + M sup 2.2
|K (2, v) K2 (x, by (E4), (ES)
KPP 600) (2, B)) MŻ L Ep(xI (lex))/la(x) - (x)
sus (z,
MT sup (z, - KP 600) (z, (e,,))le + MI [-¹ (BP(x)llo(x)) - (x) -x(x)) by Jensen's 's inequality
a.
- / l1g. halldt
MT sup | (x (x,
KP o(0)
(x. y)llr + MI /
a.y/
hrll MTsup sup oo Ko) (x, v) exp(MiT) by Gronwall's inequality (Theorem) 3)
By Proposition 2 there thus exists Pi such that for all > P1 we have 19 - hr S almost surely. We proceed nearly identically for the term |/t - frll. We need only note that for sufficiently large say > P2. we can bound K (0)) uniformly (almost surely) by a constant A > M. To see this, observe that by Proposition
2 we have that there exists almost surely sufficiently large P such that
sup,,y | (x. v)
Koco (r, v)lle A M and so by triangle inequality we have
sup KP 1l
sup 2,9
| @(0) (sr, v) (2, v)]l (x, y)llr
s sup |KP 6(0) (z, v) Koo (z, + sup
Koo (z, v)llF
z
x,y
< M =
25

---

Thereafter,
||AT frll o. |1(TT- - (fT ff)ll
T
Ep(x) KT (: ()),(( KD (X, di
EP 650) (x) (, X)xxxxxX)
Er(x) K2 (XX,
- K2 100 (3 (X,
|| (10) (xx))(Xx K ROrs) () (X, f(()))
E((x) A |10((X, (X)) - e(X,
|| 2(0) (I,8) (2) f)llF (1(
|| ((0) (2r, y) (r, (2, v)ll + Aï. {" Ep(x)]ldx(X) (1))
MT
sup 00.77
and we similarly switch from Ep(x) hx()) - f((X)l2 to the L norm lik fell above using Jensen's 's inequality, yielding
(0o) (x, K (2,00l + AŽ [" 11,e Salldl
< MT sup (0,7)
|| - frll S MT sup (.,.4
| (0o) (z. v) K (x, x exx (Ai.))
Clearly, by the same logic the above there exists P3 such that p
P3 implies
MT MT'sup su || (z, K (x, (xx((AT))))))))))ll)((()))()){)) (X)) I 6/2 by Proposition 3 Then for all p > P2, Ps). we have almost surely that ||/r- - frll S 8/2. This completes the proof, as
in this case we have by the triangle inequality that ||.r - grl and so [T(fr) L(gr) - c/2
by construction.
F Experimental Details
Our code is publicly available at https://ggithi www./declanmcnanaraaaaaaaaaoeeee http:/cccaaaaaavvviieeeeeeee gevipne 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:
θ Unif (0, 2nl]
7 ~ NN(0,2)
X (0 0, = =) (ccos() 2),
22,2),nnn sin(O + =)]
The variable a a hyperparameter in such way that E S¹ satisfies assumptions (C1) and (DI), respectively. One thousand pairs of data points x. x; 2,0000 11000 were generated independently from the model above and fixed as the dataset
model that we take to be = 0.5. The model is constructed
for which the ground-truth latent parameter values are known.
26

---

We constructed scaled, dense single hidden-layer ReLU networks of varying widths, with 2 neurons for j = 6..... ...... 12 with the same architecture as in Appendix [[and the mittialization described in condition (C2). All networks were trained to minimize the expected forward KI. objectíve; stochastic gradients were estimated using batches of 16 independent simulated (0, x)) pairs from the generative model, and SGD was performed using the Adam optimizer with a learning rate of p 0.0001. We employ a learning rate scheduler that scales the learning rate as O(1/1), 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 parameterized as " f(; Φ) + 0.0001. This small perturbation must be added because f(: :) = 0 at initialization and because the value of rj = 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 [æ Φo). Thereafter, the model in é is fixed as
f(x; Φ) f(z; Φo) + (2 ( Φ0)
where Φ, do 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 I of the manuseript are constructed by evaluating the average negative log- likelihood on the dataset at each iteration, i.e. for the fixed n = 1000 pairs of observations above, we evaluate the snite-sample loss for the expected forward KL divergence. Up fixed constants, this quantity is given by
log 9(0; fa 6))
where @ is the current iterate of the parameters (either the neural network parameters the flattened vector parameters of the same size for the linearized model). The horizontal red line in Figure
at the value 1 Eim log p(0, xz)), 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
S 50, 100?)
#1 +
2 2 = 8)
:
Mit +
X (Z 2) P2)(
We take a = 1, 0.1 and artificially fix 100 as our realization from the prior on S. Thereafter, we generate N 1000 independent, identically distributed samples {2.]] 1000 from the model above, conditional on observing 5 = 100. The other non-random hyperparameters are the cluster centers y = [-20, 100, 0, 10, 20] T_ as well as the mixture weights for the mixture of Gaussians, which are uniform, i.e. P) à for all j. These draws constitute the "data" for this problem for which inference on S is desired conditionally on the draws I1, I1000: separately, we seek to infer the random variable Z conditional on the data well. Although Z and S are closely related,
impose
a mean-field variational family (SS Z X1 = Z1..... XN = xp)) = q(S) | X1 = Ii,..., XN = I ZZ X1 ZI. XN ÆN) for ease, with both components taken to be isotropic Gaussian distributions.
We use two different parameterizations for each of (()) and (Z)). The mean parameterization fixes the variance at unity and aims to learn only the mean - this case, the natural parameter of this family of distributions is simply the mean, so the network parameterizes the variational means uss, uz for each of these distributions. We also use natural parameterization with an unknown mean and variance, and in this case, the natural parameterizations 7/s, 77 are output by the networks,
27

---

which we denote f(= and f(: 1) for S and Z, respectively. Both amortized encoder networks f(. ((:)). ()) take the entire set of points ( be permutation-invariany 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).
as input. The architecture should thus as the order of these points does not matter for inference on the latent
Parameters are fitted minimize Lp 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 g(Z) and (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 a high computational cost: using 10 parallel processes, the experiment takes about 8 hours to run. We compare L--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 (S; [x1, NN; &)) 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 La-based training perform well.
For inference on Z, we extract the mode 2 E E as the point estimate and compute the l distance to the true draw Z = z for our dataset (which is known a priori because the data were generated synthetically). ELBBO-based training illustrates label-switching behavior, converging to vector which is a permutation of the entries of the true draw 2, resulting in large l distance. Lp-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 to fit a GAN generative model of MNIST digits. The simplistic GAN uses dense networks for both the discriminator and generator and 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 g(O; fFI..... 6)) on the shared rotation angle is taken to be von Mises for minimization of Lp. We parameterize g using its natural parameterization as in Appendix F.I The encoder for θ uses three Conv2D layers to increase the number of channels to 64 (with a kernel size of 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 { 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-amortizer approach for maximizing the IWBO for this example. To com- pute the IWBO, the practitioner must posit a variational distribution on 2, for each image 2;. This because only the full joint likelihood p(0, (z+) J ( is readily available. Accordingly, for this method we construct a second network [(: 1) used to construct a variational distribu- tion q( Zz; v)) on Z; for given X, This network is amortized over all images, yielding ((44 J f((;;; v)) for any image 2y in the dataset {x:) 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 single channel (the input) up to 8 channels. After the convolutional layers, we perform a flattening layer followed 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 28-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(0, {4) (x. ) in 0 by targeting the IWBO. Conceptually, this approach is the same as placing a point mass variational family on θ, i.c., we simply initialize 6g prior to training and update its value directly to maximize the IWBO. This approach thus fits the parameters { 0o) of g(Zi; f (() and de.- the distributions on the latent representations and the angle, respectively.
28

---

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 Lp-based fitting allow it perform inference without performing inference on the latent representations Z;; thus, the distributions a(Z; fx;; 1)) 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--based training. which only fits the parameters Φ of g(:; f IN:
Minimization of Lp trains for 100,000 gradient steps, although convergence is rapid, and so trajecto- ries truncated to produce Figure 4 Runs were performed in parallel across multiple processes for both methods and completed 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 {2:} and put ELBO-based optimization on more equal footing with expected forward KL minimization.
The data are generated as follows: for i 1,..., 50, fix latents z; iid p(z), and Onee 260 degrees is set artificially. The digits z are computed via the generative adversarial network and thereafter fixed for the rest of the problem (the superscript denotes 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 {., 100 θ is thus the only latent to be inferred. In this setting, the ELBO can compute the full likelihood function p(0, _ for any value of θ A (the same Uniform(U (0, 2n) is placed on the angle). The iimulation-basee approach of expected forward KL minimization can proceed as usual: draw O Uniform(O, 2x) and simulated digits are obtained by rotating the digits according to the drawn angle and adding noise.
The variational distribution, 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 reparameteriizable 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 KI.. 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 approximations as these quantities are difficult to compute exactly. Letting Tiree denote the observed data { the forward KL divergence KL(p(0 Time) || q(O Time)) is estimated using self-normalized importance sam- pling with 1, 000 samples drawn from the prior as a proposal. The reverse KL divergence KL(g(0) ))me)) I| p(0) Tume) is computed as the difference log p(x) ELBO(q): the log-evidence is approximated via the importance weighted bound (IWBO) with = 1000
29

---

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]
The abstract and introduction outline the main contributions of our work relative preceding work, making clear that our results are primarily theoretical while still having implications for practitioners. We make clear that for our results hold, regularity conditions must be satisfied.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made 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 other settings.
is fine include aspirational goals as motivation as long 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 hold. These are enumerated even more explicitly in Appendices C D and We acknowledge in Section 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 encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results 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 the scope of the claims made, e.g., if the approach only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or 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, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
30

---

3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We include several lemmas and/or theorems in the main text. Duc to space constraints, we do not enumerate assumptions in the statements, but we do reference the assumptions numerically and list them explicitly in the relevant appendices. Our appendices contain detailed proofs all lemmas and theorems provided. As the proof of Theorem rather involved, devote 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 cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch 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 properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless 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. accepted, we will publicly release code 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 perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is 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, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make 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 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.
detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on 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
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
31

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome 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 instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: In the appendix, we include link the GitHub repository with all code necessary to reproduce experiments.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurlPS code and data submission guidelines (https://niis.cct
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
/ //nips. t/nips.cc/publl// wwtpppblic/guidese.eeeeeespeeeee heabll//guuddes/odeeuubiiiis wwweeesuummsssoooooo...
The authors should provide instructions 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. only 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 much information possible in supplemental material (appended the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
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 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 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
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
32

---

Where reasonable, we show statistical significance by aggregating runs across different random seeds. We do this in Section 5.2] our other experiments, which illustrate optimization trajectories, we omit error bars as this type of data does not readily admit uncertainty estimates.
Justification:
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, confi- dence 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 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).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful 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 the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Appendix 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 compute required for each of individual experimental runs as well 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics httpppp:/aaarrps..ccppppl/aaaaaa
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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
33

---

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 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 socictal 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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.
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 models that have 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 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 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make 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]
34

---

Justification:
We credit and cite Py 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
Guidelines:
The answer NA means that the paper does 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 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 provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, pwpesssipppeseeeeeee thcode cooddataeet has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
this information is not available online, the authors are encouraged to reach out to the asset's ss 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 released with this work.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
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 crowdsoureing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According 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
35

---

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 equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: There no human participants 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 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 the institution conducting the review.
36


Reviews:
[{"review_id": "PQkUyISNn1", "reviewer": "Reviewer_Dszp", "summary": "The 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.", "questions": "-", "limitations": "-", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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.", "weaknesses": "The 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review. We will incorporate your main suggestions as outlined below.\n\n> *The result stated in Lemma 1 is well known...please consider citing a textbook or relevant review paper*\n\nAlthough we agree that Lemma 1 follows easily from the standard properties of the exponential family, we have not been able to find the convexity of the forward KL in the natural parameter stated as such in any relevant paper, despite extensive searching. (We would welcome an exact citation if you are aware of one.) We suspect the result has not been derived before because the non-amortized objective and its gradient are not generally computable, rendering the convexity of the objective irrelevant. Similarly, the amortized objective is non-convex in the neural network parameters that are optimized, so thoughts of convexity have not been considered in-depth until now. Although the result is simple, it is fundamental to our analysis of the amortized problem.\n\nThe best we can do with what we have found in a literature search is to add a reference to Proposition 2 of Wainwright \\& Jordan's *Graphical Models, Exponential Families, and Variational Inference*, which proves that the log-partition function of an exponential family is convex in the natural parameter.\nOur lemma follows fairly easily from this proposition.\n\n> *...it is unclear how big the actual gap in performance is...a plot that reports the symmetric-KL over training for both the ELBO and FKL objective would be insightful.*\n\nThis is a great suggestion; in the global response, we have provided an analysis of the quality of the solution found by minimizing the expected forward KL via several different objective measures. By all measures, we find that the minimizer of the expected forward KL outperforms the local optimum found by minimizing the negative ELBO -- the minimizer of the expected forward KL even has a lower reverse KL divergence to the exact posterior, despite not directly optimizing this quantity.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "E9yzeckhU8", "reviewer": "Reviewer_KvcV", "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.", "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.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for the review. Below, we try to answer your main questions and concerns.\n\n> *What is the difference between the analysis in this paper and existing analysis of two layer neural network in supervised learning setting...which has been widely studied?*\n\nFirstly, let us highlight that applying the NTK to the variational inference setting is novel and a contribution in its own right. Amortization (and the use of neural networks at all) is a relatively recent advancement in variational inference, and even today the main motivation for its use is for utility or cost-saving. An analysis of this type is novel in this setting and could have a significant impact towards adoption of amortization in VI. We have shown amortization is more than just a way to save compute, but actually has tangible benefits for optimization. \n\nPurely with respect to the NTK literature, our analysis is still novel and innovative in following ways:\n- We allow for network outputs $\\eta := f(x; \\phi)$ of arbitrary dimension ($\\eta \\in \\mathbb{R}^q)$.\n- We study a general, convex loss function $\\ell$.\n- We minimize the population loss $\\mathbb{E}_{P(X)} \\ell(X; f(X;\\phi))$.\n\nExtending to the general setting above was necessary -- in the current literature, existing results could not be applied to analyze the expected forward KL objective.  Our generalizations contrast with the restrictive assumptions that are often featured in existing NTK analyses, which are often specifically focused on a particular setting consisting of i) a scalar-valued network and ii) mean-squared error loss. Additionally, virtually all works consider iii) empirical risk minimization, i.e., minimization of $\\sum_{i=1}^n \\ell(f(x_i;\\phi), y_i)$ for finite training data, rather than for an infinite population, as we require.\n\n To solve this problem, we have innovated on existing NTK results in the restrictive settings described above. One of our results that may be of particular interest to the NTK community is that of *uniform* convergence of the NTK to its limit, i.e. Proposition 2 and Proposition 3 in the appendices (l. 875 and 932, respectively). Previous works relied on pointwise convergence to analyze an empirical loss, but to analyze the population quantity $\\mathbb{E}_{P(X)} \\ell(X; f(X;\\phi))$ we required uniform convergence of kernels and proved these results. Within our proofs, several techniques may also be of interest, in particular the utilization of generalized Gronwall inequalities (e.g., line 820) to bound differences over an interval $[0,T]$. \n\nLemmas 4 and 5 (lines 900 and 921), proved using this approach, are (minor) standalone results that may aid additional work. Finally, we want to emphasize that although we write primarily for a Bayesian audience and restrict our analysis to the expected forward KL objective, our results apply to more general convex loss functions and may be extended the analysis of any objective that satisfies the convexity conditions, include those in settings beyond variational inference. \n\n> *The setting for variational inference in this paper assumes that data can be simulated during training...*\n\nWe want to emphasize that the *expected* forward KL objective that we optimize is not challenging to use; ease of implementation of this method for practically any setting is one of the strengths of this approach. A forward KL divergence without the expectation over $P(X)$, on the other hand, is infeasible to optimize -- neither the objective nor its gradient are unbiasedly estimable, as estimating either quantity requires samples from the exact posterior, which cannot be obtained. We touched on this point briefly (line 104), but have now added exposition to contrast the differences between the forward KL and the expected forward KL for clarity. \n\nOur analysis of the *amortized* objective in this work may encourage more widespread use of the forward KL. The amortized objective resolves the point above: expectations over $P(X)$ and $P(\\Theta \\mid X)$ can be estimated unbiasedly as a single expectation over $P(\\Theta, X)$ by using the ``trick'' of combining these as in Appendix B, and using ancestral sampling. The assumption we can simulate draws from $P(\\Theta, X)$ is not restrictive; in fact, this assumption is strictly looser than assuming that the likelihood function $p(\\theta, x)$ is readily available, a key assumption for any ELBO-based analysis. Modern probabilistic computing packages such as PyTorch and Pyro ensure that sampling from arbitrary distributions is straightforward via ancestral sampling of $\\Theta \\sim P(\\Theta)$, $X \\sim P(X \\mid \\Theta)$, etc. Expected forward KL minimization generalizes to standard VI settings with a single observable $x_{\\textrm{true}}$ as well with minimal additional overhead -- we discussed such settings briefly in our submission (e.g. line 68, 99, and 354), but as also requested by other reviewers, we have added more substantial discussion on practical aspects to Section 6. \n\nA counterintuitive implication of our work is that simulation-based minimization of the expected forward KL may be useful even for non-amortized VI problems such as described above. Because simulation is often computationally inexpensive, and network training rapidly converges to a global minimizer, one can obtain a unique variational approximation $q(\\theta; f(x_{\\textrm{true}}, \\phi))$. ELBO-based training, on the other hand, might yield vastly different posterior approximations to $p(\\theta \\mid x_{\\textrm{true}})$ for different initializations or random seeds. Therefore, our work could simplify VI for practitioners. However, ELBO-based optimization still has many merits. We have added discussion of additional practical considerations to our discussion section, including cases where ELBO-based training may be preferable (e.g., for fitting model parameters alongside the variational approximation; if one wants to obtain mode-seeking approximations, etc.).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Swt5yYYkfC", "reviewer": "Reviewer_grF1", "summary": "This 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.", "questions": "Minor comments. \nL45: L_p to L_P\nL771: What does \\eta over equality mean?\nL843: a,\\in - remove", "limitations": "NA", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. Below, we attempt to address your main points.\n\n> *The idea of using an expectation with P(X) also seems interesting...is this idea novel?*\n\nAlthough it's not novel (e.g. Section 2.1), we do think it is underappreciated. Our main contribution lies in the analysis of this objective function. Previous motivations for using the expected forward KL divergence come from either 1) a setting where likelihood function is not available or 2) the practitioner's desire for conservative uncertainty quantification, a property of minimizers of this objective.\n\nOur analysis suggests a much stronger motivation for using the expected forward KL objective: it behaves like a convex objective and thus yields unique solutions, regardless of random seeds, initializations, etc. The shortcomings of numerical optimization are a major trouble point for use of VI over related methods such as MCMC in practice, and our work resolves many of these concerns. Beyond popularizing this choice of objective function as an alternative to the ELBO within the VI community, we hope that our results convince other Bayesians (e.g. practitioners of MCMC) of the validity of VI. We help resolve a major concern that one may end up with a local solution of unknown suboptimality. \n\n> *Can [the authors] demonstrate [Lemma 1]'s efficacy on a toy problem that doesn't use any neural network?*\n\nLemma 1 in itself is interesting, but not practical. It illustrates the atypical paradigm of having a function that is well behaved (convex), but **not computable**. Neither $\\ell(x, \\eta)$ nor $\\nabla_\\eta \\ell(x, \\eta)$ are computable, or even unbiasedly estimable -- except in certain rare cases where the posterior itself is tractable, which of course makes VI unnecessary in the first place. \n\nOn a non-amortized problem where one has a single $x$ and wishes to minimize $\\ell(x, \\eta)$ in $\\eta$, the convexity result becomes an afterthought -- optimization of this objective cannot proceed because we cannot run SGD (we do not have access to unbiased gradients).\n\nThe corollary implying that the amortized problem $L_F$ is convex is the foundation for the rest of our work. It implies the existence of a global minimum, and the remainder of our analysis shows the gradient dynamics of $L_P$ can converge arbitrarily close to this global optimum. A distinguishing feature of $L_P(\\phi)$ is that we can estimate its gradients unbiasedly (e.g., as described in Appendix B).\n\n> *It is slightly misleading to say that the FKL objective with NTK finds global optima...the result only shows that the local solution is close to the global one.*\n\nAlthough asymptotic results such as ours will indeed only hold approximately in practice, we contend that many meaningful theoretical results take this form. Even some that might be regarded as standard results, such as the convergence of SGD with diminishing step size, are asymptotic results that are only approximated in practice with finite time and compute. Within the NTK literature in particular, asymptotic results of the nature of ours (i.e., in the limit of an arbitrarily wide network) are difficult to avoid due to the intractability of finite-width analysis. Although our results are asymptotic, we have made them as strong as possible (e.g., almost sure convergence), and have shown in our experiments that the asymptotic behavior (i.e., optimization trajectories similar to those in convex optimization) can actually be obtained in practice. \n\nYou are correct that one will ultimately obtain a local solution in the parameter space. We show, though, that the \"local solution\" can be made arbitrarily close to the global minimizer $f^*$ by increasing the network width $p$ -- this allows one to obtain the ``de facto'' global minimizer. The ability to alter the width provides the user with a large amount of control over behavior of the optimization problem. Without our analysis, understanding the results of the optimization problem seems intractable: as you correctly point out, in parameter space the objective function is highly nonconvex, and *the degree of suboptimality* of the local solution is totally unknown for any given run of the optimization routine.\n\nOur result, on the other hand, bounds the degree of suboptimality, and more importantly shows that the user can shrink this bound by increasing the network width. We can always get $\\epsilon$-close to the global optimum for any $\\epsilon$. We show this result extends to practice for even modest network widths, and $\\epsilon$ appears to be small enough on real problems that one i) converges to the same minimizer regardless of initialization and ii) this minimizer is practically indistinguishable from the true global minimum $f^*$ in the RKHS. \n\n>  *I don't understand how Lemma 1 shows that $L_F$ is strictly convex. I believe it just proves convexity.*\n\nThank you for catching this; we have added a mild condition to Corollary 1 under which strict convexity holds. Previously, any two functions $f,g$ such that $f \\neq g$ but $f = g$ almost surely would violate strict convexity because $L_F(f) = L_F(g)$ even though $f \\neq g$. We have added the condition that the domain of $L_F$ is a RKHS $\\mathcal{H}$ with respect to the measure $P(X)$. Under this assumption, two functions $f,g$ equal to each other almost surely on $P(X)$ are in fact regarded as the same element in the RKHS (they correspond to the same equivalence class of functions). With this condition, strict convexity holds as any $f \\neq g$ in the RKHS differ on a set of nonzero measure, allowing the strict convexity inequality to extend to the integral. \n\n> *L771: What does $\\eta$ over equality mean?*\n\nThis notation meant \"equality up to constants that do not depend on $\\eta$\", i.e. the right-hand side has the same gradient with respect to $\\eta$ as the left-hand side. We have updated the notation to make this more clear and now instead use a generic constant $c$ on the right-hand side.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "10Ghu4tnkN", "reviewer": "Reviewer_gbq5", "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.", "questions": "1. How restrictive are the assumptions?\n2. Can you explain the intuition more how your approach achieves convexity?\n3. 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.", "rating": 7, "confidence": 1, "soundness": 3, "presentation": 2, "contribution": 4, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your review. Based on your comments, we will add some clarifying sections to motivate the key intuition further. We aim to address your main points below.\n\n> *How restrictive are the assumptions?*\n\nWe believe our assumptions are mild enough to apply to many practical settings, and are generally less restrictive than those found in related work. Several regularity conditions such as compactness of the data space $\\mathcal{X}$, continuity of $\\sigma$ and boundedness of $\\sigma'$, and positivity of the limiting kernel $K_\\infty$ are generally standard prerequisites for proving NTK-type results. Restricting the setting to a particular neural network architecture and initialization is also a standard approach to the problem.\n\nOur work stands out from previous works in the generality of the setting: we allow for network outputs of arbitrary dimension, and consider an objective function that is an expectation over an ``infinite dataset'', i.e. $L_F(f) = \\mathbb{E}_{P(X)} \\ell(X, f(X))$. Although we consider the forward KL integrand for $\\ell$, our results actually apply to any $\\ell$ that is convex in the network outputs $f(X)$. To our knowledge, existing NTK-based analyses all restrict themselves in some way, only considering i) scalar-valued outputs, ii) mean-squared error loss, or iii) and objective function over a finite training dataset.  \n\n> *Can you explain the intuition more how your approach achieves convexity?*\n\nThe main intuition for our work is that we can make use of convexity by conducting the analysis of the forward KL optimization problem in a function space. In parameter space, there is no hope of utilizing convexity arguments -- neural networks may have millions of parameters, and the objective functions used to fit them are highly non-convex in these parameters.\n\nWhen neural network functions are viewed simply as points in a more general Hilbert space of functions, though, the forward KL objective $L_F$ is in fact convex in its function argument -- this is the result of Corollary 1. The remainder of the analysis utilizes NTK-based analysis to show that the gradient dynamics in parameter space essentially mirror those in function space, i.e. optimization behaves as if we follow a convex objective to its global optimum.\n\n> *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?*\n\nThis is a good point, and we have updated the discussion section accordingly to discuss in more detail settings where using the ELBO may be more useful for practitioners. Choosing to optimize the ELBO can be advantageous in several situations: if the likelihood function of the model is (approximately) convex in the main region of interest, optimization should behave well, and in certain situations the ELBO can even be updated with non-stochastic gradients. In the case where the generative model is unknown and needs to be learned simultaneously with the variational approximation, the ELBO can be used as an objective to fit both the model and the variational distribution, which is appealing in practice. Finally, the ELBO can be optimized in a non-amortized fashion -- this may be preferable for simplicity, but this approach is still known to struggle with the same issue of converging to local optima.\n\nWhile these particular situations do arise, minimization of the forward KL can be applied in essentially any setting, and maintains its convexity guarantees even when the model is arbitrarily complex -- this approach can even be used in some settings where ELBO-based training is impossible. One large class of methods for which the ELBO is unusable for likelihood-free inference, where the generative model typically consists of a highly complex simulator that does not admit a tractable likelihood function.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their time and helpful comments. We are encouraged by your largely positive feedback and appreciate your thoughts on areas of improvement.

We have responded to each of your points individually in the review-specific rebuttals. In this global rebuttal, we wish to emphasize the significance of our contributions with respect to two research areas: variational inference and NTK analysis. We also provide additional experimental results (as suggested by Dszp, but likely of interest to all reviewers). 

**A breakthrough in variational inference.**
In the VI community the use of posterior approximations parameterized by neural networks (i.e., amortization) is still a relatively recent phenomenon and is by no means ubiquitous. Our results are a major step towards expanding the use of amortization in variational inference. 

We show that amortization has benefits beyond reducing computational costs; in fact, targeting the amortized objective can actually be 
desirable because it admits a unique solution. Through our analysis, we resolve a significant complication for VI that previously seemed impenetrable: that of convergence only to a local optimum of the objective. The objective we consider in this work is convex, and we show that its gradient leads to the global optimum. Our results may help expand the use of VI in general. In practice, MCMC still remains more widespread than VI, in part due to potentially unreliable VI optimization -- our work is a significant step towards resolving this concern.

**A technically sophisticated adaptation of the NTK.**
 We emphasize that our results are not a mere application of existing NTK analyses to a different objective function. Existing NTK-based analyses (e.g., [1]-[5]) are restricted to specific settings that exclude variational inference. These restrictions typically include one or more of:
- a scalar-valued network output 
- a mean-squared error loss 
- a finite training set 

These assumptions are simplifying and commonly used in practice for general machine learning, but do not apply in the settings commonly seen in variational inference, which may have i) parameters of arbitrary dimension, ii) diverse loss functions, and iii) population-type objective functions (i.e., an expectation over a continuous distribution). Bridging the gap between existing literature and analysis tailored to the VI objective we consider in this work required generalizing beyond the setting outlined above. To do so, we proved results for arbitrary loss functions and introduced new machinery (e.g., uniform rather than pointwise convergence of the NTK) necessary for establishing convergence over a continuous distribution of inputs. 

[1] Generalization ability of wide residual networks, Lai et al.

[2] Gradient descent provably optimizes over-parameterized networks, Du et al.

[3] Linearized two-layer neural networks in high dimension, Ghorbani et al.

[4] Loss landscapes and optimization in over-parameterized non-linear systems and neural
    networks, Liu et al.

[5] On exact computation with an infinitely wide neural net, Arora et al.

**New experimental results**

Suggested by Dszp, we add a small case study that quantifies the *quality* of the variational approximation found by expected forward KL minimization -- in other words, how well it approximates the true posterior. We use a rotated MNIST example, with angle $\theta \sim \mathrm{Unif}[0, 2\pi]$ and for all $i=1, \dots 50$ we have $x_i \mid \theta \sim \mathcal{N}(\mathrm{Rotate}(\mu_i, \theta), \sigma^2)$, where $\mu_i$ is a fixed, synthetic MNIST digit, and $\mathrm{Rotate}(\cdot)$ applies a rotation of $\theta$ radians. The variational distribution is taken to be Gaussian with fixed $\sigma = 0.5$.

 We aim to measure the quality of the solutions found by ELBO-based optimization and expected forward KL minimization; because our approach tends toward the global optimum of the expected forward KL objective, we expect this solution to outperform local optima of other objectives by any metric used to measure the quality of the approximation.

In the attached pdf, we plot several quantities across fitting, where fitting either minimizes the negative of the ELBO or minimizes the expected forward KL. The negative ELBO objective was fit to a fixed dataset $x$ of $50$ images drawn from the model above with common rotation angle of 260 degrees. We denote these images as $x_{\textrm{true}}$, and the true rotation angle $\theta_{\textrm{true}}$. The variational approximation is denoted as $q(\theta \mid x_{\textrm{true}})$. 

For the negative ELBO objective, the estimated angle converges to about 90 degrees, while expected forward KL minimization (as expected) is centered near correct angle value (Figure 1). This translates to a better held-out negative log likelihood on the true latent angle value, as expected (Figure 2). We also display both the forward (Figure 3) and reverse (Figure 4) KL divergences across training. These quantities are difficult to estimate exactly -- we compute the forward KL using importance sampling with the prior $p(\theta)$ as the proposal and $K=1000$ importance samples, and the reverse KL is approximated using the ELBO plus an estimate of the log evidence. 

Perhaps surprisingly, expected forward KL minimization outperforms optimizing the negative ELBO *even with respect to the negative ELBO objective function*. In other words, the variational distribution fit to minimize the expected forward KL turns out to have a lower (better) reverse KL value that the distribution fit to minimize reverse KL. This arises from the intuition that a global optimum of a different objective may be preferable to a local optimum of one's original target. This is the motivation behind our manuscript: the global minimum of the expected forward KL is closer to the exact posterior than a local optimum of the ELBO.

Meta-review: This paper examines the variational inference problem and, unlike previous works, provides a global convergence guarantee. The analysis focuses on forward KL-based variational inference, where the variational family is parameterized by neural networks. The authors adapt NTK-based analysis for two-layer neural networks to achieve the desired results for their specific variational inference problem.

Almost all reviewers agree that the paper merits acceptance, and after carefully reading the rebuttal and discussion, I tend to concur. Please incorporate the reviewers’ feedback into the final version of the paper, with particular attention to the following concerns:


1- Include the empirical comparisons to the ELBO/IWAE methods (Response to  Rev DszP). 

2- Include the ref for Lemma 1 and clarify its efficiency. (Response to Rev DszP and grF1) 

3-  Include a discussion about when your method is more efficient than ELBO-based methods. (Response to Rev gbq5)
Decision: Accept (poster)