Implicit Variational Inference for High-Dimensional Posteriors

Summary

The paper proposes LIVI a variational-inference-based approach for Bayesian NNs. It hinges on implicit VI (IVI) which allows to easily sample from an arbitrary distribution, yet obtaining the gradients w.r.t its parameters or the density is hard. As the ELBO in VI decomposes to two terms, an expected likelihood term and a KL term, under this framing the KL term cannot be evaluated, and even may be ill-defined. To remedy that, the authors propose to use a Gaussian deep latent variable model (DLVM), i.e., the generator network in VAEs, as a replacement for the implicit density. And, to approximate the entropy term in the KL divergence the authors propose to linearize the output of the network. The authors also propose a lower-bound to the ELBO by approximating the log-determinant of the covariance using the minimal singular value, which can be calculated more efficiently compared to all of the values. The authors compare their method to several Bayesian methods on in-distribution and out-of-distribution tasks showing improved performance over baseline methods.

Strengths

* A novel approach for using hyper-networks to learn a Bayesian model using VI. * Good results on OOD tasks compared to baseline methods. * The paper is written clearly and addresses relevant related studies. * The results seem to be reproducible - exact experiential details were given along with the code.

Weaknesses

**Method** * Although I think the idea is nice a major shortcoming of this method is that it is challenging to apply it to large networks because of the use of Hyper-networks. Even in the experimental section, the largest network has 2.7 million parameters which is considered small these days. How do the authors propose to handle that issue and scale it to networks of 1-2 orders of magnitude larger? * To estimate the entropy term the std of the noise is taken to be zero on only one summand (the bilinear function), but not on the logdet of the covariance. How do the authors justify that? **Experiments** * The compared methods are somewhat outdated. Many Bayesian models were published in recent years, and all of them show that they beat deep ensembles in one way or another. I believe that more recent baselines should have been evaluated. * I find it a bit odd that the authors didn't compare LIVI to methods with similar pipelines, such as flow-based methods [1, 2]. * The method was evaluated on UCI benchmarks, MNIST and CIFAR-10. In my opinion, it is not enough to showcase the advantage of the method and more challenging datasets should be considered. [1] Krueger, D., Huang, C. W., Islam, R., Turner, R., Lacoste, A., & Courville, A. (2017). Bayesian hypernetworks. arXiv preprint arXiv:1710.04759. [2] Louizos, C., & Welling, M. (2017, July). Multiplicative normalizing flows for variational bayesian neural networks. In International Conference on Machine Learning (pp. 2218-2227). PMLR.

Questions

* How did your method of approximating the entropy is compared to standard MC sampling as suggested in the beginning of section 3.2? * How does your method perform compared to baseline methods in terms of memory and run-time?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The authors did not address the limitations of their method.

Summary

This work proposes one approach to implicit variational inference that avoids using adversarial training. This is achieved by first applying a small amount of Gaussian noise to the implicitly generated weights $\theta$ to ensure that they are equipped with a valid distribution (enabling valid divergences), and then the now valid ELBO is approximated through linearizing the entropy regularization term. A further lower bound to this approximation is also proposed that allows for better scaling in high-dimensional settings, which is common for applying VI to neural models.

Strengths

I felt that the approach presented was well justified, easy to follow in terms of motivation and development, and enables much more complex settings for variational inference to be applied to (in a stable manner) without having to suffer from the typical limiting mean-field assumptions. The experimental results were impressive as well, seemingly achieving much better uncertainty quantification than the other competing methods while still retaining quality predictive performance.

Weaknesses

The paper states in the contributions (line 49) that it "derive[s] a novel lower bound for variational inference"; however, I believe this is not quite the case. It is my understanding that the "novel lower bound" being mentioned here refers to $\mathcal{L}'$ and the further lower bound $\mathcal{L}''$ used for better scaling. This statement is in direct contradiction to equation 18 that states: $\log p(\mathcal{D}) \geq \mathcal{L}(\gamma) \approx \mathcal{L}'(\gamma) \geq \mathcal{L}''(\gamma)$. Strictly speaking, $\mathcal{L}'$ is not a lower bound, but rather an approximation to an actual lower bound $\mathcal{L}$. To be clear, I do not think this is bad by any means, it just needs to be communicated clearly. Aside from this, I think the experiments can be bolstered with a few additional comparisons. Namely: - The impact of using $\mathcal{L}''$ over $\mathcal{L}'$ when the latter is still eligible (i.e., when using smaller models such as the UCI dataset tasks). - How much performance (both in and out of distribution) differs between your proposed method and HMC (or other MCMC methods). This again would need to be done in low-dimensional settings, but I believe it should be feasible for at least a smaller neural network or linear model on some of the UCI dataset tasks. Lastly, one of the contributions listed cited a novel generator architecture (line 53); however, the details of this seem to be relegated to the appendix. In future revisions, I believe if you are going to cite this as one of the major contributions then the details should at least be partially included in the main paper. I understand this probably wasn't done due to space limitations, but it is worth considering in my opinion.

Questions

I do not have any direct questions, see the weaknesses for my main comments to be addressed please.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

The authors did discuss the limitations adequately. Negative societal impact is not directly applicable here in my opinion.

Summary

The paper studies the problem of approximating high-dimensional multi-modal posteriors through neural samplers specifying implicit distributions. While Bayesian methods promise a variety of benefits in terms of generalization and calibrated predictions, in practice they have seen limited success due to intractability of exact Bayesian approaches and the tradeoffs made in approximate Bayesian methods. Implicit Variational Inference provides an alternative to approximating exact Bayesian posteriors by maintaining distributions implicitly by transforming samples from simple distributions, allowing it to admit much richer distributions. Implicit Variational Inference typically requires some density ratio based adversarial objectives, which can fail on high-dimensional problems (e.g. parameters of neural networks). The authors note two major issues in existing approaches a) KL being ill-defined due to the implicit density lying on a low dimensional manifold b) intractability of the entropy of the implicit density and its gradients. To tackle these issues, the authors first introduce Gaussian noise to the output of the sampler making it a Gaussian DLVM resulting in a well-defined KL over the parameter space. Next the authors approximate the generator with a local linearization resulting in a Gaussian approximation of the output density, and obtain easy to compute approximation of the differential entropy of the output density. This results in a novel approximation to the ELBO based on the entropy approximation, which is scalable to high dimensional parameters. The authors discuss two variants based on computing the whole Jacobian or using a differentiable lower bound on the the determinant, which trade-off compute and quality of the approximation. Finally, the authors evaluate the method on a variety of tasks including impressive results on fairly large BNNs (WideResNet).

Strengths

* The paper studies the important problem of approximating expressive Bayesian posteriors on high-dimensional spaces. Due to the general applicability and promising results, the work is significant and relevant to the community. * To the best of my knowledge the main contributions of the paper namely addressing the ill-defined KL, local linearization for the entropy and the LIVI bound on the ELBO, are all novel. * Introducing Gaussian noise to the sampler to induce a Gaussian DLVM is a neat and simple way of fixing the KL with minimal additional restrictions to the model * Similarly, the local linearization of the neural sampler is a nice idea to obtain a cheaper estimator for the entropy. * The experimental results are quite impressive - in particular the results on the WideResNet. I also appreciate the authors including the code with their submission to aid reproducibility of the results. * Overall the paper is well-written with a clear exposition of ideas and most relevant details covered.

Weaknesses

* The paper proposes a local linearization to make the entropy computation tractable. However, what is not clear to me is how the local-linearization affects the expressivity of the posterior and general performance. The experiments indicate that the effect is not large, but these are still “relatively” simple tasks so a thorough study of this would be useful * The empirical comparisons also do not consider alternative approximate Bayesian methods. * Some recent work [1] proposes a closely related approach which might be worth discussing in the paper [1] Posterior Refinement Improves Sample Efficiency in Bayesian Neural Networks. Kristiadi et al., NeurIPS 2022.

Questions

* What do you think would be the critical challenges in applying the method to even larger modern networks, e.g. transformers? * What is the challenge in implementing the KIVI baseline? (since that is not included in all the experiments as the authors note) * The runtime of the method is not explicitly mentioned anywhere in the paper except for the remark relative to Deep Ensembles in the appendix. This is important information which should be included in the main paper. * Minor typos: L117 “in the following” -> “in this section” L123 “aGaussian” -> “a Gaussian”

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Despite the impressive results in the paper, there still remains a gap between the models studied in the paper and the size of models considered in practice.

Summary

## Post Rebuttal Update I have engaged with the authors for the rebuttal, and found their responses informative, prompting me to increase my score from a 7 to an 8. ## Original Review The paper addresses two issues that are common in implicit variational inference methods such as amortised neural samplers, namely - 1. The implicit density often lies on low dimensional manifolds making the KL infinite/ill-defined 2. The gradients and entropy of the implicit density are intractable The paper addresses these issues by (i) adding Gaussian noise to the output of the neural sampler, making it continuous w.r.t $\theta$, and (ii) linearising the neural sampler, resulting in a Bayesian linear model approximation to the non-linear model, giving closed-form solutions for the distributions. On different benchmarks the authors show that these methods work well on a range of small scale experiments compared to other state-of-the-art, such as last-layer Laplace and deep ensembles.

Strengths

The paper is well motivated, has clear mathematical development, and solves common issues in implicit variational inference models. I like the use of linearisation to make the entropy and gradients tractable by getting closed-form solutions. The experiments are clear, and the baselines are well-considered. The results are also quite impressive in this domain, as it is often really difficult to beat deep ensembles, and the method seems to consistently perform really well.

Weaknesses

In general I think the paper is well-written, however I have one major criticism - 1. I'm not convinced that the current approximation to $\log(J J^T)$ has been well-motivated or properly ablated. It would be nice to see some ablations on smaller-scale problems where calculating the full log determinant is tractable, and comparing it to the approximation made using just the highest singular value. Or doing a sweep over adding subsets of singular values vs the highest. Or even, plotting the eigenspectrum of $J J^T$ to show that it is true that the highest singular value is often much more dominant than the others. 2. If the paper is using the full-rank Jacobian of the neural sampler in order to estimate the entropy, I think a fairer comparison to make would be against full Laplace, not last-layer Laplace, which should be possible for UCI datasets and MNIST at least, using the Laplace library the authors cite.

Questions

I think there are some really interesting connections between the obtained linearised implicit variational model, and existing methods such as linearised Laplace and regular variational inference, and I would love for there to be a small discussion about these things potentially. For example, in linearised Laplace, you assume a fully Bayesian posterior over a NN parameterised by $\theta$, then linearise the model around the MAP estimate, resulting in a tractable Gaussian approximation to the posterior. However, this posterior is more rich compared to a mean-field variational approximation, because the covariance of the posterior is given by $J_{\theta}^T \Lambda J_{\theta} + \sigma^2 I$, where $\Lambda$ is the prior over weights, and $\sigma^2$ is the noise variance. This form of the posterior looks very close to what is obtained by implicit variational inference, where the neural sampler can be considered similar to the NN model in linearised Laplace. I would be really interested in seeing discussions about these potential connections if possible. In fact, optimising the marginal likelihood for a linearised Laplace model is akin to doing ELBO with a Gaussian approximation. I would be really interested in seeing these parallels. 2. Given that the authors are using the entirety of the Jacobian of their neural sampler, and comparing to only last-layer Laplace, where only the last layer is modelled probabilistically, I would be really interested in seeing if they can run experiments on full linearised Laplace, especially on the small UCI and MNIST datasets, where this should be tractable. In fact, there are methods that perform full probabilistic inference for linearised Laplace using samples from the posterior, such as in [1]. 3. I am also really interested in seeing what approximations to the Jacobian are best for estimating log det (J J^T), such as last-layer only, or through samples using Hutchinson's estimate, or by using more than one singular value and ablating through what the optimal number of singular values to consider is. [1] Antorán, Javier, et al. "Sampling-based inference for large linear models, with application to linearised Laplace." arXiv preprint arXiv:2210.04994 (2022).

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A.

Response to rebuttal

Thanks for the response and apologies for the delayed response! > How does the local linearisation affect the posterior in terms of expressivity and performance? Thanks for the clarification and additional results! It is indeed a bit challenging to test claims about expressivity, but I appreciate the additional experiment. It indeed appears to be the case that the performance / expressivity are not impacted considerably. > The empirical comparisons do not consider alternative approximate Bayesian methods. Thanks for the additional results. > Posterior refinement (Kristiadi et al., 2022) is closely related and should be discussed. Thanks for the explanation. I think this does merit some additional experimental validation but at the very least I hope the authors include this in the paper. > What is needed to apply LIVI to large, modern networks? Thanks for sharing these insights. On the topic of hypernetworks I would also mention recent advances (e.g. [1]) on this topic. > The runtime of the method should be included in the paper. Thanks for these results, these are indeed quite impressive and it would be great to have these in the paper. I am satisfied by the author response and encourage the authors to make the relevant changes for the camera ready version. I will maintain my score. [1] Knyazev, B., Hwang, D., & Lacoste-Julien, S. (2023). Can We Scale Transformers to Predict Parameters of Diverse ImageNet Models?. arXiv preprint arXiv:2303.04143.

Response to Authors

I thank the authors for the response. I also appreciate the effort in evaluating MNF on the UCI datasets toward this rebuttal. The main two shortcomings of this paper, in my opinion, remain. * **Scalability issues**. First, I didn't mention transformers, even standard ResNet architectures generally have tens of millions of parameters which is more than one order of magnitude compared to the largest network used in this paper. Second, in terms of "relevant to the Bayesian deep learning community" - I respect the author's opinion, but I do not feel the same. I think that in this era of deep learning, an important aspect of the model is the ability to scale it beyond the network sizes in this paper. While I do not expect the method to be readily adjusted to massive networks, scaling to network sizes such as ResNet-18 and ResNet-34 is a reasonable requirement. The authors suggested several alternatives for scaling their model which sound great. In my opinion, the submission is incomplete without showcasing that. * **Experimental section**. I stand behind my original comment that neither the datasets nor the baseline methods are strong enough. Therefore, I cannot attribute much value to the empirical evaluation in this paper when comparing LIVI to the proposed baseline methods. The authors wanted me to suggest alternatives. Well, in my opinion, this is the job of the authors as the setups in the paper are quite ubiquitous in the literature. Nevertheless, here are some suggestions: * In terms of datasets, CIFAR-100 is already more challenging than CIFAR-10, fine-grained classification datasets such as CUB, Cars, and Pets are another alternative. All appeared in the Bayesian literature before, and there are many more of course. * In terms of baselines to state a few, SWA/SWAG family which often shows good performance [1], deep kernel learning and its recent follow-up works [2], VI in function spaces and follow-up works [3], infinite-deep NNs [4, 5], and partially stochastic BNNs [6, 7]. There are other methods that I may have missed. I do not expect the authors to compare LIVI to all of these methods, but to at least some of them, I do. Overall I would like to state that I do value the proposed approach and its merits. Nevertheless, I think that currently, this paper is not ready to be published at NeurIPS. Hence, I decided to raise the score to 4 and not further. [1] Maddox, W. J., Izmailov, P., Garipov, T., Vetrov, D. P., & Wilson, A. G. (2019). A simple baseline for bayesian uncertainty in deep learning. Advances in neural information processing systems, 32. [2] Wilson, A. G., Hu, Z., Salakhutdinov, R., & Xing, E. P. (2016, May). Deep kernel learning. In Artificial intelligence and statistics (pp. 370-378). PMLR. [3] Sun, S., Zhang, G., Shi, J., & Grosse, R. (2018, September). FUNCTIONAL VARIATIONAL BAYESIAN NEURAL NETWORKS. In International Conference on Learning Representations. [4] Nazaret, A., & Blei, D. (2022, June). Variational inference for infinitely deep neural networks. In International Conference on Machine Learning (pp. 16447-16461). PMLR. [5] Xu, W., Chen, R. T., Li, X., & Duvenaud, D. (2022, May). Infinitely deep bayesian neural networks with stochastic differential equations. In International Conference on Artificial Intelligence and Statistics (pp. 721-738). PMLR. [6] Sharma, M., Farquhar, S., Nalisnick, E., & Rainforth, T. (2023, April). Do Bayesian Neural Networks Need To Be Fully Stochastic?. In International Conference on Artificial Intelligence and Statistics (pp. 7694-7722). PMLR. [7] Daxberger, E., Nalisnick, E., Allingham, J. U., Antorán, J., & Hernández-Lobato, J. M. (2021, July). Bayesian deep learning via subnetwork inference. In International Conference on Machine Learning (pp. 2510-2521). PMLR.

Response to Authors

Thank you for answering my concerns and providing more experimental results. I am satisfied by the response and maintain my original score assuming promised changes are incorporated in the camera-ready version of the paper.