SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Balhae  $\mathbf{Kim}^1$  ,Hyungi Lee',Juho Lee1,2 KAIST AI,AITRICS2
{balhaekim,lhk2708,juholee}@kaist.ac.kr
A Bayesian pseudocoreset is a compact synthetic dataset summarizing essential information of a large-scale dataset and thus can be used as a proxy dataset for scalable Bayesian inference. Typically, a Bayesian pseudocoreset is constructed by minimizing a divergence measure between the posterior conditioning on the pseudocoreset and the posterior conditioning on the full dataset. However, evaluating the divergence can be challenging, particularly for the models like deep neural networks having high-dimensional parameters. In this paper, we propose a novel Bayesian pseudocoreset construction method that operates on a function space. Unlike previous methods, which construct and match the coreset and full data posteriors in the space of model parameters (weights), our method constructs variational approximations to the coreset posterior on a function space and matches it to the full data posterior in the function space. By working directly on the function space, our method could bypass several challenges that may arise when working on a weight space, including limited scalability and multi-modality issue. Through various experiments, we demonstrate that the Bayesian pseudocoresets constructed from our method enjoys enhanced uncertainty quantification and better robustness across various model architectures.
Deep learning has achieved tremendous success, but its requirement for large amounts of data makes it often inefficient or infeasible in terms of resources and computation. To enable continuous learning like humans, it is necessary to learn from a large number of data points in a continuous manner, which requires the ability to discern and retain important information. This motivates the learning of a coreset, a small dataset that is informative enough to represent a large dataset.
On the other hand, the ability to incorporate uncertainties into predictions is essential for real-world applications, as it contributes to the safety and reliability of a model. One approach to achieving this is by adopting a Bayesian framework, where a prior distribution is established to represent our initial belief about the models. This belief is then updated through inference of posterior distributions based on the acquired knowledge. Although this approach shows promise, scalability becomes a concern when working with large-scale datasets during Bayesian inference. To address this issue, a potential solution is to employ a Bayesian coreset. A Bayesian coreset is a small subset of the original dataset where the posterior conditioning on it closely approximates the original posterior conditioning on the full dataset. Once the Bayesian coreset is trained, it can be utilized as a lightweight proxy dataset for subsequent Bayesian inference or as a replay buffer for continual learning or transfer learning.
A Bayesian coreset is constructed by selecting a subset from a large dataset. However, recent research suggests that this approach may not be effective, especially in high-dimensional settings [19]. Instead, an alternative method of synthesizing a coreset, wherein the coreset is learned as trainable parameters, has been found to significantly enhance the quality of the approximation. This synthesized coreset is
referred to as a Bayesian pseudocoreset. The process of learning a Bayesian pseudocoreset involves minimizing a divergence measure between the posterior of the full dataset and the posterior of the pseudocoreset. However, learning Bayesian pseudocoresets is generally challenging due to the intractability of constructing both the full dataset posterior and the pseudocoreset posterior, as well as the computation of the divergence between them, which necessitates approximation. Consequently, existing works on Bayesian pseudocoresets have primarily focused on small-scale problems [19, 8, 20, 21]. Recently, Kim et al. [15] introduced a scalable method for constructing Bayesian pseudocoresets using variational Gaussian approximation for the posteriors and minimizing forward KL divergence. Although their method shows promise, it still demands substantial computational resources for high-dimensional models like deep neural networks.
In this paper, we present a novel approach to enhance the scalability of Bayesian pseudocoreset construction, particularly for Bayesian neural networks (BNNs) with a large number of parameters. Our proposed method operates in function space. When working with BNNs, it is common to define a prior distribution on the weight space and infer the corresponding weight posterior distribution, which also applies to Bayesian pseudocoreset construction. However, previous studies [31, 28] have highlighted the challenge of interpreting weights in high-dimensional neural networks, making it difficult to elicit meaningful prior distributions. Additionally, in high-dimensional networks, the loss surfaces often exhibit a complex multimodal structure, which means that proximity in the weight space does not necessarily imply proximity in the desired prediction variable [24, 30]. This same argument can be applied to Bayesian pseudocoreset construction, as matching the full data and pseudocoreset posteriors in the weight space may not result in an optimal pseudocoreset in terms of representation power and computational scalability.
To be more specific, our method constructs a variational approximation to the pseudocoreset posteriors in function space by linearization and variational approximation to the true posterior. Then we learn Bayesian pseudocoreset by minimizing a divergence measure between the full data posterior and the pseudocoreset posterior in the function space. Compared to the previous weight space approaches, our method readily scales to the large models for which the weight space approaches were not able to compute. Another benefit of function space matching is that it does not constrain the architectures of the neural networks to be matched, provided that their inherited function space posteriors are likely to be similar. So for instance, we can train with multiple neural network architectures simultaneously with varying numbers of neurons or types of normalization layers, and we empirically observe that this improves the architectural robustness of the learned pseudocoresets. Moreover, it has another advantage that the posteriors learned from the Bayesian pseudocoreset in function space have better out-of-distribution (OOD) robustness, similar to the previous reports showing the benefit of function space approaches in OOD robustness [28].
In summary, this paper presents a novel approach to creating a scalable and effective Bayesian pseudocoreset using function space variational inference. The resulting Bayesian pseudocoreset is capable of being generated in high-dimensional image and deep neural network settings and has better uncertainty quantification abilities compared to weight space variational inference. Additionally, it has better architectural robustness. We demonstrate the efficiency of the function space Bayesian pseudocoreset through the various experiments.
2.1 Bayesian pseudocoresets
In this paper, we focus on probabilistic models for supervised learning problem. Let  $\theta \in \Theta$  be a parameter, and let  $p(y\mid x,\theta)$  be a probabilistic model indexed by the parameter  $\theta$ . Given a set of observations  $\mathbf{x} \coloneqq (x_{i})_{i = 1}^{n}$  and the set of labels  $\mathbf{y} \coloneqq (y_{i})_{i = 1}^{n}$  with each  $x_{i} \in \mathcal{X}$  and  $y_{i} \in \mathcal{Y}$ , we are interested in updating our prior belief  $\pi_0(\theta)$  about the parameter to the posterior,
$$
\pi_ {\mathbf {x}} (\theta) := \frac {\pi_ {0} (\theta)}{Z (\mathbf {y} \mid \mathbf {x})} \prod_ {i = 1} ^ {n} p \left(y _ {i} \mid x _ {i}, \theta\right), \quad Z (\mathbf {y} \mid \mathbf {x}) := \int_ {\Theta} \prod_ {i = 1} ^ {n} p \left(y _ {i} \mid x _ {i}, \theta\right) \pi_ {0} (\mathrm {d} \theta). \tag {1}
$$
However, when the size of the dataset  $n$  is large, the computation of the posterior distribution can be computationally expensive and infeasible. To overcome this issue, Bayesian pseudocoresets are constructed as a synthetic dataset  $\mathbf{u} = (u_j)_{j=1}^m$  with  $m \ll n$  with the set of labels  $\tilde{\mathbf{y}} := (\tilde{y}_j)_{j=1}^m$ .
where the posterior conditioning on it approximates the original posterior  $\pi_{\mathbf{x}}(\theta)$ .
$$
\pi_ {\mathbf {u}} (\theta) = \frac {\pi_ {0} (\theta)}{Z (\tilde {\mathbf {y}} \mid \mathbf {u})} \prod_ {j = 1} ^ {m} p \left(\tilde {y} _ {j} \mid u _ {j}, \theta\right), \quad Z (\tilde {\mathbf {y}} \mid \mathbf {u}) := \int_ {\Theta} \prod_ {j = 1} ^ {m} p \left(\tilde {y} _ {j} \mid u _ {j}, \theta\right) \pi_ {0} (\mathrm {d} \theta). \tag {2}
$$
This approximation is made possible by solving an optimization problem that minimizes a divergence measure  $D$  between the two posterior distributions
$$
\mathbf {u} ^ {*} = \underset {\mathbf {u}} {\arg \min } D \left(\pi_ {\mathbf {x}}, \pi_ {\mathbf {u}}\right). \tag {3}
$$
In a recent paper [15], three variants of Bayesian pseudocoresets were proposed using different divergence measures, namely reverse Kullback-Leibler divergence, Wasserstein distance, and forward Kullback-Leibler divergence. However, performing both the approximation and the optimization in the parameter space can be computationally challenging, particularly for high-dimensional models such as deep neural networks.
2.2 Bayesian pseudocoresets in weight-space
Kim et al. [15] advocates using forward KL divergence as the divergence measure when constructing Bayesian pseudocoresets, with the aim of achieving a more even exploration of the posterior distribution of the full dataset when performing uncertainty quantification with the learned pseudocoreset. The forward KL objective is computed as,
$$
\begin{array}{l} D _ {\mathrm {K L}} \left[ \pi_ {\mathbf {x}} \| \pi_ {\mathbf {u}} \right] = \log Z (\tilde {\mathbf {y}} | \mathbf {u}) - \log Z (\mathbf {y} | \mathbf {x}) \\ + \mathbb {E} _ {\pi_ {\mathbf {x}}} \left[ \sum_ {i = 1} ^ {n} \log p \left(y _ {i} \mid x _ {i}, \theta\right) \right] - \mathbb {E} _ {\pi_ {\mathbf {x}}} \left[ \sum_ {j = 1} ^ {m} \log p \left(\tilde {y} _ {j} \mid u _ {j}, \theta\right) \right]. \tag {4} \\ \end{array}
$$
The derivative of the divergence with respect to the pseudocoreset  $\mathbf{u}$  is computed as
$$
\nabla_ {\mathbf {u}} D _ {\mathrm {K L}} \left[ \pi_ {\mathbf {x}} \| \pi_ {\mathbf {u}} \right] = \mathbb {E} _ {\pi_ {\mathbf {u}}} \left[ \nabla_ {\mathbf {u}} \sum_ {j = 1} ^ {m} \log p \left(\tilde {y} _ {j} \mid u _ {j}, \theta\right) \right] - \nabla_ {\mathbf {u}} \mathbb {E} _ {\pi_ {\mathbf {x}}} \left[ \sum_ {j = 1} ^ {m} \log p \left(\tilde {y} _ {j} \mid u _ {j}, \theta\right) \right] \tag {5}
$$
For the gradient, we need the expected gradients of the log posteriors that require sampling from the posteriors  $\pi_{\mathbf{x}}$  and  $\pi_{\mathbf{u}}$ . Most of the probabilistic models do not admit simple closed-form expressions for these posteriors, and it is not easy to simulate those posteriors for high-dimensional models. To address this, Kim et al. [15] proposes to use a Gaussian variational distributions  $q_{\mathbf{u}}(\theta)$  and  $q_{\mathbf{x}}(\theta)$  to approximate  $\pi_{\mathbf{x}}$  and  $\pi_{\mathbf{u}}$  whose means are set to the parameters obtained from the SGD trajectories,
$$
q _ {\mathbf {u}} (\theta) = \mathcal {N} \left(\theta ; \mu_ {\mathbf {u}}, \Sigma_ {\mathbf {u}}\right), \quad q _ {\mathbf {x}} (\theta) = \mathcal {N} \left(\theta ; \mu_ {\mathbf {x}}, \Sigma_ {\mathbf {x}}\right), \tag {6}
$$
where  $\mu_{\mathbf{u}}$  and  $\mu_{\mathbf{x}}$  are the maximum a posteriori (MAP) solutions computed for the dataset  $\mathbf{u}$  and  $\mathbf{x}$ , respectively.  $\Sigma_{\mathbf{u}}$  and  $\Sigma_{\mathbf{x}}$  are covariances. The gradient, with the stop gradient applied to  $\mu_{\mathbf{u}}$ , is approximated as,
$$
\frac {\nabla_ {\mathbf {u}}}{S} \sum_ {s = 1} ^ {S} \left(\sum_ {j = 1} ^ {m} \log p \left(\tilde {y} _ {j} \mid u _ {j}, \mathbf {s g} \left(\mu_ {\mathbf {u}}\right) + \Sigma_ {\mathbf {u}} ^ {1 / 2} \varepsilon_ {\mathbf {u}} ^ {(s)}\right) - \sum_ {j = 1} ^ {m} \log p \left(\tilde {y} _ {j} \mid u _ {j}, \mu_ {\mathbf {x}} + \Sigma_ {\mathbf {x}} ^ {1 / 2} \varepsilon_ {\mathbf {x}} ^ {(s)}\right)\right). \tag {7}
$$
Here,  $\varepsilon_{\mathbf{u}}^{(s)}$  and  $\varepsilon_{\mathbf{x}}^{(s)}$  are i.i.d. standard Gaussian noises and  $S$  is the number of Monte-Carlo samples.
Expert trajectories Approximating the full data and coreset posteriors with variational distributions as specified above requires  $\mu_{\mathbf{u}}$  and  $\mu_{\mathbf{x}}$  as consequences of running optimization algorithms until convergence. While this may be feasible for small datasets, for large-scale setting of our interest, obtaining  $\mu_{\mathbf{u}}$  and  $\mu_{\mathbf{x}}$  from scratch at each iteration for updating  $\mathbf{u}$  can be time-consuming. To alleviate this, in the dataset distillation literature, Cazenavette et al. [7] proposed to use the expert trajectories, the set of pretrained optimization trajectories constructed in advance to the coreset learning. Kim et al. [15] brought this idea to Bayesian pseudocoresets, where a pool of pretrained trajectories are assume to be given before pseudocoreset learning. At each step of pseudocoreset update, a checkpoint  $\theta_0$  from an expert trajectory is randomly drawn from the pool, and then  $\mu_{\mathbf{u}}$  and  $\mu_{\mathbf{x}}$  are quickly constructed by taking few optimization steps from  $\theta_0$ .
In principle, we should learn the (pseudo)labels  $\tilde{\mathbf{y}}$  as well, but for classification problem, we can simply fix it as a constant set containing equal proportion of all possible classes. We assume this throughout the paper.
3 Function space Bayesian pseudocoreset
3.1 Function space Bayesian neural networks
We follow the framework presented in Rudner et al. [29, 28] to define a Function-space Bayesian Neural Network (FBNN). Let  $\pi_0(\theta)$  be a prior distribution on the parameter and  $g_{\theta}:\mathcal{X}\to \mathbb{R}^{d}$  be a neural network index by  $\theta$ . Let  $h:\Theta \rightarrow (\mathcal{X}\rightarrow \mathbb{R}^d)$  be a deterministic mapping from a parameter  $\theta$  to a neural network  $g_{\theta}$ . Then a function-space prior is simply defined as a pushforward  $\nu_{0}(f) = h_{*}\pi_{0}(f)\coloneqq \pi_{0}(h^{-1}(f))$ . The corresponding posterior is also defined as a pushforward  $\nu_{\mathbf{x}}(f) = h_{*}\pi_{\mathbf{x}}(f)$  and so is the pseudocoreset posterior  $\nu_{\mathbf{u}}(f) = h_{*}\pi_{\mathbf{u}}(f)$ .
3.2 Learning function space Bayesian pseudocoresets
Given the function space priors and posteriors, a Function space Bayesian PseudoCoreset (FBPC) is obtained by minimizing a divergence measure between the function space posteriors,
$$
\mathbf {u} ^ {*} = \underset {\mathbf {u}} {\arg \min } D \left(\nu_ {\mathbf {x}}, \nu_ {\mathbf {u}}\right). \tag {8}
$$
We follow Kim et al. [15] suggesting to use the forward KL divergence, so our goal is to solve
$$
\mathbf {u} ^ {*} = \underset {\mathbf {u}} {\arg \min } D _ {\mathrm {K L}} \left[ \nu_ {\mathbf {x}} \| \nu_ {\mathbf {u}} \right]. \tag {9}
$$
The following proposition provides an expression for the gradient to minimize the divergence, whose proof is given Appendix A.
Proposition 3.1. The gradient of the forward KL divergence with respect to the coreset  $\mathbf{u}$  is
$$
\nabla_ {\mathbf {u}} D _ {\mathrm {K L}} \left[ \nu_ {\mathbf {x}} \| \nu_ {\mathbf {u}} \right] = - \nabla_ {\mathbf {u}} \mathbb {E} _ {[ \nu_ {\mathbf {x}} ] _ {\mathbf {u}}} \left[ \log p \left(\tilde {\mathbf {y}} \mid \mathbf {f} _ {\mathbf {u}}\right) \right] + \mathbb {E} _ {[ \nu_ {\mathbf {u}} ] _ {\mathbf {u}}} \left[ \nabla_ {\mathbf {u}} \log p \left(\tilde {\mathbf {y}} \mid \mathbf {f} _ {\mathbf {u}}\right) \right], \tag {10}
$$
where  $[\nu_{\mathbf{x}}]_{\mathbf{u}}$  and  $[\nu_{\mathbf{u}}]_{\mathbf{u}}$  are finite-dimensional distributions of the stochastic processes  $\nu_{\mathbf{x}}$  and  $\nu_{\mathbf{u}}$ , respectively,  $\mathbf{f}_{\mathbf{u}} \coloneqq (f(u_j))_{j=1}^{m}$ , and  $p(\tilde{\mathbf{y}} \mid \mathbf{f}_{\mathbf{u}}) = \prod_{j=1}^{m} p(\tilde{y}_j \mid f(u_j))$ .
To evaluate the gradient Eq. 10, we should identify the finite-dimensional functional posterior distributions  $[\nu_{\mathbf{x}}]_{\mathbf{u}}$  and  $[\nu_{\mathbf{u}}]_{\mathbf{u}}$ . While this is generally intractable, as proposed in Rudner et al. [29, 28], we can instead consider a linearized approximation of the neural network  $g_{\theta}$ ,
$$
\tilde {g} _ {\theta} (\cdot) = g _ {\mu_ {\mathbf {x}}} (\cdot) + \mathcal {J} _ {\mu_ {\mathbf {x}}} (\cdot) (\theta - \mu_ {\mathbf {x}}), \tag {11}
$$
where  $\mu_{\mathbf{x}} = \mathbb{E}_{\pi_{\mathbf{x}}}[\theta]$  and  $\mathcal{J}_{\mu_{\mathbf{x}}}(\cdot)$  is the Jacobian of  $g_{\theta}$  evaluated at  $\mu_{\mathbf{x}}$ . Then we approximate the function space posterior  $\nu_{\mathbf{x}}$  with  $\tilde{\nu}_{\mathbf{x}} := \tilde{h}_{*}\pi_{\mathbf{x}}$  where  $\tilde{h}(\theta) = \tilde{g}_{\theta}$ , and as shown in Rudner et al. [29, 28], the finite dimensional distribution  $[\tilde{\nu}_{\mathbf{x}}]_{\mathbf{u}}$  is a multivariate Gaussian distribution,
$$
\left[ \tilde {\nu} _ {\mathbf {x}} \right] _ {\mathbf {u}} \left(\mathbf {f} _ {\mathbf {u}}\right) = \mathcal {N} \left(\mathbf {f} _ {\mathbf {u}} \mid g _ {\mu_ {\mathbf {x}}} (\mathbf {u}), \mathcal {J} _ {\mu_ {\mathbf {x}}} (\mathbf {u}) \Sigma_ {\mathbf {x}} \mathcal {J} _ {\mu_ {\mathbf {x}}} (\mathbf {u}) ^ {\top}\right), \tag {12}
$$
with  $\Sigma_{\mathbf{x}} = \mathrm{Cov}_{\pi_{\mathbf{x}}}(\theta)$ . Similarly, we obtain
$$
\left[ \tilde {\nu} _ {\mathbf {u}} \right] _ {\mathbf {u}} \left(\mathbf {f} _ {\mathbf {u}}\right) = \mathcal {N} \left(\mathbf {f} _ {\mathbf {u}} \mid g _ {\mu_ {\mathbf {u}}} (\mathbf {u}), \mathcal {J} _ {\mu_ {\mathbf {u}}} (\mathbf {u}) \Sigma_ {\mathbf {u}} \mathcal {J} _ {\mu_ {\mathbf {u}}} (\mathbf {u}) ^ {\top}\right), \tag {13}
$$
with  $\mu_{\mathbf{u}}\coloneqq \mathbb{E}_{\pi_{\mathbf{u}}}[\theta ]$  and  $\Sigma_{\mathbf{u}}\coloneqq \mathrm{Cov}_{\pi_{\mathbf{u}}}(\theta)$ . Using these linearized finite-dimensional distribution, we can approximate
$$
\nabla_ {\mathbf {u}} D _ {\mathrm {K L}} \left[ \nu_ {\mathbf {x}} \| \nu_ {\mathbf {u}} \right] = - \nabla_ {\mathbf {u}} \mathbb {E} _ {[ \tilde {\nu} _ {\mathbf {x}} ] _ {\mathbf {u}}} \left[ \log p (\tilde {\mathbf {y}} \mid \mathbf {f} _ {\mathbf {u}}) \right] + \mathbb {E} _ {[ \tilde {\nu} _ {\mathbf {u}} ] _ {\mathbf {u}}} \left[ \nabla_ {\mathbf {u}} \log p (\tilde {\mathbf {y}} \mid \mathbf {f} _ {\mathbf {u}}) \right], \tag {14}
$$
3.3 Tractable approximation to the gradient
Even with the linearization, evaluating Eq. 14 is still challenging because it requires obtaining  $\mu_{\mathbf{x}}$  and  $\Sigma_{\mathbf{x}}$  which are the statistics of the weight-space posterior  $\pi_{\mathbf{x}}$ . Rudner et al. [28] proposes to learn a variational approximation  $q_{\mathbf{x}}(\theta)$  in the weight-space, and use the linearized pushforward of the variational distribution  $\tilde{h}_{*}q_{\mathbf{x}}$  as a proxy to the function space posterior. Still, this approach requires computing the heavy Jacobian matrix  $\mathcal{I}_{\mathbb{E}_{q_{\mathbf{x}}}[\theta]}(\mathbf{u})$ , so may not be feasible for our scenario where we have to compute such variational approximations at each update of the pseudocoreset  $\mathbf{u}$ .
Figure 1: The conceptual overview of our proposed training procedure.
Instead, we choose to directly construct a variational approximations to the finite-dimensional distributions of the function space posteriors, that is,
$$
\left[ \tilde {\nu} _ {\mathbf {x}} \right] _ {\mathbf {u}} \left(\mathbf {f} _ {\mathbf {u}}\right) \approx q _ {\mathbf {x}} \left(\mathbf {f} _ {\mathbf {u}}\right) = \mathcal {N} \left(\mathbf {f} _ {\mathbf {u}} \mid g _ {\hat {\mu} _ {\mathbf {x}}} (\mathbf {u}), \hat {\Psi} _ {\mathbf {x}}\right), \tag {15}
$$
$$
[ \tilde {\nu} _ {\mathbf {u}} ] _ {\mathbf {u}} (\mathbf {f} _ {\mathbf {u}}) \approx q _ {\mathbf {u}} (\mathbf {f} _ {\mathbf {u}}) = \mathcal {N} (\mathbf {f} _ {\mathbf {u}} | g _ {\hat {\mu} _ {\mathbf {u}}} (\mathbf {u}), \hat {\Psi} _ {\mathbf {u}}),
$$
where  $(\hat{\mu}_{\mathbf{x}},\hat{\Psi}_{\mathbf{x}})$  and  $(\hat{\mu}_{\mathbf{u}},\hat{\Psi}_{\mathbf{u}})$  are variational parameters for the full data and coreset posteriors. Inspired by Kim et al. [15], we construct the variational parameters using expert trajectories. Unlike [15], we simply let the MAP solution computed for  $\mathbf{x}$ ,  $\theta_{\mathbf{x}}$ , by sampling a checkpoint from the later part of the expert trajectories, and obtain the MAP solution of  $\mathbf{u}$ ,  $\theta_{\mathbf{u}}$ , by directly optimizing an initial random parameter. Then we obtain  $\hat{\mu}_{\mathbf{x}}$  and  $\hat{\mu}_{\mathbf{u}}$  using  $\mathbf{u}$ . For the covariance matrices  $\hat{\Psi}_{\mathbf{x}}$  and  $\hat{\Psi}_{\mathbf{u}}$  while Kim et al. [15] proposed to use spherical Gaussian noises, we instead set them as an empirical covariance matrices of the samples collected from the optimization trajectory. Specifically, we take additional  $K$  steps from each MAP solution to compute the empirical covariance.
$$
\theta_ {\mathbf {x}} ^ {(0)} = \theta_ {\mathbf {x}}, \quad \theta_ {\mathbf {x}} ^ {(t)} = \mathrm {o p t} (\theta_ {\mathbf {x}} ^ {(t - 1)}, (\mathbf {x}, \mathbf {y})), \quad \hat {\mu} _ {\mathbf {x}} = g _ {\mathsf {s g} (\theta_ {\mathbf {x}} ^ {(0)})} (\mathbf {u}),
$$
$$
\hat {\Psi} _ {\mathbf {x}} := \operatorname {s g} \left(\operatorname {d i a g} \left(\frac {1}{K} \sum_ {k = 1} ^ {K} g _ {\theta_ {\mathbf {x}} ^ {(k)}} ^ {2} (\mathbf {u}) - \left(\frac {1}{K} \sum_ {k = 1} ^ {K} g _ {\theta_ {\mathbf {x}} ^ {(k)}} (\mathbf {u})\right) ^ {2}\right)\right), \tag {16}
$$
where  $\mathrm{opt}(\theta, \mathbf{x})$  is a step of SGD optimization applied to  $\theta$  with data  $\mathbf{x}$  and the squares in the  $\mathrm{diag}(\cdot)$  are applied in element-wise manner. Note also that we are applying the stop-gradient operations for to block the gradient flow that might lead to complication in the backpropagation procedure. The variational parameters  $(\hat{\mu}_{\mathbf{u}}, \hat{\Psi}_{\mathbf{u}})$  are constructed in a similar fashion, but using the pseudocoreset  $(\mathbf{u}, \tilde{\mathbf{y}})$  instead of the original data  $(\mathbf{x}, \mathbf{y})$ . It is noteworthy that our approach is similar to one of the methods in Bayesian learning, SWAG [18]. However, while SWAG focuses on collecting statistics on weight space trajectories, our method constructs statistics in function spaces. This distinction makes our approach more suitable and scalable for pseudocoreset construction. The overview of proposed method is provided in Fig. 1.
With the variational approximations constructed as described, we obtain a Monte-Carlo estimator of Eq. 14,