Function Space Bayesian Pseudocoreset for Bayesian Neural Networks
Authors: Balhae Kim
Hyungi Lee
Juho Lee
Venue: NeurIPS 2023 poster (2023)
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
Balhae Kim Hyungi Lee', Juho Lee
KAIST AIt, AITRICS"
(balhackim, 1hk2708, juholee) (uholee)kkais.eee Ckaist. ac. .r
Abstract
A Bayesian pseudocoreset compact synthetic dataset summarizing essential information of a large-scale dataset and thus can be used as 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, eval- uating the divergence can be challenging, particularly for the models like deep neural networks having high-dimensional parameters. In this paper, we propose novel Bayesian pseudocoreset construction method that operates on function space. Unlike previous methods, which construct and match the coreset and full data posteriors in the space of model parameters (weights), variational approximations to the coreset posterior on 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 weight space, including limited scalability and multi- modality issue. Through various experiments, we demonstrate that the Bayesian pseudocoresets constructed from method enjoys enhanced uncertainty quantification and better robustness across various model architectures.
method constructs
Introduction
Deep learning has achieved tremendous success, but its requirement for large amounts of data makes it often inefficient 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 model. One approach to achieving this by adopting a Bayesian framework, where prior distribution is established 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 concern when working with large-scale datasets during Bayesian inference. To address this issue, potential solution employ Bayesian coreset. A Bayesian coreset is 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 lightweight proxy dataset for subsequent Bayesian inference or as replay buffer for continual learning or transfer learning.
A Bayesian coreset is constructed by selecting subset from large dataset. However, recent research suggests that this approach may not effective, especially high-dimensional settings [19]. Instead, alternative method of synthesizing a coreset, wherein the coreset is learned trainable parameters, has been found to significantly enhance the quality of the approximation. This synthesized coreset is
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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, 21J. Recently, Kim et al. [15] introduced 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 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, 30J. 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 optimal pseudocoreset in terms of representation power and computational scalability.
To be more specific, our method constructs 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 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 types of normalization layers, and 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-distribun (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 probabilistic models for supervised learning problem. Let 0 E θ be parameter, and let p(y x. 6) be probabilistic model indexed by the parameter f. Given a set of observations x (zx)²= and the set of labels y := (u4)=1 with each 2; E X and y V, we are interested in updating our prior belief TO (0) about the parameter to the posterior,
(8)))
P(u I 0), xx) = / II ply: x, (d))
(0) = Z X
(1)
However, when the size of the dataset 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 datasct (u,)i with m < with the set of labels ý = (ū))

---

where the posterior conditioning on it approximates the original posterior Tx (0).
To y
II p(ū nj 0). ( u) :- / II
Tu (0)
p(ū, / ug (d)).
j=1
This approximation is made possible by solving an optimization problem that minimizes divergence measure D between the two posterior distributions
u arg min ( Tu)-
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 Setlback-Leibbee 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 the divergence measure when constructing Bayesian pseudocoresets. with the aim of achieving more even exploration the posterior distribu- tion of the full dataset when performing uncertainty quantification with the learned pseudocoreset. The forward KL objective is computed as,
DKL [TTl log Z( u) log (y |x)
+ Exg log p(yn Ziz 0)
- EiT
log p(ū; ugj 0)
(4)
The derivative of the divergence with respect to the pseudocoreset u is computed as
DuL Exr
V log plūy u,. V
Ex
log logplüzlug 0))
(5)
For the gradient, we need the expected gradients the log posteriors that require sampling from the posteriors Txx and Tu. 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 use Gaussian variational distributions qu (0) and gx (0) to approximate xx and Tu whose means are set to the parameters obtained from the SGD trajectories,
(0) = ((00; pn, En), 4x (f) = -N(A: pxx Ex),
(6)
where iiu and #x are the maximum a posteriori (MAP) solutions computed for the dataset u and x, respectively. En and Ex are covariances. The gradient, with the stop gradient applied juu is approximated as,
V
log ūn uy sg(pa) + 21/(2((0))
Σ log ūj Ujz Mx + 22/B2(6)) )4P2(())
5
Here,
(s) and (0) 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 Yu and lx as consequences of running optimization algorithms untill convergence. While this may be feasible for small datasets, for large-scale setting of our interest, obtaining Nu and 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 al. [15] brought this idea to Bayesian pseudocoresets, where pool of pretrained trajectories are assume to be given before pseudocoreset learning. At each step of pseudocoreset update, a checkpoint Oo from an expert trajectory is randomly drawn from the pool, and then Pn and fx are quickly constructed by taking few optimization steps from Op-
In principle, we should leam the (pseudo)labels ý as well, but for classification problem, we can simply fix it constant set containing equal proportion of all possible classes. We assume this throughout the paper.

---

Function space Bayesian pseudocoreset
3.1 Function space Bayesian neural networks
We follow the framework presented in Rudner et al. [29, 28] to define Function-space Bayesian Neural Network (FBNN). Let To(0) be prior distribution on the parameter and 9e : X + R
be neural network index by 6. Let h è (X' + Id be deterministic mapping from parameter 0 to neural network ge. Then function- space prior is simply defined as pushforward vo()) h. To (f) To (h 1 (f)). The corresponding posterior is also defined as a pushforward (x( Tx (f) and so is the pseudocoreset posterior (,) h.
3.2 Learning function space Bayesian pseudocoresets
Given the function space priors and posteriors, Function space Bayesian PseudoCoreset (FBPC) obtained by minimizing divergence measure between the function space posteriors,
u = min (Vvx, Mu).
(8)
We follow Kim et al. [15] suggesting to use the forward KI. divergence, so our goal is to solve
u arg min DKL [xx][a]]
The following proposition provides an expression for the gradient 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
V VuDKLlVslleu = vuBBpy. D(gg [u]] - -Eja, Vulogply [w]],
(10)
where [ux]u and [Va]u are finite-dimensional distributions of the stochastic processes Vx and Va. respectively, fu (m and plý f.,) TT=1p(ū5 f(u;)).
To evaluate the gradient Eq. 10, should identify the distributions [u and [ru]u- While this is generally intractable, as proposed in Rudner et al. [29, 28). can instead consider linearized approximation of the neural network go.
functional posterior
ge (-) 9n (-) + Jun (-) -x),),
(11)
where Px
En.. [0] and J.x is the Jacobian of ge evaluated at Px. Then we approximate the function space posterior Vx with 7x h. where h(() = Go. and shown in Rudner al. 28], the finite dimensional distribution [7 is multivariate Gaussian distribution,
(xlu(fu) N (fu 9px Jnn (u)Ex (u) (
(12)
with Σx Cov (0). Similarly, we obtain
[Tulu(fu) = N N(a gr. (g), Jy. (u.(u))nnn.(u))
(13)
with Pn = E.., (0] and En = Cov (0). Using these linearized finite-dimensional distribution, we approximate
V.DKIt[ux(lall -V V. p. plo fu]] Ej,, V lo f.)],
(14)
3.3 Tractable approximation to the gradient
Even with the lincarization, evaluating Eq. 14 is still challenging because it requires obtaining and Vix which are the statisties of the weight-space posterior xx- Rudner al. [28] proposes to learn a variational approximation &x (0) in the weight-space, and use the linearized pushforward of the variational distribution h as a proxy to the function space posterior. Still, this approach requires computing the heavy Jacobian matrix Jz., [o) (u), so may not feasible for our scenario where we have compute such variational approximations each update of the pseudocoreset

---

Byg
00
e
8a,
Expert Iraectory
solution
Minimize divergence between posteriors
8pp((0) (1)
Troin isaudocoresetu
8o_(t) (n))
0
MAP solsion
perameeter
(x))
N Nfu. 9.)
Figure The conceptual overview of our proposed training procedure.
Instead, we choose to directly construct variational approximations to the finite-dimensional distributions of the function space posteriors, that is.
(T)u((fu) R lx (fu) = Nfu I ga.= (u), Mx),
(15)
(ful(ff)) = qu qu((F) = Niu Pa. (u), 1 Wu),
where (ix, Vx) and (iu, V.)) are variational parameters for the full data and coreset posteriors. Inspired by Kim al. [15], we construct the variational parameters using expert trajectories. Unlike [15], we simply let the MAP solution computed for x, 8x. by sampling a checkpoint from the later part of the expert trajectories, and obtain the MAP solution of u, ou. by directly optimizing an initial random parameter. Then we obtain Px and Au using u. For the covariance matricies V and V 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.
((0) 0x:
((t)
(f x
1), (x,y)),
Px= = Pagia (u),
2
(16)
i = sg
diag
1 K
BLR (u)
2 9 (u)
where opt(0, x) is a step of SGD optimization applied to with data x and the squares in the diag(-) are applied in element-wise wise manner. Note also that we are applying the stop- gradientt operations for to block the gradient flow that might lead to complication in the backpropagation procedure.
The variational parameters (Pu: V (.) are constructed in similar fashion, but using the psedocoreset (u, y) instead of the original data (x, y). It is noteworthy that our approach is similar 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.
With the variational approximations constructed as described, we obtain a Monte-Carlo estimator of Eq. 14,
v, D...R(Lx(/mn (c.) (l.g f..) + R. (R.)) VuBplex)loggpe Aax + 1/2 5x)2 + Ep(ew) Vu Vologp( Au + $ 1/2 Em)
Vn logp(ý (f.)
(17)
Vulogp y Px + $ 1/2 (1(22(()) + Vulogp ÿ A $1122(1))
S 1=1
where p(Ex) and p(Eu) standard Gaussians, (sx ) and (E (e) (= are i.i.d. samples from them.
5

---

Algorithm Multi-architecture Function space Bayesian Pseudocoreset
Require: Set architectures A, expert trajectories (E(") a A), prior distributions of parameters {mh a E A). an optimizer opt.
Initialize with random minibatch of coreset size mt.
for i ,....,N N do
Initialize the gradient of pseudocoresct, g for a E A do Sample the MAP solution computed for x, 0x E E(n),
Sample an initial random parameter Og ~ (a) (0). repeat
0, + opt(O:-1, (u,
until converges to obtain the MAP solution computed for u, An.
Obtain Ax: Au 4x V. by Eq. 16.
Compute the pseudocoreset gradient g (a) using Eq. 17.
g 4 g g(() end for
Update the pseudocoreset u by using the gradient g-
end for
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 propose 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 update. This approach allows 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
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 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 issuc.
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-Carke samples of weights, which requires (Sp) where S and p represent the number of Monte- Carlo samples and the dimensionality of the weights, respectively. This dependence on Monte-Carke sampling poses a limitation for large-scale networks when memory resources are constrained. In contrast, our proposed method requires significantly less memory, ((Sd) where d represents the dimensionality of the functions. Indeed, all the results presented in this paper were obtained using a single NVIDIA RTX- X090 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,2. aims to approximate the energy function of the entire dataset using a weighted sum a small subset. However for high- sigh-dimensiona 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 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 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 ca- pabilities in performing variational inference, defining meaningful priors or efficiently infer- ring the posterior overcome this issue, researchers have increasingly focused on function space variational infer- ence [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 marginal KI. divergences over finite sets of inputs. Wang et al. [34] utilizes massecle-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.
weight space is still challenging due to their over-paramettriztion. To
Experiments
5.1 Experimental Setup
In our study, we employed the CIFARIO, CIFARIO0 and Tiny-ImageNet datasets to create Bayesian pseudocoresets of coreset size m (1, 10,50) 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 likelihodoo 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 CIFAR 10 and CIFAR 100, and 4-yayer ConvNet for Tiny-ImageNet.
We employed three baseline methods compare the effectiveness of function space Bayesian pseudocoresets. The first baseline is the random coresets, which involves selecting random mini- batch of the coreset size. The others two baseline methods, BPC-rKL [15, 19] and BPC-IKL [15], are Bayesian pseudocoresets on weight space. BPC-rKL and BPC-IKL employ reverse KL divergence and forward KL divergence as the divergence measures for their training, respectively.
5.2 Main Results
Table I 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, compare these two in Table 1, as we also presented the results for
7

---

Table 1: Averaged test accuracy and negative log-likelihoods of models trained on each Bayesian pseudocoreset from scratch using SGHMC on the CIFARIO dataset. Bold is the best and underline is the second best. These values are averaged over random seeds.
ipc
SGHMC
Random
BPC-KKL [15, 19]
BPC-FKL [15]
FBPC-isotropic (Ours)
FBPC (Ours)
Acc (1)
16.30=n.74
2..44:1.06
34.50-1.62
32.00-0.78
35.45
NLL (L)
4.6620.03
4.5
3.8620.13
3.40
3.79
Acc (f)
32.4
37. 37.9210.
56.192a.a
61.43-0.35
62.33 10.34
10
NLL (})
2.98
2.47:0.04
1.48=0.02
1.35.+0.02
1.002
Acc (t)
49.68=n.46
51.86+0.38
64.74-0.32
71.33=0.19
71.23+0.17
50
NLL (4)
2.06+0.02
1.95=0.02
1.26+0.01
1.03+0.00
1.03+0.05
Table 2: Averaged test accuracy and negative log- ligeliiklssoods of models trained on each Bayesian pseudocoreset from scratch using SGHMC on the CIFARIO0 and Tiny-ImageNer datasets. These values are averaged over random seeds.
CIFARI00
Tiny-ImageNo
ipc
10
50
10
50
Acc (f)
4.844
18.0...31
35.1...23
1,90 08
1.20
19.1510.12
Random
NLL (L)
5.55
4.57
3.35
6.18
5.77.00.02
4.88-10.01
Acc (1)
14.7+0.16
28.1=0.00
37.140.35
3,98+0.13
11.4+0.45
22.18++0.2
BPC-IKL
NLL (4)
4.17=t0.05
3.53=0.05
3.28++0.24
5.63+00.03
5.08
4.65=0.02
Acc (t)
21.010.76
39.7.0:31
44.47.1035
10.14.10.c as
19.42.0.51
26.43.1031
FBPC (Ours)
NLL (L)
3.7610.111
2.6710..0
2.63.10.0
4.69
1.14
1.30.10.05
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 and Table 2 demonstrate that method also outperforms the baseline approaches, including random coresets, BPC-rKL and BPC-fKL, in terms of both accuracy and negative oogllikelihood especially on the large-scale datasets in Table 2.
Furthermore, the Bayesian pseudocoreset can be leveraged enhance robustness against distributional shifts when combined with Bayesian model averaging. To assess the robustness of our function space Bayesian pscudocoresets on out-oo-distribution inputs, we also conducted experiments using the CIFARIO-C dataset [[]]. This dataset involves the insertion of image corruptions into the CIFARIO 10 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 'ss 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 anttiterc 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 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.
To assess this, we compare single architecture trained pseudocoreset and 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. Subse- quently, 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, mentioned in previous section.

---

-00
sigle
rerS mLhh
Pain
(a) Performance ovaluation with different normaliza- tion layers. Color represents the test architectures. The multiple architecture FBPC training enhances
(b) sample image and its corresponding function values. Despite variations in the normalization layers of different architectures, the function val- ues exhibit similarity across the architectures.
generalization ability to other architectures.
Figure 2: Results for multiple architecture FBPC training.
Table 3: Test accuracy and degradation scores of models trained each Bayesian pseudocoreset from scratch using SGHMC on the CIFARIO-C. Degradation refers to the extent by which a model's accuracy decreases when evaluated on the CIFARIO-C dataset compared to the CIFARIO test dataset.
corruption
DB
ET
JPEG
MB
PIX
Avg.
Acc (f)
33.5
34
35.9
25.1
33.7
39.1
32.7
38.3
28.9
41.2
34.3
BPC-IKL
Degradation (1)
40.3
39.4
36
55.2
39.9
30.3
41.6
31.6
48.4
26.5
38.9
Acc (f)
48.9
46.4
47.6
41.7
44.0
52.0
44.3
51.0
47.1
52.3
47.5
FBPC
Degradation (4)
21.5
25.7
23.7
33.0
29.3
16.4
28.8
18.1
24.4
23.7
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 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, 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, conduct examination of the logit values for each sample across different architectures. As an illustration, we provide an example pseudocoreset image belonging 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 high degree of similarity. indicating strong likelihood of matching function posterior distributions. Our analysis confirms that, despite architcctural 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 minimiz- ing forward KL divergence between the function posteriors of pseudocoreset and the entire dataset. To optimize the divergence, we proposed 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, terms of test performance, uncertainty quatification, OOD robustness, and architectural robustness.
Limitation Despite showing promising results on function space Bayesian pseudocoresets, there still exist few limitations in the training procedure. Our posterior approximation strategy requires

---

the MAP solutions, which necessitates training them prior to each update step preparing ex- pert trajectories. This can be time-consuming and requires additional memory to store the expert trajectories.
Societal Impacts
Our work hardly likely to bring any negative societal impacts.
Acknowledgments
This work was supported by KAIST-NAVER Hypercreative Al Center, Institute of Information & communications Technology Planning Evaluation (ITTP) grant funded by the Korea government (MSIT) (No.2019-0-00075, Artificial Intelligence Graduate School Program (KAIST), No.2022-00 00184, Development and Study of AI Technologies to Inexpensively Conform to Evolving Policy on Ethics, and No. No.2022-0-0077 Meta-learning Applicable to Real-world Problems) and the Na- tional Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. 2022RAA5A70$308122
References
[1] J. Baker, P. Fearnhead, E. B. Fox, and C. Nemeth. Control variates for stochastic gradient meme. Statistics and Computing, 29:599-615 615, 2019.
[2] D. R. Burt, S. W. Ober, Garriga-Alonso, and M. van der Wilk. Understanding variational inference in function-space. In Third Symposium on Advances in Approximate Bayesian Inference, 2020.
[3] Camphell and B. Beronov. Sparse variational inference: Bayesian coresets from scratch, In Advances in Neural Information Processing Systems 32 (NeurIPS 2019), 2019.
[4] T. Campbell and T. Broderick. Bayesian coreset construction via greedy iterative geodesic ascent. Proceedings of The 35th International Conference on Machine Learning (ICML 2018), 2018.
[5] T. Campbell and T. Broderick. Automated scalable bayesian inference via hilbert coresets. The Journal of Machine Learning Research, 20(1):551 2019.
[6] 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 Computer Vision and Pattern Recognition, pages
2020.
[7] 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 Computer Vision and Pattern Recognition, 2022.
[8] N. Chen, Z. Xu, and T. Camphell. Bayesian inference via sparse hamiltonian flows. In Advances in Neural Information Processing Systems 35 (NeurIPS 2022), 2022.
[9] Т. 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.
[10] A. G. de G. Matthews, J. Hensman, R. E. Turner, and Z. Ghahramani. On sparse variational methods and the Kullback-Leible divergence between stochastic processes. Proceedings of Machine Learning Research. 51:231 2016.
[11] D. Hendryeks and T. Dietterich. Benchmarking neural network robustness to common corrup- tions and perturbations. Proceedings of the International Conference on Learning Representa- tions, 2019.
[12] J. Hensman, N. Fusi, and N. D. Lawrence. Gaussian processes for big data. In Uncertainty in Artificial Intelligence, 2013.
[13] Hensman, A. Matthews, and Ghahramani. Scalable variational gaussian process classifica- tion. In Artificial Intelligence and Statistics, 2015.
10

---

[14] J. Huggins, T. Campbell, and T. Broderick. Coresets for bayesian logistic regression. In Advances Neural Information Processing Systems 29 (NIPS 2016). 2016.
[15] B. Kim, J. Choi, S. Lee, Lee, J.-W. Ha, and J. Lee. On divergence measures for bayesian pseudocoresets. Advances in Neural Information Processing Systems 35 (NeurIPS 2022), 2022.
[16] C. Ma and J. M. Hernández-Lobato, Functional variational inference based on stochastic process generators. In Advances in Neural Information Processing Systems, 2021.
[17] C. Ma, Y. Li, and J. M. Hernández-Loband Variational implicit processes. In Proceedings of The 36th International Conference on Machine Learning (ICML 20/9), 2019.
[18] 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.
[19] D. Manousakas, Z. Xu, C. Mascolo, and T. Campbell. Bayesian pseudocoresets, In Advances in Neural Information Processing Systems 33 (NeurIPS 2020), 2020.
[20] D. Manousakas, H. Ritter, and T. Karaletsos. Black-box coreset variational inference. In Advances in Neural Information Processing Systems 35 (NeurIPS 2022), 2022.
[21] C. Naik, Rousseau, and T. Campbell, Fast bayesian coresets via subsampling and quasi-newton refinement. Advances in Neural Information Processing Systems, 35:70-83, 2022.
[22] T. Nguyen, Z. Chen, and J. Lee. Dataset meta-learning from kernel ridge-regression. In International Conference on Learning Representations (ICLR), 2020.
[23] 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.
[24] 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.
[25] 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.
[26] M. Quiroz, R. Kohn, M. Villani, and M.-N. Tran. Speeding up meme by efficient data subsam- pling. Journal of the American Statistical Association, 2018.
[27] S. RRorrguuz-Saataan -aantan,, B. Zaldivar, and Hernandez-Lobato Function-see space inference with sparse implicit processes. In Proceedings of The 38th International Conference on Machine Learning (ICML 2022), 2022.
[28] 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.
[29] 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.
[30] т. G. J. Rudner, F. Bickford Smith, Q Feng, Y. W. Tch, and Y. Gal. Continual learning via sequential function-space variational inference. In Proceedings of the 39th International Conference Machine Learning, pages 18871-188877 2022.
[31] S. Sun, G. Zhang, Shi, and Grosse. Functional variational bayesian neural networks. In International Conference on Learning Representations, 2018
[32] M. K. Titsias, J. Schwarz, A. G. d. G. Matthews, R. Pascanu, and Y. W. Teh. Functional regularisation for continual learning with gaussian processes. International Conference on Learning Representations, 2019.

---

[33] т. Wang. J.-Y. Zhu, A. Torralba, and A. A. Efros. Dataset distillation. arXiv:1811.10955 2018.
arXiv preprint
[34] Z. Wang, T. Ren, J. Zhu, and B. Zhang. Function space particle optimization for Bayesian neural networks. In International Conference Learning Representations, 2019.
[35] R. Yu, S. Liu, and X. Wang. Dataset distillation: A comprehensive review. arXiv preprint arXiv:2301.07014 2023.
[36] B. Zhao and H. Bilen. Dataset condensation with differentiable siamese augmentation. In Proceedings of The 38th International Conference on Machine Learning (ICML 2021). 2021.
[37] B. Zhao and Bilen. Dataset condensation with differentiable siamese augmentation. In International Conference on Machine Learning, pages 12674-12685. PMLR, 2021.
[38] 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.
[39] B. Zhao, K. R. Mopuri, and H. Bilen. Dataset condensation with gradient matching. In International Conference Learning Representations, 2021.
[40] 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.
12

---

A Proofs
Proposition A.1. The gradient of the forward KL divergence with respect to the coreset is
VuDKL(Usil = p.B.p [log p(ý fu)] Ep.l.VVu Vulog p(ý fu)l,
(10)
where [ix]u and [Vulu are finite-dimensinal distributions of the stochastic processes Vx and Vu- respectively, fu := (f;),)) and p(y fu) -I;=1 P(U f(;)))
Proof. We follow the arguments in de G. Matthews et al. [10], Rudner et al. [29]. The forward KL divergence is defined as,
DRL|Uxlleul = / log dv' (f)dxx()) du
(18)
By the chain rule for the Radon-Nikodym derivative, we have
DKL|ss|bul = / log dı/ (f)dxx(f)) - deo
log dvn (f)dvx(f).
(19)
duo
The first term does not depend on u, so we investigate the second term. By the measure theoretic Bayes' rule,
dua
p(S S.
(f) = p(ý lu)
(20)
du
where p(ý Iu = L My 1) and,
p(p/u)= ( = p(ilf. u)dvo(f).
(21)
Now let PA (X' Rdd) (A Fd)) be projection function that takes function f and returns its restriction on a set A C X. Assuming that the likelihood depends only on the finite index set u. we can write
p( f.) dv/ (f) divulu (pu(f)) p(y u)
(22)
deg d[
where H), denotes the finite-dimensional distribution of stochastic process evalauted at u and Pa(s) := fu = (f(u)))f= are corresponding function values at u. Putting this back into the above equation,
f log duu (F)dz(()=
f log d[vuju (Fu)d[())u(fu)) dlvoju
deo
p( fu) d(vx)u(fuu)
(23)
log
p(ý u)
Biuch loo ply fn)) - logp(y Iu).
Now taking the gradient w.r.t. u, we get
Note also that VuD = -V.B llo p( f.)l + Va log p( u).
(24)
V Vulogp(ý u) = Vu log / p(y u)du(())
V Vup(ý f.u) dvo(f)
p(ý u)
/ V. log p( If, 1)
pl J, u) de
p(ý u)
(25)
/ V. log p(ý If. u)
dua duo
(f)dvo(f)
V. log p( f. (f)
V. log p f, (dldddlllu = Eieni logp(? fu)].
As result, we conclude that
V DKL (xLlvs//uul
= V V.E Pog f.)) Epa.,. Vu log f.)].
(26)
13

---

B Inducing points in Stochastic Variational Gaussian Processes
Stochastic Variational Gaussian Processes ISVGP: 12, 13] were introduced solution address the significant computational complexity, characterized by a cubic computational complexity of Onn)) and memory cost of O(n²), associated with performing inference using Gaussian Processes. In this context, represents the total number of training data points. SVGP effectively leverages concept called inducing points, which serves to reduce the computational complexity O(m)) and memory requirements to O(m²) 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. & n. The above description clearly shows that the inducing points have 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 = {=i,.... Zm} } involves maximizing the ELBO in order to make variational Gaussian distribution (f.,, fa) well approximate the posterior distribution p(fu Yy) This variational distribution composed of two parts: 1) p(ful()) [f),, which represents the Gaussian posterior distribution, and 2) q(fz). 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,). The goal of this optimization is to create good approximation of the posterior distribution p(ee [ie Dir) during the inference process, all while keeping the computational cost low. On the other hand, outlined in Section 3.2, the formulation of FBPC involves directly minimizing the divergence between function space posterior, specifically DKL [0x|| Vu] To sum up, while they do share some similarities in that they introduce 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 make the pseudocoreset posterior as close as possible to the full data posterior.
C Experimetal Details
The code for our experiments will available soon.
C.1 Expert trajectory
For expert trajectory, we trained the network with the entire dataset and saved their snapshot parame- ters at every epoch, following the setup described in [7]. For training, we used an SGD optimizer with 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 6x. 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 use samples from T epoch onwards as the MAP solution samples. When obtaining the MAP solution 0n. 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 learning rate of 0.001 to train the model until the training loss reached a threshold of y or below.
Next, to approximate our Gaussian variational function posterior, we employed the empirical co- variance 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 Xx and A, 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 a for pseudocoresets, and set the batch size for pseudocoresets to B. The hyperparameters used in our paper are summarized in Table 4.
14

---

Evaluation To implement the SGHMC algorithm, discussed in [9] and following the recommen- dations of [9], we employed the SGD with momentum along with an auxiliary noise term.
A = v Au = (V)(() av (O
(27)
set 17- 0.03, 0.1, and d = 0.01/m, where represents the coreset size. We perform 1000
epochs 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, trained pseudo- corests 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, also evaluated the performance on sophisticated architectures such as ResNet18, 8, ResNet34, VGG, and NlexNet. 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 ResNer18, has proven be challenging and has been explored only few previous works [38]. This This is primarily due to the lack of scalable training methods and the tendency for overfitting when training large networks. As result, even when evaluating ResNet after training on smaller network like ConvNet, the performance tends to suffer. Furthermore, has been reported that coreset training directly on ResNet initially yields lower performance compared to training on ConvNet [40, 35].
Our experiments also revealed similar trend in 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 observed to be lower compared to ConvNet. On the other hand, the second column, ResNet18 + ConvNet, refers to training both ResNet18 8 and ConvNet simultaneously using the FBPC-multi training approach. In this case, surprisingly, the test performance of ResNer18 actually improved when trained in conjunction with ConvNet using the FBPC-multi training approach. In this case, the ConvNet accuracy was recorded 60.03, which did not significantly compromise the ConvNet's performance while enhancing the performance of ResNet18. This suggests that training ConvNet acted as regularizer, preventing overfitting ResNet18 and enabling achieve better performance.
Table 4: Hyperparameter for each experiment setting.
ipc
T
Xx
A
N
1
1
0.01
0.05
0.01
1000
100
10
CIFARIO
10
2
0.1
0.05
0.01
1000
1000
100
50
10
0.2
0.05
0.01
1000
1000
500
2
0.1
5000
1000
CIFARI00
10
40
0.1
0.1
0.1
5000
1000
1000
50
20
0.2
0.01
0.01
300
1000
5000
1
2
0.1
0.1
3
5000
1000
100
Tmage
10
40
1
0.1
500
1000
100
50
40
0.1
500
1000
100
15

---

Table 5: Averaged test accuracy and negative log- sieeiikelioodd for each architecture of the pseudocore- set trained using the ConvNet architecture with CIFARIO dataset.
ConvNet
ResNet18
ResNet34
VGG
AlexNet
Ace (f)
30.7
27.720.98
22.46=0.12
26.33
21.05.:044
NLL (L)
3.44
2.87
3.0620.02
5.38
3.13:10.60
Ace (f)
22.5
47.51
25.48
47.87
32.27
10
NLL (L)
1.31
1.82#0.02
2.41
3.72=t0.09
2.96=t0.04
Acc (f)
71.20
62.02
47.97 23.37
58.24
52.42
50
NLL (L)
+0.00
1..00
2.10
3.03
2.08
Table 6: Test performance of FBPC (CIFAR1O, ipc 10) ResNet18. FBPC is trained with ResNet18 and ResNet18 + ConvNet (multiple architecture training).
ResNet18
ResNet18 ConvNet
Acc (1) 50.00
54.90
NLL (4) 1.75
1.56
D.3
Computational cost
To compare how scalable approach is compared posterior matching in the weight space, 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 the number of parameters grows, while FBPC operates very efficiently. Additionally, the coreset ipc increases memory usage proportionally its size. In terms memory considerations, FBPC excels. However, shown in Table 9, in terms of time, our method requires slightly more time because more SGD steps are needed acquire the empirical covariance. However, when using FBPC-isotropic, these steps can be reduced, trading off 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 augmen- tation during training. Interestingly, for the ipc case in BPC-IKL, performance improved when DSA was not applied. However, in all other cases, is evident that not using DSA Ieads to 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.
Figure 3: Example images of FBPC for CIFARIO with ipc 1.
16

---

Table 7: GPU memory usage (GB) for training CIFARIO FBPC with ipc 10.
LeNet
ConvNct
ResNer18
# parameters
6.2 x 10*
3.2 10
1.1 10"
FBPC
0.02
0.32
2.56
BPC-KKL
0.11
3.17
12.18
Table 8: GPU memory usage (GB) for training CIFARIO FBPC according to the ipc.
50
FBPC BPC-IKL
0.04
0.32
1.59
0.41
3.17
15.59
Table 9: Wall-clock time (sec) for step update for training CIFARIO pseudocoreset according to the ipc.
10
BPC-fKL FBPC
1.0440.10
2.59 3.20
1.37=0.13
1.5+0.13
8.38
Table 10: BPC Performances with and without DSA.
10
50
BPC-fKL (no DSA)
37.26 1.66
50.48=1.39
60.75 00.75:20.26
FBPC (no DSA)
33.69 23.73
1.30
66.03.1.22
FBPC (DSA)
35.4522.31
62.33_0.34
71.23=0.17
(a) Examples images of FBPC for CIFAR10 with ipc 10. image per class.
(b) Examples images of FBPC for CIFARIO with ipc 50. images per class.
Figure 4: Examples images of FBPCs for CIFAR10 dataset.
17

---

(a) Examples images of FBPC for CIFARI00 with ipc 1. image per class.
(b) Examples images of FBPC for CIFARI00 with ipc 10. 1 image per class.
Figure 5: Examples images of FBPCs for CIFARI00 dataset.
Figure Example images of FBPC for Tiny-ImageNet with ipc image per class.
Figure Example images FBPC for ImageNet with ipe 10. image per class.
18


Reviews:
[{"review_id": "z8MXYa9ls4", "reviewer": "Reviewer_sWjp", "summary": "This 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.", "questions": "* 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", "limitations": "I would have liked to see some discussion on the large gap to training on the full dataset.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* 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", "weaknesses": "* 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your constructive comments. We respond to the individual comments below:\n\n**[W1]** We agree that our method extends the discussion in Kim et al. [1] to function space posterior. However, we want to emphasize additional significant contributions in our work. \n* Firstly, we have demonstrated that obtaining a lower-dimensional function space matching for Bayesian pseudocoreset is easier than complex weight posterior matching, which is a novel finding. \n* Secondly, we propose scalable methods (Section 3.2, 3.3) for matching function space posterior, enabling computations in high-dimensional BNNs. Our algorithm that constructs variational posteriors directly in function spaces via trajectory statistics is significantly different from the previous function-space VI methods [2, 3] where they construct variational posteriors as pushforwards of weight-space variational posteriors. \n* Furthermore, by leveraging the similarity of function space across different architectures, our method can create architecture-robust BPC. \n\nThese unique contributions distinguish our paper from others in the field. We sincerely appreciate your consideration of these points.\n\n**[W2]** Thank you for the valuable feedback. We will carefully consider the replacement regarding the terminology 'variational distribution' in our paper during the revision period. Additionally, we have identified the similarities between our method for creating posterior distributions and the SWAG method [4]. We will add this discussion to the related work section of the paper. But again, please note that while SWAG considers collecting statistics on weight space trajectories, ours constructs statistics in **function spaces**, which makes ours more suitable and scalable for pseudocoreset construction.\n\n**[W3]** Thank you for the thoughtful comment. It is true that there remains a gap between the performance achievable by training large-size models on full data and the results obtained through our approach. We have also recognized and contemplated this issue, and that led us to propose the FBPC method, which addresses the scalability problem of Bayesian pseudocoresets and allows us to extend the models and datasets to sizes that were not easily achievable with existing Bayesian pseudocoreset methods.\nThrough FBPC, we successfully scaled our approach to handle larger models and datasets, demonstrating its potential in bridging the performance gap further. We agree that future research should continue in this direction, and we believe that our method serves as a valuable starting point for such endeavors.\n\n---\n**[Q1]** Yes, that's correct. When we increased the coreset size to 100, we obtained an accuracy of 73.45, and when increased to 200, the accuracy reached 74.80. It is evident that the performance increases as the size is expanded.\n\n**[Q2]** As the coreset size, ipc, and network size increase, memory usage also increases, and the memory burden may make conducting experiments in our current environment impractical. Specifically, BPC-fKL requires Monte Carlo samples corresponding to the parameter dimension size, whereas FBPC only needs samples from the function space dimension, resulting in a significant difference in memory usage.\n\nHowever, during the rebuttal period, we discovered that by reducing the pseudocoreset batch size during training, it becomes feasible to run experiments on a 3090 GPU. With a batch size of 100 during training, we achieved performances of accuracy 22.18 $\\pm$ 0.32 and NLL 4.65 $\\pm$ 0.02 where our FBPC achieved performances of accuracy 26.43 $\\pm$ 0.31 and NLL 4.30 $\\pm$ 0.05. \n\nFor scalability study of our method, please refer to our general response [G1].\n\n**Typos/minor** Thank you for finding them and we will correct them.\n\n---\n\n**References**\n\n[1] Balhae Kim, Jungwon Choi, Seanie Lee, Yoonho Lee, Jung-Woo Ha and Juho Lee. “On Divergence Measures for Bayesian Pseudocoresets.”  In Advances in Neural Information Processing Systems 36 (NeurIPS 2022), 2022.\n\n[2] T. G. J. Rudner, G. Chen, and Y. Gal. “Rethinking function space variational inference in\nBayesian neural networks.” 3rd Symposium on Advances in Approximate Bayesian Inference, 2020.\n\n[3] Tim G. J. Rudner, Zonghao Chen, Yee Whye Teh and Yarin Gal. “Tractable Function-Space Variational Inference in Bayesian Neural Networks.”  In Advances in Neural Information Processing Systems 36 (NeurIPS 2022), 2022.\n\n[4] Wesley J. Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov and Andrew Gordon Wilson. “A Simple Baseline for Bayesian Uncertainty in Deep Learning.”  In Advances in Neural Information Processing Systems 33 (NeurIPS 2019), 2019.\n"}, {"review_id": "9kHQ2uRP6c", "reviewer": "Reviewer_9hB5", "summary": "This 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.", "questions": "* 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]?", "limitations": "\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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* 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.", "weaknesses": "* 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your constructive comments. We respond to the individual comments below:\n\n**[W1]** Thanks for your insightful comment. We will include the relevant discussion in the paper. As you pointed out, inducing points in Stochastic Variational Gaussian Processes and the functional Bayesian pseudo coresets in BNN are similar in their purpose.\t\n* SVGP : In the context of Stochastic Variational Gaussian Processes, the process of optimizing the inducing points denoted as $Z=${$z_1,...,z_n$} 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: $p(f_{tr}|f_z)$, which represents the Gaussian posterior distribution, and $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.\n* FBPC : As outlined in section 3.2, the formulation of FBPC involves directly minimizing the divergence between function space posteriors, specifically $D(\\nu_x, \\nu_u)$.\n \nTo sum up, while they do share some similarity 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. \n\n**[W2]** Please refer to our general response [G5].\n\n**[W3]** We appreciate the valuable recommendations, and we will include the mentioned papers in the related work section during the revision.\n\n**[W4]** Please refer to our general response [G2] for the regression experiment. As you mentioned, a comparison with other methods like HMC would indeed be interesting. Analyzing the differences between the predictive distributions obtained from a-posteriori methods and those obtained from BPC could be an engaging avenue for future work.\n\n**[W5]** Please refer to our general response [G1].\n\n---\n**[Q1]** There are two major advantages in using the forward KL divergence. \n* Firstly, when compared to the reverse KL divergence, which has a mode capturing property, forward KL minimization favors solutions covering the entire target distribution. This characteristic makes it more suitable for Bayesian pseudocoreset construction, as evaluating model performance through Bayesian model averaging (BMA) allows better consideration of diversity over the entire distribution rather than focusing on individual modes. This advantage is also acknowledged in [1], where they compared three divergence measures—forward KL, reverse KL, and Wasserstein distance—and demonstrated the superiority of forward KL.\n* Secondly, another advantage of using forward KL is the ease of computing the objective. While there are many other well-defined metrics available, for BPC optimization, it is crucial to have the gradient of the divergence between two posteriors with respect to BPC. We were able to derive the derivative for forward KL through proposition 3.1, which enabled us to compute FBPC. \n\nExploring the possibility of similar derivations for other divergences could be a promising topic for future work.\n\n**[Q2]** Please refer to our general response [G4].\n\n**[Q3]** Thank you for the insightful comment. While it may raise similar concerns as you mentioned, Table 6 of the Appendix demonstrates the simultaneous learning of pseudocoresets for two entirely different architectures: CNN and ResNet. This result illustrates that our method is capable of accommodating various architectures, and it is not solely reliant on the linearization and Gaussian variational form. Instead, we attribute this capability to the effect of similar function posteriors mentioned in Section 3.4.\n\nIn addition, we explored using a diagonal covariance to achieve a more refined approximation of the posterior in weight space BPC. However, as shown below, this approach still did not yield significant benefits in terms of architecture generalization. We believe that matching the function posterior distribution is more helpful for enhancing architecture robustness, rather than solely focusing on how well we approximate the posterior in weight space.\n\n| BPC-fKL  | LN | BN | GN | IN |\n|----------|-----------|-----------|-----------|-----------|\n| Accuracy | 40.06     | 46.74     | 40.96     | 53.07     |\n\n**[Q4]** As you mentioned, there are indeed various methods to obtain predictive distributions, including a-posteriori approximations [2,3]. It is worth noting that our proposed method, utilizing the pseudocoreset, can also be applied in conjunction with a-posteriori methods to further reduce computational burden while obtaining predictive distributions.\n\nAdditionally, our approach offers the advantage of obtaining a **small condensed dataset** which can have versatile applications. This condensed dataset can be leveraged as memory data in transfer or continual learning settings, and it may hold value as a learnable prior in Bayesian settings. Utilizing a pseudocoreset as a learnable prior, rather than storing weights directly, opens up interesting possibilities for future work. Investigating the feasibility of using small datasets as a learnable prior represents a promising avenue for further research.\n\n---\n**References**\n\n[1] Kim et al. “On Divergence Measures for Bayesian Pseudocoresets.” NeurIPS 2022.\n\n[2] Daxberger et al. “Laplace Redux -- Effortless Bayesian Deep Learning.”  NeurIPS 2021.\n\n[3] Deng, Z., Zhou, F., & Zhu, J. “Accelerated Linearized Laplace Approximation for Bayesian Deep Learning.” NeurIPS 2021.\n\n\n"}, {"review_id": "yrELztNefq", "reviewer": "Reviewer_d2wU", "summary": "The 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", "questions": "1.\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).", "limitations": "NA", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\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.", "weaknesses": "1. Ambiguous writing. \n\n2. Experiments can be better. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your constructive comments. We respond to the individual comments below:\n\n**[W1, Q1]** Thank you for the valuable suggestion. We will consider reorganizing the paragraphs during the paper revision.\n\n**[Q2]** I agree with your point that the statement 'function space itself has a lower dimension compared to weight space' can be confusing. It seems better to rephrase it as 'functions have a lower dimension than weights,' as you suggested.\n\n**[Q3]** In fact, in the process of calculating $\\hat\\Psi$ in Equation 16, we have indicated the stop-gradient operator as **'sg.'** This implies that back-propagation is no longer applied to the parameters computed during SGD steps in our approach. The gradient of the pseudocoreset flows through the input $u$ in the calculation of $\\hat\\mu$. Moreover, it requires several tens of SGD steps to reach the MAP solution of the pseudocoreset, and an additional 30 SGD steps are used to obtain empirical variance. Storing gradients throughout all these processes would be computationally very expensive. This does not align with our intention to create a scalable BPC algorithm.\n\n**[Q4]** We will make sure to do that during the revision.\n\n**[Q5]** We agree with the statement that DSA [1] impacts the performance. Following the convention in the existing literature, we employed DSA in our approach and will make sure to mention it in the paper.\n\n**[Q6]** To the best of our knowledge, dataset distillation primarily focuses on achieving test performance matching, which can be seen as making the **point estimation** performance of parameters align as an objective. Accordingly, the evaluation of the distilled dataset is compared using SGD performance. On the other hand, methods like ours BPC aim to address how well **posterior distributions align.** Consequently, in our paper, we evaluate using the SGHMC sampling technique. Nevertheless, for the purpose of comparing with the relevant field, a comparison against the SGD performance of a few Dataset condensation methods yields the following [1,2,3].  \n\n**< Test accuracies for CIFAR10 dataset >**\n| ipc | GM (SGD)       | DSA (SGD)      | DM (SGD)       | Ours (SGHMC)   |\n|-----|----------------|----------------|----------------|----------------|\n| 1   | 28.3 $\\pm$ 0.5 | 28.8 $\\pm$ 0.7 | 26.0 $\\pm$ 0.8 | **35.5 $\\pm$ 0.3** |\n| 10  | 44.9 $\\pm$ 0.5 | 52.1 $\\pm$ 0.5 | 48.9 $\\pm$ 0.6 | **62.3 $\\pm$ 0.3** |\n| 50  | 53.9 $\\pm$ 0.5 | 60.6 $\\pm$ 0.5 | 63.0 $\\pm$ 0.4 | **71.2 $\\pm$ 0.2** |\n\n**[Q7]** In the appendix, we have indicated that there is typically a performance degradation of the pseudocoreset in larger architectures such as ResNet. We believe this is primarily due to the limited dataset size used to evaluate the trained pseudocoreset, leading to increased overfitting in these larger architectures. This phenomenon is also commonly observed in existing dataset distillation methods. Our contribution, utilizing the function space, enables more scalable use of larger architectures like ResNet for the pseudocoreset training, and our method's performance on ResNet is superior to that of existing literature [4]. This reaffirms the validity of our approach.\n\n**[Q8]** We appreciate the valuable recommendations, and we will include the mentioned papers in the related work section during the revision.\n\n**[W2]** Please refer to our general responses [G2, G3].\n\n---\n**References**\n\n[1] 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\n[2] G. Cazenavette, T. Wang, A. Torralba, A.A. Efros, J.Y. Zhu. “Dataset Distillation by Matching Training Trajectories.” CVPR, 2022.\n\n[3] Bo Zhao, Hakan Bilen. “Dataset Condensation with Distribution Matching.” WACV 2023, 2023.\n\n[4] R. Yu, S. Liu, and X. Wang. Dataset distillation: A comprehensive review, 2023.\n"}, {"review_id": "MyqlhI4yjF", "reviewer": "Reviewer_jYGC", "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.", "questions": "N/A.", "limitations": "N/A.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "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. \n* 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.\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your constructive comments. For the weaknesses, please refer to our general responses [G3, G4, G5]."}]

Author rebuttal:
We express our gratitude to all the reviewers for their valuable and insightful feedback. They have acknowledged the clarity in our method description (R-sWjp, R-9hB5) as well as our coverage of relevant prior works (R-9hB5). The reviewers have also recognized the paper's strong motivation and the significance of addressing an engaging problem (R-sWjp, R-9hB5, R-d2wU, R-jYGC). Moreover, they have emphasized that our proposed FBPC method demonstrates superior performance compared to other BPC approaches (R-sWjp, R-d2wU, R-jYGC), highlighting its practicality in Bayesian neural network inference (R-9hB5, R-jYGC). The reviewers have also highlighted our contributions, specifically our achievements in a multi-architecture setting (R-d2wU) and the empirical estimation of function space posterior (R-jYGC).

We sincerely appreciate all the constructive comments and we provide general responses addressing questions that were commonly raised by the reviewers below:

**[G1] (Scalability study)** In our paper, we proposed a more scalable BPC construction algorithm by focusing on distribution matching in the function space rather than the weight space. Therefore, we can effectively and efficiently create BPC even as the number of parameters increases. 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, and the results are as follows. As shown in the table, 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. We believe the same principle applies to dataset size as well.

**<GPU memory usage (GB) (CIFAR10)>**
|   ipc10           | LeNet           | ConvNet         | ResNet18        |
|--------------|-----------------|-----------------|-----------------|
| # Parameters | $6.2 \times 10^4$ | $3.2 \times 10^5$ | $1.1 \times 10^7$ |
| FBPC         | 0.02            | 0.32            | 2.56            |
| BPC-FKL      | 0.11            | 3.17            | 12.18           |

|  ipc       | 1    | 10   | 50    |
|---------|------|------|-------|
| FBPC    | 0.04 | 0.32 | 1.59  |
| BPC-fKL | 0.41 | 3.17 | 15.59 |

**[G2] (More experiments)** Our paper introduces a scalable BPC construction method for BNN classification. We extended our exploration to regression, using the yacht dataset with (308, 6) dimensions and an MLP with one hidden layer. The test mean squared error (MSE) for different coreset sizes is as follows. This confirms our method's effectiveness in regression.

The 'Random' method selected samples randomly from real datasets, potentially including outliers. As coreset size increases, the trained FBPC's test MSE approaches 0.0647, matching the test MSE from training on the entire dataset.

|        | 10     | 20     | 30     | 50     | 100    |
|--------|--------|--------|--------|--------|--------|
| Random | 0.3424 | 0.5335 | 0.2757 | 0.1230 | 0.0931 |
| FBPC   | 0.3148 | 0.2541 | 0.1434 | 0.0854 | 0.0829 |
| FBPC (Jacobian) | 0.3156 | 0.2460 | 0.1413 | 0.0872 | 0.0825 |

**[G3] (A lot of approximations)** To understand the impact of these approximations, we conducted regression task experiments on small BNNs, where it was possible to obtain the full Jacobian." As seen in the table from [G2], the results were not significantly different from ours. Analyzing the effects of each approximation in larger models and researching more refined and scalable approximation methods would be a valuable direction for future research.


**[G4] (Possibility of the usage of the full Jacobian)** 
Incorporating the full Jacobian in computations is challenging, especially for our pseudocoreset updates that involve frequent calculations of equations 12 and 13. Computing $\Sigma_x$ and $\Sigma_u$ requires inverting the Hessian matrix, often approximated using the Fisher information matrix. For the full dataset, this becomes computationally demanding, and for the coreset, small size leads to high variance with the standard Fisher approximation. These challenges persist even for small BNNs.

Initially, we attempted several approximations for FBPC, especially for diagonal approximations of $\Sigma_x$ and $\Sigma_u$, which resulted in instability and hindered convergence. Employing empirical variance proved more stable for FBPC computation. For instance, using Jacobian and Fisher Information Matrix for a diagonal covariance in the CIFAR-10 ipc10 setting resulted in 56.97% accuracy for FBPC, which was less favorable than our method. This approach introduces additional hyperparameters, complicating optimal settings identification.

**[G5] (Concerns of Gaussian approximation)**
As highlighted by Reviewer 9hB5, there's no assurance that the actual posterior distribution would perfectly align with the variational Gaussian distribution used in function space approximation. Nonetheless, it's crucial to recognize that approximations are required for effective objective computation.

Employing linearization and variational Gaussian distributions in function space is a practical strategy to manage scalability and enable feasible computations [1]. Although this approximation introduces some bias, it's a necessary compromise to address the challenges of high-dimensional models.

In the context of our work, exploring more sophisticated techniques, such as normalizing flows [2,3], to design more complex posterior distributions is indeed a promising direction for future research. 

---
**References**

[1] Tim G. J. Rudner, Zonghao Chen, Yee Whye Teh and Yarin Gal. “Tractable Function-Space Variational Inference in Bayesian Neural Networks.”  NeurIPS 2022.

[2] Kristiadi et al. “Posterior Refinement Improves Sample Efficiency in Bayesian Neural Networks.” NeurIPS 2022.

[3] Chen et al. “Bayesian inference via sparse Hamiltonian flows.”  NeurIPS 2022.


Meta-review: This paper presents a novel pseudocoreset construction method for a relevant pushforward of the posterior in certain models, instead of the raw parameter posterior as usual. The reviewers found the writing generally clear, the method interesting, and the experiments compelling. The paper is slightly incremental given Rudner et al, and involves quite a few approximations. Note that there are some minor technical issues in a few spots (e.g. lines 126/127, where densities on function spaces are used but not well-defined; 172/173 unclear claim about the dimension of function spaces, etc). Please address these points as well as the points raised during the discussion for the camera ready.
Decision: Accept (poster)