Variational Bayesian Pseudo-Coreset
Authors: Hyungi Lee
Seungyoo Lee
Juho Lee
Venue: ICLR 2025 Poster (2025)
Keywords: Bayesian Pseudo-Coreset
Variational Inference
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: In this paper, we propose a new Bayesian Pseudo-Coresets method which enables the efficient variational inference.

Abstract: The success of deep learning requires large datasets and extensive training, which can create significant computational challenges. To address these challenges, pseudo-coresets, small learnable datasets that mimic the entire data, have been proposed. Bayesian Neural Networks, which offer predictive uncertainty and probabilistic interpretation for deep neural networks, also face issues with large-scale datasets due to their high-dimensional parameter space. Prior works on Bayesian Pseudo-Coresets (BPC) attempt to reduce the computational load for computing weight posterior distribution by a small number of pseudo-coresets but suffer from memory inefficiency during BPC training and sub-optimal results. To overcome these limitations, we propose Variational Bayesian Pseudo-Coreset (VBPC), a novel approach that utilizes variational inference to efficiently approximate the posterior distribution, reducing memory usage and computational costs while improving performance across benchmark datasets.

Published as a conference paper at ICLR 2025
VARIATIONAL BAYESIAN PSEUDO-CORESET
Hyungi Lee KAIST
Seungyoo Lee KAIST punctrats@kaist.ac.kt
Juho Lee
KAIST
1bx2708ekaist.ac..
juholee@kaist.ac. kr
ABSTRACT
The success of deep learning requires large datasets and extensive training, which can create significant computational challenges. To address these challenges. pseudo-coresets, small Iearnable datasets that mimic the entire data, have been proposed. Bayesian Neural Networks, which offer predictive uncertainty and probabilistic interpretation for deep neural networks, also face issues with large- scale datasets due their high-dimensional parameter space. Prior works on Bayesian Pseudo-Coresets (BPC) attempt to reduce the computational load for computing weight posterior distribution by a small number of pseudo-coresets but suffer from memory inefficiency during BPC training and sub-optimal results. To overcome these limitations, we propose Variational Bayesian Pscudo-Corese (VBPC), novel approach that utilizes variational inference to efficiently approx- imate the posterior distribution, reducing memory usage and computational costs while improving performance across benchmark datasets.
INTRODUCTION
While deep learning has shown remarkable performance across various fields, its success requires large amounts of data storage and extensive training. However, handling such large datasets can impose significant computational burden, especially when training new models or updating ex- isting ones with new data. In settings like continual learning, where the model must be trained continuously on new data, this challenge becomes more pronounced due to the risk of catastrophic forgetting. To mitigate this, small subset of representative data, called coreset, is needed to pre- serve knowledge from previously learned data. Instead of creating small dataset as subset of the entire data to represent it, the approach of treating the small dataset itself as learnable parameters and training it to mimic the entire dataset is known as dataset distillation or pseudo-coreset (Nguyen et al., 2020; 2021; Zhou et al., 2022; Loo et al., 2023).
On the other hand, Bayesian Neural Networks (BNNs) have gained attention in fields like health- care (Abdullah et al., 2022: Lopez et al., 2023) and climate analysis (Vandal et al., 2018) because they provide posterior distribution over the weights of a deep neural network, enabling the mea- surement of predictive uncertainty and allowing for a probabilistic interpretation of parameters (Pa- pamarkou al., 2024). While this method is promising for enabling various types of statistical analysis, BNNs face significant challenges when applied to real-world scenarios that involve large- scale datasets. The high-dimensional parameter space and structure of BNNs often lead to posterior landscapes with multiple modes, which complicates efficient and straightforward computation of predictive uncertainty. To overcome this, BNNs typically rely on indirect methods such as Stochas- tic Gradient Markov Chain Monte Carlo (SGMCMC; Welling & Teh, 2011; Chen et al., 2014; Ma al., 2015) or variational inference (VI: Blei al., 2017: Fiedler & Lucia, 2023; Harrison al., 2024b) instead of directly calculating the posterior distribution in closed form. However, these approaches still depend on gradient-baad -aased updates of model weights for large-scale datasets. In par- ticular, SGMCMC-based methods face the challenge of increased computational load, as the amount of training grows linearly with the number of weight samples needed.
To overcome these issues, prior works Bayesian
Coreeet (BPC; Manousakas al., 2020; Kim et al., 2022; 2023; Tiwary et al., 2024) aim to learn a small synthetic dataset that helps effi- ciently compute the posterior distribution of BNNs' weights. These studies train the pseudo-cores by minimizing the divergence between the posterior obtained using the full dataset and the posterior obtained using the pseudo-coreset. However, these studies face three major problems: 1) require

---

Published as a conference paper at ICLR 2025
expert trajectories for training, 2) use stop-gradd stop-gradient during training, and 3) still rely SGMCMC sampling for weight space posterior computation. First, expert trajectories refer to the trajectories of model weights trained using the full dataset. In previous studies, these trajectories epoch with multiple different seeds, and they are used to approximate and match the posterior dis- tribution. This creates the problem of needing to store the model weights for the number of epochs multiplied by the number of seeds in order to train the pseudo-coreset. Secondly, when training BPC, the posterior distribution is computed using the BPC for loss computation via gradient-based methods. As the updates progress, the computational graph required to update the pseudo-coress based on the loss becomes significantly larger, resulting increased memory demands. To address this memory issue, prior works have used the siop-gradient method to reduce memory consumption. However, this approach leads to sub-optimal results because it prevents accurate updates. Finally, even after training the pseudo-coreset, the weight posterior distribution remains multi-modal, mean- ing that while the training cost is reduced, sequential training through SGMCMC sampling is still required for each sample. Additionally, after obtaining the samples, forward computation is needed for each sample to calculate the predictive distribution during Bayesian inference.
saved for every
To address these issues, propose novel BPC approach called Variational Bayesian Pseudo- Coreset (VBPC). In learning VBPC, unlike previous works, we employ VI, specifically last-layer VI (Fiedler & Lucia, 2023; Harrison et al., 2024b), to approximate the posterior distribution. During the VBPC training and inference process, we demonstrate that this variational formulation allows us obtain the closed-form posterior distribution of the last layer weights, which frees our method from relying on stop-gradient. This resolves the issuc of suboptimal performance seen in previous approaches. And, we propose a memory-efficient method approximate the predictive distribution with only single forward pass instead of multiple forwards, making the approach computationally and memory -ffooreefficient. Furthermore, we empirically show that VBPC achieves better performance compared to other baselines on various benchmark datasets.
2
PRELIMINARIES
2.1 BAYESIAN NEURAL NETWORKS AND BAYESIAN MODEL AVERAGING
In Bayesian Neural Network frameworks (Papamarkou et al., 2024; Lee al., 2024), the main objective is to compute the predictive distribution for given input x, while accounting for model uncertainty (i.e., epistemic uncertainty), as shown below:
plule, D) plula, ())(O
(1)
where represents the observed data, and 0 denotes the model parameters. This process is known as Bayesian Model Averaging (BMA). To perform BMA, we need to compute the posterior distri- bution D)) and evaluate the integral. However, due to the complexity of the model and the high- dimensional parameter space, directly computing a closed-fomn solution for p(0 is impractical. Therefore, in practice, we typically rely on posterior sampling methods such as SGMCMC (Welling & Teh, 2011; ;; Chen et al., 2014; Ma et al., 2015) VI (Blei et al., 2017; Fiedler & Lucia, 2023) to approximate the posterior distribution.
2.2 BAYESIAN PSEUDO-CORESEE
As mentioned in Section the large size of the training dataset makes computationally intensive to perform SGMCMC or VI for approximating the posterior distribution of BNNs. To address these challenges and efficiently compute the posterior distribution in terms of both computation and memory, previous works (Kim et al., 2022; 2023; Tiwary et al., 2024) introduced BPC within the SGMCMC framework. Specifically, BPC S is optimized using the following objective:
S arg min p(()))
(2)
s
where D can be various divergences between the two distributions (Kim et al., 2022). The optimiza- tion poses a challenge, as the posteriors D))[D) and p(0[S) are intractable for most of the cases. Previous works (Kim et al., 2022; 2023; Tiwary al., 2024) attempt to approximate them using weight checkpoints obtained from training trajectories based on the dataset D (i.c., expert trajecto- ries) which requires expensive computation and memory cost.
2

---

Published as a conference paper at ICLR 2025
2.3 NATURAL GRADIENT VARIATIONAL INFERENCE WITH EXPONENTIAL FAMILIES
Although several methods exist for approximating the posterior pD) in this paper, we focus on VI (Bishop, 2006; Blundell al., 2015; Blei et al., 2017). In VI, we approximate the target posterior with a variational distribution that is easier to handle and optimize the parameters of the variational distribution minimize the Kullback-Leibler (KL) divergence between the approximate and target posterior distributions. Among the many possible choices for variational distributions, we focus on the exponential family. We assume that both the prior pxo (0) and the variational distribution (x()) belong to the same class of exponential family distributions:
(3) where vh()) represents the sufficient statistics, A(-) is the log partition function, and and λo the natural parameters for QA and PAo- respectively. We further assume that the exponential family minimal, meaning that there is no non-zero vector x such that (x, +00) evaluates to constant. Under this setting, we can optimize the variational parameter À by minimizing the following loss:
QA (0) x exp((1(0); X) A())), PAo (0) x exp((v(0)) An) - A(Ao)),
(4) where B > 0 is temperature controlling the strength of the KL regularization (Blundell et al., 2015; Wenzel et al., 2020). When B = 1, minimizing Eq. is equivalent to minimizing DRL. 1x (0)llp(( Optimizing equation 4 with natural gradient descent (Amari, 1998) has been shown to be effective, especially for large-scale deep neural networks (Khan et al., 2018; Osawa al., 2019; Shen et al., 2024). The optimal solution of Eq. 4 must satisfy the following equation,
CD()) := E. I- log log p(D))) + 3D 1 (0) (0)),
= + V [log (D))))
(5)
where μ = ,,, [0)] = V (A) is the mean parumeter corresponding to the natural parameter X. Except for some cases, Eq. 5 does not admit closed-form expression for X". Therefore, one must rely on iterative algorithms to obtain it. This approach, which solves the variational inference using iterative natural gradient descent steps, covers a broad spectrum of machine learning algorithms and is commonly referred to as the Bayesian Learning Rule (BLR) (Khan & Rue, 2023).
VARIATIONAL BAYESIAN PSEUDOCORESET
In this section, we propose novel method called Variational Bayesian Pseudo-Coreset (VBPC) which effectively learns S and thereby well approximates the variational posterior distribution with full dataset distribution. Several recent studies (Fiedler & Lucia, 2023; Harrison et al., 2024b) have shown that using only last layer for variational inference is simple and computationally cheap, yet performs comparably to more complex methods. Motivated by these findings, we seek to learn a seeudo-coresed S that effectively approximates the last layer variational posterior for the classification task, all while ensuring computational and memory efficiency.
3.1 PROBLEM SETUP
Consider a supervised learning problem with dataset D (xi, U:) While our discussion
can easily extended to more general problems, in this paper, we focus on k-way classification tasks, where y/ E {0, 1} is one-hot vector representing a category. Given D and model fa parameterized by 8, we aim to learn a synthetic dataset (pseudocoreset) S == (Ā,, )t) solving Eq. 2 under constraint ń < n. We approximate the pseudocoreset posterior p(0) S) by solving the following variational inference problem,
Ls (A) = (A()) 3s DKI, (ax (0)llsx.
A/3 arg min Ls(A), A
(6)
==1 log ps (y: Ir: 0)] is the expected sum of negative log- log-likelshoods over given a choice of likelihood Ps(y xz 6). Throughout the paper, we call Eq. 6 as coreset VI prob- lem. Ideally, we would like to match the optimal solution of the coreset VI problem to the optimal variational distribution computed with the original dataset D,
where (A( := - - Σ
(A( (x( + 3D (xx (0) (0)), XD arg min (A))
X
where (A) Ey, Ei log PD (v Iz, 0)] for likelihood po(y)æ, 0). We call Eq. 7 as dataset VI problem. After obtaining >3 and Xp. to learn S, we can minimize D(a 9xz.) for some pre- defined divergence D.
3

---

Published as a conference paper at ICLR 2025
3.2 BILEVEL OPTIMIZATION
It is often challenging first compute the approximate solutions of Eqs. and 7 and then backprop- agate through the divergence D(9)3 9xz)) Instead, considering the optimization nature of the VI, cast the problem of coreset learning bilevel optimization as follows:
s = argmin min LD()'s where X = =arg min Ls(A).
(8)
Note that similar approaches have also been considered in the dataset distillation literature (Loo et al., 2023). Under the bilevel optimization formulation, learning S requires the derivative
V&LD(A3) (Vsu3)V,.C((X))
where uts = V)A(X)) is the mean parameter corresponding to A To obtain VsH.s may apply the implicit function theorem (Bengio, 2000; Krantz Parks, 2002) to Eq. 5. Specifically, if we let:
F)S, À Ao + $5 V)(s(A
(10)
With F3) = 0, applying the implicit function theorem.
VgF(S,P3) + (VsH)) = Vsr. -V,FF(S. Vsns = (V (V, + $5 (X53))
(11)
Plugging this back into the above equation, we get the expression for the gradient
VsLD (VsV (V + B3 (S(V)))()3) (v
(12)
Unfortunately, the term involving the inverse is usually intractable, one needs an approximation (e.g.. Lorraine et al. (2020)). In the next section, we describe a case where the derivatives can be computed in closed form, and develop Bayesian pseudo-coreset algorithm based on it.
3.3 LAST LAYER VARIATIONAL BAYESIAN PSEUDOCORESET
Recently, there has been growing interest in subspace Bayesian neural networks (BNNs), where only a subset of the network's 's parameters are treated as random, while the remaining parameters are kept deterministic (Sharm et al., 2023; Shen et al., 2024). An extreme form of subspace BNN would be the last layer randomization, where a neural network fo(x) E IR& is decomposed as a feature extractor d(x:) E R followed by linear layer W Raxk Denoting the j column of W wy have [=) = Φ(x) "w for j [K]. Adapting the last layer randomization scheme, we treat only the parameter W of the linear layer as random while keeping the feature extractor (xx) deterministic. From below, we describe our model more in detail.
and the j output from fe(x) as [fe(=)];
Variational distributions.
We assume the Gaussian priors and variational posteriors for W,
Px, (W) = II -w (W) II ((wylm,
(13)
with the natural parameters and the corresponding mean parameters are given as,
Ag = concat( (1o - (p- (p=112)(m) Po = concat( (0, p- p 1n1
(14)
X concal ((A,) / ((g))
where X; i (1 (1/2)V,], and μy myj Vj + mym Here, we denote Ia as the d x d identity matrix and is dre-ned precision hyperparameter of the prior. Note that the block wise approximation qx (W) reduces the space complexity of the variance parameter V = (V,) from (K2222) to O(kh²) while keeping flexibility compare mean field approximation.
4

---

Published as a conference paper at ICLR 2025
Likelihoods. For classification problem, it is common to softmax categorical likelihood, and we follow that convention for the dataset VI problem with pp. However, for the coreset VI problem, the softmax categorical likelihoods would not allow closed-form solution, which would necessitate approximations involving iterative computations to solve the bilevel optimization Eq. 8. This would, for instance, require storing the unrolled computation graph (Vicol al., 2021) of the iterative updates and performing backpropagation through it, leading to significant computational and memory overhead (Werbos, 1990). As a detour, we use the Gaussian likelihood for the ps, as it allows us to obtain a closed-form solution. While using Gaussian likelihoods may seem coun- terintuitive for a classification problem, it widely used the literature on infinitely-wide neural networks (Lee et al., 2017; 2019; 2022), and one can also interpret it as solving the classifica- tion problem as a regression, using one-hot labels as the target vector. More specifically, we set ps(wiz, 0) = (y/W >
lx) where 1 is the precision hyperparameter for the likelihood.
With our choices for pD and ps we can expand the bilevel optimization problem as follows.
1 arg min Eg
Log W ((b.), y (1)))
+ DKL (axl1xx01
(15)
n
exp(o(e;)) T (wy) Σk exp(p(x) Twi)
8" arg min Bugg
Ki log
+ BDDDL 1943 I
(16)
3.4 SOLVING CORESET VI PROBLEM
Based on our choices described in the previous section, we show how we can obtain closed-form expressions for the coreset VI problem. The likelihood term for the coreset VI problem
E
log. Nog.N T 7
Egg. [(i
((i)) w,) w )].
(17)
where Vt. indicates jth element of ÿJ for all i E [i]] and - denotes equality up to constant. Then we can further elaborate Eq. 17 as follows:
Egg
(Bir - ((t.) wo)
)"
27 e 7
20.J (1) Tr TT (TTn)))))
(18)
2
(1) = mg and (1) = V, + m,m for all j E [4]. Then by Eq. 18, the gradient of the likelihood with respect to μj can be computed as:
where Ú.jj : Ii....... Da.ll : ( ((+))]
T, M4
Y $ 4. 2
(s) - 0.g, V (
(19)
Then from Eq. 5, we obtain the closed-form solution for the coreset VI problem as follows:
7 T U. - L 2,35 + T
2
E
(20)
Bs
with Woodbury formula (Woodbury, 1950) which leads to
mj- = TT
(Bs Ia + T
V.gi VJ -
Φ p²3s
7
Ia
7
ΦΦ
(21)
p
pBs
For all; j E [k), the values V identical, meaning the full covariance calculation, though Okh?), only requires computing and storing the variance once, Ohh)). We will refer to this shared variance as V* See Appendix A. and Appendix A.2 for detailed derivations in this section.
Bilevel optimization
inffuence maximization. Before proceeding to the dataset VI prob- lem, let us describe how the last- variational model simplifies the coreset gradient Eq. 12. From Eq. 19, we have Vs (A3) = 0, leading to (s Vudss. Using this, we can show that
V
VsLo = = -VCcss((( (,3)( (V "V,
(22)
Here, (A3))( (})))())) "V, ()g3)) is the variant (in a sense that it is defined w.r.t the gradient of the variational objective by the mean parameters) of the influence function (Koh & Liang, 2017), measuring the influence of the coreset on the dataset VI loss computed with D.
5

---

Published as a conference paper at ICLR 2025
3.5 COMPUTATION FOR DATASET VI PROBLEM
Now with these coreset VI problem solutions, we have to find the optimal 5 by solving Eq. 16. However, unlike the coreset VI problem, since we use categorical likelihood with a softmax output, a closed-form solution cannot be obtained from Eq. 16. Thus we have to use iterative updates, such as Stochastic Gradient Descent (SGD), for the outer optimization problem. Then because ((x) wy ~ (o((2)) mj. ((x) (=)()) for all j E [ks], the dataset VI problem changed into
CD(A5) = Σ i=1
Ne M..BErnN (a.)
log
exp 2j
DKD [923 Ilpxo] (23)
Σ i= exp 2j
where z 121...., 2x/. mn
(x) (z)
[mi())..... (x)) Σ' xz) diag (2)(()) (x)) and
(mn (z), (x), (x)) (0(z) m. x(x) V (x)) for all j [k] and æ. For simpler notation, we will
denote (m) (x)) (m,, Σÿ) Then have to approximate Ez-N( 55+) log Σ exp
exp2;
to compute the loss LD (Xts) analytically. To compute approximate expectation for the likelihood,
we first change the form as follows:
exp(>3) -i=1
- / log K
Eg log
Σ o(2g -2i) igf
Nzzlm², "alm"","""d E')dz,
(24)
exp 2:
where (-) is the sigmoid function. Then utilize mean-field approximation (Lu et al., 2020) to the 2,8 to approximately compute the Eq. 24:
exp (2₁)
m
(m E+) log Ei=1 exp(21)
log softmax
(25)
VI +
where = K and E = oxx) T V (x). Refer to Appendix A.3 for the complete derivation of Eq. 23, Eq. 24, and Eq. 25. By Eq. 25, our outer optimization loss has changed form as follows:
(o()s) = ΣΣ vi. log softmax ==
n
aZ (()
m m (xx))
]
3D GDDKI.lugg IPx.)-
(26)
Here, since n is large, we need to employ the SGD method to optimize S. Thus, using the training batch B C (((1111)) .i)..... (xn, using stochastic loss as follows:
(Zn,,,)) }n)),,
we compute approximate loss Lo for the batch and update S
(o(A3)
n
Yi log softmax
[ ax:
m (xx)
]
+ BB BKDlala (27)
B
iEB
3.6 TRAINING AND INFERENCE
Memory Efficient Loss computing If we naïvely compute the gradient of S by directly eval- uating Eq. 27, calculating Σx and DKL[Uxg [Po] will require computations involving V*, which demands 2 memory. However, the quadratic memory requirements with respect to the feature di- mension pose challenge when training S for large-scale models. address this issue, we propose memory-efficient approach for computing loss during training in this paragraph. We will address the efficient computation of Σ* in the below paragraph Variational Inference and Memory Effi- cient Bayesian Model Averaging. Here, will focus efficiently computing the KL divergence. Since both 7xg and px, are Gaussian distributions, the KI. divergence can be expressed as follows:
lxt [xxo] _ [-k log det((V)) kp Tr(V*) + pl|m" [1]]
(28)
Thus we have to efficiently compute det V and Tr(V*). For the det V*, we use Weinstein- Aronszajn identity (Pozrikidis, 2014) which results as follows:
del V* pr deteetT +
(29)
7 σΦ
pās

---

Published as a conference paper at ICLR 2025
And for the Tr(V*), easily change the form with a property of matrix trace computation:
Tr "I:
Tr(V*)
Is ?
7h pDs
2 - pās
ΦΦ (
(30)
p55
By these formula, we can calculate the KL divergence without directly computing V reducing the memory from (hh)) to (ñ²). Refer to Appendix A.4 for the derivation of Eq. 28 and Eq. 29.
Model Pool If we train S based on only one Φ, it may overfit to that single é, resulting in an inability to properly generate the variational posterior for other Ø's. This overfitting issue is common not only in Bayesian pseudo-coresets but also in the field of dataset distillation (Zhou et al., 2022). While several prior studies (Wang et al., 2018; 2022) tackle this overfitting problem, we address it by employing a model pool during training, following the approach of Zhou et al. (2022); Loo et al. (2023). This model pool method involves generating P different 8, through random initialization during the training of S and storing them in a set M = [0:]]=- At each step, one is sampled from M, and & is constructed using this 0. Then, S is trained for one step using SGD with this 6. Afterward, A is updated by training it for one step using S and the Gaussian likelihood, and the original in M is replaced with this updated version. Once each 0, has been trained for pre-defined number of T steps, it is replaced with a new 0 generated through random initialization. Through this process, S is trained with a new & every step, allowing it to generalize better across different 6's 's and become more robust to various initialization. See Algorithm I for a summary of the whole VBPC training procedure.
Variational Inference and Memory Efficient Bayesian Model Averaging After training S, we use for variational inference. During variational inference, to improve the quality of the model's feature map Ø, we first train the randomly initialized 0 using data sampled from S for small number of steps T" with a Gaussian likelihood. Then, using the trained feature map Φ we compute the variational posterior by finding the optimal mean m and variance V for each 0f as determined in the inner optimization. However, the variance V we computed corresponds to a full covariance matrix, leading memory cost of h?. To address this, rather than calculating V explicitly, need a memory-efficieen approach for conducting BMA on test points. This can be done easily by
E
As
pBs
D TT
2 pss
DT
(I 1+++7 T TΦT pBs
Φ
te
(31)
where Pie Ru denotes the feature matrix of The number of test points. Then by storing & gaxh and (In + pps 2 -- E Ráxå instead of V*, we can reduce the memory requirements âh + n², which is much smaller than h², Refer to Algorithm for an overview of variational inference and BMA. This procedure does not require multiple forwards for BMA.
4 RELATED WORKS
Bayesian Pseudo-Coreset As discussed in Section 1 and Section 2, the large scale of modern real-world datasets leads to significant computational costs when performing SGMCMC or varia- tional inference to approximate posterior distributions. To address this issue, previous works, such as Bayesian Coreset (BC; Campbell & Broderick, 2018; 2019; Campbell & Beronov, 2019), have proposed selecting a small subset from the full training dataset so that the posterior distribution built from this subset closely approximates the posterior from the full dataset. However, Manousakas et al. (2020) highlighted that simply selecting a subset of the training data insufficient to aceu- rately approximate nigh-dimensiooal posterior distributions, and introduced BPC for simple logistic regression tasks. Later, Kim et al. (2022) extended BPC to BNNs, using reverse KL divergence, forward KL divergence, and Wasserstein distance as measures for D in Eq. 2 to assess the difference between the full posterior and the BPC posterior. Subsequent works have used contrastive diver- gence (Tiwary et al., 2024) calculated divergence in function space (Kim et al., 2023). However, as discussed in Section 1, computational and memory overhead remains an issue when training BPC and during inference using BMA. For the additional related works, refer Appendix C.
7

---

Published as a conference paper at ICLR 2025
Figure 1:
Learned VBPC images for the Fashion-MNIST (ipc=10; left), CIFARIO (ipc=10; middle) and CIFARI00 (ipc=1; right) cases. These images construct trained mean for the distribution S*.
Table 1: Comparison of the VBPC with BPC baselines for the benchmark datasets. We report ACC and NLL for the VBPC and BPC baselines.
BPC-rKL
BPC-IKL
FBPC
BPC-CD
VBPC (Ours) ACC(1)
Dataset
ipc
ACC(1) NLL()) ACC(1) NLL()) ACC(1) NLL() ACC(1)
NLL()
NLL(U)
74.8 1.90
83.0a2.
1.87 92.5
1.68
93.4
1.53
9..7 0.11
MNIST
10
95.3
1.53
92. 9
1.51
97.1
1.3
97.7=a.2
1.57
= 0.03
94.2.... 1.36
93.6#188
1.36
98.6.0.1
1.39
98.9+0.2
1.3640.01
99.4
0.02
70.5 2.47=0.02
72.5
2.30
74.7=1.4
1.81
77.3
1.90
82.9 0.47
FMNIST
78.8
1.64=0.01
83.3
1.54+0.03
85.2=0.1
1.61
88.4+0.2
1.56+0.01
89.4
0.30=0.01
50
77.0z0.
1.48=0.02
74.8
1.47
76.7
1,4620.02
89.5
1.30+0.02
91.0+0.2
0.25+0001
21.6a0.8
2.57=0.01
29.3=1.1
2,10s0.08
35.5
3.79
46.9
1.87
55, 1.34
CIFAR10
10
3. 2.13
49.9=1.4
1.73=0.01
02.3+0.3
1.31+0.02
56.4
1.72+00.3
69.8+0.7
0.89+0.02
37.5
1.93g0.0s 42.3
1.54
71.220.2
1.03
71.9
1.57.30.0
0001: 76.720.5
3.61u.1 2.69
14.7.o.2 4, 4,20
21.0a
3.76
24.0.:.11
1:.02
38.4 2.47:10.04
CIEARIO0
23.620.7
3.99
28.1
3.5329.05
39.7
2.67
28.4
3.1419.02
49.4
2.07
30.8 50
3,57
37, 1200
3,28
44.5
20.4
2.63..m1
39.6191
3.02
52.4
2.02.10:02
Tiny- Niae-emageNet
3.240.1
5.9120.77
4.0
5.63
4.69
8.420.1
4.72+a..1
23.1 3.65
10
9.8
5,2620.05
11,420.
5.08
19.4
4.14
17.8
3.64
25.8 2...
2.45
EXPERIMENT
In this section, we present empirical results that demonstrate the effectiveness of posterior approx- imation using VBPC across various datasets and scenarios. We compare VBPC with four BPC algorithms that use SGMCMC to perform Bayesian Model Averaging (BMA) with posterior sam- ples: BPC-rKL. (Kim et al., 2022), BPC-fKL. (Kim et al., 2022), FBPC (Kim et al., 2023), and BPC-CD (Tiwary et al., 2024). BPC-rKL and BPC-IKL employ reverse KL divergence and forward KL divergence, respectively, for the divergence term in Eq. 2. BPC-CD uses more complex diver- gence called contrastive divergence, while FBPC also applies forward KL divergence but matches the posterior distribution in function space rather than weight space. Following all other prior works, we adopted a three-layer convolutional network with Batch Normalization (BN; loffe, 2015) as the base model architecture. For the target dataset, we used the MNIST (LeCun et al., 1998), Fashion- MNIST (Xiao et al., 2017), CIFAR 10/100 (Krizhevsky, 2009), and Tiny-ImageNNt (I.e & Yang, 2015). Additionally, we used image-per-class (ipc) as the unit to count the number of pseudo- coresets. For a k-way classification task, m ipc signifies that a total of mk pseudo-coresets are trained. Along with evaluating classification accuracy (ACC) for each methods, we assess the per- formance of the resulting predictive distributions using negative ligllikeldoood (NLL).
In all tables, the best performance is indicated with boldfaced underline, while the second- best value is represented with underline in each row. See Appendix E for the additional experimental details.
5.1 BAYESIAN MODEL AVERAGING COMPARISON
We begin by evaluating the effectiveness of VBPC on five benchmark datasets by comparing the BMA performance across different methods. Table clearly demonstrates that VBPC surpasses other BPC baselines across all benchmark datasets and ipes in terms of ACC and NLL. Notably, VBPC achieves significantly better NLL., with large margins, while requiring only single forward
8

---

Published as a conference paper at ICLR 2025
Table Comparison with dataset distillation baselines in terms of ACC. Here, indicates the performance drop compare to original method.
Dataset
ipe|
FRePo
FRePo VI
RCIG
RCIG VI
VBPC
AVBPC
46.8 20.7
(8,2((18.6))21)))
53.9 27.8(24.1
55.1 20.3
39.7(15.44)
CIFAR10
10
65.5
55.7(9.8))
69.1 55.6(13.8.)+1.a
=9.8
67.8(2.0L)
1
28.7 19.9(8.8))
39.3
2.1(37.2.)
38.4
31.3(7.1L)
CIFAR100
10
42.5
34.8(7.7()+0.44
44.1+0.4
2.5(41.6))=u..4
29.4
44.0(5.44)=0088
Table 3: Comparison of the VBPC with BPC baselines on the OOD setting with CIFARIO-C dataset.
The -A in the first column indicates that A type corruption is applied to the CIFARIO test dataset.
BPC-rKL
BPC-IKL
FBPC
BPC-CD
VBPC (Ours) NLL(()
Comuption
ACC(1) NLL()) ACC(1) NLL()) ACC(1) NLL(() ACC(t)
NLL(U)
ACC(1)
CIFARIO
37.9
2.13 49.9
1.73+0.01
62.3+0.
1.3
56.4
1.72
69.8=0.7
0.89
Gaussian Blur
31.0+2,7 2.13 39.7=2.7
1.94
35,8
1.85
41.4+077
1.73+0.3
59,3+0.9
+.20
JPGG Compression
30.410.9 2.13 37.3129
1.54
40.1
1.7340.0z
37.3.0.2
1.71
61.9
1.12:0.22
+Snow
26.9+1.7 2.20 35.7+2.7
2.00
38.620.4
1.7840.16
37.8
1.91
59.0.
4.20
Zoom Blur
31.7a1.2 2.09 35.1
2.04
28.920.2
2.1920.11
38.3208
1.93
58.1
1.23
+Pixelate
29.0az.3 2.19 39.1
1.93
38.02o.3
1.77
39.0s15
1.92.0.07
58.8
1.26100.4
+Defncus
27.61113
2.20 36.743.7
1.99
31.7
2.0720.19
37.2+1.0
1.87:0000
63.0=007
1.08
+Motion Blur
17.422 2.73 35.2*a.3 2.01 20.05
27.9 2.29 37.1
1.92
55,9
1.32#0.03
pass for BMA. These results empirically validate that the variational distribution trained by VBPC effectively captures epistemic uncertainty with small amount of synthetic data, while keeping performance. Refer to Fig. 1 for examples of VBPC-traines images from the FARIO, and CIFAR100 datasets. For more trained VBPC images for other settings, see Appendix G.
CI-
Comparison with dataset distillation baselines In addition to the BPC baselines, we compared VBPC with two notable dataset distillation basclines, FRePo (Zhou et al., 2022) and RCIG (Loo et al., 2023), which are recognized for their strong accuracy performance. Since FRePo and RCIG do not employ cross-entropy loss for training, we only report ACC, as comparing NLL would be unfair. As shown in Table 2, although VBPC is designed to learn pseudo-coresed to approximate the variational distribution from the training data, it outperforms these dataset distillation baselines, focused mainly on ACC, in nearly all tasks except for CIFARI00 with ipc. The results for each methods (i.c., FRePo, RCIG, and VBPC) in Table 2 were evaluated based on each baseline's 's evalu- might question whether the significant performance of VBPC is due to
ation methods. However, the trained pseudo-coreset itself or the VI method. To verify that VBPC's performance isn" solely due to the VI method, we applied our VI method to the baselines' pseudo-coresets (i.e., FRePo VI and RCIG VI) and used FRePo's 's method to evaluate VBPC's pseudo-coresets (i.c., AVBPC). Al- though all methods saw some performance decline, VBPC exhibited smaller drop, indicating that its performance is not solely due to the VI method, but to its ability to effectively learn the varia- tional distribution. Full comparisons across all benchmark datasets, available in Appendix F.1, show that VBPC maintains consistent trend over dataset distillation baselines across all the datasets.
5.2 RESULTS ON OUT OF DISTRIBUTION SCENARIOS
To further demonstrate that the predictive distribution derived from the VBPC dataset enhances robustness to distributional shifts and out-of-distribution (OOD) data, we assess the performance of VBPC and BPC baselines on a corrupted version of the CIFARIO dataset, known as CIFAR10- C (Hendrycks & Dietterich, 2019). In this case, we use the CIFAR10 10ipc BPC data trained in Section 5.1 for all methods and evaluate their performance on the corrupted dataset across 7 different types of corruption. We assess performance using all 5 levels of severity provided in the dataset. Table 3 clearly illustrates that VBPC shows strong robustness against various types of corruption and consistently outperforms other baselines across all corruption types in terms of both ACC and NLL. These findings highlight that the predictive distribution obtained from the VBPC dataset improves robustness to distributional shifts and OOD scenarios.

---

Published as a conference paper at ICLR 2025
Table Comparison of the VBPC with BPC baselines on the architecture generalization. The A B in the first column indicates that B type normalization layer is used for the A model.
BPC-rKL
BPC-KKL
FBPC
BPC-CD
VBPC (Ours)
Model
ACC(T) NLL()) ACC(1) NLL()) ACC(t) NLL()) ACC(1) NLL()
ACC(1)
NLL())
Conv-BN
37,9115 2.1310.02 49.9
1.73.0.01
62.310.3
1,3
56,410.7
1.72.0.03
69.8
0.8910.02
Conv-NN
23.1 2.22 22.9
2.12.0.44
28.61488
2.17
30.1
2.05
58,4
1.46.0.05
Cony-GN
28.5::4.5 2.85
29.1=4.4
2.
31.545.2
1,93
23.8+4.2
3.07
66.8
0,95
Conv-IN
26.7+4.a 2.81
27.7=4.7
2.
31.745.3
1,96
26.9+4.4
3.29
58.1
1.22
AlexNet-NN
24.2+a.8 2.23+0m1
21.4=4.3
2.82=0.24
32.
2.91
30.8
2.24+0.11
48.0+0.4
1.04
ResNet18-BN
9.6#2.6
3.27+0.15 10.5=4.5 3.16=0.14
46.7+1.2
1.81
41.7=1.1
2.05
54.9
+.36
VGG11-GiN
10.0=2.9 2.94 10.1 2.85=0.11 37.2 1.40+0.05 44.5+1.2 1.78
5..4
1.44
Table 5: Ablation results on memory allocation and time requirements CIFAR1O 10ipe.
Naive Training
Training (Ours)
Naive BMA BMA (Ours)
Memory (MB) sec/100 steps Memory (MB) sec/100 steps Memory (MB) Memory (MB)
542.9
54.0
272.9
9.9
542.9
268.9
5.3 ARCHITECTURE GENERALIZATION
To demonstrate that VBPC can be applicd when performing BMA on unseen architectures, we conduct BMA using different model structures with various normalization layers. Specifically, we include the identity layer (NN), Group Normalization (GN; Wu & He, 2018), and Instance Normalization (IN; Ulyanov, 2016) additional normalization methods. We also incorporate AlexNet (Krizhevsky et al., 2012), ResNet18 (He et al., 2016), and VGGI1 (Simonyan & Zis- serman, 2014) as new model architectures. Similar to Section 5.2, use the CIFARIO 10ipe BPC data. As shown in Table 4, VBPC successfully performs VI across various architectures and effec- tively constructs predictive distributions through BMA. Notably, while other baselines are sensitive to changes in normalization layers, VBPC demonstrates robust learning over diverse feature maps through the model pool, resulting strong ACC and NILL performance.
5.4 MEMORY ALLOCATION AND TIME REQUIREMENTS
In this section, perform ablation study compare memory usage and time requirements be- tween the naive computation and the efficient computation for the variance V* Σ and the loss terms during both training and inference. As discussed in Section 3.6, naive loss computa- tion requires ( space complexity and (ha)) computational complexity. However, our com- putationally efficient loss computation method only requires ( space complexity and O( computational complexity. Therefore, in the BPC setting where ñ < h tyyically holds, we can sig- nificantly reduce the space and computational complexity required for training. This difference can be observed during the actual training process. As shown in Table 5, our computationally efficient training reduces the memory requirements for loss computation by nearly half and decreases the training time to under 20%. Also, we can see the similar results during the BMA procedure. Refer to Appendix F to see the various additional ablation studies including ablation on hyperparameters, pseudo-coreset initialization, and augmentations.
CONCLUSION
In this paper, present a novel BPC method for VI, referred to as VBPC. By utilizing the Gaussian likelihood, we enable the computation of a closed-form solution for coreset VI, thereby removing the need to unroll the computation graph use stop gradients. Leveraging this closed-form solution, propose a method to approximate dataset VI without weight sampling during the training of VBPC. Additionally, we introduce a computationally efficient training and BMA inference method that sig- nificantly reduces both computational and space complexity. Finally, we empirically show that the variational distribution obtained from VBPC substantially outperforms the predictive distributions derived from other BPC basclines in BMA performance across various scenarios.
10

---

Published as a conference paper at ICLR 2025
Reproducibility Statement. We present comprehensive derivations of all equations in the paper in Appendix A. The overall algorithms can be found in Appendix B. Details regarding the datasets, model architecture, data preprocessing, and hyperparameters provided in Appendix E.
Ethics Statement. We propose method that improves the computational and memory efficiency of the variational inference method for posterior approximation Bayesian Neural Networks. Thus although our approach does not have direct positive negative impact on ethical or societal aspects, it can enhance privacy preservation. Specifically, our method facilitates Bayesian inference using private training data in neural network models by generating synthetic datasets, allowing for the computation of the posterior distribution while maintaining privacy.
Acknowledgements This work partly supported by Institute of Information & communi- cations Technology Planning & Evaluation(IITPP) grant funded by the Korea govemment(MSIIT)
NN.RSS-019-IT1900D Artificial Intelligence Graduate School Program(KAISTT), the National Research Foundation of Korea(NRF) grant funded by the Korea (ivernment(MST)) (RS-2021- NR056917), Institute of Information & communications Technology Planning & Exaluation(IIIPP) grant funded by the Korea governmen((MSIT) (No.RS-2024-00) R8-2024400559929, Global AI Frontier Lab), and Institute of Information & communications Technology Planning & Evaluation(IITP) grant funded by the Korea government(MSIT) world Problems).
Meta-learnnn Applicable to Real-
REFERENCES
Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jía, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg. Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wat- tenberg. Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learn- ing heterogeneous systems, 2015. URL https:////// /ttp::///wwwwww... tensorflow org/. Software available from tensorflow.org. 22
Abdullah A Abdullah, Masoud M Hassan, and Yaseen T Mustafa. A review on bayesian deep learning in healthcare: Applications and challenges. IEEE Access, 10:36538-36562, 2022.
AF Agarap. Deep learning using rectified linear units (relu). arXiv preprint arXiv:1803.08771 2018. 23
Sungjin Ahn, Anoop Korattikara, and Max Welling, Bayesian posterior sampling via stochastic gradient fisher scoring, arXiv preprint arXiv: 120:: 6380, 2012. 20
Amari. Natural gradient works efficiently in learning. Neural Computation, 10:251-276, 1998.
Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacei, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http: / /tt:://iitin github. com/ deepmind. 22
Yoshua Bengio. Gradient-based optimization of hyperparameters. Neural computation, 12(8):1889- 1900, 2000. 4
Christopher M Bishop. Pattern recognition and machine learning, volume 4. Springer, 2006. 3, 21
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statis- ticians. Journal of the American statistical Association, 112(518):859-8777 877, 2017. 1, 2, 3, 20, 21
11

---

Published as a conference paper at ICLR 2025
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In Proceedings of The 32nd International Conference on Machine Learning (ICML 2015), 2015. 3,21
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. com/ google/jax. 22
Trevor Campbell and Boyan Beronov. Sparse variational inference: Bayesian coresets from scratch. In Advances in Neural Information Processing Systems 32 (NeurIPS 2019), 2019. 7, 20
Trevor Campbell and Tamara Broderick. Bayesian coreset construction via greedy iterative geodesic ascent. In Proceedings The 35th International Conference on Machine Learning (ICML. 2018), 2018. 7, 20
Trevor Campbell and Tamara Broderick. Automated scalable bayesian inference via hilbert coresets. Journal of Machine Learning Research, 20(15):1-38 2019. 7, 20
George Cazenavette, Tongzhou Wang, Antonio Torralba, Alexei Efros, and Lun Zhu. Dataset distillation by matching training trajectories. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 4750-4759, 2022. 21
Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In Proceedings of The 31st International Conference on Machine Learning ([CML)014), 2014. 1, 2, 20
YL Cun, L Bottou, G Orr, and K Muller. Efficient backprop, neural networks: tricks of the trade. Lecture notes in computer sciences, 1524:5-50, 1998. 23
Michael Dusenberry, Ghassen Jerfel, Yeming Wen, Yian Ma, Jasper Snock, Katherine Heller, Balaji Lakshminarayanan, and Dustin Tran. Efficient and scalable bayesian neural nets with rank-1 factors. In International conference on machine learning, pp. 2782-2792. PMLR, 2020. 21
Felix Fiedler and Sergio Lucia. Improved uncertainty quantification for neural networks with bayesian last layer. IEEE Access, 2023. 1, 2, 20, 21
J. Harrison, J. Willes, and J. Snoek. Variational Bayesian last layers. In International Conference on Learning Representations (ICLR), 2024a. 21
James Harrison, John Willes, and Jasper Snoek. Variational bayesian last layers. arXiv preprint arXiv:2404.11599 11599, 2024b. 1, 3, 22
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016. 10, 23
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common ruptions and perturbations. In International Conference on Learning Representations (ICLR), 2019. 9, 23
Jeremy Howard. A smaller subset of 10 easily classified classes from imagenet, and a little more french, 2020. URL https://githhbs com fast wtt.://aaaaaaaaaaattttaaaee 25
Sergey loffe. Batch normalization: Accelerating deep network training by reducing internal covari- ate shift. arXiv preprint arXiv: 2015. 8, 23
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gener-
alization in neural networks. In Advances in Neural Information Processing Systems 31 (NeurIPS 2018), 2018. 21
Agnan Kessy, Alex Lewin, and Korbinian Strimmer. Optimal whitening and decorrelation. The American Statistician, 72(4):309-3) 2018. 24
12

---

Published as a conference paper at ICLR 2025
M. E. Khan and H. Rue. The Bayesian learning rule. Journal of Machine Learning Research, 24:
1-46, 2023. 3
M. E. Khan, D. Nielsen, V. Tangkaratt, W. Lin, Y. Gal, and A. Srivastava. Fast and scalable bayesian deep learning by weiggh--peturrban in Adam. In Proceedings of The 35th International Confer- ence on Machine Learning (ICML 2018), 2018. 3
Balhae Kim, Jungwon Choi, Seanie Lee, Yoonho Lee, Jung-Wop Ha, and Juho Lee. On divergence measures for bayesian pseudocoresets. In Advances in Neural Information Processing Systems 35
(NeurIPS 2022), 2022. 1, 2, 7, 8, 20, 23, 24, 27, 29
Balhae Kim, Hyungi Lee, and Juho Lee. Function space bayesian pseudocoreset for bayesian neural networks. In Advances in Neural Information Processing Systems 36 (NeurIPS 2023), 2023. 1, 2, 7, 8, 20, 23, 24, 27, 29
Diederik Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6990, 2014. 24
P. W. Koh and P. Liang. Understanding black-box predictions via influence functions. In Proceedings of The 34th International Conference on Machine Learning (ICML 2017), 2017. 5
Steven George Krantz and Harold R Parks. The implicit function theorem: history, theory, and applications. Springer Science & Business Media, 2002. 4
Alex Krizhevsky. Learning multiple layers of features from tiny images. Master's thesis, University of Tront, 2009. 8, 22
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convo- lutional neural networks, In Advances in Neural Information Processing Systems 25 (NIPS 20/2), 2012. 10, 23
Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N, 7(7):3, 2015. 22
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- learning applied to document recognition. Proceedings of the IEEE, (6(11):2278-235 22224, 1998. 8, 22
Hyungi Lee, Eunggu Yun, Hongseok Yang, and Juho Lee. Scale mixtures of neural network gaussian processes. In International Conference on Learning Representations (ICLR), 2022.
Hyungi Lee, Giung Nam, Edwin Fong, and Juho Lee. Enhancing transfer learning with flexible non-
parametric posterior sampling. In International Conference on Learning Representations (ICLR), 2024.7 2
Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks gaussian processes. arXiv preprint arXiv:171..00160 2017. 5, 21
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve linear mod- els under gradient descent. In Advances in Neural Information Processing Systems 32 (NeurIPS 20/9), 2019. 5
Noel Loo, Ramin Hasani, Mathias Lechner, and Daniela Rus. Dataset distillation with convexified implicit gradients. In Proceedings of The 39th International Conference on Machine Learning (ICML 2023), 2023. 1, 2,
L Lopez, Tim GJ Rudner, and Farah E Shamout. Informative priors improve the reliability of multimodal clinical data classification. arXiv preprint arXiv:2312.00794, 2023.
Lorraine, P. Vicol, and D. Duvenaud. Optimizing millions hyperparameters by implicit dif- ferentiation. In Proceedings of The 23rd International Conference Artificial Intelligence and
Statistics (AISTATS 2020), 2020. 4
Zhiyun Lu, Eugene Ie, and Fei Sha. Mean-field approximation to muussian-softmal integral with application to uncertainty estimation. arXiv preprint 2xXXx:2006.07584 2020. 6, 18
13

---

Published as a conference paper at ICLR 2025
Yi-An Ma, Tianqi Chen, and Emily Fox. A complete recipe for stochastic gradient meme. In Advances in Neural Information Processing Systems 28 (NIPS 2015), 2015. 1, 20
Dionysis Manousakas, Zuheng Xu, Cecilia Mascolo, and Trevor Campbell. Bayesian pseudocore- sets. In Advances in Neural Information Processing Systems 33 (NeurIPS 2020), 2020. 1, 7, 20
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Baolin Wu, Andrew Ng, et al. Reading digits in natural images with unsupervised feature learning. In NIPS workshop deep learning and unsupervised feature learning. Granada, 2011. 26
Timothy Nguyen, Zhourong Chen, and Jachoon Lee. Dataset meta-learning from kemell ridge- regression. arXiv preprint arXiv:2011.000520 2020. 21
Timothy Nguyen, Roman Novak, Lechao Xiao, and Jaehoon Lee. Dataset distillation with infinitely wide convolutional networks. In Advances Neural Information Processing Systems 34 (NeurIPS 2021), 2021. 1
K. Osawa, S. Swaroop, A. Jain, R. Eschenhagen, R. Turner, R. Yokota, and M. E. Khan. Practical deep learning with Bayesian principles. In Advances in Neural Information Processing Systems 32 (NeurIPS 2019), 2019. 3
Theodore Papamarkou, Maria Skoularidou, Konstantina Palla, Laurence Aitchison, Julyan Arbel, David Dunson, Maurizio Filippone, Vincent Fortuin, Philipp Hennig, José Miguel Hernández- Lobato, et al. Position: Bayesian deep learning is needed in the age of large-scale ai. In Interna- tional Conference on Learning Representations (ICLR), 2024. 1,
Constantine Pozrikidis. An introduction to grids, graphs, and networks. Oxford University Press, USA, 2014. 6, 19
Tim GJ Rudner, Zonghao Chen, and Yarin Gal. Rethinking fiunction-space variational inference in bayesian neural networks. Third Symposium on Advances in Approximate Bayesian Inference, 2021. 20
Tim GJ Rudner, Freddie Bickford Smith, Qixuan Feng. Yee Whye Tch, and Yarin Gal. Contin- ual learning via sequential aanction-space variational inference. In International Conference Machine Learning, pp. 18871-18887, PMLR, 2022. 20
Evgenia Rusak, Lukas Schott, Roland S Zimmermann, Julian Bitterwolf, Oliver Bringmann, Matthias Bethge, and Wieland Brendel. simple way to make neural networks robust against di- verse image corruptions. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part III 16, pp. 53-69. Springer, 2020. 30
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International journal of computer vision, 115:211-2523 2015. 25
S. Sharm, M. nd Farquhr, E. Nalisnick, and T. Rainforth. Do Bayesian neural networks need to be fully stochastic? In Proceedings of The 26th International Conference Artificial Intelligence and Statistics (AISTATS 2023), 2023.
Y. Shen, N. Daheim, B. Cong, P. Nickl, G. M. Marconi, C. Bazan, R. Yokota, I. Gurevych, D. Cre- mers, M. E. Khan, and T. Möller. Variational learning is effective for large deep networks. In Proceedings of The 40th International Conference Machine Learning (ICML 2024), 2024. 3, 4, 21
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In International Conference on Learning Representations (ICLR), 2014. 10, 23
Piyush Tiwary, Kumar Shubham, Vivek V Kashyap, and AP Prathosh. Bayesian pseudo-coresets via contrastive divergence. In The 40th Conference on Uncertainty in Artificial Intelligence, 2024. 1, 2, 7, 8, 20, 23, 24, 27, 29
14

---

Published as a conference paper at ICLR 2025
Antonio Torralba, Rob Fergus, and William T. Freeman. 80 million tiny images: A large data set for nonparametric object and scene recognition. JEEE Transactions on Pattern Analysis and Machine Intelligence, 30( 100::-9988190000 2008. doi: 10.1109/TPAM.200..10 128.. 22
Ulyanov. Instance normalization: The missing ingredient for fast stylization. arXiv preprint
arXiv:1607.08022, 2016. 10, 23
Thomas Vandal, Evan Kodra, Jennifer Dy, Sangram Ganguly, Ramakrishna Nemani, and Auroop R Ganguly. Quantifying uncertainty in discrete-continuuus and skewed data with bayesian deep learning. In Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 2377-2386, 2018. 1
Paul Vicol, Luke Metz, and Jascha Sohl-Dickstein. Unbiased gradient estimation in unrolled com- putation graphs with persistent evolution strategies. In Proceedings of The 38th International Conference on Machine Learning (ICML 2021), 2021. 5
Chong Wang and David M Blci. Variational inference in nonconjugate models. Journal of Machine Learning Research, 2013. 21
Kai Wang, Bo Zhao, Xiangyu Peng, Zheng Zhu, Shuo Yang, Shuo Wang, Guan Huang, Hakan Bilen, Xinchao Wang, and Yang You. Cafe: Learning to condense dataset by aligning features. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 12196 12205, 2022. 7
Tongzhou Wang, Jun-Yan Zhu, Antonio Torralha, and Alexei A Efros. Dataset distillation. arXiv preprint arXiv: 10999999 2018 2.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of The 28th International Conference on Machine Learning (ICMI. 201 2011. 1, 2, 20
Florian Wenzel, Kevin Roth, Bastiaan S Veeling, Jakub Swiętkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good the bayes posterior in deep neural networks really? arXiv preprint 2XXi::002200200) 2020. 3
Paul Werbos. Backpropagation through time: what does and how do it. Proceedings of the IEEE, 18510):1500-1560 1990. 5
Max A Woodbury. Inverting modified matrices. Department of Statistics, Princeton University, 1950. 5, 16, 17
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pp. 3-19, 2018. 10, 23
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmark-
ing machine learning algorithms. arXiv preprint arXiv:1708.000 2017. 8, 22
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint ArXiv:1904.0000 1904.0 2019. 28
Bo Zhao and Hakan Bilen. Dataset condensation with differentiable siamese augmentation. In Proceedings of The 38th International Conference on Machine Learning (ICML 2021), 2021. 21, 26
Bo Zhao and Hakan Bilen. Dataset condensation with distribution matching. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision pp.
2023. 21
Yongchao Zhou, Ehsan Nezhadarya, and Jimmy Ba. Dataset distillation using neural feature regres- sion. In Advances in Neural Information Processing Systems 35 (NeurIPS 2022), 2022. 7, 9, 21, 22, 23, 24, 26, 27, 28, 29
15

---

Published as a conference paper at ICLR 2025
A FULL DERIVATIONS
A.1 FULL DERIVATION FOR THE INNER OPTIMIZATION
In this section, we present the full derivation calculation for the inner optimization in Section 3.4.
[- Σ Di=1 Log.N( x It) which can be computed
Let us first examine the term By. as follows:
Eag
log NiiNN ((6;),
Σ Ba Nog:N
(32)
2110: c(t.) = [W]]]
Byx
+ constant
(33)
k
Box
((EE) Tw;)²
[WiLs - ((B.)
+ constant
1=1
(34)
k
Egg
(Bis- - ((;)) (w;))
(35)
where Di.j indicates jth element of d: for all E [m]. With this approximation, now can compute as follows:
Egg
Σi= Wog WT o(E:), 1x))
k
Eyp
Σ log W Ix)
Eyp
- ( (T; w;)²]
(36)
i=1 j=1
?
Σ
Egya [0] أ0 1 22j Dwg w, x pwy)
(37)
8
(-20] æmnj Eng [w] $ (Twwyl (38)
Σ
(-20] Φm, + Tr (0 DE (uyw]])) (39)
j=1
Σ
(-20] æmy Tr (6 + [V, + m,m,])] 1))
/=1
(40)
Σ
2Íj T (1) + Tr (T mnp
(41)
where Ú3 = '9....... Jn. = (o(2),.... ((2 14 (1) = mj. and u = V, + mym]
for all j [k). Here, & denotes equality up to a constant. Eq. 39 derived from the fact that Egy Ju Φ (wy] scalar value and the property of the Tr function.
A.2 NUMERICALLY STABLE MEAN AND VARIANCE
In this section, present the full derivation calculation for the numerically stable mean and vari- ance Section 3.4. Due to the dimension of / is x h and usually ń < h, naïve computation of m; and V lead numerically unstable results. To address this issue, transformed the formulas for m and V- into equivalent but more numerically stable forms. Specifically, when calculating Vj, we applied the Woodbury formula (Woodbury, 1950). First, we utilize the kernel trick to make
16

---

Published as a conference paper at ICLR 2025
mean my more numerically stable. The derivation is as follows:
mn - -/ /
(2)a-111
1011
(42)
- In 7 TT 10-1)11111 Y TTH 90
(43)
23g
As
1 1 l +
7
-1 7
ΦT Dj
(44)
2 2
23s
As
(pIn
Y 1)
7
(pla +
7 + T (pln +
$ )-
(45)
As
88
3s
3s
3 (phn +
2
TT) 1 (00)
- T ΦΦ
(pla +
2 20 )
(46)
3s
As
3s
3s
2 LT (pla +
7 ΦΦ T) "'1
(47)
3s
As
T ( La + 2Φ ) "
(48)
Next, utilize the Woodbury formula (Woodbury, 1950) make variance Vj more numerically stable. The derivation is as follows:
is
pBs In + T
Vj
(49)
?
pis In
p5s In
pss
pss
1 11
+
Φb
In
+
In
?
?
7
7
?
(50)
Bs
7
In
7
Φ
Ia
+
7
TΦ T
-1
(51)
?
pBs
pîs
pšs
-1
7
T
la
+
7
(52)
It is important to note that for all j E [k). the VJ values are identical. This implies that while calculating the full covariance for all E [R]] can be computationally intensive (i.e. O(kh?) )), only need to compute and store the variance once (i.e. 0(12))
A.3 FULL DERIVATION FOR OUTER OPTIMIZATION PROBLEM
In this section, we present the full derivation for the outer optimization problem. Here, we first change (D() as follows:
(D(3)) = Egginggg
log pD (yilu, 04>] + Br DKI. 1PPal
(53)
n
exp(o(x;)) (wy))
Ui. log
+ 3p DKI(ox; (W)llxx. (54)
k Σl= exp(p(x))
1=1 1=1
A
exp(<b(xt) wy) Σ exp(((xx) Tw)
Vi.y Eugg
log
+ DDDDal (W)llxx (W)). (55)
17

---

Published as a conference paper at ICLR 2025
log exp
Next, order to compute approximate expectation Em 21) the form follows:
we first change
Σ exp
(m- 520)
log exp(z,) Zi=1 exp= 24
log
exp 2j -i=1 exp
N(zimm )z
(56)
f log 1 Σijs exp( (xpp(((2)- 21]) ,
1
(57)
log(( -flog(1
-exp(- =1))) -N \V(z/m" Ed)d
(58)
if1
/ log(2 K +
Σ
+ expl (2) - (2(zm", Ed)dz
ifj
(59)
-/ log(2 K +
Σ
N , dzz
(60)
o(2y- -2) 24)
i
where (-) is the sigmoid function. Then utilize mean-field approximation (Lu et al., 2020) to the 2,8 to approximately compute the Eq. 24:
exp (2z) log Ei=1 exp(2i)
Em-N(m+ 8t))
~ log
2-k
Σ
E
21)1
(f)
(61)
z log
2 + Σ
(62)
mh m
i
I+
-log
(63)
1 1
Eijj exp(
m. m
V1+055
m
log softmax
(64)
1+
log softmax
(65)
1 a>
where a and Σx = T
A.4 FULL DERIVATION FOR TRAINING AND INFERENCE
In this section, we present the full derivation for the training and inference. Since both euz and pso are Gaussian distributions, the KL divergence can be expressed as follows:
|dt(pp '1,) |dct(V']
DEL[gx; Ilxxol
[k: log
kh KTr(pT
(mg) (pT="))
(66)
1=1
|det(p-11)) 'I,)
kk log
kh ATr(pVV) |lmm"11
(67)
[det(V)]
k log det((V)) +KTT((V*) + +| m|]
(68)
2
Here, we have to reduce the memory requirements for the det( and the Tr(V*) as they require (112) memory compute directly from V*, For the det V*, we used Weinstein-Aronszain iden-
18

---

Published as a conference paper at ICLR 2025
tity (Pozrikidis, 2014) which results as follows:
Y ib -)-1
det V det(pIn
(69)
As
(70)
det(pIn + 7 Φ ()
As
(71)
del(T
① TTΦ)
pdg
(72)
p" det(ln +
7 +
pos
Thus have:
La + ? $
log det V = - log log det
log
(73)
pšs
log det
I
7
$
(74)
pβs
Also we can compute trace follows:
Tr(V*)
Tr
Bs
7
In
?
In
?
pΦ
-1
(75)
pšs
8s
7h
7
2
Tt
T
+
ΦΦ T
(76)
7
p5s
pgs
pßs
8s
yh
?
2
Tr
la
Y
T T
(77)
7
p3s
p8s
pBs
((1,
7
Tr
(78)
3sp²
pBs
These computations allow us to reduce memory requirements during training from (hh)) to O which represents significant reduction when dealing with nigh-ddmeensioal feature space h.
Memory Efficient Bayesian Model Averaging For the variance V we computed corresponds to full covariance matrix, leading to memory cost of 122 To address this, rather than calculating V explicitly, we need a memory-efficient approach for conducting BMA on test points. This can be done easily by calculating Σ* as follows:
Σx t
V T-
(79)
i
7
2
T
(-2 p3s
La [a+ TT
7 2ΦT pBs
(80)
p
where die F denotes the feature matrix of The number of test points. Then by storing & pfg 2 )-1 E ıaxa instead of V*, we can reduce the memory requirements to
paxh and (1, + 7
âh + n2, which is much smaller than h?a
B ALGORITHM FOR TRAINING AND INFERENCE
In this section, we present algorithms for training and inference. In Algorithm 1, the overall training procedures are presented, and note that we utilize the model pool M to prevent overfitting. We also use the Gaussian likelihood to update the weights contained in the model pool. Additionally, in Algorithm 2, we present computationally and memory-effiiien variational inference and BMA methods. Here, we store & and (Jn, pds + -1 instead of directly computing V*.
19

---

Published as a conference paper at ICLR 2025
Algorithm Training Variational Bayesian Pseudo-Coreset (VBPC).
Require: Training dataset D, learning rate 6.
Ensure:
Learned synthetic dataset S
1:
Initialize: Initialize synthetic dataset distribution S with á pairs of (æs, 8t).
2: Initialize: Randomly initialize P different 0,s and construct a model pool M.
3: while not converged do
4: Sample training batch B from training distribution D.
5: Uniformly sample 0g from model pool M and construct feature map Φb
6:
Efficiently compute loss Eq. 27 with Eq. 29, Eq. 30 and Eq. 31.
7: Update 2&s and ÿ:8 using gradient descent: cद <- á 8Va, Co. in t d: oV 0: 2,20
8: Update 0, with S and the Gaussian likelihood.
9: Replace 0, in M with updated 0j.
If 0, M has been updated 7 times, reinitialize 0, and replace 0 in M.
11: end while
Algorithm Variational inference and Bayesian Model Averaging using VBPC.
Require: Learned synthetic dataset S MODE which is V1 or BMA, and test dataset T. Variational posterior Bayesian Model Averaged output prediction. 1: Initialize: Randomly initialize
Ensure:
2: while not converged do
3: Update with Gaussian likelihood and S"
4: end while
5: Compute mj. 4, and (Ia + 23 2_ T)- with (&i, U.)s.
6: if MODE == VI then
7: Compute V with Eq. 21.
8: else if MODE==BMA then
9:
Compute " with Eq. 31.
10:
Compute approximate expected predictive distribution for T similar to Eq. 25.
11: end if
C ADDITIONAL RELATED WORKS
Bayesian Pseudo-Coreset As discussed in Section 1 and Section 2, the large scale of modern real-world datasets leads to significant computational costs when performing SGMCMC (Welling & Teh, 2011; Ahn et al., 2012: Chen et al., 2014; Ma al., 2015) variational inference (Blei et al., 2017; Fiedler & Lucia, 2023) to approximate posterior distributions. To address this issue, previous works, such as Bayesian Coreset (BC; Campbell Broderick, 2018; 2019; Campbell & Beronov, 2019), have proposed selecting small subset from the full training dataset so that the posterior distribution built from this subset closely approximates the posterior from the full dataset. However, Manousakas et al. (2020) highlighted that simply selecting subset of the training data insufficient to accurately approximate high-dimensional posterior distributions, and introduced BPC for simple logistic regression tasks. Later, Kim et al. (2022) extended BPC to BNNs, using reverse KL divergence, forward KL divergence, and Wasserstein distance as measures for D in Eq, assess the difference between the full posterior and the BPC posterior. Subsequent works have used contrastive divergence (Tiwary et al., 2024) or calculated divergence in function space (Kim al.. 2023) using function-space Bayesian Neural Network (FBNN; Rudner al., 2021; 2022). However, as discussed in Section 1, computational and memory overhead remains an issue when training BPC and during inference using BMA.
Dataset Distillation Similar to but distinct from BPC, dataset distillation (Wang et al., 2018) meth- ods aim to train a pseudo-coreset that preserves the essential information contained in the full train- ing dataset. These methods ensure that the model trained on the pseudo-coreset learns information that allows it to perform similarly to model trained the full dataset. This approach enables com- putationally efficient training of new models using the pseudo-coresd and helps prevent catastrophic forgetting in continual learning scenarios, leading more stable learning.
20

---

Published as a conference paper at ICLR 2025
To train these dataset distillation methods, bilevel optimization problem must be solved, requiring the computation of meta getta-gradents through unrolled inner optimization to find the solution to the outer optimization problem. To address this challenge. various learning methods have been pro- posed in the dataset distillation field, which can be broadly categorized into three approaches: 1) using surrogate objectives. 2) closed-form approximations, and 3) employing the implicit function theorem.
Examples of works in the first category include Zhao & Bilen (2021), Zhao & Bilen (2023), and Cazenavette et al. (2022), where Zhao & Bilen (2021) uses gradient matching, Zhao & Bilen (2023) focuses on feature distribution alignment, and Cazenavette al. (2022) employs a trajectory match- ing objective. Papers in the second category, Nguyen et al. (2020) and Zhou et al. (2022), calculate closed-form solutions by using the Neural Tangent Kernel (Jacot et al., 2018) and Neural Network Gaussian Process Kernel (Lee et al., 2017), respectively. Lastly, Loo et al. (2023), representing the third category, uses the implicit function theorem to compute gradients for unrolled inner optimiza- tion, allowing for the updating of the pseudo-coreset.
Variational Inference Variational inference (Bishop, 2006; Blundell et al., 2015; Blei et al., 2017), one of the most general methods for approximating most posterior distributions, is a tech- nique that approximates the target posterior distribution using variational distribution, which has a well-known and manageable form. The parameters of the variational distribution are learned by minimizing the KL divergence between the target posterior distribution and the variational distribu- tion. Aithough using all the parameters of the variational distribution can enhance its expressiveness, allowing for more accurate approximations, two common approaches are typically employed to ad- dress the computational and memory challenges that arise when handling the large scale of BNN weights: 1) mean- mean-field approximation (Blundell al., 2015; Shen et al., 2024), and 2) computing the posterior distribution for only subset of the network parameters (Dusenberry et al., 2020; Fiedler & Lucia, 2023; Harrison et al., 2024a). In both of these cases, the parameters of the variational distribution are optimized either directly using gradient descent methods to minimize the KL diver- gence (Blundell et al., 2015; Dusenberry et al., 2020; Shen al., 2024), or a closed-form solution is found (Wang & Blei, 2013).
D ADDITIONAL DISCUSSION ON VBPC
Future work direction Here, we would like to discuss some concerns and challenges we fore- in adopting the Laplace approximation the softmax likelihood instead of using variational inference with Gaussian likelihood.
Specifically, if we switch from using Gaussian likelihood employing a softmax likelihood with Laplace approximation for variational inference, there are two cases to consider: (1) using Laplace approximation on the last-layer weights without any updates, and (2) updating the last-layer weights with some gradient descent steps before applying Laplace approximation.
In the first case- case-applying Laplace approximation to weights without updating the last layer-two main issues may arise. First, the Laplace approximation assumes that the weights are near a mini- mum, allowing for the approximation of the first-order term in Taylor expansion as zero. However, this assumption may not hold for untrained weights, leading to significant approximation error. Additionally, the computational burden of calculating the Hessian for Laplace approximation substantial, and the need to compute gradients through this Hessian during pseudo-coreset updates increases the computational load further.
In the second case-nud the last layer weights with gradient steps before applying Laplace theroximationntterr''s the advantage of reducing Taylor expansion error. However, this approach involves large computational graph, which can be problematic due to the computational expense typical in bilevel optimization settings. Additionally, the need to compute gradients through the Hessian remains challenge.
Overall, we believe that solving these issues could lead to new meaningful future work for VBPC.
Limitations of the Last-Layer Approximation There might be concerns that considering the posterior distribution of only the last layer weights, rather than the entire parameter set, could limit
21

---

Published as a conference paper at ICLR 2025
the model's ability to capture uncertainty effectively, especially the model size increases and tasks become more complex. We fully agree that this is a valid concern and would like to provide a discussion based related findings.
Specifically, Harrison et al. (2024b) provides extensive empirical evidence on the effectiveness of last-layer variational inference. Their experiments span diverse tasks, including regression with UCI datasets, image classification using a Wide ResNet model, and sentiment classification leveraging LLM features from the OPT-175B model. They compared their method with other Bayesian in- ference approaches such as Dropout, Ensemble methods, and Laplace approximation for the full model. Their results demonstrate that even though last-layer variational inference focuses solely on the final layer weights, it achieves performance comparable to other comprehensive Bayesian inference techniques across various tasks.
These findings indicate that while conducting Bayesian inference on the full set of weights in neural network could potentially lead to more precise uncertainty estimation, employing last-layer variational inference is still effective in capturing meaningful uncertainty.
We believe that extending VBPC to incorporate full-weight variational inference could be a promis- ing direction for future work, offering the potential to further enhance the method's uncertainty estimation capabilities. We will include this discussion in the final manuscript to provide balanced perspective and acknowledge possible avenues for improvement.
E EXPERIMENTAL DETAILS
Our VBPC code implementation is built the official FRePo (Zhou et al., 2022)' codebase. The im- plementation utilizes the following libraries, all available under the Apache-2.0 license JAX (Brad- bury al., 2018), Flax (Babuschkin et al., 2020), Optax (Babuschkin et al., 2020), TensorFlow Datasets (Abadi et al., 2015), and Augmax For the baseline methods, we used the official code implementations provided for each. All experiments, except those on the Tiny-ImageNet (Le & Yang, 20155) dataset, were performed on NVIDIA RTX 3090 GPU machines, while Tiny-ImageNet experiments were conducted on NVIDIA RTX A6000 GPUs.
E.1 DATASETS
Datasets for the Bayesian Model Averaging comparison For the BMA comparison experi- ments, we utilize 5 different datasets: 1) MNIST (LeCun et al., 1998), 2) Fashion-MNIST (Xiao et al., 2017), 3) CIFARIO (Krizhevsky, 2009), 4) CIFARI00 (Krizhevsky, 2009), and 5) Tiny- ImageNet (Le & Yang, 2015).
MNIST:
The MNIST dataset' contains 10 classes of handwritten digits with 60,000 train- ing images and 10,000 test images, each with dimensions of 28 x 28 x 1. All images were normalized using mean of 13.1307] and standard deviation of [0.3081].
Mashion-MNIST: The Mushion-MNIST dataset consists of 10 classes of fashion article images, with 60,000 training images and 10,000 test images, each with dimensions of 28 x 28 x 1. Images were normalized using mean of [0.2861] and a standard deviation of [0.3530].
CIFAR-10/100: The CIFAR-10/100 dataset contains 10/100 classes, with 50,000 training images and 10,000 test images sourced from the 80 Million Tiny Images dataset (Torralba et al., 2008). Each image has dimensions of 32 x 32 x 3. For CIFAR-10, images were normalized with a mcan of [0.4914, 0.4822, 0.4465] and stan- dard deviation of [0.2470, 0.2435, 0.2616], while CIFAR-100 images used a mean of [0.5071, 0.4866, 0.4466 and standard deviation of [0.2673, 0.2564, 0.2762].
hhtps://github.com https///ithub..o//yoggggg http://hhh...com/yooonn///////
hhtps:///www wwtp:/wwwwwwwwwwwwwwwwp/cooo// www/liicnnse//LICCEEE22222
hhttps://yann..ee www.n.
www.llccccccoooooooeee com/exdh/mnist/
/hups://wwww. es toronto edu/ hwwp:////d///////iiffaaa..mm
22

---

Published as a conference paper at ICLR 2025
Tiny-ImageNet: The Tiny-Imase ImageNet dataset" contains 200 classes, with 100,000 train- ing images and 10,000 test images. Each image has dimensions of 64 x 64 x 3. Im- ages were normalized using mean of [0.4759, 0.4481, 0.3926] and a standard deviation of (0.2763, 0.2687, 0.2813].
Datasets for the Out of Distribution scenarios For the distribution shift and OOD scenarios, we use CIFARIO-C (Hendrycks & Dietterich, 2019), which includes seven corruption types with five severity for each corruption type: 1) Gaussian Blur, 2) JPEG Compression, 3) Snow, 4) Zoom Blur, 5) Pixelate, 6) Defocus Blur, and 7) Motion Blur
The CIFARIO-C dataset consists of 10 classes, with 50,000 test images for each corruption type. applies various corruptions to 10,000 test images from CIFARIO, with five levels of severity, each containing 10,000 images. The images are normalized using the same mean [0.4914, 0.4822, 0.4465] and standard deviation [0.2470, 0.243.,,43555 0.2616] the CIFARIO dataset.
CIFARIO-C:
E.2 MODEL ARCHITECTURE
Model architecture utilized for the Bayesian Model Averaging and Out of Distribution tasks Following previous works (Kim et al., 2022; 2023; Tiwary et al., 2024; Zhou et al., 2022), used a convolutional neural network (CNN) for the Bayesian Model Averaging comparison experiment and the Out of Distribution experiment. This model is composed of several blocks, each consisting of a 3 x 3 convolution kernel, pre-defined normalization layer, Rectified Linear Unit (ReLUU Agarap, 2018) activation, and 2 x 2 average pooling layer with stride of 2. For datasets with resolutions of 28 x 28 x and 32 x 32 x 3, we used 3 blocks, and for datasets with a resolution of 64 x 64 3, used 4 blocks. Following Zhou et al. (2022), we increase twice the number of filters when the feature dimension was halved, to prevent the feature dimensions from becoming too small. Additionally, by default, we used the Batch Normalization (loffe, 2015) layer for normalization unless stated otherwise. For initializing model weights, we conducted experiments using the Lecun Initialization (Cun et al., 1998) method, which the default initialization method of the Flax library. This configuration was applied both during the model pool in the VBPC training process and in the evaluation phase.
Model architecture utilized for the Architecture Generalization task For the Architecture generalization experiments, we incorporate three additional normalization layers and three ad- ditional model architectures. The normalization layers include Instance Normalization (Ulyanov, 2016), Identity map, and Group Normalization (Wu & He, 80. For the model architectures, we include AlexNet (Krizhevsky al., 2012), VGG (Simonyan & Zisserman, 2014), and ResNet (He al., 2016). Initially, evaluate all baselines by replacing Batch Normalization in the convolution layers with the three alternative normalization methods, referring to these as CNN-IN, CNN-NN, the three additional model architectures for evaluation.
and CNN-GN, respectively. Next, Since AlexNet does not have normalization layers in its original design, we retain this structure and refer to it AlexNet-NN. For VGG and ResNet, we use VGG11 I with Group Normalization and ResNet18 with Batch Normalization. These models are denoted as VGGI and ResNet 8-BN.
E.3 PSEUDO-CORESET INITIALIZATION PREPROCESSING, AND AUGMENTATION
Building prior works (Kim et al., 2022; 2023; Tiwary et al., 2024; Zhou al., 2022), we initialize the pseudo-coreset by randomly sampling images and labels from the original training dataset using fixed sampling seed. For the labels, following Zhou al. (2022), we initial- ize them with scaled, mean-centered one-hot vectors corresponding to each image, where the scaling Here, we train both images and factor is determined by the number of classes h; specifically Vk/10"
Initialization
labels during training.
23

---

Published as a conference paper at ICLR 2025
Data preprocessing and Augmentation Following previous works (Kim et al., 2022; Tiwary et al.. 2024; Zhou et al., 2022), we perform standard preprocessing on each dataset. with the ad- dition of ZCA (Kessy et al., 2018) transformations for all datasets with 3 channels. Consistent with Zhou et al. (2022), we apply a regularization strength of = 0.1 across all datasets. Similar to previous works (Kim et al., 2022; 2023; Tiwary et al., 2024; Zhou al., 2022), we apply the fol- lowing augmentations to the MNIST and Fashion-MNIST datasets: 'Gaussian noise", 'brightness', "crop', 'rotate', 'translate', and 'cutout'. For all other datasets, we use 'flip', 'Gaussian noise', 'color', 'crop', 'rotate", 'translate,, and 'cutout' augmentations. These augmentations are applied both during the training of the Bayesian pseudo-coreset and during evaluation with them.
E.4 HYPERPARAMTERS
Hyperparameters during training VBPC Following previous works (Kim et al., 2022; 2023; Tiwary et al., 2024), we select 1, 10, or 50 images per class for all datasets when training VBPC for evaluation. For Ss. we use â, which corresponds the number of pseudo-coresets in each experiment. This setup is designed to control the gradient magnitude by averaging, rather than summing. the expected likelihood, while maintaining the influence of the KL divergence for stable training. For 3D. we used 1c-8 as the default value, and when adjusted, it was selected from the range 1e-6, le-7, 1e-8] across all experiments. For p and 7, we set the default values to p= 1.0 and 100.0 for the ipc 1 and ipc 10 settings, and p = 10.0 and y = 100.0 for the ipc 50 settings. Except for the CIFARI00 ipc 10 setting where we utilize p = 10.0 and y = 100.0 for the default setting. When tuning these parameters, we adjusted them on a log scale in steps of 10 within the range of [-5, 5]. Following the default settings in Zhou et al. (2022). we set the number of models stored in the model pool, P. to 10. Additionally, as per Zhou et al. (2022), training steps, T, for each model in the model pool to 100. For the model pool optimizer, we used the Adam (Kingma, 2014) optimizer with fixed learning rate of 0.0003 across all experiments. For the pseudo-corese optimizer, we also used the Adam optimizer by default, with a cosine learning rate schedule starting at 0.003 for both images and labels. Lastly, we used batch size of 1024 and trained for 0.5 million steps to ensure sufficient convergence.
set the number of
Hyperparameters during variational inference and Bayesian Model Averaging For all exper- iments, the hyperparameters y. p, and As used during evaluation were the same those used for pseudo-coreset training in the corresponding experiment. The optimizer used for training the models during evaluation
the Adam optimizer with a constant learning rate of 0.0003. The number of training steps for each model was as follows: for MNIST and MaSHIOn-MNIST. 100 steps for 1 ipc, 500 steps for 10 ipc, and 1000 steps for 50 ipc. For CIFARIO, 200 steps for I ipc, 2000 steps for 10 ipc, and 5000 steps for 50 ipc. For CIFAR (0, 2000 steps for both 1 ipc and 10 ipc, and 5000 steps for 50 ipe. Lastly, for Imng-Ima,,,et, 1000 steps were used for ipc and 2000 steps for 10 ipc.
F ADDITIONAL EXPERIMENT
F.1 FULL EXPERIMENTAL RESULTS ON BAYESIAN MODEL AVERAGING COMPARISON
Here, we report the full experimental results for Section 5.1. We report results for FRePo and RCIG across the entire benchmark dataset and varying IPC settings additional to Table 1. Table clearly demonstrates that VBPC surpasses other BPC baselines across all benchmark datascts and IPC settings in terms of ACC and NLL. Notably, VBPC achieves significantly better NLL,, with large margins, while requiring only single forward pass to conduct BMA. Although VBPC is designed to learn pseudo-coresets that approximate the variational distribution derived from the training dataset, it outperforms dataset distillation baselines, which primarily focus on achieving high ACC, in nearly all tasks, except for CIFARIOS 100 with IPC and Tiny Imageeet These results empirically validate that the variational distribution trained by VBPC effectively captures epistemic uncertainty with a small amount of synthetic data, while maintaining high performance.
Comparison with dataset distillation baselines In Section 5.1, the performance was evaluated based on the training and evaluation methods proposed by each baseline's 's original papers. However, one might question whether the significant performance of VBPC is due to the trained pseudo- coreset itself or the VI method. To address this, and to validate that the significant performance of
24

---

Published as a conference paper at ICLR 2025
Table 6: Comparison of the VBPC with BPC and additional dataset distillation baselines for the benchmark datasets. We report ACC and NLL for the BPC baselines, and ACC for the dataset dis- tillation baselines. Boldfaced blue color indicates when the performance of the dataset distillation baseline surpasses that of VBPC.
FRePo RCIG BPC-rKL
BPC-KKL
FBPC
BPC-CD
VBPC (Ours)
Datasct
1pc
ACC(T) ACC(T) ACC(T) NLL(() ACC(T) NLL()) ACC(T)
NLL()
ACC(T)
NLLG)
ACC(T)
NLL(J)
93.0+e.. 94.7=0. ( 1.90=..1 83.0=a
188=0.03
92.5+0.1
1.680.00
93.4+0.1
1.53+0.01
96.=0.44 0.11+0.02
MNIST
10
98.620.1
98.9200
1.53.0.01 3
1 2
1.50
120
1.3
97.7.0a
1.57.000
99.1
0.03.0.0
99.2+0.1 50
99.2+0.0
94.2+ 1.3600.02
93.6a1
1.36=0002
98.6.0.1
1.39.0.02
98.9.02
1.36+0011
99.4=0.1
0.02+0.01
1
75.6
79.8 70.5
2.47.0.02 72.5
2.30
74.7a1.4
181
10:00
77.3
1,9010.00
K2.9
t0.0
0.47
FMNIST
10
86.2.0.
1.64..0.0
83.3100
1.54_..00
85.2.101
1.61
88.4.0.
1.56110.01
89.4
1.30
50
89.62
90.2..0:
77.00
1,48
74.8
1.47.u.
76.7.0.4
1.4610.z
89.5
1.30
91.0
20.2
0.25
1
86.8
53.9 21 3
2.57:0..1
29.3
2.10=
35.5
3.79
46.9-0.2
+8820
55.1
1.34
CIFARIO
10
65.5
1 (37.911
2.13.00.2
49.911:4
1.73_
62.3.
1,31
10:00
56.4.0.7
1.72
69.8
0.89
71.7a0.2 50
73.5
(37.5
1.93
42.3+22
1.54
21.2+02
1.03
71.92022
1.57
76.7
+0.6
+.71
28.7
39.3
3.6+0.1
4.69=0.02
14.7=002
4.20=0.10
21.0=aa
3.76+0.11
24.0=0.1
4.00
38.4
2.47+0.04
CIFARIO0
10
42.
44.1
23.6+a0
2.99.0.03
180.
3,53=0.05
39.7=0.a
2.07=0.02
28.420.2
3, 3,14
49.4
2.07=0.02
44.3+0.2 50
46.7=0.3
(30.8e0.5
3,57=0.17
37.1
3.28=0.24
44.5+0.4
2.63+0.01
39.620.2
3.02+0.01
52.4
2.02+0.02
Timy-ImageNet
1
15.4
25.611..
3.2.
5.91
4.0
5.63__0.0
1.69
8.4
4.72.0.01
23.1 Loin 1..5
10
25.420.3 29.4 9.8.0.
5.26 0s 11,4 5 5.08 19.4 4.14g0.02 17.820.4 3.64 0G
25.8 3.45 02
Table 7: Ablation experiment on BMA method. Here, we conduct our variational inference method utilizing datasets trained with other baselines.
FRePo VI
RCIG VI
BPC-KKI. VI ACC(T) NLL()
BPC-IKI VI ACC(t) NLL())
FBPC VI ACC(T)
BPC-CD VI ACC(T) NLL())
Dataset ipc ACC(1) NLL()) ACC()) NLL())
NLL(()
CIFAR1O
1 28.2a0.9 2.22a0.02
27.8 2.20+0.0:
10.1 2.30 2.251002
10.1 2.32
10.0sa.1 2.37a000
10.4:0.8 2.35 10.520.7 2.322901
10 55.720.
2.0720:0
55.611.5 2.0510022
12.0g05
20.111
2.211001
10.000.
2.37g000
CIFARI00 1 19.9 4.55
2.1 2.5+9.4 5.45#0.12 2.6 4.59+0.02
5,02+0.05
1.2a011
4.60a0.0
1,4203
4.60a0.1
1.240.2 4.60g0.0
1.230.2 4.602001
10 34.8=0.4
4.50aa.01
4.0=0.2 4.59
1.620.3 4.59=0.00
11.620.4 4.54+0.02
VBPC is not solely attributable to the VI method, we collected the pseudo-coresets trained on all baselines used in Section 5.1 for the CIFARIO and CIFAR 100 datasets in the 1ipe and 10ipc settings. We then applied our proposed VI method to these baseline pseudo-coresets to measure their BMA performance and compared the results with those reported in Table 6. Results in Table 7 and Table clearly show that the performance significantly drops for all baselines compared to their original performance. This validates that the performance is not solely attributable to the VI method, and demonstrates that VBPC successfully learns to approximate the variational distribution effectively.
F.2 ADDITIONAL EXPERIMENT RESULTS ON LARGE DATASET AND CONTINUAL LEARNING
To further highlight the ability of VBPC to handle tasks that pose challenges for other BPC baselines, we conduct additional experiments on more large datasets and the continual learning setting.
Large Datasets First, to show that our method is uniquely scalable to large datasets compared to other BPC methods, we conducted additional experiments on the ImageNetWoof (128x128x3) dataset (Howard, 2020) and the ImageNetlk Ik (64x64x3) dataset(Russakorw et al., 2015). Addi- tionally, we included an experiment performs better in practical scenarios.
continual learning scenario to validate that our method
We conducted experiments on the ImageWoof (128x128x3) dataset with ipc and ipc 10 settings, as well as the resized ImageNetlk (64x64x3) dataset with ipc I and ipc 2 settings, to demonstrate the scalability of our method to high-resolution images and larger datasets. Unlike existing BPC baselines, which encountered memory issues and failed train due to out-of-memoot errors an RTX 3090 GPU as the image resolution and number of classes increased, our method successfully completed training. Table 8 clearly shows that VBPC significantly outperforms other baselines with a large margin for both the ImageWoof and resized ImageNet1 k datasets.
Continual Learning Next, we validated the practical effectiveness of our method through con- tinual learning experiments using pseudo-coreset images learned by each method. We followed the
25

---

Published as a conference paper at ICLR 2025
Table 8: Experiments on the scalability utilizing ImageWood and resized ImageNet datasets. Here indicates the training fails due to the out-of- problems.
ImageWoof ipc
ImageWoof ipc 10 ImageNet ipe ImageNet ipe ACC(t) NLL(J)
Method
ACC(1) NLL())
ACC(t) NLL(J)
ACC(t) NLL())
Random
14.2+ 3.84
27.0==19 2.83
11+0.1 8.32+0.05
1440.1
8.10
BPC-CD
18.5+0.1 2.7641..0
FBPC
14.8+0.1 3.73+0002
28.1 2.69
BPC-fKL
14.940.9 3.74+0.25
25.0=0.8
2.90+0.27
BPC-rKL
12.0=0.5 6,0740.31
VBPC
31.2++0.1 2.13+0000 39.0=0.1 1.8 10.1=0.1 5.33+0.04 11.5+0.2 5.25++0.05
Table 9: Experiments on the continual learning setting. Here, utilize the CIFARIO0 dataset with ipe 20 setting. We assume steps during training and each step contains data from new 20 classes in the CIFARI00 dataset. Here only report accuracy due the variant of the number of classes during the steps.
Number Classes
20
80
100
BPC-CD
52.5+2.4 40.41..3
35.2
33.4
29.4
FBPC
61.4+1.8
53.2+1.5
48.8
43.9
41.2=0.3
BPC-IKL
51.8+2.2
39.8+11
35.5
33.1
29.5
BPC-rKL
48.212.7
35.5+1.8
32.0=10
29.8
25.5=0.3
VBPC
75.3+2.0
65.8
+0.1
53.3
50.3
continual learning setup described in Zhou et al. (2022); Zhao & Bilen (2021), where class- balascealcced training examples are greedily stored in memory, and the model is trained from scratch using only performed 5-step class incremental learning experiment on
the latest memory. Specifically, CIFAR 100 with an ipe 20 setting, following the class splits proposed in Zhou et al. (2022); Zhao & Bilen (2021). Table 9 demonstrates that VBPC consistently outperforms other baselines across all steps, confirming its superior practicality and effectiveness real-world continual learning scenar- ios.
F.3 ADDITIONAL. EXPERIMENTS ON OUT-OF-DISTRRIUUTINN DATA
To further validate the effectiveness of VBPC, We have conducted additional Out-of-Distributio (OOD) detection experiments and reported the results. The metrics we evaluate include AUROC, AUPR-In, and AUPR-Out, where higher values indicate better performance. We used models trained with the CIFARIO IPC 10 setting and evaluated them on CIFARI00 TinylmageNet, and SVHN (Netzer al., 201 1) datasets as OOD datasets.
The results, presented in Table 10, demonstrate that the pseudo-coress learned by VBPC performs robustly in OOD detection scenarios. These findings, combined with the corruption experiments in the main paper, validate the effectiveness and robustness of VBPC under diverse and challenging evaluation conditions.
F.4 ANALYSIS ON COMPUTATIONAL COSTS AND TRAINING TIME
In this section, we performed analyses focusing on two aspects of computational cost.
Cost of training the pseudo-coreset As mentioned in the Section 1, conventional BPC methods relying on SGMCMC require the creation of expert trajectories, which are training trajectories de- rived from the full dataset. Each dataset typically involves training with different random seeds for these trajectories, making this step computationally expensive. Since all BPC baselines share and utilize these precomputed trajectories, their associated computational cost can be considered a shared overhead.
To isolate the computational cost of training the pseudo-coreset itself, we measured the wall-clock time required for pseudo-coreset optimization by each method. The results of this comparison
26

---

Published as a conference paper at ICLR 2025
Table 10: AUROC, AUPR-In, and AUPR-Out results for the OOD detection task with model trained with the learned pseudo-coresets. Note that we used the same model structure which is utilized when training pseudo-coresets.
Dataset
Model
AUROC(T) AUPR-In(1) AUPR-Out(T)
BPC-CD
49.09
52.79
45.88
BPC-fKL
48.95
51.72
47.00
Tiny ImageNet
BPC-rKL
48.34
52.71
44.49
FBPC
45.39
49.70
43.14
VBPC
52.85
56.22
49.64
BPC-CD
55.09
35.64
73.88
BPC-FKL
54,26
34.78
75.47
SVHN
BPC-rKL
42.61
28.29
67.15
FBPC
41.34
30.12
62.18
VBPC
68.50
48.49
82.91
Table 11: Wall clock time results for training pseudo-coresets with each BPC method using CI- FARIO ipc 10 settings. We used RTX3090 GPU to measure the exact training time. Here, all methods except for VBPC share the training time for expert trajectories.
Method
|PCC-D BPC-rKL FBPC BPC-IKL VBPC
Times (hr)
548.5
549
5+10.5
5+12
5.5
summarized in Table 11, providing insights into how VBPC reduces training costs compared to other baselines.
Cost of inference When performing inference, VBPC requires training only single model, whereas other BPC baselines rely on multiple SGMCMC samples. Each sample incurs significant training and inference costs, which grow linearly with the number samples.
To quantify this difference, we measured the wall-clock time for inference across methods, with results presented in Table 12. These results highlight how VBPC achieves superior efficiency during inference by avoiding the high computational costs associated with sampling-base approaches.
These analyses demonstrate VBPC's ability perform Bayesian inference efficiently, both in terms of pseudo-coreset training and inference, and further reinforce the computational advantages of our method.
F.5 ABLATION ON RANDOM INITIALIZATION
Since our method initializes the pseudo-coreset by randomly sampling images and labels from original training dataset. following previous works (Kim et al., 2022; 2023; Tiwary al., 2024; Zhou et al., 2022), we conducted an ablation experiment using random initialization for the pseudo- coreset. In this experiment, first randomly initialized the pseudo-coreset sampling pixel values from a uniform distribution Unif(0. 1). We then trained the images after normalizing them with the predefined mean and variance for each dataset reported Appendix E. We conducted this ablation experiment on the CIFAR 10/100 1 ipc and 10 ipc settings. Fig. and Fig. 3 clearly illustrate that VBPC can effectively learn semantic information even when initialized randomly. Specifically, in the CIFAR10 ipc 1 case shown in the top figures of both Fig. 2 and Fig. 3, the images after training appear similar, whether they were initialized randomly or sampled from the training dataset. Also Table 13 shows that randomly initialized VBPC shows comparable performance compared to the VBPC.
F.6 ABLATION ON PSEUDO-CORESET OPTIMIZER
Since we use the Adam optimizer for training VBPC, which differs from the default choice in pre- vious work (Zhou et al., 2022), we conducted ablation experiment the optimizer. Following
27

---

Published as a conference paper at ICLR 2025
Wall clock time results for inference using learned pseudo-coresets. We measure the in- ference time for evaluating all the test data from the CIFARIO test dataset. After finishing training the pseudo-coresets the inference cost for the all baselines are same because they only need SGM- CMC and BMA with same number of datasets and weight samples.
Table 12:
Method BPC-CD BPC-rKL FBPC BPC-IKL VBPC
Times (s)
165
165
165
165
20
Figure Learned VBPC images from the random initialization for the CIFARIO ipc (above) and ipc 10 (below) cases. The left figure shows the random images sampled from the uniform distribution and the right figure shows the trained VBPC images starting from the left images. Training from random initialization successfully learns semantic information from the full dataset.
Zhou al. (2022), used the LAMB (You al., 2019) optimizer with cosine learning rate schedule for this ablation. We conduct this ablation experiment on the CIFAR1O I ipc and 10 ipc settings. As seen in Fig. 4, although there minor differences, the images trained with the LAMB and Adam optimizers are largely similar when starting from the same pseudo-coreset initial images. Additionally, Table 14 demonstrates that our method effectively learns pseudo-coreset with varying optimizers, closely approximating the variational distribution of the full training dataset.
F.7 ABLATION MODEL POOL MAXIMUM UPDATE STEPS
As mentioned in Appendix E.4, set T 100 the maximum update step for the weights in the model pool M across all experiments. The model pool was introduced to address VBPC's overfitting issue, the weights in the model pool trained for T steps, leading to a variety of feature maps. This prevents VBPC from learning based on a single feature map. To investigate the effect of T' we plan to conduct an ablation study to examine how changes in impact image quality and performance. We conducted an ablation experiment on the CIFAR100 ipc 10 task with
Table 13: Comparison between the random initialization and initialization with randomly sampled images. Random Initialization denotes the VBPC learned starting from the uniform random initial- ization. Here, we report ACC and NLL for both initializations.
Random Initialization
VBPC ACC(t) NLL())
Random Initialization
VBPC
Dataset ipc ACC(t)
NLL())
Datasct
ipc ACC(1)
NLL())
ACC(1) NLL())
CIFARIO
54.240.5
1.37=0.02
1.34 os 15..2 69.8 0.89
CIFARI00
27.4 2.51
38.420.2
2.47
68.940.4
40.0
10 48 4
2.20=0.03
49.4=0.1 2.07=
28

---

Published as a conference paper at ICLR 2025
Figure 3: Learned VBPC images from the randomly sampled image from the original training dataset for the CIFARIO ipc (above) and ipe (below) cases. The left figure shows the initial images sampled from the original dataset and the right figure show the final learned VBPC starting from the left images.
Figure 4: Visualization of leamed VBPC images utilizing different optimizers for the CIFARIO ipc (above) and ipc 10 (below). The left figure shows the learned VBPC images with LAMB optimizer and the right figure shows the learned VBPC images with Adam optimizer.
200 and 400. As shown in Fig. 5, the images learned with different maximum update steps appear visually similar. However, Table 15 quantitatively shows that excessive updates to the model pool weights can reduce feature diversity, potentially leading to a decline performance for unseen feature maps.
F.8 ABLATION LABEL LEARNING
Following the previous works (Kim et al., 2022; 2023; Tiwary et al., 2024; Zhou et al., 2022), learned the labels when training the pseudo-coreset. This can be crucial and effective for learning more informative BPC, the mean of the pseudo-coreset variational distribution depends the label. This dependency also impacts the loss function used in the outer optimization process. Fig. shows that without label learning, trained VBPC images significantly lost the semantic informa-
29

---

Published as a conference paper at ICLR 2025
Table 14: Comparison between the VBPC learned with LAMB optimizer and Adam optimizer. LAMB denotes the VBPC trained with LAMB optimizer. Here, we report ACC and NLL for both optimizers.
LAMB
VBPC
LAMB
VBPC
Dataset
ipc ACC(t)
NLL() ACC(1) NLL())
Dataset
ipc
ACC() NLL()) ACC(1) NLL())
54.4400.8 1.36=0.03 55,1 =0.3 1.34
38.5. 2.46#0.03 38,4 2.47 10 49.4+0.. 2.25 49.4 2.07
CIFARIO
CIFARI00
10
69.4+0.5 0.90 69.8 0 0.89
Figure 5: Learned VBPC images utilizing different maximum updates steps for the model pool
elements in the CIFARI00 ipc 10 experiment. The left figure shows the T 100 case which is the default setting for the all experiments. The middle and the right figures show the T = 200 and T = 400 cases. The learned images show minor difference in visual.
tion for each image. Also, results presented in Table 16 clearly shows that the BMA performance with VBPC variational distribution largely drops without label learning. These results validate that learning the label is important for the successful VBPC training.
F.9 ABLATION ON GAUSSIAN NOISE AUGMENTATION
During VBPC training, we apply Gaussian noise augmentation. Based on previous findings that adding Gaussian noise to images during neural network training improves robustness various image corruptions (Rusak et al., 2020), we incorporate Gaussian noise during VBPC training. This helps the learned pseudo-coreset dataset produce variational posterior that is robust to unseen model structures and corrupted test datasets. Specifically, add Gaussian noise sampled from (O 0.01) after normalizing the images using the predefined mean and standard deviation for all tasks. We conduct the ablation experiment
the existence of this Gaussian Noise during training VBPC utilizing CIFAR 10/100 ipe and 10 ipe settings. As clearly seen in the CIFAR1O ipc case in Fig. training with Gaussian noise results in much clearer and brighter images. In contrast, without Gaussian noise, the model tends to learn visually similar features in the background, unlike the cases where noise applied. Table 17 confirms that, as expected, the overall performance decreases when Gaussian noise augmentation is not applied, compared to VBPC with Gaussian noise augmentation.
Table 15: Ablation results on the model pool maximum update steps. Here, we used CIFAR 100 ipc 10 setting for the ablation. T' 200 and T 400 indicate the maximum updates for the model pool 200 and 400, respectively. Here, we report ACC and NLL for all the update steps.
VBPC ACC NLL()) ACC(t) NLL())
T 200
T 400
Datasct ipc
ACC NLL()
CIFAR100 10 49.420.1 oz 48.7 2. a.162ass 48.0 2.22-0.11
30

---

Published as a conference paper at ICLR 2025
Figure 6: Visualization of learned VBPC images with and without label learning for the CIFARIO ipe (above) and ipc (below). The left figure shows the learned VBPC images without label learning and the right figure shows the learned VBPC images with label learning.
Table 16:
Comparison between the VBPC learned with and without label learning. No Label de- notes the VBPC trained without label learing. Here, we report ACC and NLL for both results.
No Label
VBPC
No Label
VBPC
Dataset ipe ACC(t) NLL()) ACC(t) NLL())
Dataset
ipc
ACC(1)
NLL()) ACC(1) NLL()
40.641.3 1.83=0.02 55.1 1..48
10.1 5.12#0.05 38.4 2.47 10 23.5+1.2 4.9220.06 49.420.1 2.0740.02
CIFARIO
CIFAR100
10 56.610.3 1.54 69.8=0.7 0.89
F. ABLATION ON HYPERPARAMETER
In this section, conduct ablation experiments p and which the hyperparameters newly proposed in our work. We set the default values to p = 10 and ? = 100.0 for the CIFAR 100 ipc 10 settings. And, conduct ablation experiment with the CIFAR100 ipc 10 setting. And the results presented in Table 18 and Table 19 show that even when we varied our hyperparameters by orders of magnitude (in log scale, with changes up 10-fold), the performance remains consistently similar. And this concludes that our method works robustly with respect to hyperparameter changes.
F.11 ABLATION ON TRAINING STEPS DURING INFERENCE
In this section, we conduct ablation experiments on the number of training steps T" during inference. When learning the pseudo-coreset using the VBPC method, we leverage a model pool to allow the data to observe various feature maps, ensuring diverse learning. Therefore, even during inference, although the model may not perfectly fit the pseudo-coreset that was trained with the feature map, can still approximate the best variational distribution for the current last-layer weights based on
Table 17: Comparison between the VBPC learned with and without Gaussian Noise augmentation. No Noise denotes the VBPC trained without Gaussian Noise augmentation. Here, report ACC and NLL for both results.
Noise ipc ACC(t) NLL()) ACC(t) NLL())
YBPC
Noise
VBPC
Dataset
Dataset
ipc
ACC()) NLL()
ACC(t) NLL())
53.94..s 1.41 $5.1 1.34
35.440.. 2.6220.05 38.4+0.2 2.4740.04 10 48.540.4 2.22+0.06 49.4 0.07:00.2
CIFARIO
CIFAR100
68.8 840.7 0.92=0.04 69.8 0.894
31

---

Published as a conference paper at ICLR 2025
Figure 7: Visualization of learned VBPC images utilizing Gaussian noise during training for the CIFARIO ipe (above) and ipe 10 (below). The left figure shows the learned VBPC images without the Gaussian noise and the right figure shows the learned VBPC images with the Gaussian noise.
Table 18: Ablation results on the hyperparamer 7. Here, we used CIFARI00 ipc 10 setting for the ablation. y 10, 7 1000, and 7 10000 indicate that we set as 10, 1000, and 10000, respectively. Our default setting is 7 = Here, we report ACC and NLL for all the update steps.
10 NLL()) ACC(1) NLL() ACC(T) NLL() ACC(t) NLL())
VBPC
1000
10000
Dataset ipc ACC(T)
CIFARI00 10 48.5 2.30 49.410.1 2.07 49.4 2.14 49.0 2.10
the available feature map. This enables the model to achieve sufficient BMA performance even before the pseudo-coreset learning is fully completed. As shown in Table 20, the model exhibits slightly lower performance during initial steps, such as at 400 or 800 steps, compared to the best performance. However, after these early stages, the performance becomes nearly identical to the final convergence step at 2000 steps. These results further demonstrate that our VBPC approach allows for fast and efficient posterior approximation.
G TRAINED VBPC IMAGES
In this section, present the images learned under the ipc 10, and 50 settings for MNIST, Fashion-MNNNn CIFARIO, CIFARI00, and ImageNe.. To avoid overwhelming the report with too many images, have limited the number of reported images to maximum of 100 per task.
G.1 TAKE-HOME MESSAGE FROM LEARNED IMAGES
Regarding the learned pseudo-coreset images for CIFARIO. the results can be found Fig. 12 and left figure of Fig. 13, showing the outcomes for ipe values of and 10. These images reveal several interesting aspects of how VBPC captures information.
First, both ipc and ipc 10 images show that VBPC effectively learns features associated with specific classes, such as "horse" or "automobile," as can be visually confirmed. This indicates that the pseudo-corese images retain class-relevant information necessary for approximating the original dataset's 's posterior distribution. When comparing ipc and ipc 10, there are notable differences. In the case of ipc 1, where only single image per class is available, VBPC attempts to encode as many class saesiic features as possible into a single image. As a result, the learned image appears incorporate multiple discriminative features from the class symmetrically. In contrast, with ipc
32

---

Published as a conference paper at ICLR 2025
Table 19: Ablation results on the hyperparamer p. Here, we used CIFARI00 ipc 10 setting for the ablation. p 1, p 100, and p = 1000 indicate that we set p as 1, 100, and 1000, respectively. Our default setting p 10. Here, we report ACC and NLL for all the update steps.
p= 1000 Dataset ipe ACC(t) NLL()) ACC(1) NLL() ACC(1) NLL() ACC(1) NLL()
==
VBPC
100
CIFARIO0 10 49.0 2.1020.02 49.4 2.07 49 2. 2.20=0.08 47.5=a.4 2.35=0.04
Table 20:
Ablation results on the training step T during inference. Here, we used CIFARI00 ipc 10 setting for the ablation. 7" 400, T 800, T" 1200, and T" 1600 indicate that the intermediate performance at step 400, 800, 1200, and 1600, respectively. Our default setting is T' 2000. Here, we report ACC and NLL for all the update steps.
T' 400 Dataset ipc ACC(t) NLL()) ACC(1) NLL(1) ACC(1) NLL()) ACC(t) NLL()) ACC(f) NLL(L)
T 800
T" = 1200 T = 1600 VBPC
CIFARI00 48.9 2.1610.01 49.2.0.3 2.0910.01 49.4 2.07.0.01 49.510.2 2.07.0.01 49.410.1 2.07
10, where more images per class are available, VBPC distributes the class-relevant features across multiple images. This leads greater diversity of features being captured across the pseudo- coreset, enabling more comprehensive representation of the class.
Additionally, both ipc and ipe 10 images often include low-level features beyond the main class- relevant ones. These features likely help capture the dataset's variability and ensure the learned pseudo-coreset maintains close approximation the original data distribution.
These observations suggest that VBPC is effective in compressing the dataset while retaining essen- tial information. The learned images illustrate how VBPC balances feature extraction and informa- tion retention to ensure that the variational posterior distribution learned using the pscudo-coreset closely approximates the one learned using the full dataset. This further validates the interpretability and utility of VBPC in various tasks.
33

---

Published as a conference paper at ICLR 2025
Q
5
Figure 8: Visualization of learned VBPC images for the MNIST ipc 1.
4
7
5
0
0
2
0
6
0
0
0
0
6
1
7
1
7
1
1
1
/ 2
/
/
1
7
1
1
1
2
2
2
2
2
3
à
3
2 $ 4
3
2
2
2
2
2
8
2
2
2
2 3 9 5
3
3
3
3
3
3
3
8
3
3
3
3
3
3
3
3
3
0 4 5 6
4
9
/
4
4
0
4 5
4
4
4 4
7
4
9
4
4
4
5
5
$
5
5 5
5 6
8 6
5
5
5
5
5
$
5
5
6 1
6
6
6
4
6
6
6
6
6
6
6
6
6
7
7 7
7
7
7
7
2 8
7 7 8
7
7
7
7
7
7
0
7
1 8
2 8
&
8
8
8
8
8
8
8
8
8
8
8
9
6
7
9
9
9
8
4
9
9
9
9
9
7
9
Figure 9: Visualization of learned VBPC images for the MNIST 10 (left) and ipc50 (right).
Figure 10: Visualization of learned VBPC images for the Fashion-MNIST ipc
Figure 11: Visualization of learned VBPC images for the Fashion-MNIST ipc10 (left) and ipc50 (right).
Figure 12: Learned VBPC images for the CIFAR10 ipe case.
34

---

Published as a conference paper at ICLR 2025
Figure 13: Visualization of learned VBPC images for the CIFARIO ipe (left) and ipc50 (right).
Figure 14: Visualization of learned VBPC images for the CIFAR 100 ipcl.
35

---

Published as a conference paper at ICLR 2025
Figure 15: Visualization of learned VBPC images for the CIFAR100 ipc10 (left) and ipc50 (right).
Figure 16: Visualization of learned VBPC images for the Tiny-ImageNe ipel (left) and ipc10 (right).
36


Reviews:
[{"review_id": "CIB5GTyCvL", "reviewer": "Reviewer_hhFq", "summary": "The paper presents the Variational Bayesian Pseudo-Coreset (VBPC) method, aimed at efficiently approximating the posterior distribution in Bayesian Neural Networks (BNNs). Given that BNNs face substantial computational challenges when dealing with large datasets due to their high-dimensional parameter spaces, VBPC provides a promising method. Traditional Bayesian Pseudo-Coreset (BPC) techniques have been proposed to alleviate these issues, yet they often struggle with memory inefficiencies. VBPC addresses this by leveraging variational inference (VI) to approximate the posterior distribution. This method achieves a memory-efficient approximation of the predictive distribution using only a single forward pass, which makes it appealing for computationally intensive applications.", "questions": "- While VBPC demonstrates improvement over some BPC baselines, its classification accuracy on benchmark datasets like CIFAR-10, CIFAR-100, and Tiny-ImageNet remains notably lower than that of state-of-the-art classifiers. This raises concerns about the practical competitiveness of VBPC in real-world applications, particularly in image classification tasks where accuracy is crucial. Could additional optimizations or refinements to the VBPC approach improve performance?\n- The current experiments may not adequately showcase the strengths of VBPC in relevant scenarios. To enhance the paper’s impact and applicability, I suggest conducting additional experiments in settings where VBPC’s efficiency gains could be more convincingly demonstrated.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper effectively utilizes variational inference to derive a closed-form posterior distribution for the weights of the last layer, thereby addressing some of the performance limitations observed in prior BPC approaches.\nVBPC’s capability to approximate the predictive distribution in a single forward pass enhances both computational and memory efficiency, positioning it as a potentially valuable method for large-scale applications.", "weaknesses": "The experimental validation on practical application is limited.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "M5PXeia2AP", "reviewer": "Reviewer_8mt5", "summary": "This paper proposes Variational Bayesian Pseudo-Coreset (VBPC), a novel approach to efficiently approximate posterior distribution in Bayesian Neural Networks (BNNs). Bayesian Neural Networks often face issues with largescale datasets due to their high-dimensional parameter space. To reduce the computational load, many Bayesian Pseudo-Coreset (BPC) methods have been proposed, but they suffer from memory inefficiencies. VBPC addresses these limitations by using variational inference (VI) to approximate the posterior distribution. Moreover, this paper provides a memory-efficient method to approximate the predictive distribution with only a single forward pass instead of multiple forwards, making the approach computationally and memory-efficient.", "questions": "- The accuracy of these algorithms on cifar10, cifar100 and Tiny-Imagenet is too low. VBPC is effective relative to several existing BPC baselines, but the performance is significantly lower compared to existing state-of-the-art classifiers.\n- In the field of image classification, at least the scenes you choose for classification, these experiments do not seem to show the advantages of your method convincingly. \n- Please try to provide some new experiments, in more convincing scenarios, to illustrate the practical application value of your method.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This paper leverages the variational formulation to obtain the closed-form posterior distribution of the last layer weights, which resolves the issue of suboptimal performance seen in previous approaches.\n- And, the method approximates the predictive distribution with only a single forward pass instead of multiple forwards, making the approach computationally and memory-efficient.", "weaknesses": "- The experiment is not enough to illustrate the function of the algorithm.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "L8nLEBXLLf", "reviewer": "Reviewer_sqcH", "summary": "The paper studies the problem of core set extraction using Bayesian inference. The proposed solution builds on a two-stage variational inference scheme where the first stage is responsible for inferring the optimal core set while the second is to fit this core set to the full-scale data set at hand. The developed solution overarches the whole family of the distributions that belong to the exponential family.", "questions": "- Would Laplace approximation on the softmax likelihood be a better option than first choosing a variational inference scheme and then using a Gaussian likelihood? Laplace proves to be a powerful approach in Gaussian process classification.\n- Is the claim of Manousakas et al. 2020 contrary to the main message of the paper? If correct, would this not undermine the significance of the proposed solution? If incorrect, why?\n- Do we really lack an analytical solution or at least an EM-like algorithm where the E-step has an analytical solution when only the last layer of a neural net is probabilistic?\n- What is the take-home message of Figure 1? I am missing to see a particular pattern there that helps motivate the proposed solution.\n\nMy initial score is a borderline as the paper has both certain merits and clear question marks. I am happy to consider significant score updates based on a convincing rebuttal discussion.\n\n---\nPOST-REBUTTAL: The authors gave convincing answers to the above questions and they published additional results that demonstrate the advantages of the proposed method more clearly than the experiments reported in the original submission. I update my score to an accept.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is particularly well-written with a clearly defined problem scope and a solid solution methodology that follows a well-justified sequence of development steps,\n- The proposed bilevel variational inference formulation is neat and sensible.\n- The computational complexity analysis is indeed helpful to see the merit of the devised solution.\n- The reported results are strong on the chosen group of data sets.", "weaknesses": "- The paper motivates the core set extraction problem with use cases such as processing big data and addressing continual learning setups. However, the presented results are on data sets that can be considered in the present technological landscape as toy problems. I do symphathize the idea of prototyping. But given the strong applied component of the present work, I am still unconvinced about the generalizability of the observed scores to a case where coreset extraction is an actual necessity. The issue may be addressed during the rebuttal by showing results on a large enough data set used as a standard coreset extraction benchmark or a continual learning application.\n- The need to use the Gaussian likelihood to avoid the need for an approximation stage is only partially convincing. It is an artiffact of choosing variational Bayes as the inference scheme, which is exogeneous to the problem anyway. Maybe this issue, linked to my first question below, will be clarified during the rebuttal.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BabmzoVfId", "reviewer": "Reviewer_oE39", "summary": "The paper titled “Variational Bayesian Pseudo-Coreset” introduces a novel method aimed at reducing the computational and memory challenges associated with large datasets in deep learning, particularly within the context of Bayesian Neural Networks (BNNs). The authors address limitations in prior methods of Bayesian Pseudo-Coresets (BPCs), which often face inefficiencies in memory usage and suboptimal results during training. They propose a new approach called Variational Bayesian Pseudo-Coreset (VBPC), which leverages variational inference to approximate the posterior distribution of model weights. The key innovation of VBPC is the use of a closed-form solution to compute the posterior for the last layer of BNNs, eliminating the need for complex gradient-stopping techniques used in previous BPC methods. This significantly reduces memory usage and computational load. Additionally, VBPC allows for more efficient training and inference by using a single forward pass for predictive distribution computation. Empirical evaluations demonstrate that VBPC outperforms existing BPC methods on benchmark datasets, showing improvements in both accuracy and negative log-likelihood metrics across various datasets, such as MNIST, CIFAR10, and CIFAR100. The paper contributes to the field by enhancing the efficiency and scalability of BNNs, particularly in environments that require handling large-scale data while maintaining the benefits of Bayesian inference.", "questions": "1. Scalability to Larger Datasets: How does VBPC perform when applied to much larger datasets, such as ImageNet or larger text datasets? Does the method scale well in terms of both computational efficiency and accuracy, or does it encounter bottlenecks?\n\n2. Hyperparameter Sensitivity: How sensitive is VBPC to hyperparameters such as the number of pseudo-coresets, coreset size, and model initialization? Do suboptimal hyperparameter settings lead to significant performance degradation?\n\n3. Computational Costs and Training Time: Can you provide a detailed comparison of training times and wall-clock time between VBPC and other methods, particularly Bayesian Pseudo-Coreset methods using SGMCMC? How does the computational time scale with increasing dataset size?\n\n4. Limitations of the Last-Layer Approximation: Does the last-layer variational approximation hold up in deeper networks or more complex tasks such as regression? Have you observed any failure cases where this approximation does not capture enough uncertainty?\n\n5. Interpretability of Pseudo-Coresets: What do the learned pseudo-coresets represent? Are they capturing key features of the original dataset, and if so, how do they evolve during training? Is there a way to interpret or visualize the coreset to provide better intuition about what is being distilled?\n\n6. Generalization to Different Tasks: Can the VBPC method be applied effectively to tasks beyond classification, such as regression or other types of Bayesian inference? If so, how does the method adapt to these different problem types?\n\n7. Robustness to Out-of-Distribution (OOD) Data and Adversarial Attacks: Does VBPC provide any robustness to adversarial attacks or strong distribution shifts beyond what is demonstrated with CIFAR10-C? How does the method perform in more severe OOD scenarios?\n\n8. Memory-Efficient Loss Computation: How significant is the impact of memory-efficient loss computation during training in terms of accuracy or stability? Does it introduce any trade-offs in performance, particularly in very high-dimensional settings?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The paper introduces a novel approach to improving the efficiency of Bayesian Neural Networks (BNNs) by combining variational inference with pseudo-coresets. This innovation is notable because it addresses the longstanding challenges in the field related to computational and memory inefficiencies when scaling BNNs to large datasets. The proposal of a closed-form solution for last-layer variational inference is a significant departure from prior methods that relied on more memory-intensive sampling-based approaches. By focusing on variational inference and pseudo-coresets, the authors provide an original contribution that builds on existing methods but removes key limitations such as memory usage and the reliance on gradient stopping.\nThe technical rigor of the paper is high, with well-founded theoretical development and comprehensive empirical evaluations. The authors derive closed-form solutions for coreset variational inference, addressing a critical computational bottleneck in Bayesian model averaging. Their empirical results, demonstrated across multiple datasets (e.g., MNIST, CIFAR10, CIFAR100), show significant improvements over existing BPC methods in terms of both accuracy and negative log-likelihood, which strengthens the quality of their proposed method. The paper also includes a variety of comparisons with competitive baselines, reinforcing the robustness and effectiveness of their approach.\nThe paper is clearly structured and provides sufficient background for readers to understand both the motivation and the details of the proposed method. The explanation of the problem, the limitations of prior work, and the step-by-step presentation of the VBPC approach are clear and easy to follow. The use of mathematical derivations is well-supported by intuitive explanations, making the complex variational inference approach more accessible. The inclusion of visual results and performance tables also contributes to clarity, helping readers visualize the practical benefits of VBPC.\nThe significance of the work lies in its potential to influence how BNNs are applied to large-scale data problems. By significantly reducing the memory and computational burdens of Bayesian model averaging, the proposed VBPC method makes BNNs more feasible for real-world applications, such as those in healthcare and climate analysis, where uncertainty estimation is critical. The method could have broad implications for other fields requiring scalable, probabilistic neural networks. Additionally, the ability to perform Bayesian inference with less computational overhead enhances the practicality of deploying BNNs in production environments.", "weaknesses": "The paper focuses heavily on benchmark datasets like MNIST, CIFAR10, and CIFAR100, which are common in academic research but may not fully represent the complexity of real-world problems. While these datasets help establish baseline performance, the paper would benefit from exploring more challenging, domain-specific datasets, particularly those that are more representative of practical applications in fields such as healthcare or finance. Expanding the evaluation to datasets that feature more variability and noise could demonstrate the method’s robustness in real-world settings, which is especially important given the paper’s goal of making Bayesian Neural Networks more feasible for large-scale applications.\nΑlthough the paper demonstrates memory efficiency improvements, there is no extensive discussion of the scalability of the method when applied to very large datasets beyond those tested (e.g., ImageNet or even larger datasets in natural language processing). The paper could benefit from a more detailed analysis of the method’s behavior as the dataset size grows significantly. Additionally, the paper does not provide enough insight into the sensitivity of the method to hyperparameter choices such as the coreset size or the initialization of the model pool. It would be helpful to include an ablation study or sensitivity analysis that investigates how performance degrades with suboptimal hyperparameter choices and whether the method requires careful tuning to achieve competitive results.\nWhile the paper emphasizes memory savings, it does not provide a thorough comparison of training times between VBPC and existing methods, particularly in scenarios with high-dimensional datasets. A more detailed analysis of wall-clock time or computational complexity across different hardware configurations (e.g., GPUs versus CPUs) would be useful. This would help practitioners better understand the trade-offs between memory savings and potential increases in computational time, especially when scaling to larger architectures and datasets.\nThe paper relies on the last-layer variational approximation to simplify the posterior calculation, but the limitations of this approach are not thoroughly discussed. While the paper suggests that this approximation performs comparably to more complex methods, it would be valuable to include a deeper investigation of when this approximation might fail, especially in models with deep architectures or tasks requiring fine-grained uncertainty estimation. A discussion on whether the approximation is sufficient in all use cases or only certain tasks (e.g., classification versus regression) would make the paper more transparent.\nThe paper demonstrates that VBPC can learn pseudo-coresets that effectively approximate the full dataset’s posterior, but it doesn’t provide much insight into the interpretability of these pseudo-coresets. For example, what are the learned coresets capturing in terms of dataset distribution or feature representation? A qualitative analysis, such as visualizing the pseudo-coresets or interpreting what aspects of the data they retain, would help reinforce the method’s effectiveness. Additionally, further explanation of how these pseudo-coresets evolve during training and contribute to Bayesian uncertainty could strengthen the narrative.\nWhile the paper briefly touches on robustness to distributional shifts using CIFAR10-C, the evaluation of predictive uncertainty in real-world settings is somewhat lacking. It would be useful to see how VBPC handles more complex out-of-distribution (OOD) detection tasks or how well it captures uncertainty under adversarial conditions, which are critical aspects of Bayesian inference in high-stakes applications like healthcare. A more thorough evaluation in these contexts could elevate the practical relevance of the method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)