Variational Bayesian Last Layers

We thank the reviewer for their very positive feedback.

To increase the impact of the paper you need to make sure that people that are not too familiar with VI are able to easily implement the paper. This means: Make the code publicly available, especially to show how to best implement the "mixed parameterization" discussed in appendix D set good default hyperparameters

We are currently finalizing public releases of our code in both pytorch and jax, which will be released in time for the camera-ready version of the paper.

It would be useful to draw the graphical models of the models presented in Section 2, to help the reader visualize the random variables in play and their (hyper)priors/parameters In case you need more space in the paper, I would move to the appendix some of the details on the generative classification model, especially considering the poorer performances.

We thank the reviewer for these points. We have moved some details to the appendix, as noted.

We thank the reviewer for their very positive comments.

Visualizations of how tight or loose the bounds in the main text could help build more intuition; comparisons in terms of speed or efficiency to variational inference algorithms that do require sampling (such as Monte-Carlo objectives like VIMCO) could also help guide practitioners in making the correct trade-off depending on FLOPs of compute available versus the required accuracy of posterior approximation/uncertainty quantification.

We conducted experiments to investigate tighter bounds on the softmax (based on Knowles and Minka, 2011), which did not meaningfully change performance. Thus, we believe that (surprisingly) the tightness of the bound does not meaningfully impact performance, and do not emphasize it. We believe complexity/run time versus performance is a more meaningful metric which we emphasize.

For the Resnet image recognition and sentiment analysis experiments, what was the additional compute required (or time taken per iteration, if available)? The sample-efficiency is great, and understanding the practical overhead rather than theoretical complexity would be great for larger models that are in broad use.

We have included the runtime for our model compared to a standard model in for the ResNet.

We thank the reviewer for their comments.

The abstract promises "improve[d] predictive accuracy, calibration, and out of distribution detection over baselines.", similarly in the contributions, and conclusion parts of the paper. Even more, the method not only improves, but it also performs "extremely strong" and "exceptionally strong". These are some exceptionally strong statements given the actual performance. Focusing on each of the experiments in turn. The first problem is the presentation, e.g., what does a bold number mean (see question below)?

We understand that you feel the claimed contributions of the method are stated too strongly. We have revised our stated contributions to be more moderate.

Regression Experiments. Of the six data sets (see question below on this number) the proposal improves on two, slightly on one, equally on two (although better than neural net-based baselines), and worse than most of its baselines in the final one. Calling this "strong performance" is rather misleading. Two additional, though potentially minor, problems are that all of the baselines are simply cited from prior work (Watson et al., 2021). Given the wide performance variations between different train/test splits that can be observed for various UCI data sets the results are not entirely trustworthy. (Note that the reported error intervals are most likely standard errors, as is common on UCI, instead of standard deviations. But which they are is never specified.) Secondly, the authors acknowledge in the appendix that there might be differences in the way the training data is normalized compared to the cited results. (Whether these problems strongly influence the results, or bias them in favor or against VBLL is unclear.)

We partially disagree with the reviewer about the strength of the method on regression experiments. The cases where our approach outperforms baselines, the outperformance is substantial. Otherwise, the method typically performs on par with the previous state of the art. Indeed, we believe our method is one of the largest jumps in performance on this task reported. Moreover, we note competitive baselines such as RBF GP and Ensembles are dramatically more expensive than our method.

While we report baselines from (Watson et al. 2021), we note several things. First, part of the reason we built upon their results is the robustness of their experimental procedure. They perform 20 trials in which the train/val/test sets are chosen randomly, and the val set is used to automatically select the termination point of training. We follow this procedure in our experiments. On the topic of normalization, we believe our normalization matches that used in (Watson et al. 2021). In particular, they say that they use a zero-mean, unit-covariance prior “in whitened data space”. It is unclear if this denotes a direct transformation of their data (which would potentially skew reported NLLs) or indicates that prior terms are transformed to match the data. We center the data to be zero-mean, and normalize features to unit covariance in the input to the network, but do not re-scale the labels. We combine our re-scaling of the data with a zero-mean prior. Thus, comparing these strategies, we believe the practical effects of our data processing are nearly identical, although the effect of normalizing feature inputs to the network may vary.

Classification. While "Extremely strong accuracy results" are mentioned, it just performs as well or worse than competitive baselines like Laplace or Ensembles. The same for ECE, NLL, OOD detection. where "exceptionally strong performance" is claimed. The method is somewhat simpler than baselines, but it lacks a convincing argument for why this should matter. As the authors advertise this simplicity, there should be additional results on practical runtime improvements compared to the baselines to provide some evidence for the claim that a reader should use this approach.

We have weakened our claims. While we do not strictly dominate in all metrics, the performance of our approach is highly competitive, which is especially notable given the simplicity of the approach.

Direct comparison of runtimes for each model is difficult, because some methods are post-hoc only, some require a per-epoch step, and some require a change to each training step. Thus, timing depends on other training details such as batch size. We have included runtimes for our method in comparison to standard (non-Bayesian) network evaluation, and expanded our discussion of complexity for baseline methods.

We thank the reviewer for their comments, which we address point-by-point.

Overall: My main concern with the paper is the weak empirical evaluation and limited novelty of the work, that is, it seems it is essentially an application of known techniques to the special case of last-layer posteriors.

We agree with the review that the methodological novelty may appear limited compared to other papers on Bayesian deep learning. However, our fundamental motivation is that the recent history of machine learning has made clear the importance of building inexpensive but effective methods, and the importance of nailing the details in methods. We believe that the depth of our work–derivation of novel deterministic training objectives, experimental analysis, treatment of complexity and parameterization, and analysis of hyperparameters–is critical to building useful systems that may be deployed at scale.

We have added several elements to the experimental evaluation. In particular, we have added an evaluation of post-training methods (in which the VBLL is added to a model with frozen features) and added a bandit problem, in which VBLLs show extremely strong performance.

Section 2.4 lists various related works, which I believe the author claims to optimize the log marginal via gradient descent. I have not checked every citation, but it appears to me that this statement is false for at least a subset of the cited papers. It might be good to revise the exposition.

We emphasize that the list of papers in section 2.4 are papers in which Bayesian last layer models have been used. Not all of these papers used the marginal likelihood objective for training these models; we have clarified this in the text.

Eq 12 is some weighted ELBO, weighted with T for the purpose of generality, according to the authors. However, T never seems to be used later and makes the connection to the common ELBO less transparent. I believe the paper would improve in clarity if T is dropped.

The goal of including the factor of 1/T is to yield an objective that is more amenable to standard minibatch optimization, as is used in standard supervised learning. This is especially important to highlight when aggregation across minibatches should be done by averaging (as opposed to summation) and to weight the regularization terms. We have highlighted the importance of this.

Section 3.4 is very dense and it could help the reader if this section is improved in its presentation.

We have made a collection of changes to the writing in this section to improve presentation, including splitting this section into two sections: one on training objectives and one on prediction and OOD.

For the experiments, I would have expected assessments under distributional shift, a comparison to recent deterministic approaches (e.g., Zeng et al 2023 or Dhawan et al 2023), and a large-scale application of the approach as it acts on the last-layer only and should be applicable in more realistic scenarios (e.g., ImageNet).

We emphasize that, in terms of model parameter count, the wide ResNets are relatively large (~40M) parameters. Indeed, they are much larger compared to non-wide ResNets. We prioritized these experiments to enable comparison to established baselines such as SNGP. To include experiments on scalability, we also included the experiments in which we trained models on LLM features.

To include experiments that are more similar to distribution shift (and also include elements of active learning to investigate the usefulness of the uncertainty quantification) we included bandit experiments, as described in the response to all reviewers.

Page 3, Eq 11 cites "Harrison et al 2018", which I looked up but didn't find any relevant content that would discuss the use of the marginal in Bayesian deep learning as a standard objective. What is the reason for the citation?

[Harrison ea 2018] trains on the marginal likelihood, computed analytically in the regression case by iterating over the data, in the multi-task/few-shot learning setting. It also uses a Bayesian last layer approach that is structurally the same as our regression model.