Model and Feature Diversity for Bayesian Neural Networks in Mutual Learning

Summary

This paper proposed a mutual learning approach to learn a pair of Bayesian Neural Network(BNN). The posterior of BNN is approximated by Variational Inference using a Gaussian distribution with a diagonal covariance matrix. To make the BNN learn different perspective of the data, the author proposed to increase the diversity in parameter space and intermediate feature space by adding the an estimate of distance between parameter distribution and fused feature distribution of two BNN models into the objective function. Empirically, the proposed method outperform existing mutual learning method and vanilla BNN model in terms of accuracy, negative log likelihood loss and expected calibration error. An ablation study is also provided to investigate the usefulness of each component.

Strengths

The paper is well written and easy to follow. Increasing the diversity of parameter distribution and intermediate feature distribution of peer BNN models to boost performance is an interesting idea. Experiments and detailed ablation study demonstrate the effectiveness of proposed method.

Weaknesses

1. It is mentioned in the abstract and introduction that the BNN model with variational inference may underperform deterministic model or BNN obtained by MCMC, the baseline only involves BNN model trained with(DML) or without(vanilla) mutual learning. Would the proposed method close the gap to some extent? Data augmentation, optimizer may all affect performance, so it is still helpful to include deterministic model results follow with same training setup. I would expect the BNN model to outperform deterministic model at least in NLL and ECE, and with the 50 ensemble, it can outperform the accuracy. 2. Continue with last point, for MCMC method (e.g. in line 81 of the paper), I agree that traditional MCMC method(e.g. Metropolis Hasting) may not be feasible for large model, and memory storage can be an issue for MCMC method. But I don't think the stochastic gradient MCMC cited in line 81 would require prohibitive computational cost, it behaves like adding a noise to at each step of standard SGD training. 3. The code is not provided so it may hurt the reproducibility of the paper.

Questions

1. To my knowledge, it is not very clear if variational distribution(e.g. Gaussian with diagonal covariance matrix) can approximate the true posterior very well, can the author comment a bit on this, e.g. how would different choice of variational family affect the model? 2. In line 264 and line 6 of algorithm 1, it is mentioned that one BNN model is initialized with a trained model and this lead to better results empirically. Can the author discuss more on why this happened? It is a bit wired for me as it seems in the implementation detail, the pre-trained model and the model from scratch are trained with same optimizer and learning rate schedule. 3. Seems like $\alpha$ $\beta$ are set to 1,2 for CIFAR and 1,1 for Imagenet, these two parameters controls the strength of proposed penalty to the model, can the author comments a bit more on how sensitive are the model to those parameters, It can help to illustrate how diversity helps model performance. 4. As mentioned in line 268, results are average of 3 trials, I think it would be better to include the standard deviation as well to boost the significance of the results. 5. In figure A.3 in supplementary material, looks like a sharp increase of KL divergence between the fused feature distributions at around 30 epochs, but the penalty for feature is only added for last 100 epochs, can the author explain more on this?

Rating

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

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central 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

2 fair

Limitations

The authors addressed the limitations.

Summary

The paper proposes a method to combine deep mutual learning with BNN to diversify the weight distributions of each BNN networks in a pair or ensemble, to improve performance.

Strengths

1. AFAIK this is the first work combining mutual learning with BNN, so the authors can claim this point. 2. The paper is in general written clearly and easy to follow. 3. Experiments are adequate with ablation studies on individual features impact on diversity.

Weaknesses

1. Some design choices are found to be "empirically" working well without too much discussion or hypothesis. 2. Would be interesting to see how the model performs for o.o.d test data, especially uncertainty performance.

Questions

Line 178-179: The authors said adding the D(...) term will rapidly increase of this term and impact training. Wouldn't putting a smaller scaling factor for this term fix this issue?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, 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

2 fair

Limitations

None.

Summary

The paper titled addresses the challenge of improving the performance of Bayesian Neural Networks (BNNs) by leveraging the concept of mutual learning. BNNs provide a means for quantifying uncertainty in predictions through probability distributions of model parameters. However, BNNs often fall short in performance compared to their deterministic counterparts. The authors propose a novel approach that employs deep mutual learning to enhance the capabilities of BNNs.

Strengths

1. Innovative Approach: The paper introduces a novel method that combines deep mutual learning with Bayesian Neural Networks. By promoting diversity in both network parameter distributions and feature distributions, the proposed approach enables peer networks to acquire distinct features, capturing different characteristics of the input data. This innovative technique enhances the effectiveness of mutual learning in BNNs. 2. Detailed algorithm description: The paper provides a thorough and detailed description of the proposed algorithm for improving the performance of Bayesian Neural Networks (BNNs) through deep mutual learning. 3. Comprehensive Experiments: The authors conduct extensive experiments to evaluate the proposed approach thoroughly. The experimental results are statistically sound and demonstrate significant improvements in classification accuracy, negative log-likelihood, and expected calibration error compared to traditional mutual learning methods for BNNs.

Weaknesses

1. Limited variety in experimental validation: One weakness of the paper is that the proposed approach and its effectiveness are only verified through experiments conducted on Residual Neural Networks (ResNets). It would have been beneficial to include experiments on a diverse set of network architectures to demonstrate the approach's effectiveness across different model types and complexities. 2. Lack of detailed explanation for temperature, α, and β: One weakness of the paper is the limited explanation provided for the temperature parameter (T), α, and β, which are crucial components of the proposed approach. These parameters play a significant role in controlling the diversity of network parameter distributions and feature distributions, but their specific effects and optimal values are not thoroughly discussed. 3. Weakness in the conclusion: The current conclusion merely restates the experimental results and does not highlight the broader implications of the proposed approach or its potential impact on the field.

Questions

The author should supplement more experiments to prove its effectiveness.

Rating

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

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

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

The author should supplement more experiments to prove its effectiveness and strengthen the conclusion.

Summary

The paper focuses on improving the accuracy of BNNs by promoting diversity in both parameter space and feature space while training two peer BNNs with mutual learning between them. More specifically, they train two variational BNNs with a mean-field Gaussian variational loss for each along with a KL divergence term between the (temperature-scaled) predictive distributions of the two models, a Wasserstein distance term between the corresponding approximate posterior distributions across the two models (added as a softplus(-distance) term), and a KL divergence term between corresponding feature distributions. On the latter term, instead of directly maximizing the distance between corresponding feature distributions, they instead do so on "fused feature distributions". To do so, they use learned cross-attention to fuse the features from multiple feature levels in a model (two at a time). Then, they use the KL divergence between the distributions of the fused feature distributions of the two peer networks. To derive the distributions, they use the conditional probability density defined as $p_{i|j} = \frac{K(F'_i, F'_j)}{sum_{k=1, k \noteq i}^n K(F'_k, F'_j)}$, where $K(F'_a, F'_b)$ is a kernel function between two fused feature representations. Given those conditional probs, they compute a KL divergence term. Similar to the parameter space diversity term, they add this term to the loss as softplus(-divergence). The paper claims to be the first to propose maximizing the distance between feature distributions to promote diversity. In terms of experiments, the paper includes results for ResNet models on CIFAR-10/100 and ImageNet, measuring accuracy, NLL, and ECE as metrics, and comparing different approaches.

Strengths

The paper does a great job of precisely articulating the modeling approach, and discussing the relevant background info. More specifically, the proposed approach of adding terms to promote diversity in parameter and feature space is clear and would be easy to reimplement.

Weaknesses

My main concern is with the experiment section. More specifically, a few key details are unclear in the text, and importantly a deterministic baseline is missing that I believe should be present given the framing of the paper and relevant literature. Please see the Questions below. Given updates, I believe the paper would be great and I would gladly update my rating.

Questions

Main: - In the experiments, a few details are currently unclear. The following points are on Table 1, but generalize to all three tables. Please clarify these details here and in the paper. - Consider the ResNet20 section of Table 1. Is my understanding correct that the "ResNet20" results are for a pair of BNNs trained from scratch, while the "ResNet20*" results are for a pair of BNNs trained with the approximate posterior means set to the values from a deterministic model? - Is it correct that all results (all three metrics across all three approaches) are computed after averaging the predicted probs from the pair of models? - For the experiments, a deterministic baseline is missing. Given the intro that discusses how BNNs can lag behind deterministic models in acc (though not always), the experiments lack a comparison. It would be helpful to understand how the proposed approach compares to a deterministic baselines, specifically a single deterministic model and a size-2 deep ensemble. Could you add this as a baseline? I would consider this to be a blocker for the paper given the framing and relevant literature. Other: - The KL divergence term is scaled by the square of the temperature -- why? - How did you choose the values for temp, alpha, and beta? They differ between CIFAR-10/100 and ImageNet. Did you ablate values? Minor comments: - updating lines 17 & 22 of Alg 1 could be helpful for readability

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

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

No limitations are included.

Summary

This paper presents a novel method for enhancing the performance of Bayesian Neural Networks (BNNs) by employing deep mutual learning. The proposed approach aims to enhance the diversity of both network parameter distributions and feature distributions, encouraging individual networks to capture unique characteristics of the input data. The effectiveness of the proposed method is demonstrated on datasets, including CIFAR10, CIFAR100, and ImageNet.

Strengths

The proposed method improves performance and uncertainty estimation while reducing the expected calibration error (ECE). The technical approach is novel as the method introduces mutual learning in the context of BNNs and first to propose maximizing the distance between feature distributions and parameter distributions. The paper includes large scale data experiments (ImageNet) and ablation studies to demonstrate the effectiveness of each technical contribution introduced in this paper.

Weaknesses

The previous studies mentioned in the paper utilize alignments on feature maps [4] or predictions [38], rather than diversifying them. In contrast, the proposed method diversifies both feature distributions and parameter distributions which is an opposite approach to the previous works. Interestingly, both alignment-based and diversification methods improves performance over vanilla BNNs, as indicated in Table 1, 2, 3, and 5. However, the paper does not explicitly explain the reasons behind the performance improvements resulting from these contrasting approaches. Given the observed contradicting results in the experiments, where the alignment-based method (DML [38]) also enhances the performance of BNNs, an important question arises: could combining alignment-based methods with parameter diversification further improve BNN performance? Alternatively, is it necessary to diversify both feature and parameter distributions to achieve significant improvements? In the experiment section, the proposed method is only compared with [38] and not with [4]. Hyperparameters used for CIFAR experiments and ImageNet experiments are different. However, the paper does not describe details regarding the hyperparameter tuning or determination.

Questions

3 block resnet is used for CIFAR experiments while 4 block resnet is used for ImageNet experiments. Why different form of resents are used for different datasets?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited 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

3 good

Limitations

Limitations are shortly addressed in the supplementary.

Thank you for the response.

Thank you for clarifying my questions. Based on the reviews and responses, I raise my score to above borderline.

Rebuttal to tB5V

Dear tB5V Could you look at the rebuttal and check if the authors addressed your concerns? Deadline is in just 2 days. Best regards, your AC