Efficient parametric approximations of neural net function space distance
Authors: Nikita Dhawan
Sicong Huang
Juhan Bae
Roger Baker Grosse
Venue: Submitted to ICLR 2023 (2023)
Keywords: Function space distance
memory-efficiency
continual learning
influence function estimation

Abstract: It is often useful to compactly summarize important properties of a training dataset so that they can be used later without storing and/or iterating over the entire dataset. We consider a specific case of this: approximating the function space distance (FSD) over the training set, i.e. the average distance between the outputs of two neural networks. We propose an efficient approximation to FSD for ReLU neural networks based on approximating the architecture as a linear network with stochastic gating. Despite requiring only one parameter per unit of the network, our approach outcompetes other parametric approximations with larger memory requirements. Applied to continual learning, our parametric approximation is competitive with state-of-the-art nonparametric approximations which require storing many training examples. Furthermore, we show its efficacy in influence function estimation, allowing influence functions to be accurately estimated without iterating over the full dataset.

Under review as a conference paper at ICLR 2023
EFFICIENT PARAMETRIC APPROXIMATIONS OF NEURAL NETWORK FUNCTION SPACE DISTANCE
Anonymous authors Paper under double-blind review
ABSTRACT
is often useful to compactly summarize important properties of a training dataset that they can be used later without storing and/or iterating over the entire dataset. We consider a specific case of this: approximating the function space distance (FSD) over the training set, i.e. the average distance between the out- puts of two neural networks. We propose an efficient approximation to FSD for ReLU neural networks based on approximating the architecture as linear net- work with stochastic gating. Despite requiring only one parameter per unit of the network, our approach outcompetes other parametric approximations with larger memory requirements. Applied to continual learning, parametric approxi- mation is competitive with state-of-the- art nonparametric approximations which require storing many training examples. Furthermore, show its efficacy in influence function estimation, allowing influence functions to be accurately esti- mated without iterating over the full dataset.
INTRODUCTION
There are many situations in which would like to compactly summarize model's 's training data. One motivation is to reduce storage costs: in continual learning, an agent continues interacting with its environment over long time period longer than is able store explicitly but we would still like it to avoid overwriting its previously learned knowledge as it learns new tasks (Goodfellow) al. 2013). Even in cases where it is possible to store the entire training set, one might desire compact representation in order to avoid expensive iterative procedures over the full data. Examples include influence function estimation (Koh & Liang, 2017 Bae al. 2022a), model editing (De Cao et al. 2021: Mitchell et al. (02T), and unlearning (Bourtoule et al. (021). While there are many different aspects of the training data that one might like to summarize, we are often particularly interested in preventing the model from changing its predictions too much on the distribution of previously seen data.
Methods to prevent such catastrophic forgetting, especially in the field of continual learning, can be categorized at a high level into parametric and nonparametric approaches. Parametric approaches store the parameters of a previously trained network, together with additional information about how important different directions in parameter space are for preserving past knowledge; the canonical example is Elastic Weight Consolidation (Kirkpatrick et 2017 EWC), which uses diagonal approximation to the Fisher information matrix. Nonparametric approaches explicitly store a collec- tion (coreset) of training examples, often optimized directly be the most important or memorable ones (Rudner et al. 2022 Pan et al., 2020; Titsias et al., 2019). Currently, the most effective ap- proaches to prevent catastrophic forgetting are nonparametric, since is difficult to find sufficiently accurate parametric models. However, this advantage is at the expense of high storage requirements.
We focus on the problem of approximating function space distance (FSD): the amount by which the outputs of two networks differ, in expectation over the training distribution. [Benjamin et al. (2018) observed that regularizing FSD over the previous task data is an effective way to prevent catastrophic forgetting. Other tasks such as influence estimation (Bae et al.) 2022a). model edit- ing (Mitchell et al., 202I). and second-order optimization (Amari 1998; Bae et al. 2022b) have also been formulated in terms of FSD regularization equivalent locality constraints. In this pa- per, we formulate the problem of approximating neural network FSD and propose novel parametric approximations. Our methods significantly outperform previous parametric approximations despite

---

Under review as a conference paper at ICLR 2023
After Task
After Task
EWC
NTK BGLN
8
-20
Figure I: Comparison of FSD regularization on one-dimensional reggessson task. (Left) Training sequentially on tasks (blue, then yellow), results in catastrophic forgetting. (Right) BGLN retains performance on Task after training on Task more accurately than other methods. Note that the y-axis represents function space distances for each datapoint.
being much more memory-efficient, and are also competitive with nonparametric approaches to continual learning.
Several parametric approximations, like EWC, are based second-order Taylor approximation to the FSD, leading to a quadratic form involving the Fisher information matrix Fg or some other metric matrix Ge. where 0 denotes the network parameters. Second-order approximations help because one can approximate Fg or Go by sampling vectors from a distribution with these matrices the covariance. Then, tractable probabilistic models be fit these samples to approximate the corresponding distribution. Unfortunately, these tend to be inaccurate for continual learning, especially in comparison with nonparametric approaches. We believe the culprit is the second- order Taylor approximation: we show in several examples that even the exact second-order Taylor approximation can be poor match to FSD over the scales relevant continual learning, like average performance over sequentially learned tasks. This is consistent with recent line of results that find linearized upproximations of neural networks be inaccurate model of their behavior (Seleznoval
& Kutyniok 2022a b Hanin & Nica 2019
Bai et al. 2020 Huang & Yau 2020
Our contribution, the Bernoulli Gated Linear Network (BGLN), makes a parametric approxima- tion to neural network FSD which does not make a second-order Taylor approximation in param- eter space, and hence is able to capture nonlinear interactions between parameters of the network. Specifically, it linearizes each layer of the network with respect its inputs. In the case of ReLU networks, our approximation yields a linear network with stochastic gating. Linearizing the ReLU function requires computing its gradient, which can be approximated as an independent Bernoulli random variable for each unit. We derive a stochastic as well as a deterministic estimate of the FSD in this setting, both of which rely only on the first moments of the data.
To demonstrate the practical usefulness of our approximation, we evaluate its closeness to the true empirical FSD. We show that method estimates and optimizes the true FSD better than other estimators in settings that are prone to forgetting. Further, we show its application and performance in two applications. When applied to continual learning, it outcompetes state-of-the-art methods on sequential MNIST and CIFAR 100 tasks, with at least 90% lower memory requirements than nonparametric methods. When applied influence function estimation, our method achieves over 95% correlation with the ground truth, without iterating over or storing the full dataset.
BACKGROUND
Let = = f(x, 0) denote the function computed by a neural network, which takes in inputs x and pa- rameters O. Consistent with prior works, we use FSD to refer to the expected output space distance")
between the outputs of two neural networks (Benjamin et al.] 2018 Grosse 2021: [ae et al.)
2022b), with respect to the training distribution, as defined in equation When the population
Note that we use the term distance throughout since we focus on Fuclidean distance in practice. However, other metrics like KI. divergence, which not distances, also possible and commonly used.
2

---

Under review as a conference paper at ICLR 2023
distribution is inaccessible, the empirical distribution is often used as a proxy:
NV
D(0, Pdau)) = Ex--pan (e [x( 0o), f(x, 01)) =
EM(((x)) O0), f(x), (1)
N
where x is random data point sampled from the data-generating distribution Pdata. Constraining
the FSD term has been successful in preventing catastrophic forgetting (Benjamin et al. 2018).
computing influence functions (Bae et al. 2022a), training teacher-student models (Hinton et al.
2013), and fine-tune finettnnng pre-trained networks (Jiang et al. 2019 Mitchell et al. 2021).
Consider the continual learning setting as a motivating example. Common benchmarks (Normandin ct al. 2021 involve sequentially learning tasks t E {l...., ... T., using cost function J and penalty on the FSD between the learned parameters, 0 and optimal parameters of previous tasks, 0,, com- puted over the previously seen data distribution, pi, and scaled by hyperparameter, ÁPSD-
[-1
0: = arg min J(0) + AFSD Σ D(0,0..Pt)
0
Continuing with the notation in equation one way to regularize the (empirical) FSD is to store the training set and explicitly evaluate the network outputs with both 0g and @ (perhaps on random mini-batches) However, this has the drawbacks that one would need to store and access the entire training set throughout training (precisely the thing continual learning research tries to avoid) and that the FSD needs be estimated stochastically. Instead, we would like to compactly summarize information about the training set or the training distribution.
Many (but not all) practical FSD approximations are based on a sccond-order Taylor approximation:
1 D(Oa, 01, (dmaa) 2 - 0o) T Go(O: Oo):
where Go ((,, 0, P. In the case where the network outputs parametrize probability distribution and p corresponds to KL divergence, Gg reduces to the more familiar Fisher information matrix Fo Ex-pan. y- (0) Vo log p(y 0, Vo log p(y x) [] where (x) represents the model's predictive distribution over y. It is possible to sample random vectors in parameter space whose covariance is Go (Martens et al. 2012 Grosse & Martens] 2016 Grosse (021). and some parametrie FSD approximations work by fitting simple statistical models to the resulting distribu- tion. For instance, assuming all coordinates are independent gives a diagonal approximation (Kirk- patrick al.) 2017), and more fine-grained independence assumptions between network layers yield a Kronecker-factored approximation (Martens & Grosse 2015 Ritter et al., 2018). In practice, instead of sampling vectors whose covariance is Ge. many works instead use the empirical gradi-
ents during training, whose covariance is the empirical Fisher matrix. We caution the reader that the empirical Fisher matrix less well motivated theoretically and can result in different behav-
ior (Kunstner et al. 2019).
3 A DATA-BASED PARAMETRIC APPROXIMATION
Motivated by the above, we propose BGLN (Bernoulli Gated Linear Network), which approximates a given model architecture as a linear network with stochastic gating and captures more nonlineari- ties between the parameters than the previous parametric approximations discussed. is applicable to different architectures, but we first explicitly derive our approximation for multilayer perceptrons (MIPs), with T. fully-connected layers and the ReI.U activation function, ń. We then discuss and empirically evaluate its generalization to convolutional networks. In Section 3.4 we extend our
method to a class-conditioned version which takes into account different classification categorics.
For MLPs with inputs x, layer weights and biases (W(), (()) and outputs z. the computation of preactivations and activations at each layer is recursively defined as follows.
(1) ( a (-1) + (()), (1) = ((s
(4)
with (0) x. and (L z. We denote 2a and z to be the outputs obtained with parameters Oo and 0, respectively. With Euclidean distance as the output space distance, P. we can rewrite the FSD as sum of the first two moments of the difference between the output vectors Az := 21 20-
1 B | 1 ||[ + 2" Cov(Az)
Ex (p(O,, x))
(5)
3

---

Under review as a conference paper at ICLR 2023
Algorithm BGLN-S
Require: E x Cov(x), { (W, 1)1-1 ,)1-",((111 (K1)1
1: Blao]; [[a + E[x]
2: Cov(ao); Cov <- Cov(x)
3: do: d1 -V [n): Cov(ao))
sample inputs using first two moments
4: Aa f 0
5: for
1 to L do
6:
m ~ ))r(
> sample Bernoulli random variables
8g wi ag b
7:
linear layers preserve input-output relationship
8:
As + AW a
W (1) - - Ab(1)
9:
dg < m ④ 80
stochastic gating
10:
Aa - m ④ As
11:
81 + 1 + b
linear layers preserve input-output relationship
12: 01 < ag + Aa
b first-order Taylor approximation activations
13: end for
14: Az + AW ao + W Aa + Ab
return 2 A2|
3.1 LINEARIZED ACTIVATION FUNCTION
Given parameters 0g and of two networks, we linearize each step the computation with respect to its inputs. For an MLP that alternates between linear layers and non-linear activation functions, the linear transformations are unmodified (since they are already linear), while the activation functions are replaced with a first-order Taylor approximation around their current inputs. Let (W!"). (()) be
the weights and biases of layer in network é.
$u w (1) du
(1-1) + 600
(6)
(1) (Bb²)
a
(7)
8 ( wl'aa aj (1-1) + +b
(8)
(1) = + "
(0) (0)
(9)
where o' is the derivative the ReLU function defined as d'(s) 1(s > 0}. We define some additional useful notation for differences between preactivations and activations.
As (1) 8
0 = AW a + w (l-111 + A
(1)
(10)
Aa (1) ai a (1) 0'(00 As
(1)
!1)
(11)
with AW O w (1) and base cases dg
(0) and Aa 0. Observe that with Wo
held fixed, the model parametrized using W, is a linear network, i.e the network's 's exact outputs
a linear function of its inputs. There are two significant differences between our approximation and parameter space linearization, which confer an advantage to the former. First, our linearization with respect to inputs instead of parameters (Lee et al., 2019), hence capturing nonlinear inter- actions between the parameters in different layers. Second, the only computations that introduce linearization errors into our approximation are those involving activation functions, contrast with other methods which suffer linearization errors for each layer, including linear layers, where non- linear parameter dependencies exist. note that linear networks are commonly used to model nonlinear training dynamics of neural networks (Saxe et al.) 2013), hence regard our weaker form of linearity as a significant advantage over Taylor approximations in parameter space. this form, the lincarized activation function trick applies any linear (including fully-connected and convolutional) network with nonlinear activations.
3.2 BERNOULLI GATING
In the specific case of ReLU networks, our approximation only depends the training data through the signs of all the preactivations, since passing layer inputs through 6 and multiplying by the deriva- tive & both result in multiplication by 0 or 1. Hence, the maximum likelihood estimate (MLE) of
4

---

Under review as a conference paper at ICLR 2023
this multiplication factor corresponds to the MLE of a Bernoulli random variable and can be mod- eled as such. To achieve a compact approximation, we implement stochastic gating by replacing the activation signs with independent Bernoulli random variables for each unit and making a Gaussian approximation to the input distribution. We denote the vector of Bernoulli random variables for layer I as m which are summarized by their corresponding mean parameters, f( Now we can rewrite equations and 11 as (1) = m @ 8o element-we wise multiplication. Since FSD is defined as an expectation over the training distribution, are interested in the moments of and Åa, which in turn depend on the first two moments of the input data and means μu. We estimate μ with the maximum likelihood estimate (MLE) of m, which is the average number of times given unit is during training. This can be computed efficiently in the forward pass as a simple moving average during the last epoch of training on given dataset. The data moments are then propagated through the network, and multiplied by Bernoullis required, finally resulting in the FSD term in Equation 5 The steps for this computation are detailed in the following Section 3.3 and Algorithm
(9) and (a = m @ As(!) respectively, where @ denotes
3.3 BGLN-D AND BGLN-S
We observe from equation 5 that the expectation of FSD can be written in terms of the first two moments of Az, or equivalently As(") Expectations and covariances of the outputs of each layer can be computed terms of the expectations and covariances of the outputs of the previous layer. We
(1) propose computing these terms recursively by propagating the first two moments of dg and Aa(!) through the network. This can be done under our independence assumptions and using linearity of expectations. We thus obtain deterministic estimate of the expected FSD, which we refer to as BGLN-D. The exact equations for this computation are shown in Appendix B We can also compute stochastic estimate by starting with the first two moments of the inputs, sampling (")) from multivariate Guassian, and sampling Bernoulli random variables at every ReLU layer. This method, BGIN-S, shown in Algorithm [ Note that this is an unbiased estimate of the expected FSD. While we explicitly derive it for MLPs here, this estimate is not specific only MLPs, It can be easily generalized to convolutional networks where the same random variables are sampled as required in the network. We use and evaluate this generalization for continual learning task involving CIFAR-100, where the commonly used model architecture includes convolutions. We describe the corresponding computations in detail in Appendix
Note that storing moments of the data has less memory cost than storing sufficient subsets of the data itself. We present the above algorithms using the mean and co- variance of the data. However, we point out that the data covariance can be re-
S-ESVI
EWC
Number of data poins stored
3130
FROMP
BGLN
33:0
placed by its diagonal approximation, i.e.
830
the variance of each dimension of the in-
puts. This imposes an independence as- sumption on the dimensions of the in-
put data, analogous to the independence of Bernoulli random variables represent- ing activations for each layer. This further reduces the memory cost of storing mo- ments the equivalent of two data points per dataset. We empirically investigate the effect of this approximation on continual learning benchmarks.
Number tasks
Figure Comparison of number of datapoints stored for Split CIFAR 100 as number of tasks increases. EWC stores none and BGLN stores the equivalent two datapoints in the form of their moments, which is far lower than statc-of- the-art nonparametric methods.
3.4 CLASS-CONDITIONAL ESTIMATES
is also possible to a extend our method to a more fine- aneegraaddd, casss-coniitooed approximation. The lower memory cost of our method allows for this. In particular, we can fit a mixture model, with one component per class. In this case, each class has its own associated data moments and mean parameters for the Bernoulli random variables. As expected, this results in improved performance in
5

---

Under review as a conference paper at ICLR 2023
continual learning experiments, shown by the results for BGLN-D (CW) and BGLN-S (CW)
in Tables and 2
4 RELATED WORKS
Measuring distances between network parameters has been studied well in the existing literature. Here, we categorize the different approaches that are relevant to our proposed method.
Network distance estimation. Several works (Benjamin et al., 2018 Bernstein et al., 2020 Bae
jet al,, 2022b) have highlighted the importance of measuring meaningful distances between neu- ral networks. Benjamin et al. (20T8) contrast training dynamics in parameter space and function space and observe that function space distances are often more useful than, and not always corre- lated with, weight space distances. [ac al. (2022b) propose amortized proximal optimization scheme (APO) that regularizes an FSI) estimate to the previous iterate for second-order optimiza- tion. Natural gradient descent (Amari et al. 1995 Amari 1998) be interpreted
steepest descent method, using a second-order Taylor approximation to the FSD (Pascanu & Bengio 2014).
There are parametric approximations (Kirkpatrick al. 2017 Ritter al. 2018 of the FSD, which only utilize the parameters (and potentially related information) of the two networks in question. Nonparametric approximations typically store or optimize a small set of training inputs. For in- stance, Hessian-free optimization (Martens et al. 2010)) and APO (Bae et al. (002b)) can be inter- preted using batches of data to approximate the FSD. Several non-parametric methods (Pan al.] 2020 Titsias et al. [219 Rudner et al. 2022 Kapoor et al. 2021 for continual learning have been shown to give generally better empirical performance than parametric counterparts.
Continual Learning. [Parisi et al. (2019); De Lange et al. (2021); Ramasesh et al. (2020); Nor- mandin et al. (202T) have reviewed and surveyed the challenge of catastrophic forgetting in contin- learning, along with benchmarks and metrics to evaluate different methods. Parametric methods focus on different approximations to the weight space metric matrix discussed above, like diagonal (Kirkpatrick et al. 2017 EWC), Kronecker factorized (Ritter et al.) [2018, OSLA) ones. Several methods are motivated as approximations to a posterior Gaussian distribution in a Bayesian set- ting (Ebrahimi et al. 2019), for instance through variational lower bound (Nguyen al. (017) or via Gaussian process inducing points (Kapoor el al.] [0]]202]],, Non-parametric methods (Kapoor et al.)
ual T
2 Titsias et al. 2019 Pan al. 2020 Rudner et al. 2022: Kirichenko al. 2021) usually em- play some form of experience replay of stored data points. Some of these methods (Pan et al. 2020) can be related to the Neural Tangent Kernel (Jacot al.] 2018, NTK), or other words, network linearization. Doan et al.] (202T) directly study forgetting in continual learning in the infinite width NTK regime. Mirzadeh et al. (2022) further study the impact of network widths on forgetting.
5 EXPERIMENTS
We evaluate the benefits of our method in practice investigating the following questions: (1) What advantages do we gain by using our method in continual learning settings, relative existing para- metric and nonparametric FSD estimation techniques? (2) How well does our approximation esti- mate and minimize the true empirical FSD? (3) Can we use our FSD approximation to accurately compute influence functions without iterating through the full training data? In the results below, we find that (1) our methods are competitive with prior continual learning approaches despite hav- ing lower memory requirements than them, (2) it closely matches the true empirical FSI) and (3) it gives a high correlation with influence functions using only the training data moments. Influence estimates can also be applied to identify mislabeled examples, which our method is able to do as well, as detailed in Appendix
5.1 CONTINUAL LEARNING
Following the formulation described in equation we empirically evaluate BGLN and its variants.
Datasets and architectures. We visualize our method's performance on one-dimensional regression with two sequential tasks shown in Figure [T] More realistically. we test our method on image classification benchmarks used in prior works (Pan al. 2020 Rudner et al. 2022). including Split

---

Under review as a conference paper at ICLR 2023
Table 1: Average classification accuracy across tasks continual learning benchmarks and their associated memory costs. Methods are categorized as parametric, nonparametric, our BGLN methods, and ablations.
Method
Split MNIST
Permuted MNIST
Split CIFARI00
Memory Cost
EWC
63.10
84.00
71.60 + 0.40
2P P+ P+Σ PT 2P
OSLA
80.56
95.73
72.61
VCL
97.00
87.50 + 0.61
VCL (coreset)
98.40
95.50
67.40 + 0.60
2P + Nd
VAR-GP
90.57 + 1.06
97.20 0.08
2P Nd C N2
FROMP
+ 0.04
94.90 0.04
76.20 + 0.20
2P Nd C²
S-FSVI
99.54 + 0.04
95.76 + 0.02
77.60 + 0.20
2P 2 N.NN++CO C2
NTK (coreset)
99.50 + 0.09
96.46 + 0.11
78.23 + 0.20
+
BGLN-S
99.64 + 0.04
96.36 0.12
73.98 + 0.37
P+A+d+d
BGLN-D
99.72 + 0.03
96.03 + 0.20
P+A+d+d P+C(A+d+d)
BGLN-S (CW)
99.77 + 0.05
96.99 0.07
74.14 0.57
BGLN-D (CW)
99.78 + 0.02
96.85 0.02
P+C(A+d+0)
BGLN-DD-Vrr
99.64 + 0.04
94.98 + 0.18
P+ P+A+2 A 2d
BGLN-S-Var Var
99.50 + 0.03
96.36 + 0.13
72.18 + 0.26
P+ A 2d
BGLN-S (coreset)
99.50 + 0.03
96.36 0.13
72.89 0.4
P+ Nd
Table 2: Backward transfer on continual learning benchmarks (higher is better).
Method
Split MNIST
Permuted MNIST
Split CIFAR100
FROMP S-ESVI
-0.50 0.20
0.10
-2.60 0.90
-0.21
0..2
- 0.20
BGLN-S
0.00
-0.41
-7.33
BGLN-D
-0.09 0.04
-0.56 0.556.0..04
BGLN-S
-0.18
-0.37
-6.13
BGLN-D (CW)
-0.07 0.07
-1.17 0.07
MNIST, Permuted MNIST and Split CIFAR100, with standard architectures for fair comparison. More details about the datasets, architectures and hyperparameters can be found in Appendix B.2
Evaluation Metries. We evaluate the success of different methods based on average final accuracy across tasks and the backward transfer metric (Lopez-Paz & Ranzato 2017), defined as the average increase in task accuracy after training on all tasks relative to after training on that task.
Results. Figure I shows the functions learned by different methods when sequentially trained on two one-dimensional regression tasks. BGLN retains good predictions on both tasks, while EWC and exact NTK suffer catastrophic forgetting. We hypothesize that this difference in performance due to important nonlincarities between network parameters that EWC and NTK approximations unable to capture.
Table Memory cost notation.
We compare our method's performance on the above evaluation metrics image classification benchmarks existing parametric
Pt
parameters in layer
and nonparametric methods in Tables I and 2 BGLN methods out-
P
# parameters = -t # activations < data dimension
perform other parametric methods (EWC. OSLA and VCL) across
A
d
datasets and are competitive with the state-of-the-at nonparametric methods. Further. using cosss-onniiional approximations improves results. On Split CIFAR 100, BGLN methods start to bridge the gap between parametric and nonparametric methods, while the NTK approximation, which uses core- set, outperforms all previous works. We hypothesize that methods like FROMP can be viewed as approximations to the NTK method and our current result indicates that directly estimating FSD via NTK can improve continual learning performance. Table 2 compares the backward transfer metric for our methods against the current state-of-the- in continual learning. Both FROMP (Pan et al.) 2020) and S-FSVI (Rudner et al. 2022) store coresets of data, yet BGLN outperforms them the MNIST tasks, while lagging behind on Split CIFAR100. We show how the memory requirement scales in terms of number of datapoints stored in Figure 2 and analytically compare this cost for
N
coreset size classes
C

---

Under review as a conference paper at ICLR 2023
BOLN-S BGLN-D
EWC
BGLN-D BGLN- (CW)
NTK (coreset)
Log of Thme FSD from Took
NTK (coreset)
Log eatimated FSD
tree FSD hetween of networks
Task
(a) Estimating true FSD
(b) Minimizing true FSD
Figure 3: Comparison of different FSD estimators for networks trained on Permuted MNIST. Note that FSD
values plotted on log-scale. (a) Compared NTK, BGLN consistently gives closer FSD values the true empirical FSD. (b) While training on task FSD from the optimal task parameters increases with task accuracy. Optimizing BGLN-D and conditoond BGLN-D (CW) effectively minimizes the true FSD.
different methods in Table [ using notation in Table 3
This makes the memory advantage of our
methods explicit. Further details computing these costs are in Appendix E.1
5.2 COMPARING FSD ESTIMATORS
We empirically assess how well our method approximates and minimizes the true empirical FSD between two sets of parameters. Motivated by the continual learning application, we consider the Permuted MNIST task and its sequentially learned network parameters. As seen above, training on new tasks risks catastrophic forgetting predictions on old data and can significantly change the learned function space. We evaluate different FSD estimators by measuring (1) how close the estimated FSD is to the true empirical FSD computed over the full dataset and (2) how much the true empirical FSD is minimized when the estimated FSD is optimized during training. We draw attention to the NTK estimator which exactly linearizes the network using a coreset of datapoints. Figure /alplal plots the estimated FSD versus the true empirical FSD between several pairs of networks trained on different tasks, for different number of iterations and with different learning rates. Both our methods, BGLN-S and BGLN-D consistently have closer estimates to the true empirical FSD than NTK across the range of networks considered. As shown in Figure 3b our methods also minimize the true empirical FSD more effectively than EWC and NTK.
5.3 INFLUENCE FUNCTION ESTIMATION
To demonstrate that BGLN can successfully be applied to other applications involving FSD ap- proximation, we consider an influence function estimation task (Cook) 1979 Hampel 1974 Koh & Liang 2017). Given parameters 0, trained on some dataset Duain that is drawn from a distribution Prrains influence functions approximate the parameters 0_, obtained by training without partieu- lar point (x, y) E DDrain The difference in loss between 0, and A_ is an indicator for how much influence (x, y) has on the trained network.
Given data point, (æ., y). that we are interested in removing, Bae et al. (2022a) recently showed how to express influence functions in neural networks as optimizing the following proximal Breg- man response function (PBRF):
0_
= arg min (0 Oo Prrain) + 2/ 1110-001 0o11 (() GER
(12)
where N is the total number of training datapoints. Here, D33 is the Bregman divergence defined on network outputs and measures the FSD between and 0, for training distribution Prrins as defined in equation] from Section 2 For common loss functions such as squared error and cross-entropy, the Bregman divergence term is equivalent to the training error on a dataset where the original targets are changed with soft targets produced by Og. The second term is a proximity term with strength >>0 0, whose role is to prevent the parameters from moving too far in parameter space. Finally, the last term maximizes the loss of the data point are interested in removing. Intuitively, the PBRF maximizes the loss of data we would like to remove while constraining the network in both function and parameter space so that the predictions and losses of other training examples are not affected.
8

---

Under review as a conference paper at ICLR 2023
Similar to the issues encountered in continual learning, the computation of the FSD term requires storing and accessing the entire training data even after the original network is trained. In practice, Bae al. 2022a) iteratively sample batches of the training data to compute and minimize the FSD term. However, training data may not be accessible and incurs a large memory cost to store. We can employ our parametric method to estimate the FSD in PBRF without accessing the training data, but only its first two moments.
Datasets and architectures. We first train MLP with 2 hidden layers and ReLU activation for 200 epochs on regression datasets from the UCI benchmark (Dua & Graff 2017). Then, we randomly select 50 independent data points to be removed. Using the PBRF objective, we compute the differ- ence in loss after removing a single data point, which is commonly referred to as the self-influenc, score (Koh Liang 2017 Schioppa et al. 2022). Then, we apply BGLN-D to approximate the FSI) term in the PBRF objective and compare it with other approximations such as EWC and Conjugate Gradient (CG) (Martens al., 2010).
Evaluation Metric and Results. Since the PBRF can considered the ground truth for influence estimation, we compare the cor- relation of EWC, CG, and BGLN with the PBRE We used both Pearson correlation (Sedg- wick 2012) and Spearman rank-order correla- [961] to measure this align- ment. We show the results in Table 4 Across all datasets, without having over the full training dataset, BGLN-D achieves higher correlation with the PRBF objective compared to EWC and CG.
Table 4: Comparison of train loss differences computed by EWC, conjugate gradient (CG), and FSD. We show Pearson (P) and Spearman rank-order (S) correlation with the PBRF estimates.
Dataset
EWC
CG
BGLN-D
tion (Spearman
P
S
P
S
P
S
Concrete
0.78
0.57
0.92
0.94
0.96
0.97
save iterate
Energy
0.68
0.39
0.97
0.98
0.99
0.98
Housing
0.86
0.33
0,92
0.89
0.95
0.83
Kincties
0.36
0.30
0.88
0.86
0.99
0.99
Wine
0.97
0.70
0.99
0.94
0.99
0.90
5.4 ABLATIONS
We analyze the impact of the different choices our method makes by ablating them one by one. (I) BGLN-D-Var and BGLN-S-Var use a diagonal approximation to the data covariance which effectively assumes independence between features. We observe in Table thatt this approximation does not harm performance significantly. (2) BGLN-S (coreset) teases apart the effect sampling data using its first two moments by running the stochastic version of our method with small number of actual data samples. Results on backward transfer for the same can be found in Appendix E.4
6 CONCLUSIONS
We addressed the problem of compactly summarizing model's predictions given dataset, and formulated it as approximating neural network FSD. We found that BGLN methods, our novel para- metric approximations to FSD, capture nonlinearities between network parameters and are much more memory-efficient than prior works. We demonstrated the closeness of our estimate to the true FSD well its application in two use-cases. continual learning, BGLN methods outcompete existing methods on common benchmarks without storing any data samples. Furthermore, they ef- fectively estimate influence functions without iterating over or storing the full dataset. We note that while our methods make closer approximation to the true FSD in principle, they are limited by independence and probabilistic modeling assumptions. Empirically, these may explain its lower but competitive performance on Split CIFARI00 100. Interestingly, the direct NTK approximation outper- forms all previous methods, which makes optimizing its memory memory-efficienee efficiency an attractive direction for future work. Further, extending the formulation of FSD approximation to other applications like model editing or unlearning are exciting research avenues. We hope that our work inspires methods to further bridge the gap nonparametric methods' performance in memory-efficieen manner.
7 REPRODUCIBILITY STATEMENT
All details about datasets, model architectures and training can be found in the Appendix. We also plan to release code that reproduces the results in this paper once it is public.

---

Under review as a conference paper at ICLR 2023
REFERENCES
S. Amari, A. Cichocki, and H. H. Yang. A new learning algorithm for blind signal separation. In Proceedings of the 8th International Conference on Neural Information Processing Systems, 95, pp. 757-763, Cambridge, MA, USA, 1995. MIT Press.
Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Comput., 10(2):251- 276. 1998. doi: (0.11620089766633000177777 08997659830001777
URL bttps::/aci.org:
Juhan Bae, Nathan Ng, Alston Lo, Marzyeh Ghassemi, and Roger Grosse. If influence functions are the answer, then what the question? arXiv preprint (tt::2209.053644 2022a.
Juhan Bae, Paul Vicol, Jeff Z HaoChen, and Roger Grosse. Amortized proximal optimization. arXiv preprint arXiv:2203.0000 2022b.
Yu Bai, Ben Krause, Huan Wang, Caiming Xiong, and Richard Socher. Taylorized training: Towards better approximation of neural network training finite width. arXiv preprint arXiv:2002.04010 2020.
Ari S Benjamin, David Rolnick, and Konrad Kording. Measuring and regularizing networks in function space. arXiv preprint arXiv:1805.08289, 2018.
Jeremy Bernstein, Arash Vahdat, Yisong Yue, and Ming- Yu Liu. On the distance between two neural networks and the stability of learning, Advances in Neural Information Processing Systems, 33:
21370-21381, 2020.
Lucas Bourtoule, Varun Chandrasekaran, Christopher A Choquette-Choo, Hengrui Jia, Adelin Travers, Baiwu Zhang, David Lie, and Nicolas Papernot. Machine unlearning. In 2021 IEEE Symposium on Security and Privacy (SP), pp. 141-159. IEEE, 2021.
Dennis Cook. Influential observations in linear regression. Journal of the American Statistical Association, 74(365):169-174 1979.
Nicola De Cao, Wilker Aziz, and Ivan Titov. Editing factual knowledge in language models. arXiv preprint arXiv:2 08/64, 2021.
Matthias De Lange, Rahaf Aljundi, Mare Masana, Sarah Parisot, Xu Jia, Aleš Leonardis, Gregory Slabaugh, and Tinne Tuytelaars. A continual learning survey: Defying forgetting in classification tasks. IEEE transactions on pattern analysis and machine intelligence, 44(7):3366-338 2021.
Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 2916):141-142, 2012.
Thang Doan, Mehdi Abbana Bennani, Bogdan Mazoure, Guillaume Rabusseau, and Pierre Alquier. A thcoretical analysis of catastrophic forgetting through the ntk overlap matrix. In International Conference Artificial Intelligence and Statistics, pp. 1072-1080. PMLR, 2021.
Dheeru Dua and Casey Graff. UCI machine learning repository, 2017. URL httt::/t/r//iieee
1cs. uci edu/ml
Sayna Ebrahimi, Mohamed Elhoseiny, Trevor Darrell, and Mareus Rohrbach. Uncertainty-guided continual learning with bayesian neural networks. arXiv preprint arXiv: 2019.
Ian Goodfellow, Mehdi Mirza, Da Xiao, Aaron Courville, and Yoshua Bengio. An empiri- cal investigation of catastrophic forgetting in gradient- neural networks. arXiv preprint 132 2013.
Roger Grosse. University of Toronto CSC2541, Topies in Machine Learning: Neu- ral Net Training Dynamics, Chapter 4: Second-Order Optimization. Lecture Notes,
2021. URL httpp::/ww/l CS twwottoroooo... 2021/readings
ctt.aa/eeoorrse/ccc/222 csc2541.
Roger Grosse and James Martens. A kronecker factored approximate fisher matrix for convolution layers. In International Conference Machine Learning, pp. 573-582, PMLR, 2016.
10

---

Under review as a conference paper at ICLR 2023
Frank R. Hampel. The influence curve and its role in robust estimation. Journal of the American Statistical Association, 69(346):383 1974. doi: 10.10800016224459.9944 10482962. URL https:/ www.tandfonnine.. com doii
1974. 04829662
Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. In International Conference on Learning Representations, 2019.
Geoffrey Hinton, Oriol Vinyals, Jeff Dean, et al. Distilling the knowledge in neural network. arXiv preprint arXiv:1503.02533 2(7), 2015.
Jiaoyang Huang and Horng-Tzer Yau. Dynamics of deep neural networks and neural tangent hier-
arcby. In International conference on machine learming, pp. 4542-4551. PMLR, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- cralization in neural networks. Advances in neural information processing systems, 31, 2018.
Haoming Jiang, Pengcheng He. Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Tuo Zhao, Smart: Robust and efficient fine-tuning for pre-trained natural language models through principled regu- larized optimization. arXiv preprint arXiv: 191 I. 03437, 2019.
Sanyam Kapoor, Theofanis Karaletsos, and Thang D Bui. Variational auto-regressive gaussian pro- cesses for continual learning. International Conference on Machine Learning, pp. 5290-5300.
PMLR, 2021.
Polina Kirichenko, Mchrdad Farajtabar, Dushyant Rao, Balaji Lakshminarayanan, Nir Levine, Ang Li, Huiyi Hu, Andrew Gordon Wilson, and Razvan Pascanu. Task-agnostic continual learning with hybrid probabilistic models. arXiv preprint arXiv.2206. 2021.
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska. et al. Overcom- ing catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 1 114(13):351-33526 2017.
Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In International conference on machine learning, pp. 1885-1894, PMLR, 2017.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 (canadian institute for advanced search). a. URL CoLhttp:/// wwwp///wwww.//// cs. toronto. edu/ htap//rriz/cii html
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-100 (canadian institute for advanced research). b. URL http:://wwwt www.wwtwww.oooooooooooooooo.o.o.
cs toronto. wwweooottt.dd//kriiaaa .tt.larrr.mmml.............m....
Frederik Kunstner, Philipp Hennig, and Lukas Balles. Limitations of the empirical fisher approx- imation for natural gradient descent. Advances in neural information processing systems, 32, 2019.
Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
David Lopez-Paz and Mare Aurelio Ranzato. Gradient episodic memory for continual learning- Advances in neural information processing systems, 30, 2017.
James Martens and Roger Grosse. Optimizing neural networks with tronecker-factoree approximate curvature. In International conference machine learning, pp. 1408-2417. PMLR, 2015.
James Martens, Ilya Sutskever, and Kevin Swersky. Estimating the hessian by back poock-ropagating curvature. arXiv preprint arXiv: 1206.6464, 2012.
James Martens et al. Deep learning via hessian- free optimization. ICML, volume 27, pp. 735-742, 2010.
11

---

Under review as a conference paper at ICLR 2023
Seyed Iman Mirzadeh, Arslan Chaudhry, Dong Yin, Huiyi Hu, Razvan Pascanu, Dilan Gorur, and Mehrdad Farajtabar. Wide neural networks forget less catastrophically. In International Confer- on Machine Learning, pp. 15699-15771 PMLR, 2022.
Eric Mitchell, Charles Lin, Antoine Bosselut, Chelsea Finn, and Christopher D Manning. Fast model editing scale. arXiv preprint arXiv:2110.11300 2021.
Cuong V Nguyen, Yingzhen Li. Thang D Bui, and Richard E Turner. Variational continual learning. arXiv preprini 17x.iv11110.. 106.2 2017.
Fabrice Normandin, Florian Golemo, Oleksiy Ostapenko, Pau Rodriguez, Matthew D Riemer. Julio Hurtado, Khimya Khetarpal, Dominic Zhao, Ryan Lindeborg, Timothée Lesort, et al. Sequoia: A software framework unify continual learning research. arXiv preprint arXiv: 108.0 108.01005, 2021.
Pingho Pan, Siddharth Swaroop, Alexander Immer, Runa Eschenhagen, Richard Turner, and Mo- hammad Emtiyaz Khan. Continual deep learning by functional regularisation of memorable past. Advances in Neural Information Processing Systems, 33:4453-4464, 2020.
German Parisi, Ronald Kemker, Jose L. Part, Christopher Kanan, and Stefan Wermter. Contin- ual lifelong learning with neural networks: A review. Neural Netw., 113(C):54 71, may 2019. ISSN 0893-6080, doi: 10. n..00066..eunett2119.....
URL http::://ei..rr///000///
n.uumet 2019. 01.012
Razvan Pascanu and Yoshua Bengio. Revisiting natural gradient for deep networks. CoRR, abs/1301.33 33844, 2014.
Vinay V Ramasesh, Ethan Dyer, and Maithra Raghu. Anatomy of catastrophic forgetting: Hidden representations and task semantics. arXiv preprint arXiv:2007.07400, 2020.
Hippolyt Ritter, Aleksandar Botev, and David Barber. Online structured laplace approximations for overcoming catastrophic forgetting. Advances Neural Information Processing Systems, 31, 2018.
Tim G. Rudner, Freddie Bickford Smith, Qixuan Feng, Yee Whye Tch, and Yarin Gal. Continual Learning via Sequential Function-Space Variational Inference. In Proceedings of the 39th Inter- national Conference on Machine Learning, Proceedings Machine Learning Research. PMLR, 2022.
Tim GJ Rudner, Freddie Bickford Smith, Qixuan Feng, Yee Whye Teh, and Yarin Gal. Continual learning via function space variational inference.
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynam- ics learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2013.
Andrea Schioppa, Polina Zablotskaia, David Vilar, and Artem Sokolov. Scaling up influence func- tions. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 8179- 8186. 2022.
Philip Sedgwick. Pearson's correlation coefficient. Bmj, 345, 2012.
Mariia Seleznova and Gitta Kutyniok. Analyzing finite neural networks: Can we trust neural tangent kernel theory? In Mathematical and Scientific Muchine Learning, pp. 868-895 PMLR, 2022a.
Mariia Seleznova and Gitta Kutyniok. Neural tangent kernel beyond the indinite-width limit: Effects of depth and initialization, arXiv preprint arXiv:2202.00553, 2022b.
Charles Spearman. The proof and measurement of association berween two things. 1961.
Michalis Titsias, Jonathan Schwarz, Alexander G de G Matthews, Razvan Pascanu, and Yee Whye Teh. Functional regularisation for continual learning with gaussian processes. arXiv preprint arXiv:1901.11356. 2019.
12

---

Under review as a conference paper at ICLR 2023
APPENDIX
A NOTATION
f
Function corresponding to a neural network
ReLU activation function
Derivative of function ø
81
Preactivations at layer
(8))
Activations at layer /
Input data
y
Target data
(()
Input data point
Pdata
Data distribution from which sampled
d
Data dimension
N
Number of data points in a coreset
L
Number of layers in network
Number of classes in a classification task
Number of tasks
0
Parameters of neural network
0.
Parameters obtained after training on task
Number of parameters layer
P
Total number of parameters in network = Et.
f(x: 0)
Prediction of the network f x parameterized by 0 Output space distance, for instance Euclidean distance
P
D(O0. 01)
Function space divergence between networks parameterized by 0g and 0
J VoS
Network Jacobian
Go
Weight space metric matrix
Fe
Fisher information matrix
m
Bernoull random variable
ju
Bernoulli mean parameter
B RECURSION EQUATIONS FOR BGLN-D
We derive the deterministic version of our algorithm by taking expectations and covariances for the quantities in equations to II (rewriteen using Bernoulli variables). We use linearity expectations and our assumptions of independence between the Bernoulli random variables. We also assume Cov(ag, dg Aa) is close to 0 and ignore it in our computations. This assumption is tested empirically in experiments and find that it does not severely move the FSD estimate away from the true empirical FSD (see Figure Ba) The complete steps for BGLN-D computations are given in Algorithm 2
13

---

Under review as a conference paper at ICLR 2023
Algorithm 2 BGLN-D
Require: x Cov(x), { (,, b){ 1011-1 (1))
1: Blao]: [[ + E[xx
2: Cov(ae), Cov <- Cov(x)
3: E[Aa], Cov(Aa) (ov f- 0
4: for f 1 to L do
5:
[[so] + W "
E[ao] - b0
6:
E
[s)
A AwoEE(a) + W B[Aa] A60
7:
E
aol (()) @ [[so]
8:
[| + ( [[
E
Cov(#o) - w Cov W W
9:
Cov(As) Covaa))) + "/" Cor(Aa)
Cov (n (UUT)) (ov(so)
(ov - (ull), (pl (1)))) Cov(As)
13: end for
E[Az] < A AW (')
w B[Sa] + Ab
15: Cov( [z) - A AW Cov ( + w (oy(Aa)W((
16: return
E 2 + tr (ov
Algorithm 3 BGLN-S (Conv)
Require: E Cov(x), (layer] Lu]
1: E[ [[ + E[x]
2: Cov(as) (o), Cov (1) < Cov(x)
3: ag, d1 ~B((Baaa))) Cov(aa))
4: Aa 4- 0
5: for 1
L, 1
6:
if' J'Lyeer is Conv or FC then
7:
8o 1-1aer(aaa) grad-False)
8:
81 - (-1ayer((a))
9:
else layer ReLU then
10:
As = 81 8q
11:
m ~ ())(
12:
Aa <- m @ As else layer is Dropout then
13:
14:
pass
15:
else
16:
ag layer(aa)
17:
a t layer(a))
18: end if
19: end for
Az t 81 80
21: return A 1
C GENERALIZATION TO CONVOLUTIONAL NETWORKS
The generalization of BGLN-S to convolutional networks involves passing the inputs sampled us- ing the data moments through the network. convolutional networks, ReLU activation is usually applied after the convolutional and the fully connected layers. At each ReLU, use the Bernoulli mean parameters to sample activation signs and obtain the difference of activations, Aa. We treat any dropout layers as they treated at evaluation time, i.c. the identity function. Finally, the Euclidean distance between the final layer outputs leads to the stochastic estimate of the FSD. The complete procedure is shown in Algorithm
14

---

Under review as a conference paper at ICLR 2023
300
x=y
250
NTK (coreset)
BGLN-S (coreset)
Estimated FSD
200
True FSD between pairs of networks
Figure Estimating true FSD with BGLN-S (coreset) and NTK (coreset) networks trained using Split CIFAR100.
D COMPARING FSD ESTIMATORS
To directly compare the NTK approximation with the linearized activation function trick, we com- pute the estimated FSD between pairs of networks using these two kinds of linearization, when provided with the same information, and plot them against the true emprical FSD. Hence, both estimators provided with the same coreset of datapoints. In the case of BGLN-S (coreset), acti- vations are computed using this coreset directly instead of Bernoulli sampling. Figure 4 visualizes this comparison for networks trained using tasks in Split CIFAR100. BGLN-S estimates correlate better with the true FSD than NTK, hence corroborating our intuition about linearizing activation functions. To quantify this difference, we also measure the Spearman rank-order and Kendall's Tau correlation coefficients for each estimator with the true FSD. BGIN-S obtains values 96.36 and 79.19, respectively, outperforming NTK, which obtains 86.42 and 71.71, respectively.
E CONTINUAL LEARNING
E. MEMORY COST ANALYSIS
We follow the notation in Table 3 to denote as the number of parameters in layer 1, P P P as the total number of parameters in the network, as the number of activations in the network (note that A P), d the data dimension, N the number of samples in coreset, and C as the number of classes in the continual learning classification setting. We can now write analytic expressions for the memory cost incurred by the different methods considered in the continual learning experiments, as shown in Table Below we arrive at these expressions for each task that the model is continually trained on.
EWC: EWC requires storing value for each parameter of the network and one value for each diagonal element of the P x P Fisher information matrix, resulting in a cost of 2P.
OSLA: OSLA approximates the Fisher information matrix block diagonal matrix, storing p² clements for each block corresponding to layer This gives a cost of P + Σl-rp
VCL: "CLsstorss two pieces of information for the variational distribution for each param- eter, one for the mean and one for the variance in the diagonal approximation, incurring cost equal to 2P.
VCL (coreset): The "coreset" variant of VCL additionally stores N datapoints, increasing the memory cost by Nd.
15

---

Under review as a conference paper at ICLR 2023
VAR-GP, FROMP, S-ESVI: These methods all store two values per parameter, similar to VCL. Further, they require a coreset of datapoints and/or inducing points, as well as a NC NC kernel matrix.
NTK (coreset): The NTK approximation stores one value for each parameter of the net- work and one for the P-dimensional Jacobian-vector product used to linearize the net- work. It further requires a coreset of N datapoints, result in a 2P + Nd cost.
BGLN-S, BGLN-D: Our methods store each parameter once, a Bernoulli mean value for each activation and the first two moments of the data, incurring a cost of P + A A + dº,
BGLN-S (CW), BGLN-D (CW): The classwise variants our methods store separate Bernoulli means and data moments for each class, which scales those corresponding mem- ory costs by C, i.e., P + C(A + d + d
BGLN-S (Var), BGLN-D (Var): The "Var" variants of our methods make diagonal approximation to the data covariance (second moment), hence storing only d values for it. This further reduces memory cost P P+AA2 AAdd.
BGLN-S (coreset): The "coreset"" ablation of BGLN-S requires storing N datapoints instead of the data moments, giving P A + Nd cost.
E.2 EXPERIMENTAL DETAILS
Datasets. Split MNIST consists of five binary prediction tasks to classify non-overlapping pairs of MNIST digits (Deng 2012). Permuted MNIST a sequence ten tasks to classify ten dig- its, with different fixed random permutation applied to the pixels of all training images for each task. Finally, Split CIFAR100 consists of six ten-way classification tasks, with the first being CI- FARIO (Krizhevsky et al. and subsequent ones containing non-overlapping classes each from the CIFART00 dataset (Krizhevsky et al. b).
Architectures. We use standard architectures used by existing methods for fair comparison. For regression and the MNIST experiments, we use a MLP with two fully connected layers and ReLU activation. For Split CIFARI00, we use network with four convolutional layers, followed by two fully connected layers, and ReLU activation after each. Both Split MNIST and Split CIFAR100 models have a multiheaded final layer.
Hlperpaaaaeeters We have performed grid search over some key hyperparameters and used the ones that resulted in the best final average accuracy across all tasks. All hyperparameter search was done with random seed 42. We then took that best set of hyperparameters, repeated our experiments on seeds 20, 21, 22, and reported the average and standard deviation our results.
For the learning rate, we used 0.001 for all CL experiments except the BGLN-S method for Split MNIST and BGLN-D method for Permuted MNIST, where we used 0.0001 instead.
We used the same number of epochs each CL task and the exact numbers are reported in Table On the first task, all MNIST experiments used the same number of epochs as the subsequent CL tasks while CIFARI00 experiments used 200 epochs the first task.
To compute the Bernoulli mean parameters for our stochastic gating implementation, we used simple averaging the default, but also explored exponential moving averaging. While there was not much difference in performance, we report the momentum values that reproduce our results. All MNIST experiments had momentum value of 1/batch size. Note that this momentum value of 1 l/bbacc i ze corresponds to simple moving average. For CIFAR 100 experiments, we used 0.99 for BGLN-S (CW) and NTK, and 1 1/batch_s ze for BGLN-S.
For each method and dataset, the scaling factor for FSD penalty, λESD, is reported in Table Simi- larly, batch size is reported in Table 7
Evaluation Metrics. In addition to average accuracy over tasks, we measure the backward transfer metric. For T tasks, let Rag be the classification accuracy on task L after training on task 4, Then, backward transfer is given by the following formula.
R R4,
T
1=1
16

---

Under review as a conference paper at ICLR 2023
Table 5: CL tasks training epochs used in CL experiments.
Method
Split MNIST
Permuted MNIST
Split CIFARI00
NTK (coreset)
15
5
80
BGLN-S
15
BGLN-D
15
15
BGLN-S (CW)
15
BGLN-D (CW)
15
15
Table FSD scale used in CL experiments.
Method
Split MNIST
Permuted MNIST
Split CIFAR100
NTK (coreset)
1
0.005
BGLN-S
5
0.0005
BGLN-D
0.1
0.005
BGLN-S (CW)
2
0.0000001
BGLN-D (CW)
0.1
0.005
Table 7: Batch size used in CL experiments.
Method
Split MNIST
Permuted MNIST
Split CIFAR100
NTK (coreset)
256
256
512
BGLN-S
32
128
512
BGLN-D
32
128
BGLN-S (CW)
32
128
512
BGLN-D (CW)
32
128
Table 8: Ablation of hackward transfer on continual learning benchmarks (higher is better).
Method
Split MNIST
Permuted MNIST
Split CIFAR100
BGLN-D-Var
-0.18 + 0.09
-3.90 - 0.49
BGLN-S-VE
-0.13 + 0.11
-0.41 + 0.08
-9.25 + 0.06
BGLN-S (coreset)
-0.13 0.11
-0.41 + 0.08
-7.50 + 0.51
E.3 TASK-WISE CLASSIFICATION RESULTS
We show in Figure 5 the lask-wise accuracies on the Split CIFARI00 benchmark after training on all tasks is complete, drawing comparison between methods (BGLN-S and BGLN-S (CW)), NTK (with coreset) and nonparametric state-of-the-art method, FROMP.
E.4 ABLATION STUDY BACKWARD TRANSFER
We include further results on the performance of our ablations on the backward transfer metric in
Table g
F INFLUENCE FUNCTION ESTIMATION
F.1 EXPERIMENTAL DETAILS
We used Concrete, Energy, Housing, Kinetics, and Wine datasets from the UCI collection (Dua & Graff, 2017). For all datasets, we normalized the training dataset to have a mean and a standard deviation of We used 2-hidden layer ML.P with 128 hidden units and the base network was
17

---

Under review as a conference paper at ICLR 2023
0.85
0.80
Accuracy
0.75
0.70
0.65
FROMP
BGLN-S BGLN-S (CW)
NTK
0.60
Task 0
Task
Task
Task 3
Task
Task
Tasks
Figure Comparison of task-wise accuracies the Split CIFARIO0 benchmark after training all tasks complete, for our methods, NTK and nonparametric Mate-of-the-at method, FROMP.
trained for 200 epochs with SGD and batch size of 128. We performed hyperparameter searches over the learning rate in the range {0.3, 0.1, 0.03, 0.01, 0.003, 0.001] and selected the learning rate based on the validation loss.
For each random data point selected, we optimized the PBRF objective for additional 20 epochs from the base network. The FSD term computed stochastically with batch size of 128. Sim- ilarly, both EWC and BGLN were trained with the same configuration but with the corresponding approximation the FSD term.
F.2 MISLABELED EXAMPLE IDENTIFICATION
A common application of the influence function is the detection of influential or mislabeled ex- amples. Intuitively, if some fraction of the training data labels is corrupted, they would behave as outliers and have greater influence the training loss. One approach to efficiently detect and fix these examples is to prioritize and examine training inputs with higher self-influence scores.
We 10% of the MNIST dataset and corrupt 10% of the data by assigning random labels it. We train two-layer MLP with 1024 hidden units and ReLU activation using SGD with batch size of 128. Then, we use EWC and BFLN to approximate the FSD term in the PBRF objective in equation 12 for each data point. We also compare these methods against randomly sampling datapoints check for corruptions. The results are summarized in Figure Both PBRF and BGLN show significantly improved performance compared to the random baseline.
G LINEARIZATION ERROR ANALYSIS OF BGLN-S
We can analyze the error due to linearization of BGLN-S as compared to the true FSD, propagated through each step of the network's computation. Following notation in Section 3 Az and Åě denote the true difference between network outputs and the difference between network outputs estimated after applying the linearized activation functions trick. We are interested the error in this estimation,
err(z) Az 2
(13)
18

---

Under review as a conference paper at ICLR 2023
Random
Fraction of corruption fixed
0.8
PBRF
EWC
BGLN
0.6
0.4
0.2
0.0
0.4
0.5
0.6
Fraction of train data checked
Figure Effectiveness of BGLN in detecting mislabeled examples. BGLN can approximate the term the PBRF objective accurately and used in applications involving influence functions without explicitly storing or iterating over the dataset.
As errors propagate through the network, we consider the error after each step of computation. Accordingly, we define,
err(s As A and err(a = Aa A!
(14)
The key observation here that errors are only introduced due to the linearized activation computa- tions. Also note that, given two networks, linearize the activations of only one of them, leaving the other's computations accurate.
In the case of linear feulee-coneected convolutional) networks with continuously differentiable activation functions, the error introduced due to activations at any layer is the second-order Taylor error,
err(a ()) O(1 80 11)))
(15)
Any subsequent computation then has error only due to err(a())), with no further errors introduced due to linear layers. For instance, for fully-connected layers of the two networks parametrized using w and ( which have as inputs dg
(f-1) and a
(t-1), the estimation error given by
(rr(s(0) A Aw (( Aa (a(11) (1
(16)
-w w/ Aá (t)
(17)
= AW a + (1- AQ(l-1))
(18)
w ( er)(a
(19)
where the last line given by the fact that ag a since activations ag remain exact while
linearize computations for (1)
This error analysis formalizes the intuition that as errors are propagated through the network, they
only introduced due to linearized activation computations. The error incurred after linear layer a linear function of the error in its inputs, adding no new errors its own.
19


Reviews:
[{"review_id": "nbkr5hrhuI", "reviewer": "Reviewer_qY8f", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "**Post rebuttal summary**\n\nI thank the authors for their thoughtful responses. This well-written paper tackles the important problem of FSD estimation. The methodology introduced is not accompanied by a sufficiently convincing theoretical or empirical analysis. The updated empirical analysis hints that this method might be worth further exploring more thoroughly beyond isolated examples. I encourage the authors to submit an improved version of the paper in future.\n\n======\n\nSummarising the properties of a dataset (perhaps with respect to a model) finds useful application in problems such as continual learning, domain adaptation, semi-supervised learning, influence-function estimation and even regularisation. In particular, it may be used to avoid catastrophic forgetting. The authors use function space distance (FSD) for this purpose, applied to ReLU neural networks. Using a certain linear approximation and stochastic Bernoulli approximation of the FSD, they derive an efficient algorithm for approximating the FSD. They find their algorithm to be competitive with SOTA on some toy examples.", "summary_of_the_review": "While well-written, I am not convinced that this paper contains enough insight to be published. \n\nIt would be okay for the benchmark datasets to be somewhat toy *if* the mathematical formulation, motivation and analysis were complete. At the moment, I can see two unprincipled approximations: (1) The \"linearisation\" in input space, and (2) the replacement of a step function with an independent(!) Bernoulli random variable. Alternatively, if the benchmarks were more convincing, the mathematical analysis could be acceptably less complete.", "strength_and_weaknesses": "**Strengths:**\n- The problem considered is a worthy one.\n- The paper is very clearly presented, well-written, and conveys and elegant message.\n\n**Weaknesses:**\n- Section 3.1. What does linearisation (5-10) represent, and where does it come from? It is not a first-order Taylor in the inputs, since the Taylor series of the ReLU does not converge to the ReLU. \n- \"The upshot of our approximation is that, in the ReLU case, it only depends on the training data through the signs of all the activations (since passing layer inputs through $\\phi$ and $\\phi'$ both result in a multiplication by 0 or 1)\". I am not understanding something here. I thought it would depend on the sign of the *preactivations* $s$. Additionally, passing through $\\phi$ involves multiplication by zero or one, but passing through $\\phi'$ involves taking the sign of the input.\n- Beyond computational advantage, why is the step function approximated using an independent Bernoulli random variable? This changes the dynamic. Does this modification preserve the original formulation in any sense?\n- Given that this paper contains no mathematical theory, I would have hoped to see more realistic datasets for baselines beyond MNIST, CIFAR, Wine, Housing, etc. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "U1DDkVpdpw", "reviewer": "Reviewer_QzNA", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "Post rebuttal\n====\n\nThank the authors for the responses. Since the responses addressed my original concerns, I raised my score. However, I agree with reviewer i4TZ and qY8f that the theoretical and empirical studies are still somewhat simple, and could still be improved.\n\n\n\nThis paper considers how to obtain a parametric approximation to the empirical function space distance (FSD) of two networks, i.e., $\\mathbb E_{p_{emp}}[\\rho(f(x, \\theta_0), f(x, \\theta_1))]$. The key assumption here is we cannot store the training data $p_{emp}$, as it is a  infinite stream. On the other hand, the overhead of storing the model is acceptable.\n\nThe paper takes a linearization approach, BGLN, which approximates the activation of the second network $\\theta_1$ with a first-order Taylor's expansion at the activation of the first network $\\theta_0$. By the linearity of expectation, the FSD can be written with the activation gradient's moment: $\\mathbb E_{p_{emp}}[\\phi^\\prime(a_0)]$. By computing and storing such moments (i.e., sufficient statistics) in advance, the FSD can be computed without accessing the data distribution $p_{emp}$.\n\nThe proposed BGLN is empirically shown to be more accurate than competitive approaches, such as second-order approximation to the FSD on the parameter space or NTK-based approximation, which again linearizes on the parameter space. The effectiveness of the proposed approach is shown on continual learning and influence function estimation tasks.", "summary_of_the_review": "The paper proposes a reasonable solution to a practical problem of function space distance estimation. However, presentation and relationships with existing works should be improved. The significance of the paper is not clear in its current form.", "strength_and_weaknesses": "Strength:\n- The proposed approach is quite simple and accurate.\n\nWeaknesses:\n- The presentation needs to be involved. The current paper has some logical discontinuities. For example, at the end of Sec. 3.1, there misses a part on how to compute the FSD with the activation moments. There should be some derivations and pseudocode for that. Otherwise, the discussions in Sec. 3.2 and Sec. 3.3 seem somewhat random.\n\nSec. 2 might also be somewhat disruptive for the readers might not understand what is parametric and nonparametric approaches here. It might be better to put the concrete formulation of continual learning in Sec. 5.1 forward to Sec. 2 to act as a motivating example.\n\nThere is some typos: for example, in Eq. (3), some (l) should be in fact (l-1).\n\n- The advantage of linearizing the activation over competitive approaches should be discussed in more detail. Currently, there is only empirical evidences of the accuracy of BGLN, which is not strong enough. The paper would be more solid if there is some error / bias analysis and comparison. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "K1soy8W2lWT", "reviewer": "Reviewer_i4TZ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes a method to estimate the function space distance (FSD) of two networks. This method relies on transforming the ReLU layers of the network into Bernoulli stochastic gates, and measuring the distance between the networks while sampling from these gating variables. Using this method, the authors are able to estimate the function distance between two networks, and can use this technique for tasks where this function distance is required, such as in continual learning.", "summary_of_the_review": "I think this work presents an interesting idea, but given that there is not sufficient evidence that the FSD approximation proposed performs well in general, and the main results do not make a convincing argument for the usefulness of the method (other than memory requirements, which are not fully elaborated on), I believe that it needs to be improved.\n\nUpdate post-rebuttal: I have raised my score following the authors' responses, as some of my concerns have been alleviated. I believe that the paper can still be improved, by further examining why the proposed method lags behind NTK in performance (despite experimentally seeming to better approximate the FSD) and including further theoretical justification on why the gating units chosen are a good choice.", "strength_and_weaknesses": "Strengths:\n\n- The method proposed is very easy to understand and can be easily implemented. The authors measure function distance by essentially substituting the ReLU functions of the network with stochastic gating variables. Implicitly, this process samples from the various linear regions of the network, making the calculation of distance between them straightforward.\n\n- The authors examine an interesting setting for the application of their function space distance approximation in continual learning. In this setting, the authors use their FSD approximation to train models that sequentially learn new tasks, while aiming to retain accuracy in the previous ones.\n\nWeaknesses:\n\n- One major weakness of the proposed FSD approximation technique is that there is not sufficient evidence that it is a good approximation of the distance between two networks. In order to prove this theoretically, there would need to be proofs of bounds for the FSD estimated via the proposed method compared to the true FSD, or at the very least a simple toy example demonstrating the FSD estimated on a given network, on a toy dataset. To prove this experimentally, the authors would need to include further comparisons like those included in Figure 3a. Currently, this comparison between estimated and true FSD is only done on an MNIST variant, so doing so over more complicated datasets would be helpful to properly evaluate how well the proposed method performs. Including either of the two would greatly improve the paper.\n\n- In a similar vein, the experimental results on continual learning are not convincing enough for the capabilities of the proposed method. More specifically, the proposed methods only show performance benefits in the case of MNIST variants, and not on the CIFAR100 continual learning benchmark. This, combined with the previous result, makes me wonder whether the approximation derived from the proposed method is useful for the desired continual learning task. I believe the authors would need to further analyze why their method underperforms in the CIFAR100 variant.\n\n- Finally, the authors state that their method can approximate the FSD between two networks without expensive storage requirements, compared to other nonparametric methods. I believe that these storage comparisons should be included (for example, in Table 1), in order to fully understand the benefits of the proposed method, from a memory requirement perspective.\n\nUpdate post rebuttal: The authors have included extra experiments that somewhat alleviate my concerns regarding the quality of the FSD approximation. The proposed method seems to outperform NTK in Split-CIFAR100, as far as FSD approximation is concerned. Moreover, the authors have included explicit memory requirements for their proposed methods.\n\nHowever, the experimental part is still inconclusive, in the sense that while there is an experimental indication that the proposed method performs well with respect to estimating the FSD, it lags behind NTK on the same problem. Furthermore, while there is theoretical motivation with respect to the gating units used, it hinges on these being independent (as pointed out by another reviewer), which seems to be a very limiting assumption. Including a simple toy example where this choice of gating is the best one would be a good addition to the paper.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "6DBly3DGTN", "reviewer": "Reviewer_cRKz", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors propose a way to approximate the function space distance for certain broad classes of neural networks. This methods is less computationally expensive and more precise than alternatives based on the Taylor decomposition of the FSD, and are particularly suited for continual learning.", "summary_of_the_review": "I can see this paper being a useful addition to this line of literature.", "strength_and_weaknesses": "STRENGTHS:\n* The paper studies an important problem. Catastrophic forgetting is a well-known problem in the field, and the method has applications to other areas as well.\n* The algorithmic ideas proposed in the work are novel, technically sound, and non-trivial. The depth of the technical results and their derivations is significant.\n* The experiments are extensive and look at a wide range of interesting settings: influence functions, continual learning, etc.\n\nWEAKNESSES:\n* I think some of the tasks used some experiments can be made more realistic. The influence functions experiments look at a few toy UCI datasets. I'm not an expert in continual learning, but the main datasets seem versions of MNIST as well.\n* I think the presentation can be further improved by better explaining the limitations of the method relative to existing approaches. Additionally, the authors discuss extensions to non-fully-connected architectures in the paper, but I feel like most of the details are in the appendix. A clear explanation of the class of models to which this does or does not apply would help.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject