Evaluating and Improving Subspace Inference in Bayesian Deep Learning

Dear reviewer K1aU, we sincerely thank you for your detailed advice and the time you spent on our paper. Here are our replies to your comment:

Weakness:

W1: The overall methodology seems to be a straightforward combination of three existing (with minor modifications) approaches.

A1: Thank you for your comment. Our work introduces (1) a novel subspace construction method (block-averaging) that captures global trajectory efficiently, (2) new metrics (Bayes factors and evidence ratios) to evaluate subspace quality—an overlooked aspect in prior work—and (3) an efficient approach for posterior predictive computation in low-dimensional subspaces. These contributions go beyond simply combining existing methods and address limitations in prior approaches.

W2: During the subspace reconstruction, do you need to flatten the matrix from $mxn$ to $d$? What if you keep the matrix structure, but perform SVD on the matrix trajectory to get $USV^\top$, and treat $S$ as your $Z$?

A2: Thank you for your question. For all model weights, we flatten them into a $d$-dimensional vector like prior works [1, 2]. While we acknowledge that this may lose some structural information, retaining the matrix structure (e.g., performing layer-wise SVD) introduces challenges. For instance, different layers have weight matrices of varying sizes, making it unclear how to combine these layers within a unified subspace. For example, VGG-16 has about 138 million parameters across layers of varying dimensions, which complicates maintaining matrix structures for subspace construction. If there are relevant studies addressing these issues, we would be glad to explore them further.

W3: For the image classification and UCI, I am curious about the full trajectory performance. Why don't you report it?

A3: Performing PCA or randomized SVD on the full trajectory requires storing all $n$ trajectories. When the dimensionality of model weights is high, the memory and computational costs become prohibitively expensive, making full trajectory subspace construction infeasible in such scenarios. We have included this clarification in the revised manuscript (Lines 496-497). Thank you for your feedback.

[1]. Li et al. "Measuring the intrinsic dimension of objective landscapes", 2018.

[2]. Izmailov et al. "Subspace inference for Bayesian deep learning", 2020.

Thanks for the author's feedback, the response address some of my concerns. However, the block average is one simple extension (i.e. take samples with certain interval) that is not difficult to think of; the new metrics Bayes' factors and evidence ratios are also based on marginal likelihood, which has been used for model selection before. Therefore, these contribution to methodology seems to be a bit incremental. Therefore, I keep the original score.

Dear reviewer bE76, we sincerely thank you for your detailed advice and the time you spent on our paper. Here are our replies to your comment:

Weakness:

W1: The novelty of the proposed algorithm is arguable, both BA and TT subspaces are constructed from the SGD points after some initial warm-up process.

A1: While BA and TT both use the SGD points (like SWA), BA introduces a new strategy for finding the subspace direction. The TT strategy in [1] practically used the last $M$ points of the SGD trajectory, because the spacing parameter $c$ is set to 1 in all experiments. Even with $c>1$, TT still ignores the majority of the points from the SGD trajectory, and the center of these $M$ points is not $\hat{w}$, meaning the corresponding SVD decomposition no longer aligns with the interpretations of PCA. In contrast, BA divides the entire trajectory (after warm-up) into $M$ equidistant blocks and computes the block averages, capturing the global variability over the full trajectory while maintaining the same computational cost as TT.

W2: Figure 2 seems to be misleading because both BA and TT subspaces are constructed from the SGD points after some initial warm-up process that is the same in the proposed paper and in [1]. For example, for Synthetic Regression, both methods use the last 1000 points.

A2: In Figure 2, we did not include the warm-up trajectory; instead, we drew the trajectories after the warm-up phase for the online update of the mean $\hat{w}$. We believe Figure 2 precisely demonstrates the difference between FT, TT, and BA. For example, for the Synthetic experiment with $M=20$, all methods use the same trajectory from step 2000 to step 3000. BA divides this trajectory into 20 blocks, where each block represents the mean of 50 points, forming a matrix that is applied to an SVD to obtain the projection matrix $P$, as detailed in Appendix B. In contrast, TT uses only the points from step 2980 to step 3000.

W3: The provided results should contains more detailed analysis of the metrics, for example, one can provide a similar Figure as Figure 4 from [1].

A3: Figure 4 in [1] focused log-likelihoods on the UCI task, which we have reported in Tables 4 and 5 for comparison. Our evaluation extends beyond log-likelihood to include new metrics like Bayes factors and evidence ratios.

We emphasize that BNNs are designed to evaluate uncertainty in both model weights and predictions. Metrics like test log-likelihood primarily reflect accuracy on test data, which, as we observed, can often be achieved with a single model lacking proper uncertainty quantification. To address this limitation, in Section 4.2, we proposed new quality evaluation criteria.

W4: Incomplete comparisons between methods. For example, TT (RQMC) is missing from Tables 7 and 8. Also, there is no comparison between RQMC and SNIS.

A4: Thank you for highlighting this. We have included TT (RQMC) results in Tables 7 and 8 in the updated manuscript. The results show that TT (RQMC) achieves similar accuracy to TT (ESS) but with lower computational cost. Furthermore, BA (RQMC) outperforms TT (RQMC) on both the CIFAR and corrupted CIFAR datasets.

Table 7: Classification accuracy (ACC(%)) on CIFAR datasets

Table 8: Classification accuracy (ACC(%)) on corrupted CIFAR datasets using PreResNet164

Regarding RQMC-IS vs. SNIS, Theorem 2 demonstrates that RQMC consistently outperforms IS in convergence rate, supported by Table 2's RMSE results. Thus, we prioritized RQMC-IS over SNIS.

W5: There is some inconsistency between the results from the paper Figure 4 and the prior work. Is there any difference in the experiment setup that caused this difference?

A5: The key difference lies in the choice of the priors. We use a fixed prior independent of the subspace (Lines 166-168), ensuring a consistent posterior across the entire parameter space. In contrast, [1] modifies the prior within each subspace, causing different prior across different subspaces.

A fixed prior ensures the posterior reflects only subspace properties, not artifacts from varying priors. For example, if the subspace center $\hat{w}$ shifts within the same subspace, then $\mathcal{Z}$ remains unchanged and the posterior structure should not vary. We clarified this in the revised manuscript (Lines 168-170).

Dear reviewer hUx8, we sincerely thank you for your detailed advice and the time you spent on our paper. Here are our replies to your comment:

Questions:

Q1: Is it appropriate to interpret the proposed subspace construction method as a kind of downsampling scheme to the FT method?

A1: Thank you for raising this insightful point. We agree that BA can be viewed as a form of downsampling applied to the full trajectory method. We now describe this interpretation in the revised manuscript as follows (Line 183):

BA can be viewed as a structured downsampling of the full trajectory, providing a similar subspace while significantly reducing computational and memory costs, as shown in Figure 1.

Q2: How the block-averaging (BA) method differs from or improves upon SWA, TWA, and DLDR?

A2: SWA, TWA, and DLDR focus on finding a single weight $\hat{w}$ based on the SGD trajectory. For instance, SWA averages weights to achieve better generalization, while TWA identifies optimal weights within the linear space spanned by $n$ trajectories. DLDR leverages a quasi-Newton-based algorithm to find an optimal solution within the subspace.

In contrast, BA (along with TT and FT) aims to find a set of weights residing in some subspace $\{w: \hat{w} + P z, z\in \mathbb{R}^k\}$. Within such a subspace, inference algorithms (like MCMC, variational inference) are computationally more efficient, and predictions (using all points in the subspace) are combined through Bayesian model averaging. In this work, we propose BA, which constructs a high-quality subspace with limited computational cost, along with some principled subspace quality evaluation methods. We appreciate the reviewer highlighting these related works like SWA, TWA and DLDR, and we have included discussions in the Related Works paragraph (Lines 74-76).

Q3: Clarify the novel aspects of BA that go beyond simple downsampling.

A3: The novelty of BA lies in: (1) approximating FT while significantly reducing memory and computational costs; (2) building on SWA’s observation in [1] that averaging weights can improve generalization, BA partitions the trajectory into equidistant blocks and computes block averages, effectively preserving global variability; and (3) ensuring the resulting subspace aligns with PCA interpretations through SVD. Additionally, our work also introduces methods to evaluate subspace quality and efficiently perform predictive computations.

Q4: The results in Table 1 is obvious, so I don’t see much importance of taking much length showing Table 1.

A4: Thank you for the suggestion. We will adjust the discussion of Table 1 accordingly. The purpose of Table 1 is to demonstrate that BA can achieve results close to FT with lower cost, while TT shows significant angular distance in the first principal components compared to both FT and BA.

Q5: How the proposed metric (Bayes factor and evidence ratio) can be used to guide hyperparameter selection?

A5: Thank you for your suggestion. We agree that the Bayes factor and evidence ratio may provide some guidance in hyperparameter selection. For example, Figure 3 indicates that for the BA subspace, $M=5$ is a good choice as it achieves the highest Bayes factor and evidence ratio. While our current work focuses on demonstrating the utility of these metrics for evaluating subspace quality (see A6 for details), applying them for hyperparameter tuning is an important direction in future work.

Q6: Please elaborate the utility of the proposed metrics, and explain how these metrics offer insights that traditional performance metrics may not capture.

A6: Bayes Factors and evidence ratios can apply to any subspaces or subsets of original weights space, not just BA and TT. Bayes Factors evaluates subspaces using only training data, which is often overlooked in current practices for assessing fit without test data. Evidence ratios like prior predictive metrics, do not rely on the posterior derived from training data. Both metrics are computationally cheaper (see A8 for details) than the predictive performance metrics mentioned by the reviewer.