Globally Convergent Variational Inference

Summary

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

Strengths

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

Weaknesses

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

Questions

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

Rating

7

Confidence

1

Soundness

3

Presentation

2

Contribution

4

Limitations

Technical limitations are stated but it would be helpful to discuss more prominently how restrictive the assumptions are.

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.

Strengths

1. Using expected forward KL with the exponential family is simple and interesting when dealing with ELBO's non-convexity. 2. 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. 3. The connection between NTK and the gradient dynamics of the functional objective is also interesting and novel. 4. 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? 2. 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. 3. I don't understand how Lemma 1 shows that L_F is strictly convex. I believe it just proves convexity.

Questions

Minor comments. L45: L_p to L_P L771: What does \eta over equality mean? L843: a,\in - remove

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

NA

Summary

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

Strengths

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

Weaknesses

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

Questions

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

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The limitations are pointed out in the paper.

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.

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. The 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. --- Edit --- The 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).

Questions

-

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

-

Re: Rebuttal by Authors

Thank you for your efforts in addressing my concerns! - After spending about 10 minutes searching for references, I indeed found it surprisingly difficult to find more specific ones. However, I believe Wainwright & Jordan is a solid reference. - I also think the additional results are excellent for completing the overall picture and make a compelling case for why the well-posedness of the underlying optimization problem might matter, even though it is often overlooked in common practice. Overall, after considering all reviews and corresponding rebuttals, I’ve decided to raise my score to an Accept (7). I believe this is a technically solid paper that presents novel theoretical results, rigorously extending existing NTK results and applying them to the VI setting. This makes it of interest to both the variational inference community and the broader NeurIPS audience!