Title: Function Space Bayesian Pseudocoreset for Bayesian Neural Networks
Venue: NeurIPS 2023
Accepted: yes
Keywords: Bayesian pseudocoresets

Function space variational inference
Abstract: A Bayesian pseudocoreset is a compact synthetic dataset summarizing essential information of a large-scale dataset and thus can be used as a proxy dataset for scalable Bayesian inference. Typically, a Bayesian pseudocoreset is constructed by minimizing a divergence measure between the posterior conditioning on the pseudocoreset and the posterior conditioning on the full dataset. However, evaluating the divergence can be challenging, particularly for the models like deep neural networks having high-dimensional parameters. In this paper, we propose a novel Bayesian pseudocoreset construction method that operates on a function space. Unlike previous methods, which construct and match the coreset and full data posteriors in the space of model parameters (weights), our method constructs variational approximations to the coreset posterior on a function space and matches it to the full data posterior in the function space. By working directly on the function space, our method could bypass several challenges that may arise when working on a weight space, including limited scalability and multi-modality issue. Through various experiments, we demonstrate that the Bayesian pseudocoresets constructed from our method enjoys enhanced uncertainty quantification and better robustness across various model architectures.

Function Space Bayesian Pseudocoreset for Bayesian Neural Networks

Abstract

A Bayesian pseudocoreset is a compact synthetic dataset summarizing essential information of a large-scale dataset and thus can be used as a proxy dataset for scalable Bayesian inference. Typically, a Bayesian pseudocoreset is constructed by minimizing a divergence measure between the posterior conditioning on the pseudocoreset and the posterior conditioning on the full dataset. However, evaluating the divergence can be challenging, particularly for the models like deep neural networks having high-dimensional parameters. In this paper, we propose a novel Bayesian pseudocoreset construction method that operates on a function space. Unlike previous methods, which construct and match the coreset and full data posteriors in the space of model parameters (weights), our method constructs variational approximations to the coreset posterior on a function space and matches it to the full data posterior in the function space. By working directly on the function space, our method could bypass several challenges that may arise when working on a weight space, including limited scalability and multi-modality issue. Through various experiments, we demonstrate that the Bayesian pseudocoresets constructed from our method enjoys enhanced uncertainty quantification and better robustness across various model architectures.

1 Introduction

Deep learning has achieved tremendous success, but its requirement for large amounts of data makes it often inefficient or infeasible in terms of resources and computation. To enable continuous learning like humans, it is necessary to learn from a large number of data points in a continuous manner, which requires the ability to discern and retain important information. This motivates the learning of a coreset, a small dataset that is informative enough to represent a large dataset.

On the other hand, the ability to incorporate uncertainties into predictions is essential for real-world applications, as it contributes to the safety and reliability of a model. One approach to achieving this is by adopting a Bayesian framework, where a prior distribution is established to represent our initial belief about the models. This belief is then updated through inference of posterior distributions based on the acquired knowledge. Although this approach shows promise, scalability becomes a concern when working with large-scale datasets during Bayesian inference. To address this issue, a potential solution is to employ a Bayesian coreset. A Bayesian coreset is a small subset of the original dataset where the posterior conditioning on it closely approximates the original posterior conditioning on the full dataset. Once the Bayesian coreset is trained, it can be utilized as a lightweight proxy dataset for subsequent Bayesian inference or as a replay buffer for continual learning or transfer learning.

ABayesian coreset is constructed by selecting a subset from a large dataset. However, recent research suggests that this approach may not be effective, especially in high-dimensional settings [19]. Instead, an alternative method of synthesizing a coreset, wherein the coreset is learned as trainable parameters, has been found to significantly enhance the quality of the approximation. This synthesized coreset is

referred to as a Bayesian pseudocoreset. The process of learning a Bayesian pseudocoreset involves minimizing a divergence measure between the posterior of the full dataset and the posterior of the pseudocoreset. However, learning Bayesian pseudocoresets is generally challenging due to the intractability of constructing both the full dataset posterior and the pseudocoreset posterior, as well as the computation of the divergence between them, which necessitates approximation. Consequently, existing works on Bayesian pseudocoresets have primarily focused on small-scale problems [19, 8, 20, 21]. Recently, Kim et al. [15] introduced a scalable method for constructing Bayesian pseudocoresets using variational Gaussian approximation for the posteriors and minimizing forward KL divergence. Although their method shows promise, it still demands substantial computational resources for high-dimensional models like deep neural networks.

In this paper, we present a novel approach to enhance the scalability of Bayesian pseudocoreset construction, particularly for Bayesian neural networks (BNNs) with a large number of parameters. Our proposed method operates in function space . When working with BNNs, it is common to define a prior distribution on the weight space and infer the corresponding weight posterior distribution, which also applies to Bayesian pseudocoreset construction. However, previous studies [31, 28] have highlighted the challenge of interpreting weights in high-dimensional neural networks, making it difficult to elicit meaningful prior distributions. Additionally, in high-dimensional networks, the loss surfaces often exhibit a complex multimodal structure, which means that proximity in the weight space does not necessarily imply proximity in the desired prediction variable [24, 30]. This same argument can be applied to Bayesian pseudocoreset construction, as matching the full data and pseudocoreset posteriors in the weight space may not result in an optimal pseudocoreset in terms of representation power and computational scalability.

To be more specific, our method constructs a variational approximation to the pseudocoreset posteriors in function space by linearization and variational approximation to the true posterior. Then we learn Bayesian pseudocoreset by minimizing a divergence measure between the full data posterior and the pseudocoreset posterior in the function space. Compared to the previous weight space approaches, our method readily scales to the large models for which the weight space approaches were not able to compute. Another benefit of function space matching is that it does not constrain the architectures of the neural networks to be matched, provided that their inherited function space posteriors are likely to be similar. So for instance, we can train with multiple neural network architectures simultaneously with varying numbers of neurons or types of normalization layers, and we empirically observe that this improves the architectural robustness of the learned pseudocoresets. Moreover, it has another advantage that the posteriors learned from the Bayesian pseudocoreset in function space have better out-of-distribution (OOD) robustness, similar to the previous reports showing the benefit of function space approaches in OOD robustness [28].

In summary, this paper presents a novel approach to creating a scalable and effective Bayesian pseudocoreset using function space variational inference. The resulting Bayesian pseudocoreset is capable of being generated in high-dimensional image and deep neural network settings and has better uncertainty quantification abilities compared to weight space variational inference. Additionally, it has better architectural robustness. We demonstrate the efficiency of the function space Bayesian pseudocoreset through the various experiments.

2 Background

2.1 Bayesian pseudocoresets

In this paper, we focus on probabilistic models for supervised learning problem. Let θ ∈ Θ be a parameter, and let p ( y | x, θ ) be a probabilistic model indexed by the parameter θ . Given a set of observations x := ( x i ) n i =1 and the set of labels y := ( y i ) n i =1 with each x i ∈ X and y i ∈ Y , we are interested in updating our prior belief π 0 ( θ ) about the parameter to the posterior,

However, when the size of the dataset n is large, the computation of the posterior distribution can be computationally expensive and infeasible. To overcome this issue, Bayesian pseudocoresets are constructed as a synthetic dataset u = ( u j ) j m =1 with m ≪ n with the set of labels ˜ y := (˜ y j ) j m =1

where the posterior conditioning on it approximates the original posterior π x ( θ ) .

This approximation is made possible by solving an optimization problem that minimizes a divergence measure D between the two posterior distributions 1

In a recent paper [15], three variants of Bayesian pseudocoresets were proposed using different divergence measures, namely reverse Kullback-Leibler divergence, Wasserstein distance, and forward Kullback-Leibler divergence. However, performing both the approximation and the optimization in the parameter space can be computationally challenging, particularly for high-dimensional models such as deep neural networks.

2.2 Bayesian pseudocoresets in weight-space

Kim et al. [15] advocates using forward KL divergence as the divergence measure when constructing Bayesian pseudocoresets, with the aim of achieving a more even exploration of the posterior distribution of the full dataset when performing uncertainty quantification with the learned pseudocoreset. The forward KL objective is computed as,

The derivative of the divergence with respect to the pseudocoreset u is computed as

For the gradient, we need the expected gradients of the log posteriors that require sampling from the posteriors π x and π u . Most of the probabilistic models do not admit simple closed-form expressions for these posteriors, and it is not easy to simulate those posteriors for high-dimensional models. To address this, Kim et al. [15] proposes to use a Gaussian variational distributions q u ( θ ) and q x ( θ ) to approximate π x and π u whose means are set to the parameters obtained from the SGD trajectories,

where µ u and µ x are the maximum a posteriori (MAP) solutions computed for the dataset u and x , respectively. Σ u and Σ x are covariances. The gradient, with the stop gradient applied to µ u , is approximated as,

Here, ε ( s ) u and ε ( s ) x are i.i.d. standard Gaussian noises and S is the number of Monte-Carlo samples.

Expert trajectories Approximating the full data and coreset posteriors with variational distributions as specified above requires µ u and µ x as consequences of running optimization algorithms untill convergence. While this may be feasible for small datasets, for large-scale setting of our interest, obtaining µ u and µ x from scratch at each iteration for updating u can be time-consuming. To alleviate this, in the dataset distillation literature, Cazenavette et al. [7] proposed to use the expert trajectories , the set of pretrained optimization trajectories constructed in advance to the coreset learning. Kim et al. [15] brought this idea to Bayesian pseudocoresets, where a pool of pretrained trajectories are assume to be given before pseudocoreset learning. At each step of pseudocoreset update, a checkpoint θ 0 from an expert trajectory is randomly drawn from the pool, and then µ u and µ x are quickly constructed by taking few optimization steps from θ 0 .

3 Function space Bayesian pseudocoreset

3.1 Function space Bayesian neural networks

We follow the framework presented in Rudner et al. [29, 28] to define a Function-space Bayesian Neural Network (FBNN). Let π 0 ( θ ) be a prior distribution on the parameter and g θ : X → R d be a neural network index by θ . Let h : Θ → ( X → R d ) be a deterministic mapping from a parameter θ to a neural network g θ . Then a function-space prior is simply defined as a pushforward ν 0 ( f ) = h ∗ π 0 ( f ) := π 0 ( h -1 ( f )) . The corresponding posterior is also defined as a pushforward ν x ( f ) = h ∗ π x ( f ) and so is the pseudocoreset posterior ν u ( f ) = h ∗ π u ( f ) .

3.2 Learning function space Bayesian pseudocoresets

Given the function space priors and posteriors, a Function space Bayesian PseudoCoreset (FBPC) is obtained by minimizing a divergence measure between the function space posteriors,

We follow Kim et al. [15] suggesting to use the forward KL divergence, so our goal is to solve

The following proposition provides an expression for the gradient to minimize the divergence, whose proof is given Appendix A.

Proposition 3.1. The gradient of the forward KL divergence with respect to the coreset u is

where [ ν x ] u and [ ν u ] u are finite-dimensional distributions of the stochastic processes ν x and ν u , respectively, f u := ( f ( u j )) j m =1 , and p (˜ y | f u ) = ∏ m j =1 p (˜ y j | f ( u j )) .

To evaluate the gradient Eq. 10, we should identify the finite-dimensional functional posterior distributions [ ν x ] u and [ ν u ] u . While this is generally intractable, as proposed in Rudner et al. [29, 28], we can instead consider a linearized approximation of the neural network g θ ,

where µ x = E π x [ θ ] and J µ x ( · ) is the Jacobian of g θ evaluated at µ x . Then we approximate the function space posterior ν x with ˜ ν x := ˜ h ∗ π x where ˜ h ( θ ) = ˜ g θ , and as shown in Rudner et al. [29, 28], the finite dimensional distribution [˜ ν x ] u is a multivariate Gaussian distribution,

with Σ x = Cov π x ( θ ) . Similarly, we obtain

with µ u := E π u [ θ ] and Σ u := Cov π u ( θ ) . Using these linearized finite-dimensional distribution, we can approximate

3.3 Tractable approximation to the gradient

Even with the linearization, evaluating Eq. 14 is still challenging because it requires obtaining µ x and Σ x which are the statistics of the weight-space posterior π x . Rudner et al. [28] proposes to learn a variational approximation q x ( θ ) in the weight-space , and use the linearized pushforward of the variational distribution ˜ h ∗ q x as a proxy to the function space posterior. Still, this approach requires computing the heavy Jacobian matrix J E q x [ θ ] ( u ) , so may not be feasible for our scenario where we have to compute such variational approximations at each update of the pseudocoreset u .

N(μx,甲x)

h()(u)

g(2)(u)

ge(a)(u)

(

gox(u)

x

Expert trajectory

MAP

solution

Minimizedivergence

between two posteriors

g(2)(u)

Train on

gou(u)

pseudocoresetu

0o

(O)

(#(2)

gg(uy(u)

Initial

MAP

parameter

solution

h()(u)

N(pu, Pu)

Instead, we choose to directly construct a variational approximations to the finite-dimensional distributions of the function space posteriors, that is,

where (ˆ µ x , ˆ Ψ x ) and (ˆ µ u , ˆ Ψ u ) are variational parameters for the full data and coreset posteriors. Inspired by Kim et al. [15], we construct the variational parameters using expert trajectories. Unlike [15], we simply let the MAP solution computed for x , θ x , by sampling a checkpoint from the later part of the expert trajectories, and obtain the MAP solution of u , θ u , by directly optimizing an initial random parameter. Then we obtain ˆ µ x and ˆ µ u using u . For the covariance matricies ˆ Ψ x and ˆ Ψ u , while Kim et al. [15] proposed to use spherical Gaussian noises, we instead set them as an empirical covariance matrices of the samples collected from the optimization trajectory. Specifically, we take additional K steps from each MAP solution to compute the empirical covariance.

where opt ( θ, x ) is a step of SGD optimization applied to θ with data x and the squares in the diag( · ) are applied in element-wise manner. Note also that we are applying the stop-gradient operations for to block the gradient flow that might lead to complication in the backpropagation procedure. The variational parameters (ˆ µ u , ˆ Ψ u ) are constructed in a similar fashion, but using the psedocoreset ( u , ˜ y ) instead of the original data ( x , y ) . It is noteworthy that our approach is similar to one of the methods in Bayesian learning, SWAG [18]. However, while SWAG focuses on collecting statistics on weight space trajectories, our method constructs statistics in function spaces. This distinction makes our approach more suitable and scalable for pseudocoreset construction. The overview of proposed method is provided in Fig. 1.

With the variational approximations constructed as described, we obtain a Monte-Carlo estimator of Eq. 14,

where p ( ε x ) and p ( ε u ) are standard Gaussians, ( ε ( s ) x ) S s =1 and ( ε ( s ) u ) S s =1 are i.i.d. samples from them.

3.4 Multiple architectures FBPC training

One significant advantage of function space posterior matching is that the function is typically of much lower dimension compared to the weight. This makes it more likely for function spaces to exhibit similar posterior shapes in the vicinity of the MAP solutions. This characteristic of function space encourages the exploration of function space pseudocoreset training in the context of multiple architectures. Because, the task of training a coreset that matches the highly complex weight space posterior across multiple architectures is indeed challenging, while the situation becomes relatively easier when dealing with architectures that exhibit similar function posteriors.

Therefore we propose a novel multi-architecture FBPC algorithm in Algorithm 1. The training procedure involves calculating the FBPC losses for each individual architecture separately and then summing them together to update. This approach allows us to efficiently update the pseudocoreset by considering the contributions of each architecture simultaneously. We will empirically demonstrate that this methodology significantly enhances the architecture generalization ability of pseudocoresets in Section 5.

3.5 Compare to weight space Bayesian pseudocoresets

By working directly on the function space, our method could bypass several challenges that may arise when working on a weight space. Indeed, a legitimate concern arises regarding multi-modality, as the posterior distributions of deep neural networks are highly complex. It makes the optimization of pseudocoresets on weight space difficult. Moreover, minimization of weight space divergence does not necessarily guarantee proximity in the function space. Consequently, although we try to minimize the weight space divergence, there is a possibility that the obtained function posterior may significantly deviate from the true posterior. However, if we directly minimize the divergence between the function distributions, we can effectively address this issue.

On the other hand, there is an additional concern related to memory limitations. While it has been demonstrated in Kim et al. [15] that the memory usage of Bayesian pseudocoresets employing forward KL divergence is not excessively high, we can see that Eq. 7 requires Monte-Carlo samples of weights, which requires O ( Sp ) where S and p represent the number of Monte-Carlo samples and the dimensionality of the weights, respectively. This dependence on Monte-Carlo sampling poses a limitation for large-scale networks when memory resources are constrained. In contrast, our proposed method requires significantly less memory, O ( Sd ) where d represents the dimensionality of the functions. Indeed, all the results presented in this paper were obtained using a single NVIDIA RTX-3090 GPU with 24GB VRAM.

4 Related work

Bayesian coresets Bayesian Coreset [3-5, 14] is a field of research aimed at addressing the computational challenges of MCMC and VI on large datasets in terms of time and space complexity [9, 1, 26]. It aims to approximate the energy function of the entire dataset using a weighted sum of a small subset. However for high-dimensional data, Manousakas et al. [19] demonstrates that considering only subsets as Bayesian coreset is not sufficient, as the KL divergence between the approximated coreset posterior and the true posterior increases with the data dimension, and they proposed Bayesian pseudocoresets. There are recent works on constructing pseudocoreset variational posterior to be more flexible [8] or how to effectively optimize the divergences between posteriors [15, 8, 20, 21]. However, there is still a limitation in constructing high-dimensional Bayesian pseudocoresets specifically for deep neural networks.

Dataset distillation Dataset distillation also aims to synthesize the compact datasets that capture the essence of the original dataset. However, the dataset distillation places particulary on optimizing the test performance of the distilled dataset. Consequently, the primary objective in dataset distillation is to maximize the performance of models trained using the distilled dataset, and researchers provides how to effectively solve this bi-level optimization [33, 23, 22, 40, 36, 39, 7]. In recent work, Kim et al. [15] established a link between specific dataset distillation methods and optimizing certain divergence measures associated with Bayesian pseudocoresets.

Function space variational inference Although Bayesian neural networks exhibit strong capabilities in performing variational inference, defining meaningful priors or efficiently inferring the posterior on weight space is still challenging due to their over-parametrization. To overcome this issue, researchers have increasingly focused on function space variational inference [6, 2, 16, 30, 25, 32, 27, 17]. For instance, Sun et al. [31] introduced a framework that formulates the KL divergence between functions as the supremum of marginal KL divergences over finite sets of inputs. Wang et al. [34] utilizes particle-based optimization directly in the function space. Furthermore, Rudner et al. [28] recently proposed a scalable method for function space variational inference on deep neural networks.

5 Experiments

5.1 Experimental Setup

In our study, we employed the CIFAR10, CIFAR100 and Tiny-ImageNet datasets to create Bayesian pseudocoresets of coreset size m ∈ { 1 , 10 , 50 } images per class (ipc). These pseudocoresets were then evalutated by conducting the Stochastic Gradient Hamiltonian Monte Carlo (SGHMC) [9] algorithm on those pseudocoresets. We measured the top-1 accuracy and negative log-likelihood of the SGHMC algorithm on the respective test datasets. Following the experimental setup of previous works [15, 7, 39], we use the differentiable siamese augmentation [37]. For the network architectures, we use 3-layer ConvNet for CIFAR10 and CIFAR100, and 4-layer ConvNet for Tiny-ImageNet.

We employed three baseline methods to compare the effectiveness of function space Bayesian pseudocoresets. The first baseline is the random coresets, which involves selecting a random minibatch of the coreset size. The others two baseline methods, BPC-rKL [15, 19] and BPC-fKL [15], are Bayesian pseudocoresets on weight space. BPC-rKL and BPC-fKL employ reverse KL divergence and forward KL divergence as the divergence measures for their training, respectively.

5.2 Main Results

Table 1 and Table 2 show the results of each baseline and our method for each dataset. For BPC-rKL and BPC-fKL, we used the official code from [15] for training the pseudocoresets, and only difference is that we used our own SGHMC hyperparameters during evaluation. For detailed experiment setting, please refer to Appendix C.

As discussed earlier, we utilized the empirical covariance to variational posterior instead of using naïve isotropic Gaussian for the function space variational posterior. To assess the effectiveness of using sample covariance, we compare these two in Table 1, as we also presented the results for

FBPC-isotropic, which represents the FBPC trained with a unit covariance Gaussian posteriors. The results clearly demonstrate that using sample covariance captures valuable information from the posterior distribution, resulting in improved performance. Overall, the results presented in Table 1 and Table 2 demonstrate that our method also outperforms the baseline approaches, including random coresets, BPC-rKL and BPC-fKL, in terms of both accuracy and negative log-likelihood, especially on the large-scale datasets in Table 2.

Furthermore, the Bayesian pseudocoreset can be leveraged to enhance robustness against distributional shifts when combined with Bayesian model averaging. To assess the robustness of our function space Bayesian pseudocoresets on out-of-distribution inputs, we also conducted experiments using the CIFAR10-C dataset [11]. This dataset involves the insertion of image corruptions into the CIFAR10 images. By evaluating the performance of the pseudocoresets on CIFAR10-C, we can see the model's ability to handle input data that deviates from the original distribution. In Table 3, we provide the results for top-1 accuracy and degradation scores, which indicate the extent to which accuracy is reduced compared to the in-distribution's test accuracy. The result demonstrates that our FBPC consistently outperforms the weight space Bayesian pseudocoreset, BPC-fKL.

5.3 Architecture generalization

In this section, we aim to demonstrate the architecture generalizability of FBPC and emphasize the utility of multi-architecture training as we discussed in the previous section. We specifically focus on investigating the impact of varying normalization layers on the generalizability of the pseudocoreset, since it is widely recognized that a pseudocoreset trained using one architecture may struggle to generalize effectively to a model that employs different normalization layers. We have also included the results of cross-architecture experiments that involve changing the architecture itself in Appendix D.1.

To assess this, we compare a single architecture trained pseudocoreset and a multiple architecture trained pseudocoreset. For single architecture training, we initially train a pseudocoreset using one architecture with a specific normalization layer, for instance we use instance normalization. Subsequently, we evaluate the performance of this pseudocoreset on four different types of normalization layers: instance normalization, group normalization, layer normalization, and batch normalization. For multiple architecture training, we aggregate four losses for single architecture training of each architecture, and train the pseudocoreset with the sum of all losses, as mentioned in previous section.

65

LN

BN

60

GN

IN

55

50

ACC

45

40

35

30

WBPC

WBPC-multi

FBPC-single

FBPC-multi

Train

(a) Performance evaluation with different normalization layers. Color represents the test architectures. The multiple architecture FBPC training enhances the generalization ability to other architectures.

15

LN

BN

10

GN

IN

U

-10

airplane

automobile

(b) A sample image and its corresponding function values. Despite variations in the normalization layers of different architectures, the function values exhibit similarity across the architectures.

As depicted in Fig. 2a, we observe that both WBPC (Weight space Bayesian pseudocoresets) and FBPC-single (Function space Bayesian pseudocoresets trained on a single architecture) exhibit a notable trend, that they tend to not perform well when evaluated on the architecture that incorporates different normalizations, regardless of whether it is trained on weight space or function space. On the other hand, when trained with multiple architectures, both WBPC-multi and FBPC-multi perform well across the all architectures, while notably FBPC-multi significantly outperforms WBPC-multi.

As mentioned in the previous section, we hypothesize that the superior performance of FBPC compared to WBPC can be attributed to the likelihood of having similar function space posterior across architectures. To validate this, we conduct an examination of the logit values for each sample across different architectures. As an illustration, we provide an example pseudocoreset image belonging to the class label "dog" along with its corresponding logits for all four architectures. As Fig. 2b shows, it can be observed that the logits display a high degree of similarity, indicating a strong likelihood of matching function posterior distributions. Our analysis confirms that, despite architectural disparities, the function spaces generated by these architectures exhibit significant similarity and it contributes to superiority of FBPC in terms of architecture generalizability.

6 Conclusion

In this paper, we explored the function space Bayesian pseudocoreset. We constructed it by minimizing forward KL divergence between the function posteriors of pseudocoreset and the entire dataset. To optimize the divergence, we proposed a novel method to effectively approximate the function posteriors with an efficient training procedure. Finally, we empirically demonstrated the superiority of our function space Bayesian pseudocoresets compared to weight space Bayesian pseudocoresets, in terms of test performance, uncertainty quatification, OOD robustness, and architectural robustness.

Limitation Despite showing promising results on function space Bayesian pseudocoresets, there still exist a few limitations in the training procedure. Our posterior approximation strategy requires

the MAP solutions, which necessitates training them prior to each update step or preparing expert trajectories. This can be time-consuming and requires additional memory to store the expert trajectories.

Societal Impacts Our work is hardly likely to bring any negative societal impacts.

Acknowledgments

This work was supported by KAIST-NAVER Hypercreative AI Center, Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.2019-0-00075, Artificial Intelligence Graduate School Program (KAIST), No.2022-000184, Development and Study of AI Technologies to Inexpensively Conform to Evolving Policy on Ethics, and No.2022-0-00713, Meta-learning Applicable to Real-world Problems) and the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. 2022R1A5A708390812)

References

J. Baker, P. Fearnhead, E. B. Fox, and C. Nemeth. Control variates for stochastic gradient mcmc. Statistics and Computing , 29:599-615, 2019.

D. R. Burt, S. W. Ober, A. Garriga-Alonso, and M. van der Wilk. Understanding variational inference in function-space. In Third Symposium on Advances in Approximate Bayesian Inference , 2020.

T. Campbell and B. Beronov. Sparse variational inference: Bayesian coresets from scratch. In Advances in Neural Information Processing Systems 32 (NeurIPS 2019) , 2019.

T. Campbell and T. Broderick. Bayesian coreset construction via greedy iterative geodesic ascent. In Proceedings of The 35th International Conference on Machine Learning (ICML 2018) , 2018.

T. Campbell and T. Broderick. Automated scalable bayesian inference via hilbert coresets. The Journal of Machine Learning Research , 20(1):551-588, 2019.

E. D. Carvalho, R. Clark, A. Nicastro, and P. H. Kelly. Scalable uncertainty for computer vision with functional variational inference. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 12003-12013, 2020.

G. Cazenavette, T. Wang, A. Torralba, A. A. Efros, and J.-Y. Zhu. Dataset distillation by matching training trajectories. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , 2022.

N. Chen, Z. Xu, and T. Campbell. Bayesian inference via sparse hamiltonian flows. In Advances in Neural Information Processing Systems 35 (NeurIPS 2022) , 2022.

T. Chen, E. B. Fox, and C. Guestrin. Stochastic gradient Hamiltonian Monte Carlo. In Proceedings of the 31st International Conference on Machine Learning (ICML 2014) , 2014.

A. G. de G. Matthews, J. Hensman, R. E. Turner, and Z. Ghahramani. On sparse variational methods and the Kullback-Leibler divergence between stochastic processes. Proceedings of Machine Learning Research , 51:231-239, 2016.

D. Hendrycks and T. Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. Proceedings of the International Conference on Learning Representations , 2019.

J. Hensman, N. Fusi, and N. D. Lawrence. Gaussian processes for big data. In Uncertainty in Artificial Intelligence , 2013.

J. Hensman, A. Matthews, and Z. Ghahramani. Scalable variational gaussian process classification. In Artificial Intelligence and Statistics , 2015.

J. Huggins, T. Campbell, and T. Broderick. Coresets for bayesian logistic regression. In Advances in Neural Information Processing Systems 29 (NIPS 2016) , 2016.

B. Kim, J. Choi, S. Lee, Y. Lee, J.-W. Ha, and J. Lee. On divergence measures for bayesian pseudocoresets. Advances in Neural Information Processing Systems 35 (NeurIPS 2022) , 2022.

C. Ma and J. M. Hernández-Lobato. Functional variational inference based on stochastic process generators. In Advances in Neural Information Processing Systems , 2021.

C. Ma, Y. Li, and J. M. Hernández-Lobato. Variational implicit processes. In Proceedings of The 36th International Conference on Machine Learning (ICML 2019) , 2019.

W. J. Maddox, P. Izmailov, T. Garipov, D. P. Vetrov, and A. G. Wilson. A simple baseline for bayesian uncertainty in deep learning. Advances in neural information processing systems , 32, 2019.

D. Manousakas, Z. Xu, C. Mascolo, and T. Campbell. Bayesian pseudocoresets. In Advances in Neural Information Processing Systems 33 (NeurIPS 2020) , 2020.

D. Manousakas, H. Ritter, and T. Karaletsos. Black-box coreset variational inference. In Advances in Neural Information Processing Systems 35 (NeurIPS 2022) , 2022.

C. Naik, J. Rousseau, and T. Campbell. Fast bayesian coresets via subsampling and quasi-newton refinement. Advances in Neural Information Processing Systems , 35:70-83, 2022.

T. Nguyen, Z. Chen, and J. Lee. Dataset meta-learning from kernel ridge-regression. In International Conference on Learning Representations (ICLR) , 2020.

T. Nguyen, R. Novak, L. Xiao, and J. Lee. Dataset distillation with infinitely wide convolutional networks. In Advances in Neural Information Processing Systems 34 (NeurIPS 2021) , 2021.

P. Pan, S. Swaroop, A. Immer, R. Eschenhagen, R. E. Turner, and M. E. Khan. Continual deep learning by functional regularisation of memorable past. Advances in neural information processing systems , 2020.

P. Pan, S. Swaroop, A. Immer, R. Eschenhagen, R. E. Turner, and M. E. Khan. Continual deep learning by functional regularisation of memorable past. Advances in neural information processing systems , 2020.

M. Quiroz, R. Kohn, M. Villani, and M.-N. Tran. Speeding up mcmc by efficient data subsampling. Journal of the American Statistical Association , 2018.

S. Rodríguez-Santana, B. Zaldivar, and D. Hernandez-Lobato. Function-space inference with sparse implicit processes. In Proceedings of The 38th International Conference on Machine Learning (ICML 2022) , 2022.

T. G. Rudner, Z. Chen, Y. W. Teh, and Y. Gal. Tractable function-space variational inference in bayesian neural networks. In Advances in Neural Information Processing Systems 35 (NeurIPS 2022) , 2022.

T. G. J. Rudner, G. Chen, and Y. Gal. Rethinking function space variational inference in Bayesian neural networks. 3rd Symposium on Advances in Approximate Bayesian Inference , 2020.

T. G. J. Rudner, F. Bickford Smith, Q. Feng, Y. W. Teh, and Y. Gal. Continual learning via sequential function-space variational inference. In Proceedings of the 39th International Conference on Machine Learning , pages 18871-18887, 2022.

S. Sun, G. Zhang, J. Shi, and R. Grosse. Functional variational bayesian neural networks. In International Conference on Learning Representations , 2018.

M. K. Titsias, J. Schwarz, A. G. d. G. Matthews, R. Pascanu, and Y. W. Teh. Functional regularisation for continual learning with gaussian processes. In International Conference on Learning Representations , 2019.

T. Wang, J.-Y. Zhu, A. Torralba, and A. A. Efros. Dataset distillation. arXiv preprint arXiv:1811.10959 , 2018.

Z. Wang, T. Ren, J. Zhu, and B. Zhang. Function space particle optimization for Bayesian neural networks. In International Conference on Learning Representations , 2019.

R. Yu, S. Liu, and X. Wang. Dataset distillation: A comprehensive review. arXiv preprint arXiv:2301.07014 , 2023.

B. Zhao and H. Bilen. Dataset condensation with differentiable siamese augmentation. In Proceedings of The 38th International Conference on Machine Learning (ICML 2021) , 2021.

B. Zhao and H. Bilen. Dataset condensation with differentiable siamese augmentation. In International Conference on Machine Learning , pages 12674-12685. PMLR, 2021.

B. Zhao and H. Bilen. Dataset condensation with distribution matching. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision , pages 6514-6523, 2023.

B. Zhao, K. R. Mopuri, and H. Bilen. Dataset condensation with gradient matching. In International Conference on Learning Representations , 2021.

Y. Zhou, E. Nezhadarya, and J. Ba. Dataset distillation using neural feature regression. In Proceedings of the Advances in Neural Information Processing Systems (NeurIPS) , 2022.

A Proofs

where [ ν x ] u and [ ν u ] u are finite-dimensional distributions of the stochastic processes ν x and ν u , respectively, f u := ( f ( u j )) j m =1 , and p (˜ y | f u ) = ∏ m j =1 p (˜ y j | f ( u j )) .

Proof. We follow the arguments in de G. Matthews et al. [10], Rudner et al. [29]. The forward KL divergence is defined as,

By the chain rule for the Radon-Nikodym derivative, we have

The first term does not depend on u , so we investigate the second term. By the measure theoretic Bayes' rule,

where p (˜ y | f, u ) := ∏ m j =1 p (˜ y j | u j , f ) and,

Now let ρ A : ( X → R d ) → ( A → R d ) be a projection function that takes a function f and returns its restriction on a set A ⊆ X . Assuming that the likelihood depends only on the finite index set u , we can write

where [ · ] u denotes the finite-dimensional distribution of stochastic process evalauted at u and ρ u ( f ) := f u := ( f ( u j )) j m =1 are corresponding function values at u . Putting this back into the above equation,

Now taking the gradient w.r.t. u , we get

Note also that

As a result, we conclude that

B Inducing points in Stochastic Variational Gaussian Processes

Stochastic Variational Gaussian Processes [SVGP; 12, 13] were introduced as a solution to address the significant computational complexity, characterized by a cubic computational complexity of O ( n 3 ) and memory cost of O ( n 2 ) , associated with performing inference using Gaussian Processes. In this context, n represents the total number of training data points. SVGP effectively leverages a concept called inducing points , which serves to reduce the computational complexity to O ( m 3 ) and memory requirements to O ( m 2 ) during inference, while still providing a reliable approximation of the posterior distribution of the entire training dataset. Notably, m denotes the number of inducing points, which is typically much smaller than the total number of training data points i.e. m ≪ n . The above description clearly shows that the inducing points have a similar purpose to FBPC. However, there are some differences in their learning objectives. In the context of SVGP, the process of optimizing the inducing points denoted as Z = { z 1 , . . . , z m } involves maximizing the ELBO in order to make a variational Gaussian distribution q ( f tr , f z ) well approximate the posterior distribution p ( f tr , f z | y tr ) . This variational distribution is composed of two parts: 1) p ( f tr | f z ) , which represents the Gaussian posterior distribution, and 2) q ( f z ) , which is the Gaussian variational distribution. During this optimization, we focus on training the inducing points Z as well as the mean and variance of the variational distribution q ( f z ) . The goal of this optimization is to create a good approximation of the posterior distribution p ( y te | x te , D tr ) during the inference process, all while keeping the computational cost low. On the other hand, as outlined in Section 3.2, the formulation of FBPC involves directly minimizing the divergence between function space posterior, specifically D KL [ ν x ∥ ν u ] . To sum up, while they do share some similarities in that they introduce a set of learnable pseudo data points, they are fundamentally different in their learning objectives. SVGP is interested in approximating the full data posterior through the inducing points, while ours aims to make the pseudocoreset posterior as close as possible to the full data posterior.

C Experimetal Details

The code for our experiments will be available soon.

C.1 Expert trajectory

For expert trajectory, we trained the network with the entire dataset and saved their snapshot parameters at every epoch, following the setup described in [7]. For training, we used an SGD optimizer with a learning rate of 0.01. We saved 100 training trajectories, with each trajectory consisting of 50 epochs.

C.2 Hyperparmeter setting

Training In our training procedure, we have some hyperparameters. Firstly, we sampled the MAP solution of x , denoted as θ x , from the later part of the expert trajectories. The decision of how many samples from the later part to utilize was treated as a hyperparameter for each experimental setting. We chose to use samples from T epoch onwards as the MAP solution samples. When obtaining the MAP solution θ u , there are also several hyperparameters involved, the optimizer and convergence criteria for training the MAP solution from random parameters. We used an Adam optimizer with a learning rate of 0.001 to train the model until the training loss reached a threshold of γ or below.

Next, to approximate our Gaussian variational function posterior, we employed the empirical covariance of the function samples. During the process of drawing function samples, we performed an additional training steps. This step involved specifying the optimizer and the number of steps used. We used an SGD optimizer with learning rates of λ x and λ u for a total of 30 steps during the additional training step for drawing function samples for x and u , respectively.

Lastly, we used the training iteration N , a learning rate of α for pseudocoresets, and set the batch size for pseudocoresets to B . The hyperparameters used in our paper are summarized in Table 4.

Evaluation To implement the SGHMC algorithm, as discussed in [9] and following the recommendations of [9], we employed the SGD with momentum along with an auxiliary noise term.

we set η = 0 . 03 , α = 0 . 1 , and d = 0 . 01 /m , where m represents the coreset size. We perform 1000 epochs of SGHMC and collect samples every 100 epochs.

D Additional Experiments

D.1 Cross-architecture generalization

In order to assess the cross-architecture generalization performance of FBPC, we trained pseudocorests of sizes {1, 10, 50} using the ConvNet architecture and tested them on various architectures trained from scratch. In addition to the ConvNet architecture used during training, we also evaluated the performance on sophisticated architectures such as ResNet18, ResNet34, VGG, and AlexNet. The results are presented in Table 5. As evident from Table 5, FBPC demonstrates considerable performance even on architectures different from those used during training, highlighting its strong cross-architecture generalization capabilities.

D.2 Training FBPC on larger neural networks

To evaluate the scalability of our method to large networks, we trained FBPC with the ResNet18 architecture. Training coreset with larger networks, such as ResNet18, has proven to be challenging and has been explored in only a few previous works [38]. This is primarily due to the lack of scalable training methods and the tendency for overfitting when training large networks. As a result, even when evaluating ResNet after training on a smaller network like ConvNet, the performance tends to suffer. Furthermore, it has been reported that coreset training directly on ResNet initially yields lower performance compared to training on ConvNet [40, 35].

Our experiments also revealed a similar trend in our findings as shown in the first column of Table 6. Although FBPC exhibits excellent scalability, making it easy to train on ResNet18 and larger networks, its performance was observed to be lower compared to ConvNet. On the other hand, the second column, ResNet18 + ConvNet, refers to training both ResNet18 and ConvNet simultaneously using the FBPC-multi training approach. In this case, surprisingly, the test performance of ResNet18 actually improved when trained in conjunction with ConvNet using the FBPC-multi training approach. In this case, the ConvNet accuracy was recorded at 60.03, which did not significantly compromise the ConvNet's performance while enhancing the performance of ResNet18. This suggests that training ConvNet acted as a regularizer, preventing overfitting in ResNet18 and enabling it to achieve better performance.

D.3 Computational cost

To compare how scalable our approach is compared to posterior matching in the weight space, we measured GPU memory usage corresponding to the number of parameters. As shown in Table 7 and Table 8, the memory usage for weight space BPC significantly increases as the number of parameters grows, while FBPC operates very efficiently. Additionally, the coreset ipc increases memory usage proportionally to its size. In terms of memory considerations, FBPC excels. However, as shown in Table 9, in terms of time, our method requires slightly more time because more SGD steps are needed to acquire the empirical covariance. However, when using FBPC-isotropic, these steps can be reduced, trading off a slight decrease in performance for time savings.

D.4 Differentiable siamese augmentation

Table 10 shows the result for BPC-fKL and FBPC without using DSA [37] and without any augmentation during training. Interestingly, for the ipc 1 case in BPC-fKL, performance improved when DSA was not applied. However, in all other cases, it is evident that not using DSA leads to an average performance drop of approximately 4.7%. Moreover, even when training BPC without augmentation, we observe that function space BPC outperforms weight space BPC.

D.5 Visualization

In this section, we provide visualizations of the pseudocoreset examples for each dataset.

(b) Examples images of FBPC for CIFAR10 with ipc 50. 10 images per class.

Peer reviews:
comments: SUMMARY:
This paper presents an alternative approach to the construction of Bayesian pseudocoresets by considering the quality of function space approximations. Specifically, they seek to minimise the KL-divergence between function space approximations of posteriors conditioned on the Bayesian pseudocoreset, and the true posterior.

STRENGTHS:
* I consider the method very well motivated---the use of variational approximations in weight-space is known to result in pathologies, particularly with high-dimensional models such as BNNs. Function space approximations offer a solution to this by instead computing variational approximations to distributions that exhibit much less multi-modality and related behaviour that is difficult to approximate. 
* Although the authors largely follow the work of Rudner et al., they depart from this method in estimating the parameter values of the function space approximation. Instead, they empirical estimates which, although inexact, I expect result in considerable speed-up in computation time.
* Experimental results are good, with the method outperforming the benchmarks.

WEAKNESSES:
* A lot of approximations are involved in the method, and it would be nice to see some evaluation of the effects of each.

QUESTIONS:
N/A.

LIMITATIONS:
N/A.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe following work proposes to a new way to construct Bayesian Pseudo-Coresets. Particularly, the authors propose to optimise the KL-Divergence between posteriors associated with real data and synthetic  in function space rather than the parameter space of large networks. The main argument posed by the authors is that function space is typically lower in dimension as compared to the parameter space of the network. Further, the authors use the theory laid out in [1] to approximate the function space posteriors associated with real data and pseudo-coreset. Further, since the function space can be same for several architectures, the proposed method can also be used with multiple architectures.\n\nSTRENGTHS:\n1.\tThe argument behind the use of function space seems logical enough, inference in high dimensional parameter space is indeed difficult.\n2.\tThe demonstrated result shows decent improvement compared to SoTA BPC methods.\n3.\tThe “multi-architecture” setting that comes for free because of working in the function space is very intriguing. To the best of my knowledge, this is the first work that enables one to use multiple architectures.\n\nWEAKNESSES:\n1. Ambiguous writing. \n\n2. Experiments can be better.\n\nQUESTIONS:\n1.\tIt seems that the paper could have been better written and presented. Especially for someone who is not aware of function space variational inference, the paper is very hard to follow through. Perhaps sections 2.1 and 2.2 can be trimmed to accommodate some background on function-space VI.\n\t2.\tFurthermore, related to the above point, though I understand why the term “function-space” might be appropriate, it gives an impression that authors are referring to “space of functions”. Then, it is mentioned that “the function space is typical of much lower dimension compared to the weight space”. However, the space of functions has infinite dimensions. Are the authors may be referring to the dimension of the range of functions? This is actually a bit confusing. If this is not the case, then the motivation of function space VI is not strong enough.\n\t3.\tIn Eq. 16, the authors note that while calculating $\\hat{\\mu_u}$ there is a stop-gradient operator in action. The same is not said for the step calculating $\\hat{\\Phi_u}$. I assume this means that the gradients are back-propagated through the SGD steps to update the pseudo-corsets. If yes, then I would like to see the effect of removing the stop-gradient operator while calculating $\\hat{\\mu_u}$. Further,  doesn’t this make the proposed method computationally expensive too? As is the case with BPC? However, if this is not the case, then this should be highlighted more explicitly.\n\t4.\tThe authors have shown the architecture generalization by changing the type of normalization layer for a single architecture. However, in my opinion, this is not good enough. In dataset summarisation literature, architecture generalization is often shown by changing the architecture itself. Hence, I would suggest bringing Tables 5 & 6 from the appendix to the main manuscript.\n\t5.\tDo the authors make use of Differentiable Siamese Augmentation (DSA) [2]? This is not mentioned in the main text and can affect the result by large margins. If not, then I suggest the authors try it as this can improve the results as well.\n\t6.\tI think the authors can also include a few basic Dataset Condensation methods in comparison for eg. GM [3], DM [4], DSA [2].\n\t7.\tIn Table 6 of the appendix, the authors note that there is performance degradation while using big architectures like ResNet. However, doesn’t this in some sense defy the motivation noted in the proposed work? Since, the proposed method operated in function space rather than the parameter space which (as the authors say) is lower dimensional, shouldn’t the performance be agnostic to architecture choice?\n\t8.\tFew recent works might be relevant to mention in related works [5, 6, 7]\n\n\n[1] T. G. Rudner, Z. Chen, Y. W. Teh, and Y. Gal. Tractable function-space variational inference in bayesian neural networks, 2022. \n[2] B. Zhao and H. Bilen. Dataset condensation with differentiable siamese augmentation. In Proceedings of The 38th International Conference on Machine Learning (ICML 2021), 2021. \n[3] B. Zhao, K. R. Mopuri, and H. Bilen. Dataset condensation with gradient matching. In International Conference on Learning Representations, 2021. \n[4] Bo Zhao and Hakan Bilen. Dataset condensation with distribution matching. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pages 6514–6523, 2023\n[5] Xu, Yue, et al. \"Distill Gold from Massive Ores: Efficient Dataset Distillation via Critical Samples Selection.\" arXiv preprint arXiv:2305.18381 (2023).\n[6] Tiwary, Piyush, Kumar Shubham, and Vivek Kashyap. \"Constructing Bayesian Pseudo-Coresets using Contrastive Divergence.\" arXiv preprint arXiv:2303.11278 (2023).\n[7] Yin, Zeyuan, Eric Xing, and Zhiqiang Shen. \"Squeeze, Recover and Relabel: Dataset Condensation at ImageNet Scale From A New Perspective.\" arXiv preprint arXiv:2306.13092 (2023).\n\nOverall, the paper is definitely intriguing and excites the reader. However, the lack of background in the paper confuses the reader. Further, some points need to be clarified (see above).\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a novel approach called Function Space Bayesian Pseudocoreset (FBPC) for constructing Bayesian pseudocoresets for Bayesian Neural Networks. A Bayesian pseudocoreset is a compact synthetic dataset that summarizes essential information from a large-scale dataset and can be used as a proxy dataset for scalable Bayesian inference. Typically, the construction of a Bayesian pseudocoreset involves minimizing the divergence between the posterior conditioning on the pseudocoreset and the posterior conditioning on the full dataset. However, evaluating this divergence measure can be challenging, especially for models like deep neural networks with high-dimensional parameters.\n\nIn contrast to previous methods that construct and match coreset and full data posteriors in weight space (model parameter space), this proposed method operates directly in function space. By working directly in function space instead of weight space, several challenges such as limited scalability and multi-modality issues can be bypassed.\n\nThe method constructs variational approximations to the coreset posterior on function space by linearization and variational approximation to true posterior distributions. It then matches these approximations with full data posteriors also defined in function spaces.\n\nWorking directly in function spaces allows this approach to scale well even for large models where traditional weight-space approaches struggle computationally. Additionally, it does not constrain matching only specific architectures of neural networks but rather allows training with the equivalent of multiple architectures simultaneously while still achieving similar results.\n\nFurthermore, using function-space matching improves out-of-distribution robustness compared to previous approaches based on weight spaces.Overall, experiments conducted demonstrate that using Function Space Bayesian Pseudocoresets leads to enhanced uncertainty quantification abilities compared to traditional methods based solely on weights or model parameters.\n\nSTRENGTHS:\n* The authors tackle a very interesting problem, extracting reliable posterior distributions out of BNNs to perform more precise inference. The Bayesian pseudocoreset approach allows for a simple but practical solution to this task.\n\n* The paper is easy to follow and is very well written. The authors provide much or the context needed to understand the submission and the previous state-of-the-art.\n\n* The function space approach is a very promising framework to conduct inference in BNNs and other models, and advances such as this one would help spreading its usage.\n\nWEAKNESSES:\n* While reading the paper, I could not help but miss some a comparison with the concept of inducing points, so common in sparse variational approximations for Gaussian Processes. From my point of view this is missing and would be interesting, since the proposed approach seems to be conceptually very strongly related to SVGPs, so somewhat including a comparison would be helpful.\n\n* To avoid scalability issues regarding the Jacobian defined in the linearized approximation, the authors then approximate the posterior in function space using variational Gaussian distributions that then myst be fit. This can be scalable, although I fear it may induce strong bias to the final posterior estimate, since there is no guarantee that the posterior distribution would behave in this manner.  \n\n* I think that some contributions relevant to the topic at hand are left out from the related work section, specially in the function space variational inference paragraph in section 4. I suggest including [1,5] as context, maybe [6] as well for context. Moreover, in terms of approximating the posterior distribution from BNNs, all a-posteriori Laplace approximation-based approaches are missing, such as [4] and [7]. These last two are not crucial for the topic at hand, although could serve to complete further the discussion.\n\n* I would expect a more detailed experimental phase, characterizing further the properties of the model. I would suggest complementing the results with some regression results. Moreover, it would be interesting to compare the resulting predictive distribution in each case with other models such as HMC, MFVI or other function space methods s.a. [1,2,3,5]. Moreover, it would be interesting to know how the distributions obtained compare against a-posteriori methods s.a. [4].\n\n* I would include some scalability study for the method, specially in terms of the size and dimensionality of the dataset.\n\n_Note:_ I condition my final evaluation on the submission to the fact that the authors address these issues. Were this not the case, I may consider revising the score.\n\n---\n\n**References**:\n\n[1] Rodrı́guez-Santana, S., Zaldivar, B., & Hernandez-Lobato, D. (2022, June). Function-space Inference with Sparse Implicit Processes. In International Conference on Machine Learning (pp. 18723-18740). PMLR.\n\n[2] Rudner, T. G., Chen, Z., Teh, Y. W., & Gal, Y. (2022). Tractable function-space variational inference in Bayesian neural networks. Advances in Neural Information Processing Systems, 35, 22686-22698.\n\n[3] Ma, C., & Hernández-Lobato, J. M. (2021). Functional variational inference based on stochastic process generators. Advances in Neural Information Processing Systems, 34, 21795-21807.\n\n[4] Deng, Z., Zhou, F., & Zhu, J. (2022). Accelerated Linearized Laplace Approximation for Bayesian Deep Learning. Advances in Neural Information Processing Systems, 35, 2695-2708.\n\n[5] Ma, C., Li, Y., and Hernández-Lobato, J. M. (2019). “Variational implicit processes”. In: International Conference on Machine Learning, pp. 4222–4233.\n\n[6] Fortuin, V. (2022). Priors in bayesian deep learning: A review. International Statistical Review, 90(3), 563-591.\n\n[7] Antorán, J., Janz, D., Allingham, J. U., Daxberger, E., Barbano, R. R., Nalisnick, E., & Hernández-Lobato, J. M. (2022, June). Adapting the linearised laplace model evidence for modern deep learning. In International Conference on Machine Learning (pp. 796-821). PMLR.\n\nQUESTIONS:\n* What role does the forward KL divergence play in constructing the Bayesian pseudocoresets, and why would you prefer this to well-defined distance measurement or any other divergence measure? Does the assymetry of the KL divergence bias the results in any significative manner that one should be aware of when using the method?\n\n* Have you performed any checks comparing the usage of the full Jacobian [lines 151 and 152] with the case where you induce the variational approximateion to the finite-dimensional distributions of the posteriors in function space? I agree that this cannot be done in general, specially for big BNNs where the computational cost of computing the Jacobian is high. However, tests could be performed in small BNNs to check whether this assumption holds in those simpler instances. If you consider this to not be necessary, please argue why.\n\n* Robustness to the selection of BNN architecture is an important point in the submission. While I agree that the results obtained in this regard seem to point in that direction, I would also argue that differences may be observed if it were not for the approximations regarding the linearization of the system and the imposed Gaussian variational form. Is there any reason to think otherwise?\n\n* Why would this approach be preferred over those that obtain predictive distributions with a-posteriori approximations s.a. [4] and [7]?\n\nLIMITATIONS:\n* The Bayesian pseudocoreset is only constructed comparing Gaussian distributions, which may be insufficient to express different posteriors induced by the data.\n\n* The experimental part of the submission could be improved.\n\n* There are no studies on the scalability of the method.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper combines the works of Kim et al. (2022) and Rudner et al. (2022) to propose a Bayesian coreset learning method based on function space variational inference. The authors find that this leads to coresets with improved accuracy/NLL compared to (Kim et al., 2022) on CIFAR and Tiny-ImageNet benchmarks as well as improved robustness against data corruption on CIFAR10. Finally, the authors argue that their method allows for improved incorporation of different architectures in coreset learning based on an improved generalization across architectures with different types of normalization layers for CIFAR10.\n\nGenerally, I found the paper quite incremental, although I appreciate that clear credit is given to prior work. The use of terminology is also rather loose, not any ad-hoc construction of a distribution over the weights constitutes an approximation to a Bayesian posterior. Finally, while performance is improved over that of Kim et al. (2022), in absolute terms it is still really far off that of training on the full dataset. I would really want to see that the coreset performance converges to that of training on the full dataset, so at this point I would argue for rejection.\n\nSTRENGTHS:\n* The method is new and described clearly\n* Clear credit is given to prior work that the method is based on\n* Performance is better than that of the baseline\n* The architecture generalization supports/motivates the use of function space inference\n\nWEAKNESSES:\n* The method as such seems really incremental, essentially it plugs the functional variational inference of Rudner et al. (2022) into the coreset learning and posterior “approximation” method of Kim et al. (2022)\n* That being said, calling a Gaussian distribution constructed from an optimization trajectory by taking the mean and variance of the parameters and calling it a variational distribution seems like a bit of an abuse of terminology. Not every distribution approximates the posterior. Further, this approach seems mostly equivalent to SWAG (Maddox et al., 2019), so this should be cited (and any potential difference discussed if applicable).\n* The absolute accuracies even for the largest coresets are still really far off what we would get when training on the full dataset. So the method seem quite far away from being useful to practitioners.\n\nQUESTIONS:\n* Does the gap with full training close when further increasing the coreset size?\n* The experimental section states that BPC-fKL was not possible on Tiny-Imagenet with 50 coreset points: why is that exactly? It would be useful to have some more details on memory consumption/runtime and the corresponding scaling behavior with coreset size/dateset size/dimension/number of parameters for the different methods. Reporting. some empirical numbers from the experiments on VRAM use etc could be helpful.\n\n**Typos/minor:**\n* l65-66: remove “number of”?\n* l106: “the means” -> “means”\n* l238: “we” -> “We”\n* l247: “refer the appendix” -> “refer to the appendix”\n* eq (6): $\\Sigma_u, \\Sigma_x$ appear for the first time here, but are only defined after eq (12).\n* I found the name for the FBPC-random baseline really confusing vs the random coreset points baseline, perhaps calling FBPC-iso or FBPC-isotropic would be more descriptive?\n* The references could use another round of editing, please ensure that all citations have a venue (missing e.g. for 1, 2, 6, 15, 22, 27), and that title capitalization and venue names are consistent.\n* The abstract is a bit verbose while providing fairly limited specifics. I would suggest cutting back the introductory first half — unfamiliar readers won’t learn much from this anyway and experts know all of this already.\n\nLIMITATIONS:\nI would have liked to see some discussion on the large gap to training on the full dataset.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}