Paper: Lightweight uncertainty modelling using function space particle optimization (ICLR-2024)

Reviewer comment (Experiments):
most importantly, the paper sorely lacks a competitive baseline for efficient uncertainty estimation. DDU underperforms even the single net baseline on all non-MNIST benchmarks, on some of them by a large margin (CIFAR100 in Table 2, bottom panel of Fig 4). I don't know if DDU hasn't been tuned properly or simply isn't suited to tasks with a large number of classes, but in either case there are plenty of alternatives available in the literature that have reported good results in this setting, e.g. (Liu et al., Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness, Neurips 2020) or (Kristiadi et al., Posterior Refinement Improves Sample Efficiency in Bayesian Neural Networks, 2022).

Author response:
Dear Reviewer,

We appreciate your evaluation and thank you for your feedback.


We acknowledge the importance of competitive baselines and recognize that the absence of alternatives in our current version is a limitation. Due to time limitations we could not include them in this submission. However, we added ablation studies with experimental results for additional base network architectures and varying spectral normalization coefficients. We show that our method improves OOD detection performance for all analyzed base network architectures. We observe the best performance for the Resnet-18 architecture, although no strong correlation between spectral normalization strength and OOD detection performance. The appropriate choice of context points is still the most important factor in all experiments. The key point of this this work is that the use of the function space variational inference method with an appropriate choice of context points can significantly improve the performance of a single network. 

The additional ablation studies in the appendix are structured as follows:

A4) Base network architecture 

A5) Ensemble head architecture

A6) Number of particles

A7) Choice of context points 


We genuinely appreciate your feedback and the opportunity to improve our paper. Thank you for your time and consideration.

Reviewer comment (Experiments):
the paper is a bit light on ablations. In particular the use of the spectral normalization is not justified empirically. From a theoretical standpoint, as far as I am aware, prior work has justified its need based on the network being distance preserving being necessary for the use of a distance aware output head, however I do not see the proposed method relying on this property, so I would like to see a baseline without spectral normalization.

Author response:
Dear Reviewer,

We appreciate your evaluation and thank you for your feedback.


We acknowledge the importance of competitive baselines and recognize that the absence of alternatives in our current version is a limitation. Due to time limitations we could not include them in this submission. However, we added ablation studies with experimental results for additional base network architectures and varying spectral normalization coefficients. We show that our method improves OOD detection performance for all analyzed base network architectures. We observe the best performance for the Resnet-18 architecture, although no strong correlation between spectral normalization strength and OOD detection performance. The appropriate choice of context points is still the most important factor in all experiments. The key point of this this work is that the use of the function space variational inference method with an appropriate choice of context points can significantly improve the performance of a single network. 

The additional ablation studies in the appendix are structured as follows:

A4) Base network architecture 

A5) Ensemble head architecture

A6) Number of particles

A7) Choice of context points 


We genuinely appreciate your feedback and the opportunity to improve our paper. Thank you for your time and consideration.

Reviewer comment (Experiments):
Similarly, the structure of the ensemble heads appears as a magic hyperparameter without justification. I realize that a linear head without hidden layers will probably underperform due to the categorical log likelihood being a convex/unimodal objective w.r.t. the weights of the linear output layer, but this should be justified explicitly and sensitivity of the method w.r.t. this hyperparameter evaluated.

Author response:
Dear Reviewer,

We appreciate your evaluation and thank you for your feedback.


We acknowledge the importance of competitive baselines and recognize that the absence of alternatives in our current version is a limitation. Due to time limitations we could not include them in this submission. However, we added ablation studies with experimental results for additional base network architectures and varying spectral normalization coefficients. We show that our method improves OOD detection performance for all analyzed base network architectures. We observe the best performance for the Resnet-18 architecture, although no strong correlation between spectral normalization strength and OOD detection performance. The appropriate choice of context points is still the most important factor in all experiments. The key point of this this work is that the use of the function space variational inference method with an appropriate choice of context points can significantly improve the performance of a single network. 

The additional ablation studies in the appendix are structured as follows:

A4) Base network architecture 

A5) Ensemble head architecture

A6) Number of particles

A7) Choice of context points 


We genuinely appreciate your feedback and the opportunity to improve our paper. Thank you for your time and consideration.

Reviewer comment (Experiments):
evaluate the sensitivity of the method w.r.t. depth/width of the ensemble head and discuss the trade-off vs. computational efficiency.

Author response:
Dear Reviewer,

We appreciate your evaluation and thank you for your feedback.


We acknowledge the importance of competitive baselines and recognize that the absence of alternatives in our current version is a limitation. Due to time limitations we could not include them in this submission. However, we added ablation studies with experimental results for additional base network architectures and varying spectral normalization coefficients. We show that our method improves OOD detection performance for all analyzed base network architectures. We observe the best performance for the Resnet-18 architecture, although no strong correlation between spectral normalization strength and OOD detection performance. The appropriate choice of context points is still the most important factor in all experiments. The key point of this this work is that the use of the function space variational inference method with an appropriate choice of context points can significantly improve the performance of a single network. 

The additional ablation studies in the appendix are structured as follows:

A4) Base network architecture 

A5) Ensemble head architecture

A6) Number of particles

A7) Choice of context points 


We genuinely appreciate your feedback and the opportunity to improve our paper. Thank you for your time and consideration.

Reviewer comment (Experiments):
I have doubts about the fairness of the comparisons, as 20-30 particles are used for MH-POVI and variations, but only an ensemble of five networks. I believe a fairer comparison is to equalize the number of ensemble networks and particles, or to evaluate both methods with a variable number of particles and ensemble members, which also would provide information about how both methods scale with the number of particles.

Author response:
Regarding your doubts towards the fairness of the comparison: We compared with a deep ensemble consisting of 5 ensemble members in order to keep the number of additional parameters comparable. We added additional ablation studies (see appendix A5, A6), where we evaluate the influence of the ensemble head network architecture and the number of particles. There, we show that for the image classification OOD detection task an ensemble of linear layers can already achieve competitive results. However, this is only achieved for an appropriate selection of context points. By using linear layers, we only add minimal computational and memory costs for the forward pass. Thus, we argue that the comparison to the 5-member deep ensemble is fair, given the significant reduction in parameter cost. However, in our additional experiment we show the results of our ensemble head consisting of 5 particles (see Figure 11) to address your concern.

Reviewer comment (Novelty):
About multi-head architectures for uncertainty, there is also the concept of a sub-ensemble, which is very similar to the proposed method in this paper minus the particle optimization, it is meant to be an approximation to a deep ensemble, here I do not suggest a comparison but just integration into the literature for multi-head network architectures.

Author response:
Thank you for bringing the work of Valdenegro-Toro (2023) to our attention. In the revised manuscript, we will incorporate a reference to this work in the related works section.

Reviewer comment (Reproducibility):
Could you clarify how the network heads are trained using function space particle optimization? In particular how Eq 2 is evaluated.

Author response:
The network heads with the function space particle optimization are trained as follows: 
In each training iteration we load a batch of training data and context points. The context points are used to approximate the particles in function space. For these context points, the repulsion kernel (in our case a Gaussian kernel) enforces diverse predictions of each ensemble member. The other objective, the posterior term, enforces that all particles correctly predict the training data.

Reviewer comment (Reproducibility):
Is the model trained end-to-end? Please clarify the training process.

Author response:
We train our model post-hoc without changing the base network in order to maintain a fair comparison to DDU and the base network scores. Our key observation is the importance of context points for the uncertainty estimation performance (see appendix A5-A7). In all added ablation studies, we test our method for the same selection of context points as in the main paper. In each experiment we make the same observation: The largest improvements are obtained by using unlabeled data points that are related to the training data. As you correctly mentioned, using Gaussian noise does not improve the uncertainty estimates of the function space method compared to the unconstrained ensemble head. This highlights the need and usefulness of context points that carry additional information - in this case, the post-hoc extension to the single model can be significantly improved.

Reviewer comment (Theory):
Can you motivate the selection of context points? For example using Gaussian noise as a context point, how do you know that these points are supported by the data distribution?

Author response:
Our key observation is the importance of context points for the uncertainty estimation performance (see appendix A5-A7). In all added ablation studies, we test our method for the same selection of context points as in the main paper. In each experiment we make the same observation: The largest improvements are obtained by using unlabeled data points that are related to the training data. As you correctly mentioned, using Gaussian noise does not improve the uncertainty estimates of the function space method compared to the unconstrained ensemble head. This highlights the need and usefulness of context points that carry additional information - in this case, the post-hoc extension to the single model can be significantly improved.

Reviewer comment (Experiments):
The authors evaluated the proposed multi-headed architecture on several benchmark tasks. However, the reasons for the selection of these benchmark tasks were not discussed.

Author response:
2) The experimental setting has been selected to be similar to cited related works to simplify comparison.

Reviewer comment (Experiments):
The topic of the article is uncertainty estimation under lightweight, the author's description of lightweight is not sufficient, and the analysis of lightweight is lacking in the experimental part.

Author response:
Dear Reviewer,

we appreciate your thoughtful review on our paper. We have carefully considered your review, and we would like to share the revisions and clarifications we have made in response to your comments.

1) We provide additional ablation studies in the appendix of the revised version. The additional ablation studies in the appendix are structured as follows:

A4) Base network architecture 

A5) Ensemble head architecture

A6) Number of particles

A7) Choice of context points 

We show that for the image classification task, a repulsive ensemble of linear layers is sufficient to improve OOD detection given that the context points used for the repulsion term are appropriate and informative. By using unlabeled and related image data (kMNIST for DirtyMNIST, CIFAR100 for CIFAR10), we can improve the single model by a large margin with minimal computational overhead for the forward pass.

3) We show that we are able to improve the uncertainty estimates of a single network by adding only a minimum number of additional parameters. During the training of the ensemble head, we use the same batch size for training data and context points. By relying on the features extracted from the pretrained base network, we show in our additional experiments that the size of the ensemble head can be minimal and still provide informative uncertainty estimates. 

We hope that the additional experiments and discussions contribute to a better understanding of our proposed model.

Reviewer comment (Experiments):
How to preserve the bi-Lipschitz conditions and how to avoid feature collapse? What's the key component to deal with these challenges? Can you show some ablation studies to verify the effectiveness?

Author response:
Additionally, we added ablation studies regarding the influence of the base network and the use of spectral normalization for networks with residual connections (see appendix A4 of the revised version). We show that our method is not as sensitive to the choice of architecture and spectral normalization strength, *given* an appropriate choice of context points. We also show that in the image classification task, an ensemble of repulsive linear layers can already achieve competitive results to deep ensembles (see appendix A5). By using informative context points, e.g. unlabeled and related image data (kMNIST for DirtyMNIST, CIFAR100 for CIFAR10), we can improve the single model by a large margin with minimal computational overhead for the forward pass.

Reviewer comment (Theory):
How to disentangle aleatoric and epistemic uncertainty? Which features mainly handle these capabilities? What's the theoretical foundation of the proposed method or by leveraging existing works?

Author response:
We used the uncertainty decomposition proposed in Depeweg, Stefan, et al. "Decomposition of uncertainty in Bayesian deep learning for efficient and risk-sensitive learning." to distinguish between aleatoric and epistemic uncertainty. The equation is provided on page 6.

Reviewer comment (Evaluation):
How about the computational cost and memory cost comparison with the baseline methods on multiple tasks? I think it is important to show these since the paper mainly claims the high efficiency and low memory requirements.

Author response:
Regarding the choice of number of particles we added an ablation study (appendix A6). E.g, for CIFAR10 we can achieve competitive results when using 20 linear layers on top of the pre-trained feature space, which is a small overhead in terms of computational and memory costs compared to the Resnet-18 base network.