Paper: Variational Bayesian Last Layers (ICLR-2024)

Reviewer comment (Reproducibility):
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.

Author response:
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.

Reviewer comment (Presentation):
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.

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

Reviewer comment (Evaluation):
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.

Author response:
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.

Reviewer comment (Evaluation):
In the regression experiments, of the six data sets 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.

Author response:
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.

Reviewer comment (Experiments):
In the regression experiments, 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.)

Author response:
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.

Reviewer comment (Experiments):
In the regression experiments, 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.

Author response:
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.

Reviewer comment (Evaluation):
In the classification experiments, while "Extremely strong accuracy results" are mentioned, it just performs as well or worse than competitive baselines like Laplace or Ensembles. The same holds for ECE, NLL, and OOD detection, where "exceptionally strong performance" is claimed.

Author response:
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.

Reviewer comment (Presentation):
When submitting to ICLR please make sure that you follow the ICLR style guide. For example, table captions belong above tables, not below.

Author response:
We have addressed these typos.

Reviewer comment (Writing):
Sec 5.2 first paragraph has a typo: "We refer to this model **as** G-VBLL-MAP..."

Author response:
We have addressed these typos.

Reviewer comment (Writing):
In equation (6) there is a typo: $y|\rho$

Author response:
We have addressed these typos.

Reviewer comment (Novelty):
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.

Author response:
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.

Reviewer comment (Writing):
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.

Author response:
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.

Reviewer comment (Writing):
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.

Author response:
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.

Reviewer comment (Presentation):
Section 3.4 is very dense and it could help the reader if this section is improved in its presentation.

Author response:
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.

Reviewer comment (Experiments):
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).

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
[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.

Reviewer comment (Theory):
Eq. 14 uses a rather loose bound, is it possible that this is the reason the approach underperforms in the classification settings compared to the regression setting? If so, is there any way to obtain a tighter bound?

Author response:
Several tighter bounds have been proposed, which we mostly did not address in the paper for simplicity. These are best discussed in:

Non-conjugate Variational Message Passing for Multinomial and Binary Regression, Knowles and Minka, NeurIPS 2011. 

We investigated the adaptive bound from this paper, in which the bound is tightened by exploiting the shift equivariance of the LSE function. This bound introduces new variables which control the shifting of the bound, and can be trained along with other network weights. 

We found that the adaptive bound did not meaningfully change the performance of our method. Because the features are learned, it is likely that the automatic bound tightening enabled by Knowles and Minka’s bound was already happening with the simpler Jensen-only bound.

Reviewer comment (Writing):
From what I understand T is never actually used. Is this correct?

Author response:
T is an important term in the weighting of the regularization terms, and is an important feature of the algorithm. In particular, our experiments set the weighting of the KL term to be 1/T. We have clarified this in the paper.