Function-space Parameterization of Neural Networks for Sequential Learning
Authors: Aidan Scannell
Riccardo Mereu
Paul Edmund Chang
Ella Tamir
Joni Pajarinen
Arno Solin
Venue: ICLR 2024 poster (2024)
Keywords: Neural networks
Bayesian deep learning
deep learning
Gaussian processes
Laplace approximation
sequential learning
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: Converting neural networks to sparse function-space representation for sequential learning

Abstract: Sequential learning paradigms pose challenges for gradient-based deep learning due to difficulties incorporating new data and retaining prior knowledge. While Gaussian processes elegantly tackle these problems, they struggle with scalability and handling rich inputs, such as images. To address these issues, we introduce a technique that converts neural networks from weight space to function space, through a dual parameterization. Our parameterization offers: (*i*) a way to scale function-space methods to large data sets via sparsification, (*ii*) retention of prior knowledge when access to past data is limited, and (*iii*) a mechanism to incorporate new data without retraining. Our experiments demonstrate that we can retain knowledge in continual learning and incorporate new data efficiently. We further show its strengths in uncertainty quantification and guiding exploration in model-based RL. Further information and code is available on the project website.

Published as a conference paper at ICLR 2024
FUNCTION-SPACE PARAMETERIZATION OF NEURAL NETWORKS FOR SEQUENTIAL LEARNING
Aidan Scannell; Riccardo Mereu, Paul Chang, Ella Tamir, Joni Pajarinen & Arno Solin Aalto University, Espoo, Finland
(aidan. scannell, caameell,,icccc riccardo. mereu @aalto. fi
ABSTRACT
Sequential learning paradigms pose challenges for gradient-based deep learning due to difficulties incorporating new data and retaining prior knowledge, While Gaussian processes elegantly tackle these problems, they struggle with scalability and handling rich inputs, such as images. To address these issues, we introduce a technique that converts neural networks from weight space to function space, through a dual parameterization. Our parameterization offers: (i) a way to scale function-space space methods to large data sets via sparsification, (ii) retention of prior knowledge when access to past data is limited, and (iii) mechanism to incorporate new data without retraining. Our experiments demonstrate that we can retain knowledge in continual learning and incorporate new data efficiently. We further show its strengths in uncertainty quantification and guiding exploration model- based RL. Further information and code is available on the project website
INTRODUCTION
Deep learning (Goodfellow al., 2016) has become the cornerstone of contemporary artificial intelligence, proving remarkably effective with large-scale and complex data, such as images. On the other hand, Gaussian processes (GPs, Rasmussen & Williams, 2006), although limited simpler data, offer functionalities beneficial in sequential learning paradigms, such as continual learning (CL, Parisi al., 2019; De Lange al., 2021), reinforcement learning (RL. Sutton & Barto, 2018; Deisenroth Rasmussen, 2011), and Bayesian optimization (BO, Garnett, 2023). In CL, the challenge is preventing forgetting over the lifelong learning horizon when access to previous data is lost (MeCloskey & Cohen, 1989). Notably, GP's function space provides more effective representation for CL than an NN's weight space. RL and BO, GPs provide principled uncertainty estimates that balance the exploration-exxpliin
trade-off. Furthermore, GPs do not require retraining weights from scratch when incorporating new data.
While GPs could offer several advantages over NNs for sequential learning, they struggle to scale to large and complex data sets, which are common the real world. In essence, NNs and GPs have complementary strengths that aim to combine. Previous approaches for converting trained NNs to GPs (Khan et al., 2019; Immer et al., 2021b) show limited applicability to real-world sequential learning problems they (i) rely subset approximations which do not scale large data sets and (ii) can only incorporate new data by retraining the NN from scratch.
In this paper, we establish connection between trained NNs and dual parameterization of GPs (Csató& & Opper, 2002; Adam et al., 2021; Chang et al., 2023), which is favourable for sequential learning. In contrast to previous work that utilizes subsets of training data (Immer et al., 2021a), our dual parameterization allows us to sparsify the representation whilst capturing the contributions from all data points, essential for predictive uncertainty. We refer to our method Sparse Function-space Representation (SFR)-a) sparse GP derived from trained NN. A preliminary version of SFR was presented Scannell et al. (2023). We show how SFR's dual parameterization (i) maintains representation of previous data in CL, (ii) incorporates new data without needing retrain the NN (see Fig. 1), and (iii) balances the exploration-exploiie exploitation trade-off in sequential decision decision-making
"Equal contribution. hattps://aaleee http://aaalmmm..iittttt. github. github.io/sfr

---

Published as a conference paper at ICLR 2024
Data
Mean 95% interval
Data
Neural net output
Mcan
interval
Outgut
SFK
NEW DATA
Sparse inducing points
Spurse inducing points
2
Taput,
Input
Topur,
Figure 1: Regression example with an MLP: Left: Predictions from the trained neural network. Middle: Our approach summarizes all the training data at the inducing points. The model captures the predictive mean and uncertainty, and (right) incorporates new data without retraining the model.
(RL), via predictive uncertainty. Importantly, our results show that SFR scales to data sets with rich inputs (e.g., images) and millions of data points.
Related work Probabilistic methods in deep learning (Neal, 1995; Wilson, 2019) have recently gained increasing attention in the machine learning community as a means for uncertainty quantification. Calculating the posterior distribution of Bayesian neural network (BNN) is usually intractable. It calls for approximate inference techniques, such as variational inference (Blei et al., 2017), deep ensembles (Takshminarayanan et al., 2017) and MC dropout (Gal & Ghahramani, 2016) each with its own strengths and weaknesses. A common approach for uncertainty quantification in NNs is the Laplace-GGi GGN approximation (Daxberger et al., 2021), which takes a trained NN and linearizes it around the optimal weights. The Hessian is approximated using the generalized Gauss -ewtoon approximation (GGN, Botev et al., 2017). The resulting linear model, with respect to the weights, can be used to obtain uncertainty estimates and refine the NN predictions (Immer et al., 2021a). Immer et al. (2021a) extend Khan et al. (2019) to show the GP connection of the Laplace approximation. However, it does not scale well as it cannot incorporate all data points for uncertainty estimates. Recent work from Ortega al. (2023) attempts form more principled variational sparse GP but resorts to running a second optimization loop. In this paper, we show that such an optimization not necessary when taking a dual parameterization as shown in Adam et al. (2021) and how such a parameterization lends itself to sequential learning (Chang et al., 2023).
Recent methods CL, such as FRCL (Titsias et al., 2020), FROMP (Pan et al., 2020), DER (Buzzega et al., 2020), and S-FSVI (Rudner et al., 2022) (see Parisi et al., 2019; De Lange et al., 2021, for an in depth overview) rely on functional regularization methods to alleviate catastrophic forgetting. These approaches obtain state-of-the-arr performances among the objective-based techniques in several CL benchmarks compared to their weight-space counterparts: Online-EWC (Schwarz et al., 2018), and VCL (Nguyen et al., 2018). In contrast to the method presented in this paper, these functional regularization methods lack a principled way to incorporate information from all data. Instead, they rely on subset approximations (see App. B.1 for more related work on CL).
2 BACKGROUND
We consider supervised learning with inputs æ; E RDD and outputs y: E RC (e.g., regression) or yi E (0, 1 (e.g., classification), giving a data set f. : R + BC with weights E R and use a likelihood function p(y: [u (2;)) to link the function values to the output y, (e.g., categorical for classification). For notational conciseness, we stick scalar outputs yi and denote sets of inputs and outputs as X and y, respectively.
{( y:) We introduce a NN
BNNs In Bayesian deep learning, we place a prior over the weights p(w) and aim to calculate the posterior over the weights given the data p( w DD. Given the weight posterior p( w [D), BNNs make probabilistic predictions PENN(Us æ; D) Ep(w D) [p[M]] fw (=:))] The posterior p(w D) x p(y| Jw X)) p(w) is generally not available in closed form so we resort to approximations.
MAP It is common train NN weights to minimize the (regularized) empirical risk,
arg min "
((D,w) = arg min ΣN
+ R(w)
(I)
log p( ((D,w))
log plu: (mx)))
log p(w)
2

---

Published as a conference paper at ICLR 2024
where the objective L(D, w) corresponds to the negative log-joint distribution - p(p w) as the loss f( (z)), 3/) can be interpreted as a negative log-likelihood - log p f (z:)) and the
regularizer R(w) corresponds to a negative log-prior log (w). For example, weight decay
regularizer R(w) 2 w corresponds to Gaussian prior p(w) N(w 0,,8-11..0 8811))) with prior precision 6. As such, we can view Eq. (1) as the maximum a posteriori (MAP) solution.
Laplace approximation The Laplace approximation (MacKay, 1992: Daxberger et al., 2021) builds upon this objective and approximates the weight posterior around the MAP weights (w*) by setting the covariance to the Hessian of the posterior,
p(w (D = q(w) (ww w², Σ) with
V Viw log w [D) [===--]
(2)
Computing this requires calculating the Hessian of the log-likelihood from Eq. (1). The GGN approximation is used extensively ensure positive semi-definiteness,
V log ply fw (X))
GGN
Jw (x)
log(y f) and Jw [z) (3)
where f ( denotes set of function values at the training inputs. In practice, the diagonal (Khan et al., 2018; Graves, 2011) Kronecker factorization (Ritter al., 2018) of the GGN are used for scalablity. Immer et al. (2021b) highlighted that the GGN approximation corresponds to local linearization of the NN, showing predictions can be made with a generalized linear model (GLM),
PGLM (14 c; D) = (u)(w) [p(u.] I (mx)) with f (a) (x:) + (a)) w²) (4)
GPs As Gaussian distributions remain tractable under linear transformations, we can convert the linear model from weight space to function space (see Ch. 2.1 in Rasmussen & Williams, 2006). As shown in Immer et al. (2021b), the Bayesian GLM in Eq. (4) has an equivalent GP formulation,
PGp (v s, D) = B. [p(1: [+]] and ((h)) N (Jw- (ar,), ku I (K æx + A--1--1 kias)
(5) where the kernel x') J. (x) Jur (xx) is the Neural Tangent Kernel (NTK, Jacot et al., 2018), (x,) is the function output at Øi A - log p(y ff can be interpreted as per-input noise, the 27 entry of matrix Kax E NXN is x(x,, xz)), Ras is a vector where each jih element is k(x;, x;), and kü (((z;;2;).
Sparse GPs The GP formulation in Eq. (5) requires inverting an N x N matrix which has complexity O(N3). This limits its applicability to large data sets, which are common in deep learning. Sparse GPs reduce the computational complexity by representing the GP as a low-rank approximation at set of inducing inputs Z = [z1....,
E MxD with corresponding inducing variables = f (Z) (see Cilioonero-annder & Rasmussen, 2005, for an early overview). The approach by Titsias (2009) (also used in the DTC approximation) defines the marginal predictive distribution as qre (f.) = [[[[s u) q(u) du with q(u)= = N (u m, S) (as in Titsias, 2009; Hensman et al., 2013).
The sparse GP predictive posterior is
(6) where K.z and kiz are defined similarly to KE and Kni but over the inducing points Z. We have assumed a zero mean function. Note that the parameters (m and S) are usually obtained via variational inference, which requires further optimization.
[f = k K and Var = kus K [ 2=
K- kan,
The GP formulation from Immer et al. (2021b), shown in Eq. (5), struggles to scale to large data sets and it cannot incorporate new data by conditioning on it because its posterior mean is the NN fw- (x). Our method overcomes both of these limitations via a dual parameterization, which enables us to (i) sparsify the GP without further optimization, and (ii) incorporate new data without retraining.
3 SFR: SPARSE FUNCTION- REPRESENTATION OF NNS
In this section, we present our method, named SFR, which converts a trained NN into a GP (see Fig. 2 for an overview). SFR is built upon dual parameterization of the GP posterior. That is, in contrast to previous approaches, which adopt the GP formulation in Eq. (5), we use dual parameterization consisting of parameters a and B, which gives rise to the predictive posterior,
Epc [fi kT and Varp w) [fa] Kis 1TT (Kar aur + (1))(())1)-11-
3

---

Published as a conference paper at ICLR 2024
UNCERTAINTY QUANTIFICATION
INCORPORATE NEW DATA
Full GP xN
Sparse GP e.M
Punction-spas Laplace (10)
Dual param. Digs. (7) and (11)
Dual updates
Digqs (12) and (13)
bdgy (17)
Weight space Laplace bliq 12)
Computationaly infeasible
Full GP
GP subset
NN MAP weights w bEqq.
1 Function spice 1 Weight spsice
Linearize
Weight-spnce
Luplace DEy
Laplace GLM
Cosdl retraining
Figure 2: SFR overview: We linearize the trained NN around the MAP weights w² and interpret in function space, via a kernel formulation s((.-) (Eq. (9)). contrast to previous approaches, we perform Laplace approximation on the function-space space objective Eq. (10). This leads to SFR' dual parameterization, scaling to large data sets Eq. (12) and incorporating new data efficiently Eq. (17).
See App. A.1 for background. Eq. (7) states that the first two moments of the resultant posterior process (not restricted to GPs), can parameterized via the dual parameters or, 3 E R.N. defined as
Ep(s: v) V p(y f and B, -Ep(fi v) Vi log ff)l==s.]- (8)
Eq. (8) holds for generic likelihoods and involves no approximations since the expectation is under the exact posterior, given that the model can be expressed in a kernel formulation. Eq. (7) and Eq. (8) highlight that the approximate inference technique, usually viewed as posterior approximation, alternatively be interpreted an approximation of the expectation of loss (likelihood) gradients.
Linear model for dual updates Previous NN to GP methods (Khan et al., 2019; Immer et al., 202 set the GP mean be the NN's prediction (Eq. (5)). In contrast, we linearize
1 5 Ju (a) J (x)),
fa (ar) no Jwn (x) n0
u(x) = and <(x; x') =
because this leads a GP with zero mean function. Importantly, this formulation enables to incorporate new data via our dual parameterization, as we detail in Sec. 4.
Dual parameters from NN Eq. (9) gives us a way convert weight-space NN into function space. However, we go step further and convert the weight-space objective in Eq. (1) to function space, i.e. p(f). We can then approximate the function-space posterior q()) = NJ my: Sy) by applying the Laplace approximation to this function-space objective. obtain my and S from the stationary point of the objective:
C(D, w) DN i log p(s:
That is,
y log p(y: [f)l== K=lmy my and S Vi log p(u Afs=, K-1 (10)
Comparing Eq. (10) and Eq. (7), simplified the dual parameter calculation in Eq. (8), i.e. the expectation is now removed, meaning:
that approximate inference technique (Laplace) has
& = V log [ and 3, == -V; log p(u: 10)
(11)
The variables in Eq. (1I) are easily computable using the NN MAP, i.e. ft = fw (æ:). Note that if we had a weight-space posterior ((w) could use the law of unconscious statistician (pushforward measure) and replace the expectations in Eq. (8) with q( ww. ). As such, our method can be applied to any weight-space approximate inference technique. See App. A.2 for details on the dual parameters for different losses. Substituting Eq. (11) into Eq. (7), we obtain our GP based on the trained NN.
Predicting with Eq. (7) costs ((N3), which limits its applicability on large data sets.
Sparsification via dual parameters Sparse GPs reduce the computational complexity by rep- resenting the GP as a low-rank approximation induced by a sparse set of inducing points (see -uinenero-Candela & Rasmussen, 2005, for early overview). In the general setting, with non- Gaussian likelihoods, it is not obvious how to sparsify the full GP in Eq. (5). Previous approaches have resorted to (i) simply selecting subset (Immer al., 2021b) and (ti) performing variational in- ference, i.e. running a second optimization loop (Ortega et al., 2023). The dual parameter formulation in Eq. (7) enables leverage any sparsification method. Notably, we do not need to resort selecting
4

---

Published as a conference paper at ICLR 2024
subset or performing second optimization loop. In this work, we opt for the approach suggested by Titsias (2009); Hensman et al. (2013), shown in Eq. (6). However, instead of parameterizing a variational distribution q(u) N (u m, S), m, that the posterior under this model bears structure akin to Eq. (7). As such, we project the dual parameters onto the inducing points giving us sparse dual parameters. Using this sparse definition of the dual parameters, our sparse GP posterior is given by
we follow the insights from Adam et al. (2021),
(f) k_ au and Var kus with sparse dual parameters,
k
k
(K5z + Bu) Bukkk
(12)
Ou = Ei R2i O; RR and Bu = Simi lzzi 3r k E RRMMM
(13)
In our experiments, we sampled the inducing inputs Z - :.y.... 2M] inputs X. Note that the sparse dual parameters are now sum over all data points. Contrasting Eq. (12) and Eq. (7), we can see that the computational complexity went down from O to ( with M < N. Crucially, our sparse dual parameterizatton (Eq. (13)) is compact representation of the full model projected using the kernel. Alg. A1 details how we compute SFR's 's dual parameters and App. A.3 analyzes SFR's 's computational complexity.
RR from the training
We highlight that SFR differs from previous NN to GP methods (Khan et al., 2019; Immer et al., 2021b). To the best of our knowledge, we are the first to formulate dual GP from a NN. SFR's dual parameterization has two main benefits: (i) it enables us to construct a sparse GP (capturing information from all data points) which does not require further optimization, and (ii) it enables us to incorporate new data without retraining by conditioning on new data (using Eq. (13), see Eq. (17)).
4
SFR FOR SEQUENTIAL LEARNING
In this section, we show how SFR's ' sparse dual parameterization can equip NNs with important funetionalities for sequential learning: (i) maintaining representation of the NN for CL and (ii) incorporating new data without retraining from scratch.
Continual learning In the continual learning (CL) setting, training divided into T tasks, each with its own training data set D, = ((z,, y:) y.)).=. Once a task e {I..... T} is complete its data D; cannot be accessed in the subsequent tasks, ie., it is discarded. CL methods based on rehearsal and function-space regularization typically keep a subset of each task's is training data to help alleviate forgetting (Buzzega et al., 2020; Pan et al., 2020; Rudner et al., 2022). They show better performance than their weight space equivalents and report state-of-the-ard results on CL benchmarks. Previous approaches attempt to perform function-space VCL (Rudner et al., 2022; Pan et al., 2020), however, in order to scale to large data sets, they are forced to approximate the posterior at subset of the training data, ignoring information from the rest of the training data. In contrast, sparse function space VCL method from the GP literature (Chang et al., 2023) uses the sparse approximation to scale using the ELBO the current iteration,
4" (u) = arg max LVCL= = arg max ((u)
EN E (o [J]] - KL[[((u) (Qaa(()) (w)),
(14)
(
where Gola(tt) (t) is the posterior from the previous iteration. Using SFR we can approximate sparse function-space VCL. (see App. B.2 for details). For task I,, the regularized objective is
1 Ei M ||5 fw (2,) B
w = arg min L(Dr, w) +
(15)
Rea(w, Me-1)
where 7 E R is hyperparameter that scales the influence of the regularizer, w, denotes the MAP weights from task I, and Z, E IR MD denotes a set of task-specific inducing inputs selected from D,. For notational convenience, we restrict ourselves to single-output NNs and refer the reader to App. B.3 for details on the multi-output setting. The regularizer resembles a Mahalanobis distance with covariance matrix B, given by
B K B.KK MXM s.t.
B, = = kzi TT
(16)
5

---

Published as a conference paper at ICLR 2024
where === and igg are the Gram matrix and the vector computed on the task-specific inducing points Z, In practice, we randomly select a set of M task-specific inducing inputs Z₃ from the previous task's ' data set D, Given these inducing inputs,
then calculate the corresponding inducing variables u = f Z,). Finally, we calculate the regularization matrix Eq. (16) and add all of these == )) summarize previous tasks.
entries to a memory buffer, M {(Z.. Us B
Intuitively, the regularizer in Eq. (15) keeps the function values of the NN J (Z.) close to those at the It is worth noting that the covariance structure
MAP of previous tasks [u,}==1 = (fw: (2.))
in B, relaxes the regularization where there is no training data. Eq. (16) sidesteps computing the full posterior covariance for each task's dataset D, as it leverages SFR's dual parameterization to efficiently encode the information from Dz. As such, it offers a scalable solution whilst capturing information from the entire data set D,.
Incorporating new data without retraining Incorporating new data Dnew into a trained NN is not trivial. Existing approaches typically consider weight space formulation and directly update the NN's weights (Kirsch et al., 2022; Spiegelhalter & Lauritzen, 1990). In contrast, GPs defined in the function space can incorporate new data easily (Chang et al., 2023). Given our sparse dual parameters (Eq. (13)), we can incorporate new data into SFR with dual updates.
au f au + ΣiEDa k dj and B t Bu + EiED... kz 3, I
(17)
See Fig. 1 for an example of SFR incorporating new data and see Alg. A2 and App. A.3 for details of the dual updates and their computational complexity. Importantly, SFR's 's dual parameterization enabled us to incorporate new data (i) efficiently and (it) when access to previous data is lost.
EXPERIMENTS
We present a series of experiments specifically designed to showcase the power of SFR's dual parameterization. We first provide supervised learning experiments (Sec. 5.1) to highlight that SFR's s sparse dual paramcterization scales to data sets with (i) image inputs and (ii) over I million data points. As samity check, also show that SFR's 's uncertainty quantification generally matches (or is better than) other post hoc BNN methods, such as the Laplace approximation. We then provide sequential learning experiments (Sec. 5.2) which show that (i) SFR's FR's sparse dual parameterization can help retain knowledge from previous tasks in CL and (ii) SFR'S 's dual updates can incorporate new data fast. We implement all methods in PyTorch (Paszke al., 2019) and use GPU cluster. We provide full experiment details in App. D. For illustrative purposes, we show SFR on ID regression (Fig. 1) and on the 2D BANANA classification task (Fig. A5). Fig. I middle shows a trained MLP being converted into SFR and right shows new data being incorporated with the dual updates (see Sec. 4). App. E we visualize SFR's 's posterior covariance structure on the ID regression problem.
5.1 SUPERVISED LEARNING EXPERIMENTS
This section aims to demonstrate the efficacy of SFR quantifying uncertainty in several regression and classification tasks and test its performance in dealing with complex and large-scale data sets.
Experiment setup We evaluate the effecttivnness of SFR's sparse dual parameterization on eight UCI (Dua & Graff, 2017) classification tasks, two image classification tasks: Fashion-MNIST (FMNIST, Xiao al., 2017) and CIFAR-IO (Krizhevsky al., 2009), and the large-scale House- Electric data set. We used a two-layer MLP with width 50 and tanh activation functions for the UCI experiments. We used CNN architecture for the image classification tasks and a wider and deeper MI.P for Housefilectric. See App. 1.2 for full experiment details.
Baselines SFR represents information from all the training data at a set of inducing points, we compare it to making GP predictions with a subset of the training data (GP subset) on UCI and the image data sets. Note that for a given number of inducing points, the complexity of making predictions with the GP subset matches SFR. As such, the GP subset acts as baseline whose predictive performance we want to match whilst using fewer inducing points. As a sanity check, also compare SFR to the Laplace approximation (Laplace PyTorch, Daxberger al., 2021) when making predictions with (i) the nonlinear NN (BNN). (ii) the generalised linear model (GLM) in Eq. (4), and (iãi) the GP predictive from Immer al. (2021b), the generalization of DNN2GP (Khan

---

Published as a conference paper at ICLR 2024
CIFAR-10
F-MNIST
GLASS
VEHICLE
WAVEFORM
SATELLITE
06 08
NUFD
M % N
M N
M 4 of . %6 of
M % of N
.M % N
Figure 3: Effective sparsification: Comparison of convergence in number of inducing points M in NLPD (sea over 5 seeds) on classification tasks: SFR (-) vs. GP subset (-). Our SFR con- verges fast for all cases showing clear benefits of its ability to summarize all the data in a sparse model.
et al., 2019) to non-Gaussian likelihoods. As GP predictive makes predictions around the NN, also include results for SFR making predictions with the NN (SFR-NN), instead of the GP mean. Note that the LAPLACE DIAG and LAPLACE KRON experiments use the diagonal (Graves, 2011) or Kronecker factorization (Martens & Grosse, 2015) of the GGN respectively.
SFR's sparsification Fig. 3 compares how the predictive performance in terms of negative log predictive density (NLPD, lower is better) of SFR and the GP subset deteriorates as the number of inducing points lowered from M 100% of N to M = 1% of N. SFR is able to summarize the full data set more effectively than the GP subset method as it maintains a good (low) NLPD) with fewer inducing points. This demonstrates that the dual parameterization is effective for sparsification.
SFR on UCI data sets Tables A4 and A5 further demonstrate SFR's 's uncertainty quantification on the UCI classification tasks. They show that SFR outperforms the Laplace approximation (when using BNN, GLM and GP predictions) and the GP subset with M 20% of N on all eight tasks.
SFR experiments on image data sets We evaluate SFR on image data sets, which present a more challenging task and require more complicated NN architectures than UCI experiments. In particular, we report the results obtained using CNN architecture on FMNIST (Xiao et al., 2017) and CIFAR-10 (Krizhevsky et al., 2009). In Table I, we report the results obtained keep- ing the same prior precision 8 at training and inference time. We refer the reader to Apps. 1. and 1.1.3 for detailed expla- nation of the experiment. The performance of SFR using M = 2048 and M = 3200 inducing points outperforms the GP subset method and Laplace approximation when using both BNN and GLM predictions. This experiment indicates that SFR's 's sparse dual parameterization effectively captures information from all data points even when dealing with high-dimensional data.
Table 1: Image classification results using CNN: We report NLPD, accuracy and AUROC (mean-tstd over 5 seeds). SFR outperforms the GP subset and Laplace methods. The prior precision ò is not tuned post hoc.
MODEL
M
NLPD
ACC. (%)
AUROCT
NN MAP
0.23+-0.01
91.98-0.44
0. 0.83--0.00
TAPLACE DIAG
2.42:0.02
10.21
0.50
LAPLACE KRON
2.30110.01
9.871:0.66
0.511:0.02
LAPLACE GLM DIAG
1.66110.02
65.19:2.21
0.67:10.03
LAPLACE GLM KRON
1.09:10.04
84.79=1.00
0.96:10.01
PREDICTIVE
3200
0.47:0006
91.51:10.0
0.97
SER-NN (Ours)
3200
91.86 0.44
0.97
0.9720.13
71.32=8.83
0.93
GP SUBSET
3200 0.79-10.09
82.52=44110
0.95=0000
STR (Ours)
2048 0.30 01
91. 01.6810.55
0.90110.02
3200 0.09
91.74
0.47
0.96
NN MAP
0.09:10.03
1.00
0.86:10.02
LAPLACE DIAG
2.37:10.05
10.08
0.18
LAPLACE KRON
2.36
9.78
0.49-+0.01
LAPLACE GLM DIAG
1.33-:0006
1..96=1.35
0.82 03
CWAK-1
LAPLACE GLM KRON
1.04:0.08
75.56=1.65
0.64:10.04
GP PREDICTIVE
3200
0.90.10.02
7..07
0.7910 02
SER-NN (Ours)
3200
0.79
76.59
0.79 0..00
2048
1. 1.18 18.00.06
66,
40
69
0.71100.5
SUBSET
1. 4-0.05
69.75--3.
Out-of dittriiitittinn detection We follow Ritter et al. (2018); Osawa et al. (2019) and compare the predictive entropies of the in-distribution (ID) vs. -ii-tii-ittiittionnn (OOD) data. We desire low predictive entropy IDD to indicate that predictions can be made confidently and high predictive entropy OOD to indicate predictions cannot be made confidently. Fig. 4 shows predictive entropy histograms for models trained on FMNIST and evaluated OOD on MNIST data. SFR has low predictive entropy at the ID data and high predictive entropy at the OOD data, just as we desire. We further quantify SFR'S OOD detection using the area under the receiver operating characteristic curve (AUROC), a commonly used threshold-fre free evaluation metric. Table I reports the AUROC metric (higher is better), which further demonstrates that SFR has good OOD performance. See App. 1. for further details and experiments CIFAR-10.
SER (Ours)
2048
0.7 4 0.0.02
78.40
88
0.79
0.72 02
78.4820.95
0.79:20.02
7

---

Published as a conference paper at ICLR 2024
LAPLACE KRON
LAPLACE GI KRON
SURSET
SPR (Ours)
A
AUROC
AUROCASI
A
AUROC 80)
DEMBITY
AD
1OOD
PRED. ENTROPY
PRED. ENTROPY
PRED. ENTROPY
PRED. ENTROPY
PRED. ENTROPY
Histograms showing each method's predictive entropy at ID data (FMNIST, blue) where lower is better and at OOD data (MNIST, red) where higher is better.
Figure 4: OOD detection with CNN:
SFR scales to large data sets It is well known that GPs struggle to scale large data sets. Nevertheless, Wang et al. (2019) scaled GPs to million data points. We repeat their main experiment on the HouseElectric data set and recover better NLPD of -0.153++00000 SVGP model (Hensman et al., 2013) on the same data set with a 5-fold split. Furthermore, we also record a better wall-clock time of 1999 s vs. 115 s. This indicates that SFR outperforms GPs on large data sets in terms of both computation time and predictive performance.
vs. 0.024 for an
5.2 SEQUENTIAL. LEARNING EXPERIMENTS
The results in Sec. 5.1 motivate using SFR in more challenging sequential learning problems. In this section, we demonstrate the effectiveness of SFR's sparse dual parameterization in sequential learning settings. Representation: First, we show that SFR's sparse dual parameterization can help retain knowledge from previous tasks in CL. Incorporating new data: We then show that SER incorporate data fast via dual updates. Uncertainty: Finally, demonstrate that SFR's ss uncertainty estimates are good enough to help guide exploration in model-based RL.
Continual learning The reg- ularizer described in Sec. 4 can be used in CL retain com- pact representation of the NN. We report experiments in the single-head (SH) setting because it more realistic (and harder) than multi-head (MH) (see dis- cussion Van de Ven Tolias (2019)). Our regularizer is eval- uated on three CL benchmarks: MNISS-MNIST (S-MNIST), Split-
Table 2:
Continual learning experiments: sccuracy=tsd,,
bolding based on a f-test." "Methods rely weight regularization.
S-NNIST (SH)
S-MNIST (SH)
S-FMNIST (SH)
P-MNIST (SH)
Method
40 pts./ask
200 pis./usk
200 pts.Aakk
200 pts.Aask
ONLINE-EWC*
19.95
19.95-0.25 28
19.48-10.01
74.87
19.8240.09
19.82=0.09
19.80::0.21
88.39
VCL (Wlcnresets)
22.31 12.00
32.11 1.116
53.59 74
93.08
85.26-00.44
92.13-0044
$2.03-0.55
93.08
FROMP
75.21 12.05
89.54=0.77
78.8310044
94.90
S-FSVI
84.51 +1.30
92.87 00.14
77.54
0.000
95.76 0..00-0.02
SFR (Ours)
89.22:10.0
94.19
81.96=0.25
95.5810.08
L2 (Ous)) abl.
87.53 +-336
93.72 0..10
81.21 H.0.8
9-1.66 1:115
FashionMNIST (S-FMNIST), and the 10-tasks PerMUted-MNIST (P-MNIST). We adhere to the same setups outlined in Rudner al. (2022); Pan et al. (2020), which use two-layer MLP with 256 hidden units and ReLU activation for S-MNIST and S-FMNIST and a two-layer MLP with 100 hidden units for P-MNIST, ensure consistency (see App. D.2 for full details).
We compare our method against two categories of methods: (i) weiightregullarization methods: Online-EWC (Schwarz et al., 2018), SI (Zenke et al., 2017), and VCL (with coresets) (Nguyen et al., 2018), and (ūi) function-based regularization methods: DER (Buzzega et al., 2020). FROMP (Pan al., 2020), and S-FSVI (Rudner et al., 2022). We also introduce an ablation study where we replace our B, with an identity matrix IM, equivalent to having an L2 regularization between current and old function outputs. In Table 2, we use 200 points for each task and we further demonstrate our method's s ability to compress the task data set information on lower number of points for S-MNIST.
From Table 2, clear that the weight-space regularization methods fail entirely compared to the function-space methods on S-MNIST and S-FMNIST but are still achieving comparable results on P-MNIST. Considering only function-space methods, we can see that SFR achieves the best results on most data sets and is particularly effective when using fewer points per task. On S-FMNIST, our method obtains close results to DER, which regularizes the model by taking the mean squared error between current and old function outputs without accounting for the covariance structure of the loss, similar to our L2 ablation. However, DER resorts to reservoir sampling (Vitter, 1985) that continuously updates the set of points instead of selecting them at the task boundary. The SFR-based
8

---

Published as a conference paper at ICLR 2024
Table 3: SER's dual updates are fast: Comparison of incorporating new data D2 via SPR's dual updates (Eq. (17)) vs. retraining from scratch (on D1 U D2). SFR's 's dual updates improve the NLPD (val.+std as bar, lower better) whilst being significantly faster than retraining from scratch.
NLPD 4
Tine (5)
Trai D2
Updates D2
Retrain D1 U 22
Train D1
Updates D2 Retrin Dj U D2
AIRFOIL 0.60
0.50
0.47
19.65
0.041
18.22
BOSTON
0.23
0,16
0.13
11,45
0.021
7,48
PROTEIN
0.42
0.15
0.14
30.17
0.821
30.61
regularizer obtains comparable results to the best-perforiing method on P-MNIST, S-FSVI, which requires heavy variational inference computations compared to SFR.
Incorporating new data via dual updates SFR's ' dual parameterization enabled us to formulate equations for incorporating new data D2 into a NN previously trained on D2 (using Eq. (17)). We test SFR's ' so-called dual updates three UCI regression data sets. See App. D.3 for more details on our experimental set-up. Table shows that incorporating new data D2 via SFR's 's dual updates is significantly faster than retraining the NN from scratch (on Di D22) and offers improvements in the NLPD. Whilst SFR's 's dual updates show benefits in simple regression tasks, our initial experiments on image data sets did not show any improvement in predictive performance. We think this is because the dual updates require the kernel (NTK) to be accurate outside of the training data and this is unlikely because it was not enforced to be stationary. Nevertheless, incorporating new data via dual updates is very fast relative to retraining from scratch. As such, investigating SFR's 's dual updates in downstream applications, where retraining from scratch too costly, is worthwhile. For example, batch Bayesian optimization (Wu & Frazier, 2016).
Reinforcement learning under sparse rewards As final experiment, we demonstrate the capability of SFR to use its uncertainty estimates as guidance for exploration in model-based RL. We use SFR to learn a dynamics model within a model-based RL strategy that employs posterior sampling to guide exploration (Osband & Van Roy, 2017). We use the cartpole swingup task in MuJoCo (Todorov et al., 2012) (see Fig. A7); the goal is to swing the pole up and balance it. We increase the diffieulty of exploration by using sparse reward function. See App. C for overview of the RL problem, details of the algorithm, and the experiment setup.
Fig. A7 shows training curves for using SFR as the dynamics model (-), along with Laplace- GGN (Immer et al., 2021a) with GLM predictions (=), an ensemble of NNs (=), and a basic ML.P without uncertainty (-). To ensure fair comparison, we maintain the same MLP architec- ture/training scheme across all these methods and incorporate them into the same model-based RL algorithm (see App. C). We also compare our results with (DDPG, Lillicrap et al., 2016), a model-free RL algorithm (-). The training curves show that SFR's 's uncertainty estimates help exploration converge in fewer episodes, demonstrating higher sample efficiency.
6 DISCUSSION AND CONCLUSION
We introduced SFR, a novel approach for representing NNs in sparse function space. Our method is built upon dual parameterization which offers powerful mechanism for capturing predictive uncertainty, providing a compact representation suitable for CL, and incorporating new data without retraining. SFR is applicable on large data sets with rich inputs (e.g., images), where GPs are known to fail. This is because SFR learns the covariance structure so can learn non-stationary covariance functions, which would otherwise be hard to specify. These aspects were demonstrated in wide range of problems, data sets, and learning contexts. We showcased SPR's 's ability to capture uncertainty in UCI and image classification, established its potential for CL., and verified its applicability in RI..
I Limitations practical terms, SFR serves a role similar to a sparse GP. However, unlike vanilla GPs, it does not provide straightforward method for specifying the prior covariance function. This limita- tion can be addressed indirectly: the architecture of the NN and the choice of activation functions can be used to implicitly encode the prior assumptions, thereby incorporating broad range of inductive biases into the model. is important to note that we linearize the network around the MAP weights w". resulting in the function-space prior (and consequently the posterior) being only a locally linear approximation of the NN model. As such, when we incorporate new data with Eq. (17) the model becomes outdated. Nevertheless, in some cases we can still improve predictions without retraining.

---

Published as a conference paper at ICLR 2024
AUTHOR CONTRIBUTIONS
AJS (Aidan Scannell), RM, and PC jointly devised the methodological ideas used in the paper through discussion with JP and AHS. AJS and RM jointly developed the code base. RM conducted the CL experiments, whilst AJS conducted the RL experiments. AJS and RM jointly conducted the supervised learning experiments with help from ET. AJS lead the paper writing with the help of RM and PC. All authors contributed to finalizing the manuscript.
ACKNOWLEDGMENTS
AJS was supported by the Research Council of Finland from the Flagship program: Finnish Center for Artificial Intelligence (FCAI). AHS acknowledges funding from the Research Council of Finland (grant id 339730). We acknowledge CSC - IT Center for Science, Finland, for awarding this project access to the LUMI supercomputer, owned by the EuroHPC Joint Undertaking, hosted by CSC (Finland) and the LUMI consortium through CSC. We acknowledge the computational resources provided by the Aalto Science-IT project. We thank Rui Li, Severi Rissanen, and the rest of the Aalto Machine Learning research group for valuable feedback on manuscript. Finally, this research was supported by the NVIDIA AI Technology Center (NVAITC) Finland. We thank Niki Loppi of NVAITC Finland for his help with multi-GPU/node implementation and experimenting with float and half precision.
REFERENCES
Idan Achituve, Aviv Navon, Yochai Yemini, Gal Chechik, and Ethan Fetaya. GP-Tree: Gaussian process classifier for few-shot incremental learning. In Proceedings of the 38th International Conference on Machine Learning (ICML), volume 139 of Proceedings of Machine Learning Research, pp. 54 PMLR, 2021.
Vincent Adam, Paul Chang, Mohammad Emtiyaz Khan, and Arno Solin. Dual parameterization of sparse variational Gaussian processes. In Advances in Neural Information Processing Systems 34 (NeurIPS), pp. 11474-1 11476. Curran Associates, Inc., 2021.
Ari Benjamin, David Rolnick, and Konrad Kording. Measuring and regularizing networks in function space. In International Conference on Learning Representations (ICLR), 2019.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American Statistical Association, 112(518):859-877 2017.
Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical Gauss-Newton optimisation for deep learning. In Proceedings of the 34th International Conference on Machine Learning (ICML), volume 70 of Proceedings of Machine Learning Research, pp. 557-565, PMLR, 2017.
Pietro Buzzega, Matteo Boschini, Angelo Porrello, Davide Abati, and Simone Calderara. Dark experience for general continual learning: strong, simple baseline. In Advances in Neural Information Processing Systems 33 (NeurIPS), pp. 15920-15930. Curran Associates, Inc., 2020.
Paul Edmund Chang, Prakhar Verma, S. T. John, Amo Solin, and Mohammad Emtiyaz Khan. Memory-based dual Gaussian processes for sequential learning. In Proceedings of the 40th International Conference on Machine Learning (ICML), volume 202 of Proceedings of Machine Learning Research, pp. 4035-4054. PMLR, 2023.
Olivier Chapelle. Training a support vector machine in the primal. Neural Computation, 19(5): 1155-1178, 2007.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems 31 (NeurIPS), pp. 4765- Curran Associates, Inc., 2018.
Lehel Csató and Manfred Opper. Sparse on- on-line Gaussian processes. Neural Computation, 14(3):
641-668, 2002.
10

---

Published as a conference paper at ICLR 2024
Sebastian Curi, Felix Berkenkamp, and Andreas Krause. Efficient model-based reinforcement learning through optimistic policy search and planning. In Advances in Neural Information Processing Systems 33 (NeurIPS), pp. 14156-14170. Curran Associates, Inc., 2020.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Bayesian deep learning, In Advances in Neural Informa-
Philipp Hennig. Laplace fion Processing Systems 34 (NeurIPS), pp. 20089-20103. Curran Associates, Inc., 2021.
Matthias De Lange, Rahaf Aljundi, Marc Masana, Sarah Parisot, Xu Jia, Aleš Leonardis, Gregory Slabaugh, and Tinne Tuytelaars. A continual learning survey: Defying forgetting in classification tasks. IEEE Transactions Pattern Analysis and Machine Intelligence, 44(7):3366-338 2021.
Richard Dearden, Nir Friedman, and David Andre. Model based Bayesian exploration. In Proceedings of the Fifteenth Conference Uncertainty in Artificial Intelligence (UAI), pp. 150-159. Morgan Kaufmann Publishers, 1999.
Marc Deisenroth and Carl Rasmussen. PILCO: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference on Machine Learning (ICML), volume 28, pp. 465-472. ACM, 2011.
Dheeru Dua and Casey Graff. UCI machine learning repository, 2017. URL httttt///acciiie ics. uci.cci... edu/ml.
Yarin Gal and Zoubin Ghahramani. Dropout as a Bayesian approximation: Representing model uncertainty deep learning. In Proceedings of the 33rd International Conference on Machine Learning (ICML), volume 48 of Proceedings of Machine Learning Research, pp. 1050-1059 PMLR, 2016.
Yarin Gal, Rowan McAllister, and Carl Rasmussen. Improving PILCO with Bayesian Neural Network Dynamics Models. ICML Workshop on Data-Efficien Machine Learning. 2016.
Roman Garnett. Bayesian Optimization. Cambridge University Press, 2023. Cambridge, UK.
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016.
Alex Graves, Practical variational inference for neural networks. In Advances in Neural Information Processing Systems 24 (NIPS), pp. 2348-23482 2356.. Curran Associates, Inc., 2011.
James Hensman, Nicolò Fusi, and Neil D. Lawrence. Gaussian processes for big data. In Proceedings of the 29th Conference on Uncertainty Artificial Intelligence (UAI), pp. 282 290. AUAI Press, 2013.
Rein Houthooft, Xi Chen, Yan Duan, John Schulman, Filip De Turek, and Pieter Abbeel. VIME: Variational information maximizing exploration. In Advances in Neural Information Processing Systems 29 (NIPS), pp. 110-11117. Curran Associates, Inc., 2016.
Alexander Immer, Matthias Bauer, Vincent Fortuin, Gunnar Rätsch, and Mohammad Emtiyaz Khan. Scalable marginal likelihood estimation for model selection deep learning. In Proceedings of the 38th International Conference on Machine Learning (ICML), volume 139 of Proceedings of Machine Learning Research, pp. 4563-4573, PMLR, 2021a.
Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions of Bayesian neural nets via local linearization. Proceedings of The 24th International Conference on Artificial Intelligence and Statistics (AISTATS), volume 130 of Proceedings of Machine Learning Research, pp. 703-711 PMLR, 2021b.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization neural networks. In Advances in Neural Information Processing Systems 31 (NeurIPS), pp. 8571-8580. Curran Associates, Inc., 2018
Thomas Jaksch, Ronald Ornner and Peter Auer. Near-optimal regret bounds for reinforcement learning. Journal of Machine Learning Research (JMLR), 11651):1533-160 2010.
11

---

Published as a conference paper at ICLR 2024
Sanket Kamthe and Mare Deisenroth. Data-efficient reinforcement learning with probabilistic model predictive control. In Proceedings of the Twenty-Firg International Conference on Artificial Intelligence and Statistics (AISTATS), volume 84 of Proceedings of Machine Learning Research, pp. 1701-1710. PMLR, 2018.
Sanyam Kapoor, Theofanis Karaletsos, and Thang D Bui. Variational auto-regressive gaussian processes for continual learning. In International Conference on Machine Learning (ICML), pp. 5290-5300. PMLR, 2021.
Mohammad Khan, Didrik Nielsen, Voot Tangkaratt, Wu Lin, Yarin Gal, and Akash Srivastava. Fast and scalable Bayesian deep learning by sright-petuurbation Adam. In Proceedings of the 35th International Conference on Machine Learning (ICML), pp. 2611-2620. PMLR, 2018.
Mohammad Emtiyaz Khan and Wu Lin. conjugate-computatie variational inference: Converting variational inference in non-conjugate models to inferences in conjugate models. In Proceedings of the 20th International Conference Artificial Intelligence and Statistics (AISTATS), volume 54 of Proceedings of Machine Learning Research, pp. 878-887. PMLR, 2017.
Mohammad Emtiyaz Khan, Alexander Immer, Ehsan Abedi, and Maciej Korzepa. Approximate inference turns deep networks into Gaussian processes. In Advances in Neural Information Processing Systems 32 (NeurIPS). pp. 3094-3104. Curran Associates, Inc., 2019.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference Learning Representations (ICLR), 2015.
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka catastrophic forgetting in neural networks. Proceedings of the National Academy of Sciences, 114 (13):3521-3526, 2017.
et al. Overcoming
Andreas Kirsch. Jannik Kossen, and Yarin Gal. Marginal and joint cross-entropies & predictives for online Bayesian inference, active learning, and active sampling. arXiv preprint arXiv:2205.08766, 2022.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Technical report, University of Toronto, Toronto, ON, Canada, 2009.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural Information Processing Systems 30 (NIPS), pp. 6402-6413. Curran Associates, Inc., 2017.
Zhizhong Li and Derek Hoiem. Learning without forgetting. IEEE Transactions on Pattern Analysis and Machine Intelligence, 40(1)) (2)):23935-299 22229995--27,,,,- 2018.
Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In International Conference on Learning Representations (ICLR), 2016.
David J. C. MacKay. Bayesian Interpolation. Neural Computation, 4(3):415 447, May 1992.
James Martens and Roger Grosse. Optimizing Neural Networks with Kronecker-factored Approxi- mate Curvature. In Proceedings of the 32nd International Conference on Machine Learning, pp. 2408 -417. PMLR, June 2015.
Michael McCloskey and Neal Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of Learning and Motivation, volume 24, pp. 109-165. Elsevier, 1989.
Radford M Neal. Bayesian Learning for Neural Networks. PhD thesis, University of Toronto, Toronto, Canada, 1995.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Ng. Reading digits in natural images with unsupervised feature learning. In Advances in Neural Information Processing Systems, 01 2011.
12

---

Published as a conference paper at ICLR 2024
Cuong V. Nguyen, Yingzhen Li, Thang D. Bui, and Richard E. Turner. Variational continual learning. In International Conference on Learning Representations (ICLR), 2018.
Manfred Opper and Cédric Archambeau. The variational Gaussian approximation revisited. Neural Computation, 21(3):786-799 2009.
Luis A Ortega, Simón Rodríguez Santana, and Daniel Hemández-LLoato. Variational linearized Laplace approximation for Bayesian deep learning. arXiv preprint arXiv:2302. /2565, 2023.
Kazuki Osawa, Siddharth Swaroop, Mohammad Emtiyaz Khan, Anirudh Jain, Runa Eschenhagen, Richard Turner, and Rio Yokota. Practical deep learning with Bayesian principles. In Advances in Neural Information Processing Systems, volume 32, pp. 4287-4299, Curran Associates, Inc., 2019.
lan Osband and Benjamin Van Roy. Why is posterior sampling better than optimism for reinforcement learning? In Proceedings of the 34th International Conference on Machine Learning (ICML), volume 70 of Proceedings of Machine Learning Research, pp. 2701-2710. PMLR, 2017.
Ian Osband, Daniel Russo, and Benjamin Van Roy. (More) efficient reinforcement learning via posterior sampling. In Advances in Neural Information Processing Systems 26 (NeurIPS), pp. 3003-3011. Curran Associates, Inc., 2013.
Pingbo Pan, Siddharth Swaroop, Alexander Immer, Runa Eschenhagen, Richard Turner, and Mo- hammad Emtiyaz Khan. Continual deep learning by functional regularisation of memorable past. In Advances in Neural Information Processing Systems 33 (NeurIPS), pp. 4453-4464. Curran Associates, Inc., 2020.
Yunpeng Pan, Evangelos Theodorou, and Michail Kontitsis. Sample efficient path integral control under uncertainty. In Advances in Neural Information Processing Systems 28 (NIPS), pp. 2314- 2322. Curran Associates, Inc., 2015.
German Parisi, Ronald Kemker, Jose Part, Christopher Kanan, and Stefan Wermter. Continual lifelong learning with neural networks: A review. Neural Networks, 1 11::54-7,, 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32 (NeurIPS), pp. 8026-8037. Curran Associates, Inc., 2019.
Joaquin Quiñonero-Cannn Candell and Carl Edward Rasmussen. A unifying view of sparse approximate Gaussian process regression. The Journal of Machine Learning Research (JMLR), 11999-19599 2005.
Carl Edward Rasmussen and Christopher K. I. Williams. Gaussian Processes for Machine Learning. MIT Press, Cambridge, MA, 2006.
Hippolyt Ritter, Aleksandar Botev, and David Barber. A scalable Laplace approximation for neural networks. In International Conference on Learning Representations (ICLR), 2018.
Tim G. J. Rudner, Freddie Bickford Smith, Qixuan Feng, Yee Whye Teh, and Yarin Gal. Continual Learning via Sequential Function-Space Variational Inference. In Proceedings of the 39th Interna- tional Conference on Machine Learning (ICML), volume 162 of Proceedings of Machine Learning Research, pp. 18871 18887. PMLR, 2022.
Remo Sasso, Michelangelo Conserva, and Paulo Rauber. Posterior sampling for deep reinforcement learning. In Proceedings of the 40th International Conference Machine Learning (ICML), volume 202 of Proceedings of Machine Learning Research, pp. 30042-30061. PMLR, 2023.
Aidan Scannell, Riccardo Mereu, Paul Chang, Ella Tami, Joni Pajarinen, and Arno Solin. Sparse function-space representation of neural networks. In ICML 2023 Workshop Duality Principles for Modern Machine Learning, 2023.
Frank Schneider, Lukas Balles, and Philipp Hennig. DeepOBS: A deep learning optimizer benchmark suite. In International Conference on Learning Representations (ICLR), 2019.
13

---

Published as a conference paper at ICLR 2024
Jonathan Schwarz, Wojciech Czarnecki, Jelena Luketina, Agnieszka Grabska-Barwinske Barwinska, Yee Whye Teh, Razvan Pascanu, and Raia Hadsell. Progress & compress: A scalable framework for continual learning. In Proceedings of the 35th International Conference Machine Learning (ICML), volume 80 of Proceedings of Machine Learning Research, pp. 4528-4537. PMLR, 2018.
David J Spiegelhalter and Steffen L Lauritzen. Sequential updating of conditional probabilities on directed graphical structures. Networks, 20(5):579-60) 1990.
Richard S Sutton and Andrew G Barto. Reinforcement Learning: An Introduction. Adaptive Computation and Machine Learning Series. MIT Press, second edition, 2018.
Michalis K Titsias. Variational learning of inducing variables in sparse Gaussian processes. In Pro- ceedings of the Twelth International Conference on Artificial Intelligence and Statistics (AISTATS), volume 5 of Proceedings of Machine Learning Research, pp. 567-574. PMLR, 2009.
Michalis Titsias, Jonathan Schwarz, Alexander G de G Matthews, Razvan Pascanu, and Yee Whye Teh. Functional regularisation for continual learning with Gaussian processes. In International Conference on Learning Representations (ICLR), 2020.
Emanuel Todorov, Tom Erez, and Yuval Tassa. MuloCo: physics engine for model-based control. In International Conference Intelligent Robots and Systems (IROS), pp. 5026-5033, IEEE, 2012.
Gido M Van de Ven and Andreas S Tollias. Three scenarios for continual learning. arXiv preprint arXiv: 1904:9904.774 2019.
Jeffrey S Vitter. Random sampling with a reservoir. ACM Transactions on Mathematical Software (TOMS), 11(1):37-57, 1985.
Ke Wang, Geoff Pleiss, Jacob Gardner, Stephen Tyree, Kilian Q Weinberger, and Andrew Gordon Wilson. Exact Gaussian processes on million data points. In Advances in Neural Information Processing Systems 32 (NeurIPS), pp.
Curran Associates, Inc., 2019.
Grady Williams, Andrew Aldrich, and Evangelos A. Theodorou. Model predictive path integral control: From theory to parallel computation. Journal of Guidance, Control, and Dynamics, 40(2): 344-357, 2017.
Andrew Gordon Wilson. The case for Bayesian deep learning. https://cimm nyu. edu/ 2019. Technical Report.
-aadreewg/caaee wttp://wwwaa.foo.......ooooo
Jian Wu and Peter Frazier. The parallel knowledge gradient method for batch Bayesian optimiza- tion. In Advances in Neural Information Processing Systems 29 (NIPS), pp. 3126-3134. Curran Associates, Inc., 2016.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST: A novel image dataset for bench- marking machine learning algorithms. arXiv preprint 1XXiv:7708.0744: 2017.
Friedemann Zenke, Ben Poole, and Surya Ganguli. Continual learning through synaptic intelligence. In Proceedings of the 34th International Conference on Machine Learning (ICML), volume 70 of Proceedings of Machine Learning Research, pp. 3987-3995. PMLR, 2017.
14

---

Published as a conference paper at ICLR 2024
APPENDIX
This appendix is organized as follows. App. A provides a more extensive overview of the technical details of SPR. App. B covers additional details on the SFR-based regularizer used for CL. App. C provides more extensive write up of the reinforcement learning setup used in the experiments. App. D provides the full details of each individual experiment in the main paper. Finally, App. provides a visualisation of SFR's 's covariance matrix.
METHOD DETAILS
16
A.1 DUAL PARAMETERIZATTION BACKGROUND AND DERIVATION
16
A.2 COMPUTATION THE DUAL PARAMETERS FOR DIFFERENT LOSSES
17
A.3 COMPUTATIONAL COMPLEXITY
18
A.4 PRACTICAL CONSIDERATIONS
18
A.5 ALGORITHMS
19
SFR FOR CONTINUAL LEARNING
B.1 RELATED WORK IN THE SCOPE OF CL
19
B.2 DERIVATION THE REGULARIZER TERM
20
B.3 EXTENDING THE CL REGULARIZER TO MULTI-CLASS SETTINGS
21
SFR FOR MODEL-BASED REINFORCEMENT LEARNING
C.I RELATED WORK IN THE SCOPE OF MODEL-BASE RL
21
C.2 BACKGROUND
22
C.3 ALGORITIIM
22
D EXPERIMENT DETAILS
22
D. SUPERVISED LEARNING EXPERIMENTS
23
D.1.1 PRIOR PRECISION TUNING
23
D.1.2 UCI EXPERIMENT CONFIGURATION
23
................................
D.1.3 IMAGE EXPERIMENT CONFIGURATION
24
D.1.4 OUT-OF-DISTRIBUTTI DETECTION EXPERIMENTS
26
D.2 CONTINUAL LEARNING EXPERIMENT DETAILS
26
D.3 INCORPORATING NEW DATA VIA DUAL UPDATES EXPERIMENT DETAILS
27
D.4 REINFORCEMENT LEARNING EXPERIMENT DETAILS
E COVARIANCE VISUALIZATION
29
15

---

Published as a conference paper at ICLR 2024
A METHOD DETAILS
In this section, we provide further details on our method. App. A. provides fürther background and derivations of the dual parameters. App. A.2 shows how SFR's dual parameters are computed for the different losses used in our experiments. We then discuss the computational complexity of our method in App. A.3. In App. A.4 we detail considerations for using SFR in practice. In App. A.5 we provide algorithms for (i) the computation of SFR's ' sparse dual parameters and (ii) SFR 's dual updates.
A.1 DUAL PARAMETERIZATION BACKGROUND AND DERIVATION
Any empirical risk minimization objective (that assumes likelihood terms are factorized across data) can be reparameterized in to what Adam et al. (2021) call the dual parameterization In the support vector machine (SVM) literature, the connection to dual parameterization was pointed out in Chapelle (2007). In the GP literature, they first appeared in Csató & Opper (2002) and then Opper & Archambeau (2009) showed them for the variational GP objective. More recently, Adam et al. (2021) coined the dual parameterization for sparse variational GPs and showed they correspond to the dual Lagrangian optimization objective.
Dual parameters take slightly different forms depending the types of objective. In particular, whether the objective is Bayesian (e.g. ELBO) or non-Bayesian (e.g. MAP objective). They are identified from the stationary point of the objective. As parameters (my and S)), given the dual parameters, we follow Adam et al. (2021) and call them dual parameters. We derive the dual parameters for these two type of objectives.
determine the posterior's original
Dual parameters for ELBO objective seeking maximize the ELBO, given by
Variational inference casts inference to an optimization
mÿ. S = arg max my Sy
Σ
B.ccyloooppe (() - KL(g(f) p(f).,
(18)
with variational posterior s(f) = m,, Sy) and prior = = ( K). Note that my and
Sy the variational parameters to be optimized. The stationary point of the ELBO in Eq. (18) is given by
(5 where (= my, St) and a² [ag ................................ ayy] and 9 [Bā..... 3817 are vectors with entries
Ka = my.
diag (B*) K
(19)
a, Eg (V, log (1)
(20)
8- = Eg (Vol fulft log [+)
(21)
If we have access to a and ,3 then we can determine the posterior's original parameters (m and Sy). Hence, we follow Adam et al. (2021) and call them dual parameters. Eq. (19) follows because the KL Eq. (18) has the following form,
1 THKK Syp
1 1
KLig (Li(f))llP())
In SSyl + const.
(22)
As such, the derivatives of the KI. with respect the variational parameters (m, and Sy) given by,
m, V ()) - is, -
= K-,m,,
(23)
1
(24)
Similarly, the derivatives of the expected log likelihood under the variational posterior are given by, [fa] Egg.a log [.)]ss
my E
(25)
s, Bogsr [.)]] = BeCr /V V.. logp(us [[]].
(26)
16

---

Published as a conference paper at ICLR 2024
Dual parameters for MAP objective The Laplace approximation approximates the posterior using the curvature around the MAP solution fe, which given by
fr arg max CLA (f) arg max Σ
log logp(y: f() p(f)
(27)
Note that the objective does not contain the posterior QLA (f). Taking the derivative of the objective in Eq. (27) w.r.t. f we obtain
[fCi(()) = V logp(y f) K- V2 Lia = V log ff K K- The Laplace approximation then approximates the posterior with.
(28)
(29)
( S.A) where
(30)
where the covariance SLA is set to the inverse of the Hessian of the objective (in Eq. (27)) at the MAP solution ft. Similar to before, the stationary point of Eq. (27) exhibits dual parameter decomposition,
Ko ff; ST diag( K-1
(31)
where the dual parameters (a and (T,)) are vectors with entries
a V f: log p(y: fi)l==ff--
(32)
i - -V f..f log p(ye [f));=;;--
(33)
Comparing the dual parameters for the ELBO in Eq. (18) to those for the MAP in Eq. (27), we can see that the only difference is the expectation over the posterior.
A.2 COMPUTATION OF THE DUAL PARAMETERS FOR DIFFERENT LOSSES
As shown in Eq. (8) in the main paper, our method relies on the computation of the dual parameters. We derive them in general for the family of exponential likelihoods, by rewriting the likelihood in the natural form,
{(h()).)) log p(v (f) A()) (1w)) f),
(34)
where A(f) is the log partition function, h(-) is the inverse link function, t(y) are the sufficient statisties which for all our likelihoods is simply t(y) = y. and f denotes the NN output before the inverse link function. Then, taking the first and second derivatives of the loss in this form become:
V Vy((h()) A':) yi and
(.)(()()))))l) Afa:).
(35)
In this section, we provide the derivatives needed for the losses used in this work, i.e., the mean squared error (MSE) likelihood, used for the regression tasks, the Bernoulli likelihood, for binary classification, and the categorical likelihood, used for multi-class classification.
For regression, the negative log-likelihood comes from the MSE and for single input it can be written as
1 ((((f)) - [ 9;)
(36)
where the inverse link function is simply h(f) = Thus, the first and second derivatives are Vyl(o()) = fi- - and (if(((()) w.)(k(())))l==.... =1
(37)
In binary classification, we have a Bernoulli likelihood. The inverse link function is the sigmoid function h(f)) o(f) I and we can then write the derivatives as follows
v,((h(f).u)lp o(fi) yi and (gf((b(f).
-ff)(( o(f)).
(38)
In multi-class classification, the function outputs are no longer scalar a vector of logits (ax;)
fa E RC. The inverse link function for the categorical loss is the softmax function softmax(f)
CC [o, Ij defined as h(f) softmax (f) _e f.) E [0, 1] If we write the targets as one-hot encoded vector Yis we can write the negative likellhodd as
exp(f)
((h(f,))) y.) - log (y. softmax(f))
(39)
17

---

Published as a conference paper at ICLR 2024
Finally, we can write the first and second derivatives as follows
(V s=)s= = [Softmax (f) yol,
(40)
(diag (V ( ,:)1,=;)>}, [Softma(ff)), ( [Softmax (f)))
where we only need the diag given our independence assumption.
A.3 COMPUTATIONAL COMPLEXITY
In this section, we analyze the computational complexity of SFR. We first need to train a NN fw and then compute the posterior covariance matrix that characterizes our sparse GP representation of the model. The first step then requires the usual complexity cost for training NN, where defining the cost of a forward pass (and backward pass ) as [FP], this can be roughly sketched as O(2EN[FP)) where E' is the number of training epochs and N the number of training samples. Thus, there is no additional cost involved in the training phase of the NN.
Dual parameter computation Given the trained network, construct the covariance function, which is defined through the NTK (Jacot et al., 2018). We relied on the Jacobian contraction formulation implemented using PyTorch (Paszke al., 2019). Computing the kernel Gram matrix for a batch of B samples with this implementation has a computational complexity of (BC[[FP) B²C²p), where C the number of outputs and P is the number of parameters of the NN. In Eq. (13), we compute B., iterating through all the N training points to compute all the kzi. Then, we use them compute the outer product, which has a cost of O(M2), with M being the number of inducing points. The total cost of this operation becomes O which is also the dominating cost of our approach.
Dual updates The issue with adding new data with standard GP equations is that it's computationally infeasible for large data sets, which are common in deep learning. This due to the matrix inversion of (K + --1 which is of complexity O((N + Nnew)3), There are some linear algebra tricks you can employ that would then make the update O((N + )²) but it would still grow with data set size. Incorporating new data into the sparse dual parameters (see Alg, A2) includes the computation of ki at Nnew data points which has complexity ne + N C) Computing the outer products k takes MM which is the dominating cost in this step.
A.4 PRACTICAL CONSIDERATIONS
Numerical stability In practical terms, some care needs to be taken during the implementation, both regarding memory consumption and numerical stability. For the former, batched over the number of training samples, N, which keeps the memory consumption in check without the need for any further approximations. The computation of the terms in Eq. (12) requires the inversion of an M x M matrix, that we implemented in numerically stable manner using Cholesky factorization. We also introduce small amount of jitter and clip the GGN matrix for numerical stability where needed (see discussions in App. D). It is also worth noting that we train the neural network in floating point precision (float32) and then switch to double precision (float64) for calculating SFR's dual parameters and making predictions. This is because we ran into issues when using single precision, which we believe was due to the Cholesky decompositions used implement the matrix inversions. However, we have been testing different approaches for implementing the matrix inversions e.g. using the Woodbury matrix identity and IUU and SVD -based direct solves. In some small scale experiments these enabled us to remain in single precision throughout. However, we leave further investigation for future work.
Hardware We ran our experiments on a cluster and used single GPU. The cluster is equipped with four AMD MI20XX GPUs based on the 2nd Gen AMD CDNA architecture. A MI250x GPU is a multi-chip module (MCM) with two GPU dies named by AMD Graphics Compute Die (GCD). Each of these dies features 110 compute units (CU) and have access to 64 GB slice of HBM memory for total of 220 CUs and 128 GB total memory per MI250x module.
18

---

Published as a conference paper at ICLR 2024
Algorithm A1 Compute SFR's sparse dual parameters
Input: D ((r,, y: |= prior precision 6, number of inducing points M, trained NN Sw, negative log-likelihood function P(-), batch size b, kernel function k(.;"
Output: Duals uu,, B..
Sample inducing points Z from X
ou = E MM Bu = 0 E R2 RMxM
for i in [0..... len (en(D)// /b) do
Select rows indices SLICE = (t.b:i.b+
x SLICE,:
y YSLICE
Compute a 1b), (0(00 for X (b) y
Eq. (11)
Ou c- au + *x (()) ((6)
Eq. (13)
Bu t Bu + k Z ( 3 k: ( x())
Eq. (13)
end for
Algorithm A2 SFR's dual updates
Input:
Duals Ou. Bu. Dnew ( / Now, inducing points Z, kernel function ko(..))
Output: New duals new Bhew for = to Neew do
Compute à;, B3 at Yi
Eq. (11)
Qu au did xz,
Eq. (17)
B,, + Bu N² Z, 8c 3
Eq. (17)
end for
A.5 ALGORITHMS
Alg. A1 details how we compute the sparse dual parameters in Eq. (13). Important considerations include batching over the training data to keep the memory in check. Alg. A2 then details how SFR updates the sparse dual parameters to incorporate information from new data.
B SFR FOR CONTINUAL LEARNING
This section provides additional details about the related works in the CL literature, the detailed derivation of the SFR-based regularizer described in Sec. 4, and how we extended it for dealing with the oulti- setting in our experiments.
B.1 RELATED WORK IN THE SCOPE OF CL
Continual learning (CL) hinges on how to deal with a non-stationary training distribution, wherein the training process may overwrite previously learned parameters, causing catastrophic forgetting (Mc- Closkey & Cohen, 1989). CL approaches fall into inference-based methods (see (Parisi et al.. 2019: De Lange al., 2021) for a detailed survey). The methods in the first category usually tackle this problem with weight-space regularization, such as EWC (Kirkpatrick al., 2017), 7), SI (Zenke et al., 2017)). VCL (Nguyen al., 2018), 8), that induce retention of previously learned information in the weights alone. However, regularizing the weights does not guarantee good quality predictions. Function-space regularization techniques (Li & Hoiem, 2018; Benjamin al., 2019; Titsias et al., 2020; Buzzega al., 2020; Pan et al., 2020; Rudner et al., 2022; Achituve et al., 2021) address this by relying on training data subsets for each task to compute a regularization term. These methods can be seen hybrid between objective and rehearsal-based methods, since they also store a subset of training data to construct the regularization term. Recent function- based methods, e.g., DER (Buzzega al., 2020), FROMP (Pan al., 2020), and S-FSVI (Rudner et al.,
ctheersal--bssed and model-based
19

---

Published as a conference paper at ICLR 2024
2022), achieve satee-o--hhe-a performance among the objective-based techniques in several CL benchmarks.
B.2 DERIVATION OF THE REGULARIZER TERM
For CL, the model cannot access the training data from previous tasks up to task 1, denoted here as Dou = - D., and for each new task i receives a new chunk of data Dnew D,. In contrast, in the supervised learning setting (described in Sec. 2), the NN model has access to all the training data at once, and then the objective we are trying to optimize
arg min C(Dola U D Dnew, w) arg min
Σ
((fro(a))) ()) R(w).
(41)
iEDadUDie
Our goal in the continual setting is to replace the missing data coming from Dota by resorting to a functional regularizer that accounts for that by forcing the predictions of the NNs over the inducing points Ze to come from (2;) ~ g (u), Vz; Zt. As shown in (Khan & Lin, 2017; Adam al., 2021), terms. Using this insight, we can rewrite the sparse GP posterior for q(u) in terms of the following approximate normal distributions:
view the dual parameters derived in Sec. 3 as linked approximate likelihood
B - a, u2)² 2
(ota(es) & NN K2z) II iEDaa
exp
(42)
x p(u) N u. Bu);
(43)
where a. k and D; a; fw (3. The sparse GP posterior g(u) is proportional the product between the prior over the inducing points p(u) and Gaussian term N(;; u, Bu), which arises from the sparse approximation to the likelihood. Replacing a with x w. we can see the similarity between the sparse GP and linear model, for which the kernel is the linear kernel.
Now that we have transformed SFR's dual parameters in to multivariate form, we can detail the motivation for our function-space spacee CL regularizer and also detail how it differs from previous methods. Using a NN model for function-space VCL (Rudner al., 2022) has previously been formulated through converted full GP model for CL. We now show the equivalent for sparse GP model (Kapoor et al., 2021),
N
("(u) = argmax LVCL. argmax ((u)
Σ
E Beutt)llogpt [A] KL(g(u) (o)(())),
(44)
q(u)
where q(u) is parameterized by N(m,, S.). Substituting Eq. (42) into LvC. we get the following expression for sparse VCI.,
N
LVCL = Σ Bn p(y: [f)] + Bogg u, B.)l KL(g(u) (I
(45)
So adding in the expectation of the dual parameters in MVN form is equivalent to doing sparse VCL in function space. Expanding the expectation gives
By(() [og.N(ü u, B.)] = m
(2g)
log Bul
2
1
1
1
mm))
[Bul-1 m.)
Trace Trace(B,;15) BB, S.). (46)
2
2
Therefore, we use the quadratic term from N (ū; u, Ba) as our regularizer in place of the missing data term, the regularizer becomes (ÿ u.) Bu 111 -ū u). Given that ÿ are in the function space, we can instead use the previous NN predictions fold (2:) and save computations.
Our objective is similar to both Rudner et al. (2022) and Pan al. (2020) they both approximate the KL in function-space VCL. However, for scalability we use the sparse function-space VCL formulation, whereas they resort subset approach.
20

---

Published as a conference paper at ICLR 2024
B.3 EXTENDING THE CL REGULARIZER TO MULTI-CLASS SETTINGS
In a multi-class classification setting, the NN can be seen as a function approximator that maps inputs R to the logits S IR, i.e., f RDD R As a consequence, the Jacobian matrix Jw (z: = fw k(æ, x²) is tensor in RCPxxP This means that we need to compute the dual parameters and the matrix B,C for all the function's outputs. If we denote ke( x') E IR, as the kernel related to the et output function, Lee, the logit for class e, we can write the dual parameters as
E RRCP and the kernel of the associated sparse GP derived in Sec. 3,
Ouuc = Σm kzi,c di,e E RRM and Bu,c = Ei kzi,e Br. I E C &RMX
(47)
where lizi. is the kernel :)):;-) computed for the inducing points z in 2 and the input point æ; for class c. In this setting, the dual variables are
áGc- = | = v y p(u ff)l=slle E R, _ic= = diag( (3.) = 1- v log ffl== ER,
(48)
where the operator [e is taking the ethh element of vector. In fact, in the multi-class setting O; is vector in IR and B, a matrix in ROxC For the SFR-regularizer, at the end of each task, we compute the matrix B as 22,c K2 Bu,
Pu. =z,c E MXM, s.t. Bu, - kzzi Bi,c K
(49)
The kernel computation is based the relationship between the observed data and the parametrized function. For this reason, computing the kernel K. for classes e, for which the model has not observed any data yet, is not meaningful. In practice, for our method we define B be B IM Ve & C.. where C2 is the set of classes observed to task t. For example, for S-MNIST, after task 1 the set of observed classes for which we will compute the kernel is C1 = {0, 1}. then after task it becomes C1 = {0, 1, 3), et cetera. Finally, we can rewrite REFR for task as follows
1
REFR M
M [[[fw; (Z.)] - [o (2.) B
(50)
2 s=1
C SFR FOR MODEL-BASEL REINFORCEMENT LEARNING
As a final experiment in the main paper (see Sec. 5.2), we provided demonstration of using SFR in model-based reinforcement learning. As such, this is a direct extension of the uncertainty quantification and representation capabilities already demonstrated in the supervised (and continual) learning examples in the main paper. However, as the RL setting differs from the supervised and CL settings, present additional background material on it and thus App. C has been separated out to separate section in the appendix.
In this appendix, we provide background of the reinforcement learning problem and details of our model-based RL algorithm. See App. D.4 for an overview of the cart pole swing up environment and all details required for reproducing our experiments.
C.1 RELATED WORK IN THE SCOPE OF MODEL-BASER RL
A key challenge in RL is balancing the trade-off between exploration and exploitation (Sutton & Barto, 2018). One promising direction to balancing this trade-off is to model the uncertainty associated with a learned transition dynamics model and use it to guide exploration. Prior work has taken an expectation over the dynamics model's posterior (Deisenroth Rasmussen, 2011; Kamthe & Deisenroth, 2018; Chua et al., 2018), sampled from it (akin to Thompson sampling but referred to as posterior sampling in RL) (Dearden et al., 1999; Sasso et al., 2023; Osband et al., 2013), and used it to implement optimism in the face of uncertainty using upper confidence bounds (Curi et al., 2020; Jaksch et al., 2010). No single strategy has emerged as a go-to method and we highlight that these strategies are only good as their uncertainty estimates. Previous approaches have used GPs (Deisenroth & Rasmussen, 2011: Kamthe Deisenroth, 2018), ensembles of NNs (Curi et al., 2020; Chua et al., 2018) and variational inference (Gal et al., 2016; Houthooft et al., 2016). However, each method has its pros and cons. In this paper, we present a method which combines some of the pros from NNs with the benefits of GPs.
21

---

Published as a conference paper at ICLR 2024
C.2 BACKGROUND
We consider environments with states 8 E S C P= actions a E A C R. and transition dynamics f: x A S, such that 81+-1 = [(s, az) + ét, where e, is i.i.d. transition noise. We consider the episodic setting where the system is reset to an initial state 8y at each episode and we assume that there is known reward function S x A + R. The goal of RL is to find the policy π : S A (from set of policies II) that maximises the sum of discounted rewards in expectation over the transition noise,
DO J( = Eeg Exrr(se t=0
s.t. 8141 f(81, at) + 6t
(51)
where y [0, 1) is discount factor. In this work, we consider model-based RL where a model of the transition dynamics is learned fw- ~ f and then used by a planning algorithm. A simple approach is to use the learned dynamics model fw. and maximise the objective in Eq. (51),
Greedy arg max J(fw=, xell
(52)
However, we can leverage the method in Sec. to obtain function- space posterior over the learned dynamics (u [D), where D represents the state transition data set DD = {(s;, az), 8 J Im- portantly, the uncertainty represented by this posterior distribution can be used balance the exploration-explon exploitation trade-off, using approaches such as posterior sampling (Osband & Van Roy, 2017; Osband et al., 2013),
PS arg max (F,T) s.t. f~ (2) JEll
(53)
where a function / is sampled from the (approximate) posterior qu (D) and used find a policy as in Eq. (52). Intuitively, this strategy will explore where the model has high uncertainty, which in turn will reduce the model's uncertainty as data is collected and used to train the model. This strategy should perform well in environments which are hard to explore, for example, those with sparse rewards.
C.3 ALGORITHM
In practice, it is common to approximate the infinite sum in Eq. (53) by planning over a finite horizon and bootstrapping with a learned value function. Following this approach, our planning algorithim is given by,
H Σ
7'r(s,, |80 s Qa(sn, an):
PS (8) arg max max E B
(54)
t=
such that 8 [(0, a:) + ex, where f ~ Qu ( is a sample from the dynamics posterior and a. = al is a learned value function with parameters 0. We use deep deterministic policy gradient (DDPG, Lillicrap al., 2016) to learn the action value function Qo. Note that we also learn a policy but its sole purpose is for learning the value function.
Q0(s, a) xe E[ " 7'r(3;, a.)
Model predictive path integral (MPPI) We adopt model predictive path integral (MPPI) control (Pan et al., 2015; Williams et al., 2017) to plan online over a H step horizon and then bootstrap with learned value funetion. MPPI is an online planning algorithim which iteratively improves the action trajectory 0t:t и using samples. At each iteration N trajectories are sampled according to the currect action trajectory d' The top K trajectories with highest returns Σh=
i+1 ""r a/=..) + 70 ah) are selected. The next action trajectory 4 is then computed by taking the weighted average of the top K trajectories with weights from the softmax over returns from top K trajectories. After J iterations the first action is executed in the environment.
See App. D.4 for details of our experiments and results.
D EXPERIMENT DETAILS
In this section, we provide further details for the experiments presented in the main paper.
22

---

Published as a conference paper at ICLR 2024
SFR
Neural network prediction
Sparse function-space representation
HMC result as baseline
Figure A5: Uncertainty quantification for binary classification (E vs. @). We convert the trained neural network (left) to a sparse GP model that summarizes all data onto a sparse set of inducing points (middle). This gives similar behaviour as running full Hamiltonian Monte Carlo (HMC) on the original neural network model weights (right). Marginal uncertainty depicted by colour intensity.
Table bolding In tables throughout the paper, we mark the best-performing method/setting in bold. Following this, we check with a two-tailed t-test whether the bolded entry significantly differs from the rest. Those not differing from the best performing one with significance level of 5% are also bolded.
Toy example The illustrative examples in Figs. 1 and A5 include toy problems for regression and classification. The ID regression problem (Fig. 1) uses an MLP with two hidden layers (64 hidden units each, with tanh activation), whereas the 2D banana binary classification problem uses MLP with two hidden layers (64 hidden units each, with sigmoid activation). Both examples use Adam the optimizer. We include Python notebooks for trying baseline in Fig. A5 (right) was implemented with the help of hamiltorch². We used ten chains with 10,000 samples each and remove 2000 samples as burn-in. The HMC result more expressive but the results are still similar in terms of quantifying the uncertainty in low-data regions.
the models. The illustrative HMC
D.1 SUPERVISED LEARNING EXPERIMENTS
In this section, we provide details of the supervised learning experiments reported in the main table. In particular, we discuss our choice of reporting the results without any prior precision tuning in App. D.1.1, the details of the UCI data set experiments App. D.1.2 and the ones for image data sets in App. D.1.3.
D.1.1 PRIOR PRECISION TUNING
It is common practice tune the prior precision 8 after training the NN. This post hoc step is additional tuning to find a prior precision dune. which has a better NLPD on the validation set than the ó used to train the NN. Note that this leads different prior precision being used for training and inference. We highlight that this technically invalidates the Laplace/SFR methods, as the NN weights no longer the MAP weights. Nevertheless, in the following sections, App. D. and App. D.1.3, we report results both when the prior precision is tuned and when it is not. When tuning the prior precision 6, SFR matches the performance of the Laplace methods (BNN/GLM), In contrast the other methods, SFR already performs well without tuning the prior precision 5, whereas the Laplace methods require this additional post hoc optimization step to achieve good results.
1.1.2 UCI EXPERIMENT CONFIGURATION
In Sec. 5.1 in the main paper, we showed experiments on 8 UCI data sets (Dua & Graff, 2017). We used two-layer MLP with width 50, tanh activation functions and
https://gitt thub. com hwtp//aa/moomoaaammoomaa Itorch
23

---

Published as a conference paper at ICLR 2024
Table A4: UCI classification with prior precision 8 tuning Comparisons on UCI data with negative log predictive density (NLPD sstd, lower better). SFR, with M 20% of N, is on par with the Laplace approximation (BNN/GLM) and outperforms the GP subset when the prior precision (6) is tuned. SFR-NN uses the NN as its posterior mean function. GP subset, GP predictive, and SFR use the same inducing points.
DC C
NN MAP
LAPLACE full
LAPLACE GLM full
GP prodictive
GP subsct
SFR 20% of N M
SFR-KN
M 20% of N M
20% of N M
20% of N
AUSTRALIAN
690
14
2
0.33#0.00
0.33 :3.00
0.33 0.00
0.42 0.155
0.43+0,04
0.33 00007
0.33:00.07
CANCER
0.1010.04
0.1210.03
0.10
0.10
DIGITS
351
34
2
0.09=0001
0.14+0.02
0.10+0.01
0.09-0.01
0.22++0.09
0,11=0.01
0.09::0.01
GLASS
214
6
0.91
0.92 10.15
0.92
0.92.0.16
1.1810.16
0.92.10.05
0.93-10.17
IONOSPHERE
0.33==0.05
0.36
0.32-0005
0.34-0.08
0.41+0.05
0.32-00.07
0.33 06
SATELLITE
1000
21
3 0.27:10.02
0.30
0.27
0.51-0.57
0.34:00.01
0.26:10.00
0.27::0:02
VEHICLE
1797
64
10 0.35
0.37
00.04
0.35=0.05
0.34=0.05
0.46:20.04
0.35=0.05
0.34:0005
WAVEFORM
6135 35
6 0.33+50.0
0.001
0.33-0.01
0.33
0.40-+0.05
0.32 0.002
0.33
Table A5: UCT classification without prior precision á tuning: Comparisons on UCT data with negative log predictive density (NLPDeEstd, lower better). SPR, with M 20% of N, outperforms the Laplace approximation (BNN/GLM) (Daxberger et al., 2021), GP predictive from Immer et al. (2021b) and the GP subset when the prior precision (d) not tuned. SFR-NN uses the NN as the posterior mean function. GP subset, GP predictive and, SFR use the same inducing points.
N
D C
NX MAP
LAPLACE Full
LAPLACE GLM
GP predictive
GP subset
SFR
SFR-NX
full
M 20% of N
M
20% of N
M
201 of N
M
20% of IV
AUSTRALIAN
690
14
0.33==..0
0.71#0002
0.41#0.03
0.34=0.06
0.42=0.05
0.33=0.07
0. 0.33+0.00 0000
BREAST CANCER
683
10
0.10.::004
0.7320000
0.36.20.00 05
0. 0.22:0001
0.20 07
0.16::000 02
0.1610002
DIGITS
351
34
0.08..0.01 01
2.34.0.02
3.12
1.0910.00
1.10.00
1.07:0000 05
1 1.06.10.0 06
0.9420018
1.79.0.05
1.66:20.30
1.12:10.09
1.21 2:0.15
0.86:0.11
0.97:10:1
IONOSPHERI
18
0.33::0.08
0.70:0002
0.34:0004
0.34:0.05
0.40:20.04
0.32:0.06
0.33:20.05
SAVELLITE
1000
21
0.27--0.02
1.80++0.22
0.80+0005
0.29:0.01
9.34-0.01
0.26=0.00
0.27::0.02
VEDNEL:
1797
64
10
0.35.20000 05
1.30.20.02
1.5210.02
0.77.00.03
0.78.2000 0.
0.72:0004
0. 0.76:10.03
WAYUFORM
6435
35
0.33::001
1.10
1.08:0003
0.62:0002
0.82:00.05
0.52:0005
0. 0.65:20.03
70% (train): 15% (validation) 15% (test) data split. We trained the NN using Adam (Kingma & Ba, 2015) with a learning rate of 10 and batch size of 128. Training was stopped when the validation loss stopped decreasing after 1000 steps. The checkpoint with the lowest validation loss was used as the NN MAP. Each experiment was run for 5 seeds. For the HouseElectric experiment in Sec. 5.1 of the main paper, we keep the same 70% (train) 15% (validation) 15% (test) data split, but we switched to a deeper and wider MLP with 3 layers and 512 units per layer. with tanh activation functions. The MLP is trained with learning rate of 10 3 for 100 epochs and batch size of 128. To adhere with the number of inducing points reported in the paper of Wang et al. (2019), we used M 1024. We ran the experiment for five random seeds.
D.1.3 IMAGE EXPERIMENT CONFIGURATION
Sec. 5.1 in the main paper contains the experiments two image data sets, i.e., FMNIST (Xiao et al., 2017) and CIFAR-10 (Krizhevsky et al., 2009). For all the experiments, we used a CNN composed of three convolutional layers with ReLU activations, followed by three linear layers with tanh activation functions. This model follows the implementation standardized in Schneider et al. (2019) without using dropout layers, as also done in Immer et al. (2021b). We trained the NNs using Adam (Kingma& & Ba, 2015) for 500 epochs with a batch size of 128. For FMNIST we used a learning rate of 10 10-4 and for CI we increased it to 10-3, The training was stopped when the validation loss stopped decreasing after 15 steps. We performed a grid search to select the best value of the prior precision. We reported the runs performed with 8 = 0.0018 for CIFAR-I 10 and 8 = 0.0013 for F-MNIST. The runs are repeated five times with different random seeds. As discussed in Sec. 5.1 and in more detail App. D. we reported only the results obtained without any post hoc tuning for the prior precision in the main paper. However, here in Table A6, also report the results obtained by tuning the prior precision for SFR, GP subset and Laplace GLM/BNN For SFR and GP subset, we implemented the 5 tuning as a BO with 20 trials, while we used the built-in methods provided by the Laplace Redux library (Daxberger et al., 2021), with which we implemented the Laplace baselines.
The improvement in performance for the Laplace methods is notable for both FMNIST and CIFAR- 10. SFR matches their performance after tuning, with the advantage of performing better than them without the need of post hoc prior precision tuning. It is also important to note that in Table A6, we
24

---

Published as a conference paper at ICLR 2024
Table A6: SFR results on image data sets Classification results on image data using CNN. We report NLPD, accuracy, expected calibration error (ECE) and AUROC (mean over 5 seeds). SPR outperforms the SUBSET method and the LAPLACE methods without any post hoc tuning, and matches them after the post hoc prior precision tuning. We report the results for SFR using different M values, with and without a post hoc tuning for 8.
MODEL
M
NIPD
ACC. (Si f
ECF 4
MAP
0.23=0001
91.98-20.44
0.01
0.83:0.05
LAPLACE
2.42
10.21:0000
0.10
0.50=0003
LAPLACE KRON
2.39
9.87-40066
0.07=#0.01
0.51-+00.0
LAPLACE GLM DIAG
1.66:0.02
65.19:12
0.44::0.02
0.67:10.08
GLM KRON
1.09
84.79:11.00
0.47:10.01
0.96:10.01
GP PREDICTIVE
3200
0.47=00.6
91.51 4-0.45
0.23:40.03
0.97 00000
512
1.83:0.68
18.18114493
0.26::0.08
0.77: 10
1024
0.300
66.90:114.95
0.33:10.06
0.88:10.4
SUBSET
2048
0.97=0.13
77.32 8:8 8.888
0.30:10.08
0.93
3200
0.79=0009
82.52
0.29:40.05
0. 0
312
0.38
90.9910.55
0.15:10:01
0.95
1024
0.34
91.41 00046
0.12::0.02
0.95::0.00
SFR (Ours)
2048
0.30
91. 68 0.551
0.09::0.00
0.95++0.02
3200
0.29=0002
91. 74:10.00.4
0.08:10.01
0.96::0.0
P-NNIST
512
0.37=00.03
91.31 000.64
0.15:::0.0
0.97::00
SFR-NN (Ours)
3200
0.31=0003
91.86-00.40
0.10=:0.02
0.97 01
LAPLACE
1.88
41.022:4.00
0.32210:08
0.04
LAPLACE KRON
1.39:00.01
82.31 4:3.8
0.55
0.92++0.00
LAPLACE GLM DIAG
0.52:20.22
90. 00..97200.55
0.26:
0.96 02
GLM KRON
0.25=2001
91.78:0..8
0.05=10.00
0.91
GP PREDICTIVE
3200
0.23 0..00
92.09 ::0 00.20
0.01
1.87=+0.02
512
1.61:0.20
50,26 1..::13.00
0.22:10.09
0.80=0.07
cuning
1024
1.30=0.27
65.70:116.. 1.144
0.30
0.82:10.
SUBSET
2048
0.96
77.49 8..6
0.31:10007
0.93
3200
0.79=0009
82, 32.76
0.29
0. 9.95
312
0.37=0.01
99.6911..8
0.11::0:33
0.88:00
1024
0.34
91.38 0..00.000
0.12::0.02
0.95:00.02
SFR (Ours)
2048
0.00
91.57 0..44
0.09
0.95++0002
3200
0.29=0.02
01.82:00550
0.08:10.01
0.96:10.01
512
0. 0.23:00.0
91.87 0.55
O. 0.01+00.00
0.85-10005
SFR-NN (Ours)
3200
0.23:0001
92.14::.22
0.01:
0.86::0002
0.69=0.03
71.001
0.04
0.85
LAPLACE DIAG
2.37=0008
10.08
0.002
0..6:
LAPLACE KRON
2.36=2001
9.78:10.41
0.06::0.01
0.49:
GLM
1.33=0.05
71.96
0.39:10.02
0.82
LAPLACE GLM KRON
1.04=0.08
75.56
0.30+40.08
0.64 0.44
GP PREDICTIVE
3200
0.90=0002
76.07
0.23:10.02
0.79:10.02
312
1.5620.08
50.89155.30
0.20
0.64:10.66
1024
1.38
58.30
0.22:10.07
0.64::0.1 10
SUBSET
2048
1.08
66.40+13.00
0.24-10.04
0.71:1:0.05
3200
1.08=0.05
69.75113.22
0.24:10.03
0.75:10.00
512
0.86=0.02
77.58=0.90
0.24:10.01
0.80=00.2
SFR (Ours)
1024
0.79
78. 39 0..89
0.20:#0.01
0.80++0.02
2048
0.74=0002
78.40=0.68
0.16:10.00
0.79:10.00
3200
0.72 0.002
78 0.99
0.14:#0.01
0.79++0.00
CIFAR-I0 10
312
0.89=0.02
75.9721440
0.22:1002
0.7920.02
SFR-NY (Ours)
3200
0.79=0002
76.59
0.16:10.02
0.79-10.03
LAPLACE DLAG
2.13
28.89114.58
0.14::0.05
0.39 0..0
LAPLACE KRON
2.13:0001
28. 00:0000333.0
0.14:10.00
0.35: 06
0.7420.04
76.90
0.11::0.0 1:0002
0.81-10.00
LAPLACE GLM KRON
0.69
78. $3.95-.. 1.122
0.05-10.02
0.82- 0.03
GP PREDICTIVE
3200
0.68
76.91
0.02
0. 0.84:::0.0
312
1.5620.08
50.62
0.202:0:05
0.64:00.6
6 tuning
1024
1.37
17.95--61
0.22-10.06
0.65-10.10
SUBBSTT
2048
1.08
66.24133.8
0.24-10.04
0.71:+0005
3200
1.08=00.0
69.88.12.70
0.24:10.03
0.75:10.00
512
0.78=0.06
75.12
0.11 0.11::0.06
0.82-L0.02
SFR (Ours)
1024
0.74
77.19 11112
0.12::0.06
0.82 0.00
2048
0.71::0.03
77.48
0.11:10.05
0.82:10.0
3200
0.70+0..
77.67
11001
0.10+#0.04
0.82 0.00
312
0.68
76. 06.92211
0.02:0000
0.84
SFR NN (Ours)
3200
0.68
76.94
0.02:10.01
0.84:10.02
report the expected calibration enror (ECE) metric, which is problematic as a standalone metric and should be carefully assessed considering the accuracy. For instance, in the CIFAR-10 without 5
25

---

Published as a conference paper at ICLR 2024
SPR (Ours)
AUROC
AUROC
AUROC
AUROCAIS
AUROC a
AD 1OOD
PRED. ENTROPY
PRED, ENTROPY
PRED. ENTROPY
PRED. ENTROPY
PRED. ENTROPY
Figure A6: OOD detection with CNN: Histograms showing each method's predictive entropy at in-distribution (ID) data (CIFAR-10 10 test set, blue) where lower is better and at out-of-distribution (OOD) data (SVHN test set, red) where higher is better. The prior precision à was tuned post hoc.
tuning, the ECE is very low for Laplace DIAG and KRON even if the accuracy is very poor, meaning that the model is well calibrated but producing inaccurate predictions.
Table A6 results show how SFR consistently outperforms the GP subset method, with both lower and higher numbers M of inducing points. This shows that dual parameterization better captures the information from all training data.
D.1.4 OUT-OF-DISTRIBUTION DETECTION EXPERIMENTS
In Sec. 5.1 we evaluated SFR'S OOD detection. In all experiments, we tuned the prior precision post hoc. Fig. A6 provides further OOD detection results when training on CIFAR-10 and using SVHN (Netzer et al., 2011) as the OOD data set.
D.2 CONTINUAL. LEARNING EXPERIMENT DETAILS
In Sec. 5.2 in the main paper, showed how the regularizer described in Sec. can be used in CL to retain a compact representation of the NN. We conducted our experiments on three common CL data sets (De Lange et al., 2021; Pan al., 2020; Rudner et al., 2022):
Split-MNIST (S-MNIST), that consists of five binary classification tasks, each on pair of MNIST classes, i.e., @ and @
@ and @. ..... @ and D..
Split-FashhonMMNISTT (S-FMNIST), composed of five binary classification tasks, similar to S-MNIST, but using Fashion-MNIST.
Permuted-MNISS (P-MNIST), that consists of ten tasks, where each one is ten-way classification task on the entire MNIST data set, where the images' pixels are randomly permuted.
Following Pan al. (2020); Rudner et al. (2022), used two-layer MLP with 256 hidden units and ReLU activations for C.MIITT/SFMNIISS and a two-layer MLP with 100 hidden units and ReLU activations for P-MNIST, to ensure consistency with the other methods. In our experiments we only consider the single-head (SH) setting. This means that for S-MNIST and S-FMNIST, we are dealing with Innse-nacreettall Learning (CIL) scenario, where for each task pair of dijoint classes are presented to the model. The P-MNIST experiments fall into the Domain- Inciiiinnnneeeentalll Learning (DIL) setting, where the model deals with a distribittiittion shift, by progressively observing tasks with different permutations of the image pixels that need to be classified into ten different classes.
In our experiments, we compared the performance of the based regularizer against two families of methods: weight-space regularization methods and function- spaooe regularization ones. For the weight-space methods, we considered Online-EWCC (Schwarz et al., 2018), SI (Zenke al., 2017), and VCL (Nguyen et al., 2018). For the function-space methods, we considered DER (Buzzega et al., 2020), FROMP (Pan et al., 2020), and S-FSVI (Rudner et al., 2022).
For method, DER, OnWC and SI, relied on the Mammoth framework proposed by Buzzega et al. (2020). We used the available open-source FROMP codebase to conduct the experiments with FROMP, and the S-FSVI codebase to run the experiments S-FSVI and VCL. We selected the hyperparameters to use with all these methods by trying to adapt the ones pointed out in the original papers and reference implementations. For this reason, relied on SGD for DER,
26

---

Published as a conference paper at ICLR 2024
Online-EWC. and SI, which requires higher learning rates, compared to the others methods for which we used Adam (Kingma & Ba, 2015), i.e., our SFR regularizer, FROMP, S-FSVI, and VCL.
For the methods relying on stored subset of training data for each task, we make sure that the number of data points stored are the same for each method. For DER, which does not select the rehearsal training samples at task boundaries, cannot directly apply the definition of points per task (M), and thus we set the dimension of the reservoir buffer to be T- M, i.e., total number of tasks times the points per task.
The results reported in Table 2 are obtained using the following hyperparameters and repeating the runs five times with different seeds. The results in-line with what has been reported in previous work, notably even better for some of the baseline methods. Thus, we consider the presented results to give realistic picture of the relative performance of the different methods, demonstrating a clear benefit of using SFR for CL.
S-MNIST (40 pts./task): Online-EWCC: SGD, learning rate 0.03, batch size 10. number of epochs
1. A = 90, = 1.0. SI: SGD, learning rate 0.1, batch size 10, number of epochs 1, c 1.0, { - 0.9. VCL (with coresets): Adam, learning rate 0.001, batch size 256, number of epochs 100, selection method for coresets random choice. DER: SGD with learning rate 0.03, batch size 10, buffer batch size 10, number epochs 1, a = 0.3, reservoir buffer size 200. FROMP: Adam, learning rate 10-1, batch size 32, number of epochs 10, 10, random selection for memorable points. S-FSVI: Adam, learning rate 5 10 batch size 128, number of epochs 10. SFR: Adam, learning rate 3 10 10-4, 1, 8 0.0001, number of epochs 1.
batch size 32,
S-MNIST (200 pts./task):
Online-EWCC: SGD, learning rate 0.03, batch size 10, number of epochs 1 90, 7 = 1.0. SI: SGD, learning rate 0.1. batch size 10, number of epochs I, c 1.0, { = 0.9. VCL (with coresets): Adam, learning rate 0.001, batch size 256, number of epochs 100, selection method for coresets random choice. DER: SGD, learning rate 0.03, batch size 10, buffer batch size 10, number epochs 1, a 0.2, reservoir buffer size 1000. FROMP: Adam, learning rate - 10 batch size 32, number of epochs 10, 10, random selection for memorable points. S-FSVI: Adam, learning rate 5 10 batch size 128, number of epochs 80. SFR: Adam, learning rate 1 10 batch size 32, T = 0.5,
10- number of epochs 5.
S-FMNIST (200 pts./task): Online-EWCC: SGD, learning rate 0.03, batch size 10, number of epochs 5, 90, ? 1.0. SI: SGD, learning rate 0.1, batch size 10. number of epochs 5, c 1.0. <= 0.9. VCL (with coresets): Adam, learning rate 1 10-3 batch size 256, number of epochs 100, selection method for coresets random choice. DER: SGD, learning rate 0.03, batch size 10, buffer batch size 10, number epochs 5, 0.2, reservoir buffer size 200. FROMP: Adam, learning rate 1 10-4, batch size 32, number of epochs 10, T = 10, random selection for memorable points. S-FSVI: Adam, learning rate 5 10-4 batch size 128, number of epochs 60. SFR: Adam, learning rate 3- 10 batch size 32, T = 1.0, 5 = 0.0001, number of epochs 5.
P-MNIST (200 pts./task): Online-EWC: SGD, learning rate 0.1, batch size 128, number of epochs 10, A = 0.7. = 1.0. SI: SGD, learning rate 0. batch size 138, number of epochs 10, = 0.5, { 1.0. VCL (with coresets): Adam, learning rate 0.001, batch size 256, number of epochs 100, selection method for coresets random choice. DER: SGD, learning rate 0.01, batch size 128, buffer batch size 128, number epochs 10, = 0.3, reservoir buffer size 2000. FROMP: Adam, learning rate 1- 10-3 batch size 128, number of epochs 10, r = 0.5, lambda descend selection for memorable points. S-FSVI: Adam, learning rate 10-4 batch size 128, number of epochs 10. SFR: Adam, learning rate 3 10 batch size 64, r = 1.0, 8 = 0.0001.
Note that we used the same hyperparameters for the SFR ablation experiments.
D.3 INCORPORATING NEW DATA VIA DUAL UPDATES EXPERIMENT DETAILS
This section details how we configured and ran our incorporating new data via dual conditioning
experiments.
In all experiments, used two-layer MLP with 128 hidden units and tanh activations. We trained the NN using Adam (Kingma & Ba, 2015) with a learning rate of 10-4 and a batch size of 50. Training was stopped when the validation loss stopped decreasing after 100 steps. The checkpoint
27

---

Published as a conference paper at ICLR 2024
400
Episode Retum
300
200
100
SER (Ouns) LAPLACE-GLM ENSEMBLE
0
10
20
30
40
50
MLF
DDPG
Episade
Figure A7: Cartpole swingup with sparse reward: Training curves show that SFR's uncertainty esti- mates improve sample efficiency in RL. Our method (-) converges in fewer environment steps than the baselines. The dashed lines mark the asymptotic return for the methods coverged in the plot.
with the lowest validation loss was used as the NN MAP. Each experiment was ran for 5 seeds. The Gaussian likelihood's 's noise variance was set to " and trained alongside the NN's parameters.
Data split Each data set was split 50% 50% into an in-distribution (ID) set Dj and an out-of- distribution (OOD) set DOoD by ordering the data along the input dimension with the most unique values and splitting down the middle. The in-distribution (ID) data set D2 was then split 70% train and 30% validation. The out-of-distribute (OOD) data set DOOD was then split 70% for updates D2 and 30% for calculating test metrics.
The results in Table 3 were obtained with the following procedure. We first trained the MLP on Dj and calculated SFR's 's sparse dual parameters. Table reports SFR's test NLPD as well as the time to train the NN and calculate SFR's s sparse dual parameters (Train D1). We then took the trained NN and incorproated the new data D2 using dual conditioning from Eq. (17) (Updates w. D2) . Finally, compare incorporating new data via SFR's dual conditioning retraining from scratch. That is, reinitializing the NN and training on D2 U D2 (Retrain w. D2 U D2).
D.4 REINFORCEMENT LEARNING EXPERIMENT DETAILS
This section details how we configured and ran our reinforcement learning experiments.
Environment We consider the cartpole swingup task in MuJoCo (Todorov et al., 2012). However, make exploration difficult by implementing sparse reward function which returns 0 unless the reward is over a threshold value. That is, our reward function is given by,
((( ac), rf a:) 2 0.6
f(s a.)
00
otherwise
In all experiments we collected an initial data set using a random policy for one episode and we set action repeat as two.
Dynamics model In all experiments we used an MLP dynamics model with a single hidden layer of width 64 and tanh activation functions. At each episode used Adam (Kingma & Ba, 2015) to optimize the NN parameters for 5000 iterations with a learning rate of 0.001 and batch size of 64. We reset the optimizer after each episode. As we are performing regression instantiate the loss function in Eq. (1) as the well- mean squared error. This corresponds to Gaussian likelihood with unit variance. We then set the prior precision as 0.0001.
Value function learning We use DDPG to learn the action value function. DDPG learns both policy and value function but we do not use the policy. In our experiments, we parameterized both the actor and critic MLPs with two hidden layers of width 128 with ELU activations. We train them using Adam for 500 iterations at each episode, using a learning rate 0.0003 and a batch size of 512. DDPG used f = 0.005. DDPG is an off policy algorithm where the exploration policy samples from a noise process, which here Gaussian distribution with 0.1 and clipping 0.3.
a target value function to stabilize learning and for the soft target updates
28

---

Published as a conference paper at ICLR 2024
Model predictive path integral (MPPI) In all model-based RL experiments, we used a H = 5 step horizon and a discount factor of y = 0.9. We sampled N = 256 trajectories for = iterations with K 32. We used temperature of 0.5 and momentum 0.1.
MLP As a baseline, we ran experiments using only the MLP dynamics model with no uncertainty. As be seen in Fig. A7, this strategy is not always able to solve the task, indicated by the magenta line converging at ~ 360. This is expected because this strategy has no principled exploration mechanism.
used M 128 inducing points this seemed to be sufficient. The results in Fig. A7 indicate that SFR's 's uncertainty estimates are of high quality. This is because the agent is able solve the task with less environment interaction, indicated by the cyan training curve converging quicker.
SFR In the SFR experiments
Laplace-GGN In the Laplace-GGN experiments we used the Laplace library from (Daxberger et al., 2021) with the full hessian structure over all of the network's weights. We then used the GLM predictions. As expected, the Laplace-GGi is also able to solve the sparse reward environment.
Ensemble The ensemble experiments used 5 members where each member was an MLP with one hidden layer of width 128 and tanh activations, i.e. the same architecture as the other experiments. At each episode, the ensemble was trained for 10000 iterations using the Adam optimiser. We used a learning rate 0.0001 and a batch size of 64. At each time step in the planning horizon we used the following procedure to sample from the ensemble. First, the mean and variance of the ensemble members predictions were calculated. Next, we sampled from the resulting Gaussian. Fig. A7 shows that the ensemble required more environment interaction to solve the task. This is indicated by the green line not converging within the 50 episodes shown in the figure.
DDPG The DDPG experiment serves our model-free RL baseline. In the DDPG experiment use a standard deviation of 0.6 for the actor to help it explore the environment. We increased the width of the MLPs 512 and we increased the discount factor y 0.99. As expected, the results show that DDPG requires lot more environment interaction to solve the task. Fig. A7 also shows that DDPG does not converge to the same asymptotic performance as SFR. This is due some of the random seeds not being able to successfuly explore the environment.
Summary Our results motivate further experiments looking into NN uncertainty quantification for model-bases RL. To the best of our knowledge, we are the first to leverage the Laplace-GGN for mosel-based RL. Our results indicate that both the Laplace approximation and SFR could offer benefits over ensembles of NNs, which are typical in model-based RL. Although further experiments are required, our results indicate that SFR could offer better uncertainty estimates these tasks. We believe this may be due to SFR sampling in function space as opposed to weight space like Laplace-GGN.
E COVARIANCE VISUALIZATION
(a) Full GP N 200
(b) GP subset w. M = 40
(c) SFR W. M 40
Figure A8: Comparison of covariance matrix for a full GP vs. GP subset vs. SFR. The contour plot shows the covariance matrix evaluated at the training inputs 1e x of the ID regression example from Fig. 1.
Fig. A8 compares the covariance structure of SFR and the GP subset for the ID regression problem in Fig. 1. It shows that SFR accurately recreates the full solution despite having a limited number of inducing points. In contrast, the GP subset (with the same computational complexity in terms of M) fails to accurately recreate the full solution with the same number of inducing points.
29


Reviews:
[{"review_id": "jENOiQguIw", "reviewer": "Reviewer_V1B4", "summary": "This paper introduces a new method for approximating a weight-space neural network using a function-space Gaussian process. To overcome the inhibitive computational complexity it then uses two techniques: a dual parameterization using a Laplacian approximation (equation 10) and sparsification (equations 12 and 13). The resulting sparse GP doesn't require any further training, but makes it easy to incorporate new data, which makes it suitable for continual learning. The paper proceeds with a series of experiments which show that the proposed model is effective at sparsifying (figure 3 and table 2), uncertainty quantification (table 1), continual learning (table 3), and incorporating new data (table 4).", "questions": "These might be obvious to someone more familiar with the literature, but questions that I was left with were:\n\n* The final model (equation 12) seems to rely on inversing $\\\\bf{K_{zz}}$. How are the inducing points to form this matrix chosen?\n* In the paragraph \"GPs\" it says that $\\\\bf{\\Lambda}$ can be considered per-input noise. How so?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "I must preface this review by saying that this paper is not in my field of expertise. However, as someone who is unfamiliar with the topic, I found this paper well written and clear. The authors do a good job reviewing the set of results from the literature that they build on, before presenting their own model. They argue convincingly show that their model has practical advantages (i.e., for incorporating new data and continual learning), and the experiments seem to support the theory.", "weaknesses": "Perhaps the text would benefit from some pseudo-code or explanations that relate the theory to the actual practical algorithm, particularly for a reader like myself who is unfamiliar with the related literature.\n\nI am also not too familiar with the continual learning or reinforcement learning literature, so I find it difficult to judge whether table 3 and figure 4 present a competitive set of baselines, but the results seem compelling.\n\nMinor:\n\n* \"adpot\" should be \"adopt\" in section 3\n* \"subest\" shoudl be \"subset\" in \"Sparsification via dual parameters\"", "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": "N9gHZ85lJt", "reviewer": "Reviewer_7JfH", "summary": "The paper presents a new method to turn a pre-trained NN, trained with the MAP objective, into a GP model with a kernel based on the Neural tangent kernel. They then show its usefulness on many different tasks.", "questions": "Please explain how you arrive at eq. 7-11", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "- The method, to the best of my knowledge is novel.\n- The experiments cover a diverse set of problems", "weaknesses": "- My main issue with this paper is that it does not explain or derive any of its equations, not even in the supplementary. It is not clear how the authors arrived at eq. 7,8,10 & 11. It isn't even clear what the authors mean by dual parameterization (dual in what sense?) as it was never explained. While basic things like Laplace approximation and sparse GP are explained, the main points of this paper aren't.\n- I am unclear about the issue with adding points to the standard GP. Once you have your kernel, given to you by the MAP trained NN, you can use it on any number of points you want. You can also train a standard sparse-GP approximation using the fixed kernel. The authors claim that this is an issue that their work solves, but I don't see any reason why it should be so and they do not explain it properly.\n- Baselines in many of the experiments are very basic. For instance, in Table 1 you compare to Laplace approximation and using a simple subset selection sparse GP. Other Bayesian NN methods should be compared against, for example NN-GP, deep kernel learning, SWAG.\n- I would also point out another relevent work on incremental learning by GPs Achituv et al \"GP-Tree: A Gaussian process classifier for few-shot incremental learning\"", "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": "d95BwJg5yg", "reviewer": "Reviewer_C2Qu", "summary": "This paper extends recent dual parameterization approaches for sparse GPs to function space inference for neural networks. This allows for efficient uncertainty estimation and the authors further formulate a continual learning objective as well as a variation that allows for incorporating new data without retraining. The experiments report competitive performance compared to Laplace- and GP-based baselines from the literature for uncertainty estimation on MNIST and CIFAR scale datasets, continual learning on MNIST and model-based reinforcement learning for the cart pole mujoco challenge.\n\nOverall, the paper is quite rich and in terms of its technical contribution more than sufficient for a conference paper. Unfortunately, there is almost so much going on that the paper lacks in clarity and focus. The experiments tend to be somewhat small scale, so that I find it hard to argue for acceptance on a purely quantitative basis, and qualitatively the most interesting aspect - the possibility of incorporating new data without re-training - is mostly an afterthought in the empirical evaluation.\n\nAll things considered, I would lean towards rejection at this point, but with a thorough update of the manuscript would hope to increase my score over the course of the rebuttal.\n\nEDIT: in light of the rebuttal, I would now tend towards acceptance.", "questions": "1. How are the inducing points chosen? This isn’t discussed in the main text, they are simply assumed to exist around Eqs 12 & 13. The pseudocode in the appendix seems to suggest they are sampled from the real data?\n2. Similarly, the continual objective in Eq 14 appears out of nowhere. I appreciate that there is a derivation in the appendix, but there is no reference to this in the main text. Does the objective correspond to some principled approximation to the posterior as in e.g. VCL or online Laplace (although these of course accumulate approximation errors)? Or is it some ad-hoc expression to maintain close predictions on old data with previous parameters?\n3. The paper states at the end of the related work section “In contrast to the method presented in this paper, these functional regularization methods lack a principled way to incorporate information from all data. Instead, they rely on subset approximations”, but this seems to be the case for the present paper as well (assuming I'm not mistaken in Q1)? I understand that the dual parameters ultimately are determined as a sum over all data, yet the projection onto the inducing points does seem like a subset approximation to me, in particular since the regularization term is a sum over the inducing points, i.e. a subset.\n4. Would optimizing the inducing points render the method equivalent to (Ortega et al., 2023)? If so, that seems like a highly relevant baseline/ablation.\n5. The objective in Eq 14 seems quite closely related to that formulated in (Rudner et al., 2022) — which is admittedly rather generic. Would your approach fit in as a variational method or is it lacking the penalty terms related to the precision matrix?\n\nMinor:\n- I think there is a mixup in the bibtex entry for the VCL citation, this should presumably be (Nguyen et al., Variational Continual learning, ICLR 2018)\n- Table 1: MAP seems to perform by far the best on half the datasets, do you have any comments on this? It might also be worth adding a HMC baseline (at least if you think of your method as a BNN method, I don’t think this is stated anywhere explicitly).\n- Table 2: accuracy values are decimals rather than percent as indicated in the header\n- Table 2: including the Laplace results without tuning the prior precision seems fairly pointless, it is well known across the literature that sampling from the Laplace approximation in weight space does not work well without cross-validation or online tuning. Why not use the latter if rejecting the former due to breaking the MAP property of the parameters?\n- NLPD should be defined\n- Laplace Kron should be referenced/cited in the experiments\n- \"Hessian is often infeasible and in practice it is common to adopt the GGN approximation\" This seems like a slightly inaccurate statement, infeasible is calculating the full matrix, the GGN approximation is to ensure positive definiteness\n- Why the two separate for loops in Algorithm 1? It seems like the SLICE variables just add notational clutter and the updates to the sum could be done in a single loop?\n- The appendix mentions using double precision for the post-training computations, is that a precautionary choice or do e.g. the matrix inversions fail in single precision?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- there is quite a lot going on technically, extending recent dual parameterization methods for GPs to neural networks not only for uncertainty estimation, but also continual learning\n- there is a wide range of experiments, I would not have expected the model-based RL benchmark on top of the uncertainty estimation/continual learning problems\n- prior work that the paper builds on is clearly referenced\n- the possibility of incorporating new data without re-training is potentially practically impactful", "weaknesses": "- Largely due to the breadth of material covered, I find that the paper somewhat lacks depth and details in the main text. I appreciate that the authors don’t attempt to break their work into the smallest publishable units, but with the conference format being what it is, in my opinion there is also a case to be made for splitting work e.g. into a generic uncertainty estimation method and then extending it to continual learning in a separate paper in order to leave enough space for covering enough details on the method/design choices (Q1&2) and discussing the relationship with prior approaches (Q3-5).\n- The title seems to place the emphasis on the continual learning part of the paper, however the empirical evaluation here is extremely small scale (nothing beyond MNIST-scale). I would find it difficult to argue for the method on a quantitative basis and would find e.g. Split CIFAR results helpful.\n- Similarly the in-distribution uncertainty evaluation is a bit limited. I appreciate that these experiments were run as sanity checks, but to me it would be sufficient to leave them in the appendix (or run the full suite of uncertainty evaluation experiments on OOD/shifted data to make a more convincing quantitative case for the method).\n- The evaluation of the dual parameter updates is mostly an afterthought, Table 4 seems to only show that it is possible in principle. I think this is a real shame, I’m not aware of any other near-zero cost approach for incorporating new data. I wish this was a much bigger focus in the empirical evaluation and there was more of an attempt to push the limits here, e.g. by using smaller and smaller initial datasets for learning the NN parameters, longer streams of data or two distinct datasets, e.g. MNIST/FashionMNIST or CIFAR10/SVHN where the learned parameters may be less relevant for the second datasets (essentially I would like to see the boundaries being pushed and get a sense of where the dual parameter updates start to fail).\n- The continual objective seems to scale linearly in the number of tasks (both in terms of memory and computational complexity), in contrast to weight space methods that tend to add a constant overhead.", "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)