Streamlining Prediction in Bayesian Deep Learning

We thank reviewer tEnH for their thorough review. We appreciate that they found our method conceptually simple. We address the comments and questions below.

C1: The idea of linearizing the activation function has been done before in the context of continual learning by Dhawan et al., ICML 2023. The authors should cite this work and compare their method in detail.

Thank you for pointing this out. We agree Dhawan et al. (ICML 2023) is a relevant related work and we have added discussion on it in the related work.

C2: It is unclear to me why one can still maintain good performance when the network is effectively linear in its inputs. Intuitively, this should not be the case since the representation learning of the base network is negated by the linearization. Could the authors please elaborate in detail (additional, empirical investigations like measuring the degree of linearity of the resulting linearized might be one way to go about this).

In our approximation, the network is only locally linear and we do not globally linearise the network. Hence, we still capture interactions between the model parameters and globally capture the non-linear nature of the original model. This means we assume that the change of the function output behaves linearly locally around the input.

Could the reviewer elaborate on the proposed: "empirical investigations like measuring the degree of linearity of the resulting linearized"? One way to locally assess the induced error would be in terms of an $L_p$ norm between the original and the approximate network in spirit of metrics used in certified robustness, e.g., [1].

C3: The main selling point of the approach is that to obtain the distribution over outputs, only a single pass is needed. However, it's currently unclear how much speedup this method provides compared to, e.g., the linearized Laplace (GLM)...

We are currently conducting an empirical analysis of the efficiency of our method. We will post the results as soon as we have them and update the paper accordingly.

C4: To make your selling point stronger, I suggest additional experiments with larger models such as Llama-3 7B etc. [...]

We agree that applying our method to larger language models like Llama-3 [2] is an exciting direction. While extending our work to, for example, Llama-3 is possible, it is due to the architectural choices a non-trivial task. In particular, the choice of the SwiGLU activation functions [3] $\operatorname{SwiGLU}(x, W, V, b, c, \beta)=\operatorname{Swish}_\beta(x W+b) \otimes(x V+c)$, where $W, V, b, c, \beta$, which is not a scalar activation function anymore and can lead to additional computational challenges. Similar challenges are to be expected w.r.t. the RMSNorm [4]. Propagating distributions through these parts may require additional approximations to obtain an efficient algorithm. Nevertheless, in the paper, we applied our method to GPT-2 and ViT as an important initial step to show that our method is indeed scalable to large-scale models (see Section 4.2). Extending our work to larger families of models, such as Llama-3, is an interesting future direction.

Q1: Suppose the user has their model that has been Laplace-approximated via laplace-torch. What does the user need to do to enjoy the proposed method?

We plan to package our method into a library and integrate our method with the laplace-torch library for easy use after acceptance.

---

References:
[1] Changming Xu and Gagandeep Singh. "Cross-input Certified Training for Universal Perturbations", In ECCV 2024.

[2] Touvron, Hugo, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière et al. "Llama: Open and efficient foundation language models." arXiv preprint arXiv:2302.13971 (2023).

[3] Shazeer, Noam. "Glu variants improve transformer." arXiv preprint arXiv:2002.05202 (2020).

[4] Zhang Biao and Rico Sennrich. "Root mean square layer normalization." In NeurIPS 2019.

We thank reviewer oXsH for their detailed and down-to-point review. We appreciate that they found our proposed method simple yet effective and that the experiments cover significant breadth and appreciate the sensitivity analysis. We address both comments and questions below.

C1: The main weakness of the paper is that it does not consider alternative methods for approximating the predictive distribution of a BNN with a single forward pass. You do not compare the method to last-layer variants of the Laplace approximation and MFVI. This is important to determine if there is any benefit in trying to linearise the neural network to allow for more efficient predictions vs. just treating the network up to the last layer deterministically.

First, we want to mention that in the case of last layer Laplace, we recover the method by [1]. Moreover, even though last layer Laplace approximation have empirically shown to be successful, they are not universal conditional density estimators [2], which might limit their effectiveness in practice. Hence, to assess the quality of our local linearisations we compared against moment-matching, as performed in DVI [3], with results shown in Table 3. We are happy to add additional results on last layer Laplace to the paper.

C2: I think the title is too general.

We agree that our current title might not be specific enough to reflect the content, and we have changed the title as suggested.

Q1: When mentioning stable distributions, please define them informally in one sentence first.

Thank you for pointing this out, we added the definition in the paper.

Q2: Have you considered applying this method to architectures that use convolutions?

We considered a transformer architecture for vision tasks due to its popularity. Applying our method to convolution layers is an interesting avenue and we will add details into the appendix as soon as possible.

Q3: I find the word “useful” in Figure 1 confusing. What makes this “useful” outlier detection/sensitivity analysis? What is non-useful outlier detection/sensitivity analysis?

We have changed the wording in Figure 1 to avoid unnecessary confusion.

Q4: In the paragraph about ensemble methods in the related work [...] there are efficient ensembling approaches that only require a single forward pass and could be mentioned here [...].

Thank you for the suggested related works. We have added the suggested references on single forward-pass ensemble approaches to Section 2 in the revised paper.

Q5: Have you thought about how to improve the usability of your method?

We will release the code publicly and plan to package our method into a library for easy use upon acceptance.

---

References:
[1] Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. "Being Bayesian, even just a bit, fixes overconfidence in ReLU networks." In ICML 2020.

[2] Mrinank Sharma, Sebastian Farquhar, Eric Nalisnick, and Tom Rainforth. "Do Bayesian Neural Networks Need To Be Fully Stochastic?" In AISTATS 2023.

[3] Anqi Wu, Sebastian Nowozin, Edward Meeds, Richard E. Turner, José Miguel Hernández-Lobato, and Alexander L. Gaunt. "Deterministic variational inference for robust Bayesian neural networks". In ICLR 2019.

I still think it would be valuable to compare to a last-layer Laplace approximation, especially in the larger scale experiments in section 4.2. [...]

We are starting the comparison experiment for last-layer Laplace approximation for ViT Base now, and post the result as soon as possible.

an empirical analysis of the computational cost of your method vs. last-layer/GLM Laplace would be interesting here.

We have added a wall-clock time results for a Bayesian MLP (c.f., Section 4.1) and a partially Bayesian ViT base (c.f., Section 4.2) to the Appendix B.7. In the case where more layers are treated Bayesian, we have comparable wall-clock with GLM and faster than sampling. We will add wall-clock times for last-layer Laplace after the experiment is finished.

reviewer 8AJz pointed out that tuning the scaling factor requires validation data -- it might potentially be possible to use a Laplace approximation of the marginal likelihood to tune this hyperparameter without the need for validation data, at least when you use a Laplace approximation to fit the approximate posterior distribution over the weights

Thank you for the suggestion. It would indeed be interesting to explore the approximated marginal likelihood from LA to tune the scaling factor. We added details to the paper.

Have you considered applying this method to architectures that use convolutions?

We have added a derivation for convolutional layers in Appendix A.7.

We thank reviewer KRRC for the thoughtful comments and discussion and thank you for pointing out typos. We address the comments and questions below.

C1: "[...] they do NOT specify the number of samples used in the MC sampling baseline."

We thank the reviewer for pointing out the missing detail about the number of samples used in the MC sampling baselines. This was not done intentionally. We have added the information on the number of MC samples for every experiment in Section 4. For both MFVI and LA Sampling, we used $1000$ samples in regression and classification experiments in Section 4.1, and $50$ samples for ViT and GPT models in Section 4.2.

C2: "Can the author provide an ablation study on number of samples used?"

We have added an ablation study on the number of samples for the MC sampling baselines into Appendix B.4 of the revised paper. We show the ACC and NLPD evaluated on the test set, and measured the wall-clock time for evaluating the whole test set for different numbers of MC samples.

Moreover, we want to point out that MC sampling has shown to be problematic in Laplace approximations, e.g., [1,2], and results in a computational hurdle for larger models. Our work focuses on making predictions in Bayesian deep learning more practical by providing a computationally feasible approximation to the posterior predictive.

C3: "Residual connection is essentially a sum, and the elements involved here are two Gaussians [...] this step hide a big approximation: you can't store a mixture of Gaussians here [...]"

We believe that there might be a misunderstanding. In the case of linear layers and residual connections, we end up with a sum of Gaussian random vectors (RVs) rather than their distribution. For residual connections, this consists of Gaussian RVs representing the inputs and outputs of a non-linear layer due to the skip connection. To obtain the distribution over the outputs of the residual connection, we assume independence between the input RV and output RV of the non-linear layer. Hence, the output RV follows a Gaussian distribution. We agree that assuming independence is a strong assumption, especially for the residual blocks, but find that this empirically works well while remaining computationally feasible. We elaborate on this in the limitations in Section 5, and a detailed derivation for the case of a linear layer is given in Appendix A.1. Please let us know if there are any further questions.

Line 80 "recent work (2024) developed ... and have shown good performance (2016) (2019)". I get what you mean, but the way the sentence is formulated imply that something developed in 2024 was used in 2016 and 2019

To clarify, the work by [3] (2024) is the first time MFVI was successfully applied to the large scale models [4] (2016) and [5] (2019). We believe there might have been a misunderstanding. For completeness, below we are quoting line 80 from our submission:

Recent work by Shen et al. (2024) developed an optimiser to ease the use of MFVI, and have shown good performance on large-scale models such as ResNets (He et al., 2016) and GPT (Radford et al., 2019).

---

References:
[1] Lawrence, Neil David. "Variational inference in probabilistic models." PhD diss., University of Cambridge, 2001.

[2] Alexander Immer, Maciej Korzepa, and Matthias Bauer. "Improving predictions of Bayesian neural nets via local linearization". In AISTATS 2021.

[3] Yuesong Shen, Nico Daheim, Bai Cong, et al. "Variational Learning is Effective for Large Deep Networks". In ICML 2024.

[4] Saining Xie, Ross Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He. "Aggregated Residual Transformations for Deep Neural Networks". In CVPR 2016.

[5] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. "Language Models are Unsupervised Multitask Learners". OpenAI blog 2019.

We thank reviewer 8AJz for their constructive review and appreciate they find our work novel, versatile, and solidly validated. We address both comments and questions below.

C1: The treatment of attention layers (using deterministic queries/keys) needs stronger justification.

Recall that the attention score is computed as $\text{Softmax}(\mathbf{Q}\mathbf{K}^{\top})$. In case we treat queries and keys as random vectors, the softmax function will "squish" the uncertainties over queries and keys and not result in a random vector following a distribution close to a Gaussian. Henceforth, we decided to treat queries and keys deterministically, which also elevates a potential computational bottleneck. A possible remedy is to leverage an approximation to the softmax function such as [1]. We have added a discussion on this to Section 3.2.

C2: The scaling factor for predictive variance is tuned on validation data - this seems ad hoc.

The use of our scaling factor is similar in notion to the pseudo-count used in [2]. But we agree that it is somewhat ad hoc. Hence, we mention the scaling factor in the Limitations paragraph in Section 5.

C3: No discussion of potential failure modes or limitations of the local linearisation assumption.

We have added a Limitations paragraph in Section 5 of the revised paper, where we mention the induced error from local linearisation and independence assumptions, as well as the need for a validation set for fitting the scaling factor.

---

References:
[1] Jiachen Lu, Jinghan Yao, Junge Zhang, Xiatian Zhu, Hang Xu, Weiguo Gao, Chunjing Xu, Tao Xiang, Li Zhang. "SOFT: Softmax-free Transformer with Linear Complexity". In NeurIPS 2021.

[2] Hippolyt Ritter, Aleksandar Botev, and David Barber. "A scalable laplace approximation for neural networks." In ICLR 2018.