Efficient Parametric Approximations of Neural Network Function Space
  Distance
It is often useful to compactly summarize important properties of model
parameters and training data so that they can be used later without storing
and/or iterating over the entire dataset. As a specific case, we consider
estimating the Function Space Distance (FSD) over a training set, i.e. the
average discrepancy between the outputs of two neural networks. We propose a
Linearized Activation Function TRick (LAFTR) and derive an efficient
approximation to FSD for ReLU neural networks. The key idea is to approximate
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 estimating influence functions accurately
and detecting mislabeled examples without expensive iterations over the entire
dataset.
Nikita Dhawan, Sicong Huang, Juhan Bae, Roger Grosse
# Efficient Parametric Approximations of

###### Abstract

It is often useful to compactly summarize important properties of model parameters and training data so that they can be used later without storing and/or iterating over the entire dataset. As a specific case, we consider estimating the Function Space Distance (FSD) over a training set, i.e. the average discrepancy between the outputs of two neural networks. We propose a Linearized Activation Function TRick (LAFTR) and derive an efficient approximation to FSD for ReLU neural networks. The key idea is to approximate 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 estimating influence functions accurately and detecting mislabeled examples without expensive iterations over the entire dataset.

## 1 Introduction

As machine learning models are trained on increasingly large quantities of data or experience, it can be useful to compactly summarize information contained in a training set. In continual learning, an agent continues interacting with its environment over a long time period -- longer than it is able to store explicitly. A natural goal is to avoid overwriting previously learned knowledge as it learns new tasks (Goodfellow et al., 2013) while controlling storage costs. Even in cases where it is possible to store the entire training set, a compact representation circumvents the need for expensive iterative procedures over the full data.

We focus on the problem of estimating _Function Space Distance (FSD)_ for neural networks: 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., 2022), model editing (Mitchell et al., 2021), unlearning (Bourtoule et al., 2021) and second-order optimization (Amari, 1998; Bae et al., 2022) have also been formulated in terms of FSD regularization or similar locality constraints.

Methods for summarizing the training data can be categorized as parametric or nonparametric. In the context of preventing catastrophic forgetting, parametric approaches typically store the parameters of a previously trained network, along with additional information about the importance of different directions in parameter space for preserving past knowledge. The canonical example is Elastic Weight Consolidation (Kirkpatrick et al., 2017, EWC), which uses a diagonal approximation to the Fisher information matrix. Nonparametric approaches explicitly store, in addition to network parameters, a collection (coreset) of training examples, often optimized directly to be the most important or memorable ones (Rudner et al., 2022; Pan et al., 2020; Titsias et al., 2019). Currently, the most effective approaches to prevent catastrophic forgetting are nonparametric due to the lack of sufficiently accurate parametric models. However, this advantage comes at the expense of high storage requirements.
In this paper, we formally formulate neural network FSD estimation and propose novel parametric approximations. To motivate our approach, notice that several parametric approximations, like EWC, can be interpreted as a second-order Taylor approximation to the FSD. This leads to a quadratic form involving the Fisher information matrix \(\mathbf{F_{\theta}}\) or some other metric matrix \(\mathbf{G_{\theta}}\), where \(\mathbf{\theta}\) denotes the network parameters. Second-order approximations are practically useful because one can estimate \(\mathbf{F_{\theta}}\) or \(\mathbf{G_{\theta}}\) by sampling vectors from a distribution with these matrices as the covariance (Martens and Grosse, 2015). Then, tractable probabilistic models can be fit to these samples to approximate the corresponding distribution. Unfortunately, these tend to be inaccurate for continual learning compared to 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 perform poorly in terms of average classification accuracy and backward transfer in sequentially learned tasks. Since such an approximation can be interpreted as network linearization (Grosse, 2021), this finding is consistent with a recent line of results that find linearized approximations of neural networks to be an inaccurate model of their behavior (Seleznova and Kutyniok, 2022, 2022, 2022, 2019, 2020). In Section 2, we present this network linearization perspective of some existing approaches for regularization in function space.

Our method, based on a _Linearized Activation Function TRick (LAFTR)_, does _not_ make a second-order Taylor approximation in the parameter space, and hence is able to capture nonlinear interactions between parameters of the network. Specifically, it linearizes each step of the network's forward pass with respect to its inputs. In the case of ReLU networks, our approximation yields a linear network with stochastic gating, which we refer to as the _Bernoulli Gated Linear Network (BGLN)_. We derive a stochastic and a deterministic estimate of FSD, both of which rely only on the first two moments of the data. This allows the application of our methods in different scenarios where stochasticity is or isn't desirable.

We evaluate our BGLN approximation in the contexts of continual learning and influence function estimation. Our method significantly outperforms previous parametric approximations despite being much more memory-efficient. For continual learning tasks, our method is competitive with nonparametric approaches. For influence function estimation tasks, it closely matches an oracle estimate of a network's loss after a data point is removed, but without having to iterate over the whole dataset.

The key contributions and findings of this work are:

* We introduce LAFTR, an idealized FSD approximation, which improves over parameter space linearization by capturing nonlinear interactions between weights in different layers.
* We propose the Bernoulli Gated Linear Network (BGLN), an efficient parametric FSD approximation for

Figure 1: Comparison of FSD regularization on a 1-D regression task. **(Left)** Training sequentially on two tasks (blue yields \(f_{1}\), then yellow yields \(f_{2}\)) results in catastrophic forgetting. The LAFTR approximation more closely matches the true function \(f_{2}\) than its NTK approximation does. **(Right)** BGLN retains performance on task 1 after training on task 2 more accurately than EWC and NTK.
ReLU networks based on LAFTR, which stores only aggregate statistics of the data and the activations.
* In continual learning, BGLN outcompetes state-of-the-art methods on sequential MNIST and CIFAR100 tasks, with at least significantly lower memory requirements than nonparametric methods.
* For influence function estimation, BGLN accurately approximates the effect of removing a single data point without iterating over or storing the full dataset.

## 2 Background

Let \(\mathbf{z}=f(\mathbf{x},\mathbf{\theta})\) denote the function computed by a neural network, which takes in inputs \(\mathbf{x}\) and parameters \(\mathbf{\theta}\). Consistent with prior works, we use FSD to refer to the expected output space distance1\(\rho\) between the outputs of two neural networks (Benjamin et al., 2018; Grosse, 2021; Bae et al., 2022) with respect to the training distribution, as defined in equation 1. When the population distribution is inaccessible, the empirical distribution is often used as a proxy:

Footnote 1: Note that we use the term _distance_ throughout since we focus on Euclidean distance in our derivation. However, other metrics like KL divergence can also be used, as shown in Section 5.

\[D(\mathbf{\theta}_{0},\mathbf{\theta}_{1},p_{\text{data}}) =\mathbb{E}_{\mathbf{x}\sim p_{\text{data}}}[\rho(f(\mathbf{x}, \mathbf{\theta}_{0}),f(\mathbf{x},\mathbf{\theta}_{1}))] \tag{1}\] \[\approx\frac{1}{N}\sum_{i=1}^{N}\rho(f(\mathbf{x}^{(i)},\mathbf{ \theta}_{0}),f(\mathbf{x}^{(i)},\mathbf{\theta}_{1})), \tag{2}\]

where \(p_{\text{data}}\) is the data-generating distribution. Constraining the FSD term has been successful in preventing catastrophic forgetting (Benjamin et al., 2018), computing influence functions (Bae et al., 2022), training teacher-student models (Hinton et al., 2015), and fine-tuning pre-trained networks (Jiang et al., 2019; Mitchell et al., 2021). Natural choices for \(\rho\) are Euclidean distance for networks trained using mean-squared error (e.g. regression) and KL divergence for those trained with cross-entropy loss (e.g. classification).

Consider the continual learning setting as a motivating example. Common benchmarks (Normandin et al., 2021) involve sequentially learning tasks \(t\in\{1,\dots,T\}\), using loss function \(\mathcal{L}\) and a penalty on the FSD between the parameters \(\mathbf{\theta}\), and the parameters \(\{\mathbf{\theta}_{i}\}\) fit to previous tasks. The penalty is computed over the previously seen data distribution \(p_{i}\) and then scaled by a hyperparameter \(\lambda_{\text{FSD}}\):

\[\mathbf{\theta}_{t}=\arg\min_{\mathbf{\theta}}\mathcal{L}(\mathbf{\theta})+\lambda_{ \text{FSD}}\sum_{i=1}^{t-1}D(\mathbf{\theta},\mathbf{\theta}_{i},p_{i}). \tag{3}\]

Continuing with the notation in equation 2, one way to regularize the FSD is to store the training set and explicitly evaluate the network outputs using both \(\mathbf{\theta}_{0}\) and \(\mathbf{\theta}_{1}\) (perhaps on random mini-batches). However, this has the drawbacks of having to store and access the entire training set throughout training (precisely the thing continual learning research tries to avoid) and necessarily estimating FSD stochastically. Instead, we would like to compactly summarize information about the training set or distribution.

Many (but not all) practical FSD approximations are based on a second-order Taylor approximation:

\[D(\mathbf{\theta}_{0},\mathbf{\theta}_{1},p_{\text{data}})\approx\frac{1}{2}(\mathbf{ \theta}_{1}-\mathbf{\theta}_{0})^{\top}\mathbf{G}_{\mathbf{\theta}}(\mathbf{\theta}_{1}-\mathbf{ \theta}_{0}), \tag{4}\]

where \(\mathbf{G}_{\mathbf{\theta}}=\nabla_{\mathbf{\theta}}^{2}D(\mathbf{\theta}_{0},\mathbf{\theta},p_{ \text{data}})\) is the corresponding Hessian. In the case where the network outputs parametrize a probability distribution and \(\rho\) corresponds to KL divergence, \(\mathbf{G}_{\mathbf{\theta}}\) reduces to the more familiar Fisher information matrix \(F_{\mathbf{\theta}}=\mathbb{E}_{\mathbf{x}\sim p_{\text{data}},\mathbf{y}\sim P_{ \mathbf{y}|\mathbf{x}}(\mathbf{\theta})}[\nabla_{\mathbf{\theta}}\log p(\mathbf{y}|\bm {\theta},\mathbf{x})\nabla_{\mathbf{\theta}}\log p(\mathbf{y}|\mathbf{\theta},\mathbf{ x})^{\top}]\), where \(P_{\mathbf{y}|\mathbf{x}}(\mathbf{\theta})\) represents the model's predictive distribution over \(\mathbf{y}\). It is possible to sample random vectors in parameter space whose covariance is \(\mathbf{G}_{\mathbf{\theta}}\)(Martens et al., 2012; Grosse and Martens, 2016; Grosse, 2021) and some parametric FSD approximations work by fitting simple statistical models to the resulting distribution. For
instance, assuming all coordinates are independent gives a diagonal approximation (Kirkpatrick et al., 2017), and more fine-grained independence assumptions between network layers yield a Kronecker-factored approximation (Martens and Grosse, 2015; Ritter et al., 2018). In practice, instead of sampling vectors whose covariance is \(\mathbf{G_{\theta}}\), many works use the empirical gradients during training, whose covariance is the empirical Fisher matrix. We caution the reader that the empirical Fisher matrix is less well motivated theoretically and can result in different behavior (Kunstner et al., 2019).

## 3 A Parametric Estimate with LAFTR

We introduce and apply **LAFTR** (Linearized Activation Function TRick) to linear ReLU networks and propose **BGLN** (Bernoulli Gated Linear Network) which approximates a given model architecture as a linear network with stochastic gating. While it is applicable to different architectures, we first explicitly derive our approximation for multilayer perceptrons (MLPs) with \(L\) fully-connected layers and ReLU activation function \(\phi\). We also discuss its generalization to convolutional networks and empirically evaluate the same in Section 5.

For MLPs with inputs \(\mathbf{x}\) drawn from \(p_{\text{data}}\), layer \(l\) weights and biases \((\mathbf{W}^{(l)},\mathbf{b}^{(l)})\), and outputs \(\mathbf{z}\), the computation of preactivations and activations at each layer is recursively defined as follows:

\[\mathbf{s}^{(l)}=\mathbf{W}^{(l)}\mathbf{a}^{(l-1)}+\mathbf{b}^{(l)},\ \mathbf{a}^{(l)}=\phi(\mathbf{s}^{(l)}) \tag{5}\]

with \(\mathbf{a}^{(0)}=\mathbf{x}\), and \(\mathbf{s}^{(L)}=\mathbf{z}\). We denote \(\mathbf{z}_{0}\) and \(\mathbf{z}_{1}\) to be samples of the output distribution obtained with parameters \(\mathbf{\theta}_{0}\) and \(\mathbf{\theta}_{1}\), respectively.

### Linearized Activation Function TRick

Given parameters \(\mathbf{\theta}_{0}\) and \(\mathbf{\theta}_{1}\) of two networks, we linearize _each step of the forward pass_ around its value under \(\mathbf{\theta}_{0}\). For an MLP that alternates between linear layers and non-linear activation functions, the linear transformations are unmodified while the activation functions are replaced with a first-order Taylor approximation around their inputs. Hence, the network's computation becomes linear in \(\mathbf{x}\) (but, importantly, remains nonlinear in \(\mathbf{\theta}\)). Let \((\mathbf{W}_{i}^{(l)},\mathbf{b}_{i}^{(l)})\) denote the weights and biases of layer \(l\) in network \(i\).

\[\mathbf{s}_{0}^{(l)} =\mathbf{W}_{0}^{(l)}\mathbf{a}_{0}^{(l-1)}+\mathbf{b}_{0}^{(l)} \tag{6}\] \[\mathbf{a}_{0}^{(l)} =\phi(\mathbf{s}_{0}^{(l)})\] (7) \[\mathbf{s}_{1}^{(l)} =\mathbf{W}_{1}^{(l)}\mathbf{a}_{1}^{(l-1)}+\mathbf{b}_{1}^{(l)}\] (8) \[\mathbf{a}_{1}^{(l)} =\phi(\mathbf{s}_{0}^{(l)})+\phi^{\prime}(\mathbf{s}_{0}^{(l)})\odot(\bm {s}_{1}^{(l)}-\mathbf{s}_{0}^{(l)}), \tag{9}\]

where \(\phi^{\prime}\) is the derivative of the activation function.

We define some additional notation for differences between preactivations and activations. For \(\Delta\mathbf{s}^{(l)}=\mathbf{s}_{1}^{(l)}-\mathbf{s}_{0}^{(l)}\), \(\Delta\mathbf{a}^{(l)}=\mathbf{a}_{1}^{(l)}-\mathbf{a}_{0}^{(l)}\), and \(\Delta\mathbf{W}^{(l)}=\mathbf{W}_{1}^{(l)}-\mathbf{W}_{0}^{(l)}\), we have the following formulae:

\[\Delta\mathbf{s}^{(l)} =\Delta\mathbf{W}^{(l)}\mathbf{a}_{0}^{(l-1)}+\mathbf{W}_{1}^{(l)}\Delta\mathbf{ a}^{(l-1)}+\Delta\mathbf{b}^{(l)} \tag{10}\] \[\Delta\mathbf{a}^{(l)} =\phi^{\prime}(\mathbf{s}_{0}^{(l)})\odot\Delta\mathbf{s}^{(l)}. \tag{11}\]

Here, base cases are \(\mathbf{a}_{0}^{(0)}=\mathbf{x}\) and \(\Delta\mathbf{a}^{(0)}=0\). Written this way, the \(\Delta\) terms at each step of computation rely linearly on corresponding terms from the immediately preceding step, making LAFTR conducive to implementation via propagation of the base cases through the network. Observe that with \(\mathbf{W}_{0}\) held fixed, the model parametrized using \(\mathbf{W}_{1}\) is a linear network, i.e. the network's exact outputs are a linear function of its inputs.
There are two significant differences between LAFTR and parameter space linearization (referred to as NTK (Jacot et al., 2018) henceforth), which confer an advantage to the former. First, our linearization is with respect to inputs instead of parameters (Lee et al., 2019), hence capturing nonlinear interactions between the parameters in different layers. Second, the only computations that introduce linearization errors into our approximation are those involving activation functions, in contrast with other methods which suffer linearization errors for each layer, including linear layers, where nonlinear parameter dependencies exist. In fact, our method is exact for linear networks, whereas NTK is only approximate. We note that linear networks are commonly used to model nonlinear training dynamics of neural networks (Saxe et al., 2013). Hence, we regard our weaker form of linearity as a significant advantage over Taylor approximations in parameter space.

LAFTR applies to any linear (including fully-connected and convolutional) network with nonlinear activations. Next, we use the intuition above to motivate two probabilistic approximations which enable a memory-efficient implementation of this algorithm.

### Bernoulli Gating

In the specific case of ReLU networks, our approximation depends on the training data through the signs of preactivations, which we denote as a mask \(\mathbf{m}=\mathds{1}\{\boldsymbol{s}>0\}\), since \(\phi(\boldsymbol{s})=\mathbf{m}\odot\boldsymbol{s}\) and \(\phi^{\prime}(\boldsymbol{s})\odot\Delta\boldsymbol{s}=\mathbf{m}\odot\Delta \boldsymbol{s}\). Here, \(\phi^{\prime}\) is the piecewise derivative of the ReLU function given by \(\phi^{\prime}(\boldsymbol{s})=\mathds{1}\{\boldsymbol{s}>0\}\). Given this structure, we model \(\mathbf{m}\) as a vector of independent Bernoulli random variables and fit its mean vector \(\boldsymbol{\mu}\) using maximum likelihood estimation (i.e., computing the fraction of times the unit is activated). We can accordingly rewrite equations 7 and 11 as \(\boldsymbol{a}_{0}^{(l)}=\mathbf{m}^{(l)}\odot\boldsymbol{s}_{0}^{(l)}\) and \(\Delta\boldsymbol{a}^{(l)}=\mathbf{m}^{(l)}\odot\Delta\boldsymbol{s}^{(l)}\), respectively, where \(\odot\) denotes element-wise multiplication and \(\mathbf{m}^{(l)}\sim Ber(\boldsymbol{\mu}^{(l)})\). For efficiency, we compute this average over the last epoch of training. We call this approximation the Bernoulli Gated Linear Network (BGLN).

Note that this gating technique is not specific only to MLPs. It can be implemented in ReLU convolutional networks by replacing activations with Bernoulli random variables.

### Propagating Moments of the Activations

A key insight enabling efficient computation is that when the output distance \(\rho\) is chosen to be (squared) Euclidean distance, the FSD depends only on the first and second moments of the output difference \(\Delta\boldsymbol{z}:=\boldsymbol{z}_{1}-\boldsymbol{z}_{0}\):

\[\mathbb{E}\left[\tfrac{1}{2}||\Delta\boldsymbol{z}||^{2}\right]=\tfrac{1}{2}|| \mathbb{E}[\Delta\mathbf{z}]||^{2}+\tfrac{1}{2}\mathit{tr}\;(\mathrm{Cov}( \Delta\mathbf{z})). \tag{12}\]

We compute these terms recursively by propagating the first two moments of \(\boldsymbol{a}_{0}^{(l)}\) and \(\Delta\boldsymbol{a}^{(l)}\) through the network. Using equations 10 and 11, we obtain the following equations:

\[\mathbb{E}[\boldsymbol{s}_{0}^{(l)}] =\boldsymbol{W}_{0}^{(l)}\mathbb{E}[\boldsymbol{a}_{0}^{(l-1)}]+ \boldsymbol{b}_{0}^{(l)}\] \[\mathbb{E}[\boldsymbol{a}_{0}^{(l)}] =\boldsymbol{\mu}^{(l)}\odot\mathbb{E}[\boldsymbol{s}_{0}^{(l)}]\] \[\mathbb{E}[\Delta\boldsymbol{s}^{(l)}] =\Delta\boldsymbol{W}^{(l)}\mathbb{E}[\boldsymbol{a}_{0}^{(l-1)}] +\boldsymbol{W}_{1}^{(l)}\mathbb{E}[\Delta\boldsymbol{a}^{(l-1)}]+\Delta \boldsymbol{b}^{(l)}\] \[\mathbb{E}[\Delta\boldsymbol{a}^{(l)}] =\boldsymbol{\mu}^{(l)}\odot\mathbb{E}[\Delta\boldsymbol{s}^{(l)}]\] \[\mathrm{Cov}(\boldsymbol{s}_{0}^{(l)}) =\boldsymbol{W}_{0}^{(l)}\mathrm{Cov}(\boldsymbol{a}_{0}^{(l-1)}) \boldsymbol{W}_{0}^{(l)T}\] \[\mathrm{Cov}(\Delta\boldsymbol{s}^{(l)}) \approx\Delta\boldsymbol{W}^{(l)}\mathrm{Cov}(\boldsymbol{a}_{0}^ {(l-1)})\Delta\boldsymbol{W}^{(l)T}+\boldsymbol{W}_{1}^{(l)}\mathrm{Cov}( \Delta\boldsymbol{a}^{(l-1)})\boldsymbol{W}_{1}^{(l)T}\] \[\mathrm{Cov}(\boldsymbol{a}_{0}^{(l)}) =(\boldsymbol{\mu}^{(l)}\boldsymbol{\mu}^{(l)T})\odot\mathrm{ Cov}(\boldsymbol{s}_{0}^{(l)})\] \[\mathrm{Cov}(\Delta\boldsymbol{a}^{(l)}) =(\boldsymbol{\mu}^{(l)}\boldsymbol{\mu}^{(l)T})\odot\mathrm{ Cov}(\Delta\boldsymbol{s}^{(l)})\]
We assume that \(\text{Cov}(\mathbf{a}_{0},\Delta\mathbf{a})\) is close to \(0\). This assumption is tested empirically in our experiments and we find that it does not severely move the FSD estimate away from the true empirical FSD (see Figure 2). Hence, the only information that needs to be stored about the data itself are the first and second moments. While the choice to store only the moments is justified only for squared Euclidean distance, we find that it also works well empirically for other output space metrics such as KL divergence.

### BGLN-D and BGLN-S

The most straightforward way to use the BGLN approximation is to draw Monte Carlo samples of the random variables. When only the first and second moments matter, we are free to assume Gaussianity of the inputs. We denote this method BGLN-S (S for "stochastic"). This is sufficient in situations where FSD is used as a regularization term in stochastic gradient-based optimization (as we do in our continual learning experiments). In other situations, it is advantageous to have a deterministic computation; for instance, optimization with nonlinear conjugate gradient requires a deterministic objective. In the case of Euclidean distance as the output space metric, we can exactly compute the BGLN approximation by propagating the first and second moments of all random variables through the forward pass. We call the deterministic estimator BGLN-D. **BGLN-S** is outlined in Algorithm 1 and **BGLN-D** in Algorithm 2. We describe the analogous BGLN-S computations for convolutional networks in Appendix C.

We present the above algorithms using the mean and covariance of the data. Note that storing moments of the data usually has less memory cost than storing sufficient subsets of the data itself. We can further reduce the memory requirement by approximating the covariance matrix as a diagonal matrix, i.e., using only the variance of each dimension of the inputs. This is equivalent in cost to storing two data points per task. We empirically investigate the effect of this approximation on continual learning benchmarks. Furthermore, capturing the expected FSD over the training distribution in a single term, via a single forward pass, is far less computationally expensive than iterative alternatives.

### Class-conditional Estimates

In the classification setting, it is also possible to a extend our method to a more fine-grained, class-conditional approximation. In particular, we can fit a mixture model for our probabilistic approximations, with one component per class. In this case, each class has its own associated input moments and Bernoulli mean parameters. The lower memory cost of our method allows for this when number of classes is not too large. We refer to their variant as BGLN-CW. As expected, it boosts performance in our continual learning experiments at the expense of a slightly higher memory requirement, as shown in Tables 1, 2 and 3.

## 4 Related Works

Several works (Benjamin et al., 2018; Bernstein et al., 2020; Bae et al., 2022) have highlighted the importance of measuring meaningful distances between neural networks. Benjamin et al. (2018) contrast training dynamics in parameter space and function space and observe that function space distances are often more useful than, and not always correlated with, parameter space distances. Bae et al. (2022) propose an Amortized Proximal Optimization (APO) scheme that regularizes an FSD estimate to the previous iterate for second-order optimization. Natural gradient descent (Amari et al., 1995; Amari, 1998) can also be interpreted as a steepest descent method, using a second-order Taylor approximation to the FSD (Pascanu and Bengio, 2014).

Parisi et al. (2019); De Lange et al. (2021); Ramasesh et al. (2020); Normandin et al. (2021) have reviewed and surveyed the challenge of catastrophic forgetting in continual learning, along with benchmarks and metrics to evaluate different methods. Parametric methods focus on different approximations to the weight
space metric matrix, like diagonal (Kirkpatrick et al., 2017, EWC) or Kronecker-factored (Ritter et al., 2018, OSLA). As described in Section 2, we interpret these as second-order Taylor approximations to the FSD with further structured approximations to the Hessian. Several methods are motivated as approximations to a posterior Gaussian distribution in a Bayesian setting (Ebrahimi et al., 2019), for instance through a variational lower bound (Nguyen et al., 2017) or via Gaussian process inducing points (Kapoor et al., 2021). Non-parametric methods (Kapoor et al., 2021; Titsias et al., 2019; Pan et al., 2020; Rudner et al., 2022; Kirichenko et al., 2021) usually employ some form of experience replay of stored or optimized data points. Some of these methods (Pan et al., 2020) can also be related to the Neural Tangent Kernel (Jacot et al., 2018, NTK), or in other words, network linearization. Doan et al. (2021) 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.

In this paper, we also examine influence functions (Cook, 1979; Hampel, 1974) which is another application that involves the FSD between networks. Influence functions are a classical robust statistics technique that has since been used in machine learning (Koh and Liang, 2017). Bae et al. (2022) formally study influence functions in neural networks and show that they approximate an objective called the proximal Bregman response function (PBRF). This approximation depends on a FSD term that is typically computed by iterating through the full training dataset.

## 5 Experiments

We empirically assess the effectiveness of LAFTR (our idealized method) and the BGLN (our practical algorithm) in approximating FSD as well as their usefulness for downstream tasks: continual learning and influence function estimation. The experiments investigate the following questions:

* Can LAFTR outperform NTK in approximating FSD?
* Does LAFTR improve performance and memory cost on continual learning benchmarks relative to existing methods?
* How do choice of output space metric \(\rho\), or the use of the Gaussian input and the Bernoulli activation approximations, impact empirical performance?
* Can the BGLN perform competitively with iteration-based influence function estimators without requiring iteration over the dataset?

### Comparing FSD Estimators

To conduct further empirical analysis of our methods' estimation and minimization of the true (empirical) FSD, we use tasks and models from standard continual learning settings which are prone to forgetting. These include Split MNIST, Permuted MNIST and Split CIFAR100 (Pan et al., 2020; Rudner et al., 2022). In addition to directly evaluating the continual learning performance, we also use a collection of networks trained in the course of this experiment (with varying hyperparameter settings) to directly evaluate the accuracy of the FSD estimates. Specifically, we vary the learning rate and the number of training iterations, and consider the set of trained networks that result; on each pair of these networks, we compare the FSD estimates against the true empirical FSD computed using the full training set.

Figure 2 (Left) shows that BGLN-S and BGLN-D consistently estimate the true FSD more accurately

Figure 3: **(Left)** While training on task 2, FSD from the optimal task 1 parameters increases with task 2 accuracy. Optimizing BGLN-D and class-conditioned BGLN-D-CW effectively minimizes the true FSD. **(Right)** LAFTR has a higher correlation with true FSD than NTK, with a more significant advantage as network depth (and hence a number of nonlinear interactions) increases.

Figure 2: Compared to NTK, LAFTR-based approximations consistently give closer FSD values to the true empirical FSD when evaluated on Permuted MNIST **(Left)** and Split CIFAR100 **(Right)**. Quantified in terms of Spearman rank-order and Kendall’s Tau correlation coefficients, LAFTR (96.36 and 79.19, respectively) outperforms NTK (86.42 and 71.71, respectively).
than NTK. Analogous analysis for LAFTR using CIFAR100 in Figure 2 (Right) shows a similar trend when the estimators are given access to the same coreset of inputs. We can also measure how the true FSD changes when different FSD estimates are optimized during training, as in Figure 3 (Left), where BGLN methods more effectively minimize true FSD as new task accuracy increases. Finally, we train networks of varying depths on CIFAR100 tasks and measure correlation (Spearman rank-order Spearman, 1961) and Kendall's Tau (Kendall, 1938)) with true FSD. Figure 3 (Right) shows that LAFTR has a higher correlation using both metrics, and its advantage over NTK increases with network depth, and hence with a number of nonlinear interactions between parameters. This corroborates our intuition that LAFTR captures nonlinearities that NTK is unable to account for.

### Continual Learning

Recall the formulation of continual learning in terms of FSD, as described in equation 3. We visualize our method's comparative performance on 1-D regression with two sequential tasks shown in Figure 1. More realistically, we test our methods on standard benchmarks used in prior works (Pan et al., 2020; Rudner et al., 2022), with standard architectures for a fair comparison. See Appendix D.2 for details on the datasets, architectures, and hyperparameters. We evaluate average final accuracy across tasks, backward transfer (Lopez-Paz and Ranzato, 2017) and memory cost.

**Toy Regression.** Figure 1 shows the functions learned by different methods when sequentially trained on two one-dimensional regression tasks. LAFTR gives a better approximation of the learned function than NTK. When used to regularize the network, BGLN retains good predictions on both tasks, while EWC and exact parameter space linearization (NTK) suffer catastrophic forgetting. We hypothesize that this difference in performance is due to important nonlinearities between network parameters that EWC and NTK approximations are unable to capture.

**Split and Permuted MNIST.** As shown in Tables 1 and 2, our LAFTR-based methods outperform other parametric methods (EWC, OSLA, and VCL) on Split and Permuted MNIST tasks and are competitive with the state-of-the-art (SOTA) nonparametric methods, in terms of average accuracy. Class-conditional approximations further boost performance and the diagonal approximation to input covariance (BGLN-S-Var, BGLN-D-Var) does not harm it significantly. With respect to the backward transfer, a more direct measure of forgetting, BGLN methods significantly outperform the SOTA. Finally, they are also amenable to successful adaptation to the nonparametric setting when a coreset is available for use in place of Gaussian samples.

**Split CIFAR100.** We consider the much more challenging Split CIFAR100 task to compare our method to existing approaches and tease apart the effects of our algorithmic choices and approximations. Table 3 summarizes these results and analytically compares the memory costs associated with each method (see Appendix D.1 for details). **Coreset** refers to a coreset of real inputs vs. Gaussian samples, **Bernoulli** refers

\begin{table}
\begin{tabular}{l|c c} \hline \hline
**Method** & **Split MNIST** & **Permuted MNIST** \\ \hline Nonparametric & & \\ \hline
**VCL (coreset)** & 98.40 & 95.50 \\
**VAR-GP (coreset)** & \(90.57\pm 1.06\) & \(\mathbf{97.20\pm 0.08}\) \\
**FROMP (coreset)** & \(99.00\pm 0.04\) & \(94.90\pm 0.04\) \\
**-FSVI (coreset)** & \(99.54\pm 0.04\) & \(95.76\pm 0.02\) \\
**NTK (coreset)** & \(99.50\pm 0.09\) & \(96.46\pm 0.11\) \\
**BGLN-S (coreset)** & \(99.50\pm 0.03\) & \(96.36\pm 0.13\) \\ \hline Parametric & & \\ \hline
**EWC** & 63.10 & 84.00 \\
**OSLA** & 80.56 & 95.73 \\
**VCL** & 97.00 & \(87.50\pm 0.61\) \\
**BGLN-D** & \(99.72\pm 0.03\) & \(96.03\pm 0.20\) \\
**BGLN-D-CW** & \(99.78\pm 0.02\) & \(96.85\pm 0.02\) \\
**BGLN-S** & \(99.64\pm 0.04\) & \(96.36\pm 0.12\) \\
**BGLN-S-CW** & \(99.77\pm 0.05\) & \(96.99\pm 0.07\) \\
**BGLN-D-Var** & \(99.64\pm 0.04\) & \(94.98\pm 0.18\) \\
**BGLN-S-Var** & \(99.50\pm 0.03\) & \(96.36\pm 0.13\) \\ \hline \hline \end{tabular}
\end{table}
Table 1: Average accuracies of nonparametric and parametric approaches on Split and Permuted MNIST datasets.

\begin{table}
\begin{tabular}{l|c|c} \hline \hline
**Method** & **Split MNIST** & **Permuted MNIST** \\ \hline
**FROMP** & \(-0.50\pm 0.20\) & \(-1.00\pm 0.10\) \\
**S-FSVI** & \(-0.21\pm 0.06\) & \(-0.65\pm 0.21\) \\ \hline
**BGLN-S** & \(\mathbf{-0.04\pm 0.03}\) & \(\mathbf{-0.41\pm 0.08}\) \\
**BGLN-D** & \(\mathbf{-0.09\pm 0.04}\) & \(-0.56\pm 0.04\) \\
**BGLN-S-CW** & \(-0.18\pm 0.06\) & \(\mathbf{-0.37\pm 0.04}\) \\
**BGLN-D-CW** & \(\mathbf{-0.07\pm 0.07}\) & \(-1.17\pm 0.07\) \\ \hline \hline \end{tabular}
\end{table}
Table 2: Backward transfer on Split and Permuted MNIST. Lower is better.
to Bernoulli activations vs. simply passing preactivations through the ReLU function (termed LAFTR here) and **CW** refers to the class-conditonal estimate. Our empirical analysis shows that given a random coreset of the same size as comparable methods, LAFTR outperforms the SOTA as well as the NTK baseline on average accuracy and backward transfer significantly. Other LAFTR and BGLN variants remain competitive with prior methods while incurring lower memory costs. We also observe that performance is hurt to some extent by the Gaussian and Bernoulli modeling assumptions, while it is improved by class-conditioning. We present a more fine-grained task-wise comparison of accuracies in Figure 5 of Appendix D.3.

### Influence Function Estimation

To further assess BGLN's applicability to other settings involving FSD estimation and regularization, we consider influence function estimation (Cook, 1979; Hampel, 1974; Koh and Liang, 2017). Given parameters \(\mathbf{\theta}_{0}\) trained on dataset \(\mathcal{D}_{\text{train}}\) of size \(N\), influence functions approximate the parameters \(\mathbf{\theta}_{-}\) that would be obtained by training without a particular point \((\mathbf{x},\mathbf{y})\in\mathcal{D}_{\text{train}}\). The difference in loss between \(\mathbf{\theta}_{0}\) and \(\mathbf{\theta}_{-}\) is an indicator of the influence of \((\mathbf{x},\mathbf{y})\) on the trained network.

Bae et al. (2022a) show that influence functions in neural networks can be formulated as solving for the proximal Bregman response function (PBRF):

\[\mathbf{\theta}_{-}=\operatorname*{arg\,min}_{\mathbf{\theta}\in\mathbb{R}^{d}}- \frac{1}{N}\mathcal{L}(f(\mathbf{x},\mathbf{\theta}),\mathbf{y})+D_{B}(\mathbf{\theta},\mathbf{ \theta}_{0},p_{\text{train}})+\frac{\lambda}{2}\|\mathbf{\theta}-\mathbf{\theta}_{0} \|^{2}. \tag{13}\]

Here, the first term maximizes the loss of the data point we are interested in removing. The second term is the Bregman divergence defined on network outputs and measures the FSD between \(\mathbf{\theta}\) and \(\mathbf{\theta}_{0}\) over training distribution \(p_{\text{train}}\) similar to the FSD term as defined in equation 1. For standard loss functions like squared error and cross-entropy, the Bregman divergence term is equivalent to the soft training error where the original targets are replaced with soft targets produced by \(\mathbf{\theta}_{0}\). Finally, the last term is a proximity term with strength \(\lambda>0\), which prevents large changes in weight space. Intuitively, the PBRF maximizes the loss of data we would like to remove while constraining the network in both function and weight space so that the predictions and losses of other training examples remain unaffected.

Existing approaches for this optimization face two key challenges: (1) the entire training dataset must be stored and iterated over, requiring as many forward passes as there are mini-batches and (2) techniques like

\begin{table}
\begin{tabular}{l c c c c c c} \hline \hline
**Method** & \(\mathbf{\rho}\) & **Corset** & **Bernoulli** & **CW** & **Average Accuracy \(\uparrow\)** & **Backward Transfor \(\downarrow\)** & **Memory Cost** \\ \hline Nonparametric & & & & & & & \\
**VCL (coreset)** & & & – & & \(67.40\pm 0.60\) & – & \(2P+Nd\) \\
**VAR-GP (coreset)** & & & – & & \(-\) & \(2P+Nd+C^{2}N^{2}\) \\
**FROMP (coreset)** & & & – & & \(76.20\pm 0.20\) & \(-2.60\pm 0.90\) & \(2P+Nd+C^{2}N^{2}\) \\
**S-PSVI (coreset)** & & & – & & \(77.60\pm 0.20\) & \(-2.50\pm 0.20\) & \(2P+Nd+C^{2}N^{2}\) \\
**NTK (coreset)** & KL & – & & & \(77.61\pm 0.20\) & \(-2.03\pm 0.04\) & \(2P+Nd\) \\
**LAFTR (coreset)** & KL & ✓ & ✗ & ✗ & \(78.33\pm 0.01\) & \(-0.78\pm 0.10\) & \(P+Nd\) \\
**BGLN-S (coreset)** & KL & ✓ & ✓ & ✗ & \(73.27\pm 0.01\) & \(-4.99\pm 0.28\) & \(P+Nd\) \\
**LAFTR (coreset)** & Euclidean & ✓ & ✗ & ✗ & \(76.22\pm 0.01\) & \(-2.64\pm 0.40\) & \(P+Nd\) \\ \hline Parametric & & & – & & \(71.60\pm 0.40\) & – & \(2P\) \\
**EWC** & & & – & & \(72.61\) & – & \(P+\frac{7L_{+}}{2}p_{1}^{2}\) \\
**OSLA** & & & – & & \(-\) & \(2P\) \\
**VCL** & & & – & & \(-\) & \(2P\) \\
**LAFTR** & KL & ✗ & ✗ & \(75.61\pm 0.01\) & \(-19.03\pm 0.59\) & \(P+d+d^{2}\) \\
**LAFTR-CW** & KL & ✗ & ✓ & \(76.22\pm 0.01\) & \(-1.45\pm 0.63\) & \(P+C(d+d^{2})\) \\
**BGLN-S** & KL & ✗ & ✓ & ✗ & \(72.37\pm 0.01\) & \(-8.20\pm 0.04\) & \(P+d+d^{2}\) \\
**BGLN-S-CW** & KL & ✗ & ✓ & ✓ & \(74.02\pm 0.01\) & \(-2.44\pm 0.15\) & \(P+C(A+d+d^{2})\) \\
**LAFTR** & Euclidean & ✗ & ✗ & \(75.51\pm 0.01\) & \(-3.12\pm 0.44\) & \(P+d+d^{2}\) \\
**BGLN-S** & Euclidean & ✗ & ✓ & ✗ & \(74.29\pm 0.01\) & \(-5.49\pm 0.05\) & \(P+4+d+d^{2}\) \\
**BGLN-S-CW** & Euclidean & ✗ & ✓ & ✓ & \(77.78\pm 0.01\) & \(-1.75\pm 0.50\) & \(P+C(A+d+d^{2})\) \\ \hline \hline \end{tabular}
\end{table}
Table 3: Split CIFAR100: Average Accuracy and Backward Transfer. Notation for memory cost: \(p_{l}=\#\) parameters in layer \(l\), \(P=\#\) parameters \(=\sum_{l=1}^{L}p_{l}\), \(A=\#\) activations \(<P\), \(d=\) data dimension, \(N=\) coreset size, \(C=\#\) classes.
nonlinear Conjugate Gradient (CG) (Hager and Zhang, 2006) do not work well with the stochastic gradients produced by sampling batches of data. LAFTR enables estimating the PBRF (or FSD) by storing only the first two data moments, requires just a single forward pass to compute it and provides a deterministic function to optimize, implemented as BGLN-D.

**Regression.** We first train a MLP with two hidden layers and ReLU activations for 200 epochs on regression datasets from the UCI benchmark (Dua and Graff, 2017). Then, we randomly select 50 independent data points to be removed. For each removed point, we sample batches and use a Stochastic Gradient Descent (SGD) optimizer to minimize the PBRF objective and compute the difference in loss after removing that data point, commonly referred to as the self-influence score (Koh and Liang, 2017; Schioppa et al., 2022). Next, we follow the same procedure as above but approximate the FSD term in the PBRF objective with EWC, CG (Koh and Liang, 2017) and BGLN-D. Since the direct minimization of PBRF can be considered as the ground truth for influence estimation, we compare the alignment of these methods' estimates with that of PBRF via Pearson correlation (Sedgwick, 2012) and Spearman rank-order correlation (Spearman, 1961). The results are shown in Table 4. Without having to iterate over or store the entire dataset, BGLN-D correlates with PBRF more strongly than EWC and CG (Koh and Liang, 2017), which can be seen as minimizing a linearized version of the PBRF objective.

**Mislabeled Example Detection.** Influence function estimators are commonly evaluated in terms of their ability to identify mislabeled examples. Intuitively, if some fraction of the training labels is corrupted, they would behave as outliers and have a more significant influence on the training loss (self-influence score). One approach to efficiently detect and correct these examples is to prioritize and examine training inputs with higher self-influence scores. Following the evaluation setup from Bae et al. (2022), we use 10% of the MNIST dataset and corrupt 10% of it by assigning random labels to it. We train a two layer MLP with 1024 hidden units and ReLU activations using SGD with a batch size of 128. Then, we use EWC, CG and BGLN-D to approximate the FSD term in equation 13 and compute individual self-influence scores. We also compare these methods against a baseline of randomly sampling data points to check for corruption. The results are summarized in Figure 4. BGLN-D significantly outperforms the random baseline and EWC and closely matches the oracle PBRF and CG, while being much faster, cheaper and more memory-efficient.

\begin{table}
\begin{tabular}{l c c c c c c} \hline \hline
**Dataset** & \multicolumn{2}{c}{**EWC**} & \multicolumn{2}{c}{**CG**} & \multicolumn{2}{c}{**BGLN-D**} \\ \cline{2-7}  & P & S & P & S & P & S \\ \hline Concrete & 0.78 & 0.57 & 0.92 & 0.94 & **0.96** & **0.97** \\ 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 \\ Kinetics & 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 \\ \hline \hline \end{tabular}
\end{table}
Table 4: Evaluation of training loss differences computed by EWC, CG and BGLN-D. We compare Pearson (P) and Spearman rank-order (S) correlations with the PBRF estimates.

Figure 4: Effectiveness of BGLN in detecting mislabeled examples. BGLN can approximate the FSD term in the PBRF objective accurately and be used in applications involving influence functions without explicitly storing or iterating over the dataset.
Conclusions

In this work, we addressed the problem of compactly summarizing a model's predictions on a given dataset, and formulated it as approximating neural network FSD. We developed the Linearized Activation Function TRick as an improvement over network linearization in parameter space and proposed novel parametric methods, BGLN, to estimate FSD. Our methods capture nonlinearities between network parameters, are much more memory-efficient than prior works and are amenable to adaptation to the nonparametric setting when a coreset of data is available.

We empirically show that LAFTR-based estimates are highly correlated with the true FSD across several settings. In continual learning, our methods outcompete existing methods without storing any data samples. Further, in influence function estimation, they estimate influence-scores with high correlation and can efficiently detect mislabeled examples without expensive iteration over the whole dataset.

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 enhance memory and computational efficiency in settings where estimating or constraining FSD is relevant.

## Acknowledgements

We would like to thank Florian Shkurti for useful discussions, Cem Anil for feedback on the draft, and Gerald Shen for assistance with the compute environment. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute (www.vectorinstitute.ai/partners).
