Title: Non-Stationary Learning of Neural Networks with Automatic Soft Parameter Reset
Venue: NeurIPS 2024
Accepted: yes
Keywords: Non-stationarity

plasticity loss

online learning

deep learning
Abstract: Neural networks are most often trained under the assumption that data come from a stationary distribution. However, settings in which this assumption is violated are of increasing importance; examples include supervised learning with distributional shifts, reinforcement learning, continual learning and non-stationary contextual bandits. Here, we introduce a novel learning approach that automatically models and adapts to non-stationarity by linking parameters through an Ornstein-Uhlenbeck process with an adaptive drift parameter. The adaptive drift draws the parameters towards the distribution used at initialisation, so the approach can be understood as a form of soft parameter reset. We show empirically that our approach performs well in non-stationary supervised, and off-policy reinforcement learning settings.

Non-Stationary Learning of Neural Networks with Automatic Soft Parameter Reset

Abstract

Neural networks are most often trained under the assumption that data come from a stationary distribution. However, settings in which this assumption is violated are of increasing importance; examples include supervised learning with distributional shifts, reinforcement learning, continual learning and non-stationary contextual bandits. Here, we introduce a novel learning approach that automatically models and adapts to non-stationarity by linking parameters through an OrnsteinUhlenbeck process with an adaptive drift parameter. The adaptive drift draws the parameters towards the distribution used at initialisation, so the approach can be understood as a form of soft parameter reset. We show empirically that our approach performs well in non-stationary supervised, and off-policy reinforcement learning settings.

1 Introduction

Neural networks (NNs) are typically trained using algorithms such as stochastic gradient descent (SGD), which implicitly assume that the data come from a stationary distribution. This assumption is incorrect for scenarios such as continual learning, reinforcement learning, non-stationary contextual bandits, and supervised learning with distribution shifts [20, 53]. Although the parameters can be updated online as new data are encountered, this approach often leads to a loss of plasticity [12, 2, 13], manifesting either as a failure to generalise to new data despite reduced training loss [4, 2], or as an inability to reduce training error as the data distribution changes [13, 37, 1, 42, 34].

In [38], the authors argue for two factors that lead to loss of plasticity: shifts in the distribution of preactivations, leading to dead or dormant neurons [47], and growth in the parameter norm, causing training instabilities. These problems are often addressed using hard resets based on heuristics such as detecting dormant units [47], assessing unit utility [13, 12], or simply after a fixed number of steps has elapsed[43]. Although effective at increasing plasticity, hard resets can be inefficient as they can discard valuable knowledge captured by the parameters.

We propose an algorithm that instead implements a form of soft parameter reset, avoiding the pitfalls associated with hard resets. A soft reset moves the parameters in the direction of the initial-value distribution, while maintaining dependence on their previous values. It also increases the learning rate applied to gradient-based updates, allowing new parameters to adapt faster to the changing data. The magnitude of the soft reset is governed by an Ornstein-Uhlenbeck drift process linking the parameters in time, with the scale of the drift itself chosen adaptively. In effect, the model approximates the action of a dynamical Bayesian prior over the NN parameters, which is adapted online to new data.

Our contributions can be summarised as follows. We provide an explicit formalisation of the model for the drift in NN parameters as sketched above, and derive a procedure to estimate the parameters of the drift model online. Second, we use the learnt drift model to modify the NN parameter update algorithm. Third, we explore the effectiveness of the approach in supervised learning experiments, showing that it avoids the challenge of plasticity loss, as well as in an off-policy reinforcement learning setting.

2 Non-stationary learning with Online SGD

Consider a non-stationary supervised learning setting with changing data distributions p t ( x, y ) , where x ∈ R L and y ∈ R K ( y may be restricted to "one-hot" categorical indicator vectors). Given a single-point loss L ( θ, x, y ) , we define the time-dependent expected loss function for parameters θ ∈ R D to be

Our goal is to find a parameter sequence Θ = ( θ 1 , . . . , θ T ) that reduces the dynamic regret

relative to an oracular reference sequence Θ ⋆ = ( θ ⋆ 1 , . . . , θ ⋆ T ) satisfying θ ⋆ t ∈ arg min θ L t ( θ ) . One common approach to online learning employs online stochastic gradient descent (SGD) [23]. An initial parameter value θ 0 is updated sequentially for each batch of data B t = { ( x i t , y i t ) } B i =1 s.t. ( x i t , y i t ) ∼ p t ( x t , y t ) . The update rule is:

where ̂ L t ( θ ) = 1 B ∑ B i =1 L ( θ, x i t , y i t ) is the empirical loss on the batch, and α t is a learning rate. See also Appendix G for the connection of SGD to proximal optimization.

Convex settings. For convex losses, online SGD with a fixed learning rate α can track nonstationarity [56]. By selecting α appropriately-potentially using additional knowledge about the reference sequence-we can optimise the dynamic regret in (2). In general, algorithms that adapt to the observed level of non-stationarity can outperform standard online SGD. For example, in [29], the authors propose to adjust the learning rate α t , while in [21] and in [29], the authors suggest modifying the starting point of SGD from θ t to an adjusted θ ′ t proportional to the level of non-stationarity.

Non-convex settings. Non-stationary learning with NNs is more complex, since now there is a changing set of local minima as the data distribution changes. Such changes can lead to a loss of plasticity and other pathologies. Alternative optimization methods like Adam [30], do not fully resolve this issue [13, 37, 1, 42, 34]. Parameter resets [13, 48, 12] partially mitigate the problem, but may be too aggressive if the data distributions are similar.

3 Online non-stationary learning with learned soft parameter resets

Notation. We denote by N ( θ ; µ, σ 2 ) a Gaussian distribution on θ with mean µ and variance σ 2 . We denote by θ j the j -th component of the vector θ = ( θ 1 , . . . , θ D ) . We assume that the NN outputs define a distribution over targets y , so that the single-point loss is the negative log likelihood L ( θ, x, y ) = -log p ( y | x, θ ) . As above, B t = { ( x i t , y i t ) } B i =1 is a batch of data sampled from p t ( x, y ) with a fixed batch size B (which could be 1 ) and ̂ L t ( θ ) = log p ( B t | θ ) = 1 B ∑ B i =1 L ( θ, x i t , y i t ) is the average negative log likelihood on batch B t . Element-wise product is denoted by ◦ .

Our proposed soft reset scheme tracks changes in the distribution of non-stationary data using an explicit drift model for the parameters p ( θ t +1 | θ t , γ t ) . Specifically, the drift model we adopt assumes

Graphical model

(d) Non-stationary case

without dynamical model

(e) Non-stationary case with

dynamical model

(a) i.i.d. assumption

(b) non i.i.d. assumption

Toy Bayesian Inference example

that the parameters change in the direction of the initialization distribution at each time t . The amount of change, and thus the strength of reset, is controlled by non-stationary hyperparameters γ t . These hyperparameters are themselves estimated online from the data.

From a probabilistic standpoint, the drift model implies an empirical Bayesian prior on parameters θ t +1 , which is displaced from the posterior on θ t towards the prior distribution, and has increased variance (controlled by γ t ). As we argue below, the effect can also be interpreted in the context of SGD learning, as an adjustment in the starting point for batch B t +1 to a point ˜ θ t ( γ t ) between the previous estimate θ t and the mean of initializing distribution, with a corresponding increase in the learning rate. This approach is inspired by prior work in online convex optimization for non-stationary environments [e.g., 25, 21, 8, 18, 29].

3.1 Toy illustration of the advantage of drift models

Consider online Bayesian inference with 2-D observations y t = θ ⋆ + ϵ t , where θ ⋆ ∈ R 2 are unknown true parameters and ϵ t ∼ N (0; σ 2 I ) is Gaussian noise with variance σ 2 . Starting from a Gaussian prior p 0 ( θ ) = N ( θ ; µ 0 , Σ 0 ) , the posterior distribution p t +1 ( θ ) = p ( θ | y 1 , . . . , y t +1 ) = N ( θ ; µ t +1 , Σ t +1 ) is updated using Bayes' rule

The posterior update (3) arises from assumption that data are i.i.d. (Figure 1a), since then p t +1 ( θ ) ∝ p 0 ( θ ) ∏ t +1 s =1 p ( y s | θ ) . By the Central Limit Theorem, the posterior mean µ t converges to θ ⋆ and the covariance matrix Σ t shrinks to zero (the radius of red circle in Figure 1c).

Suppose now that the true parameters θ ⋆ t (which were fixed up to time t ) change to new parameters θ ⋆ t +1 at time t +1 . The i.i.d. assumption is thus violated and the update (3) becomes problematic because the low uncertainty (small radius of red dashed circle in Figure 1d) in p t ( θ ) causes the posterior p t +1 ( θ ) (see blue circle) to adjust slowly towards θ ⋆ t +1 (blue dot) as illustrated in Figure 1d.

The issue is addressed by allowing explicitly for the possibility that, before observing new data, the parameters drift according to p ( θ t +1 | θ t , γ t ) . The corresponding conditional independence structure is shown in Figure 1b. The posterior update then becomes:

For a suitable choice of drift model p ( θ t +1 | θ t , γ t ) , this modification allows p t +1 ( θ ) (blue circle) to adjust more rapidly towards the new θ ⋆ t +1 (blue dot), see Figure 1e. This is because the new prior ∫ p ( θ | θ ′ t , γ t ) p t ( θ ′ t ) dθ ′ t has larger variance (green circle) than p t ( θ ) and its mean is closer to the center of the circle. Ideally, the parameter γ t should capture the underlying non-stationarity in the data distribution in order to control the impact of the prior ∫ p ( θ | θ ′ t , γ t ) p t ( θ ′ t ) dθ ′ t . For example, if at some point the non-stationarity disappears, we want the drift model to adaptively eliminate changes and recover the stationary posterior update (3). This highlights the importance of the adaptive nature of the drift model.

3.2 Ornstein-Uhlenbeck parameter drift model

We motivate a specific choice of drift model which is useful for maintaining plasticity. Assume that the NN is flexible enough to capture any stationary dataset in a fixed number of iterations starting from a good initialization θ 0 ∼ p 0 ( θ ) [see, e.g., 24, 16]. Informally, we refer to the region of high probability under p 0 as a plastic region .

In a non-stationary data setting, learning for the batch B t +1 is initialized based on the learnt parameter at the previous time step, θ t . Empirically, if this parameter is far from the initial plastic region, the NN may suffer from a loss of plasticity with SGD and similar online methods failing to adapt to the new data. In this case, a hard reset of the parameters to θ 0 (or an alternative draw from p 0 ) may be helpful. However, if changes in the data distribution are gradual then performance may be improved by retaining information about θ t . Furthermore, since θ is high-dimensional, each dimension might benefit from different updates.

Drift model. A drift model p ( θ t +1 | θ t , γ t ) which captures these desiderata is given by

defined independently for each parameter dimension θ j , where p 0 ( θ j 0 ) = N ( θ j 0 ; µ j 0 ; σ j 0 2 ) is the perparameter prior distribution and γ t = ( γ 1 t , . . . , γ D t ) controls drift in each parameter separately. The model is a discretized Ornstein-Uhlenbeck (OU) process [50] (see Appendix B for the derivation).

The parameters γ j t ∈ [0 , 1] control the degree to which each parameter is reset. When γ j t = 1 , the model has θ j t +1 and θ j t equal, and the data in B t +1 are used to refine the estimate of the common parameter. When γ j t = 0 , θ j t +1 is drawn from the prior p 0 independently of θ j t , and so B t +1 is used independently of previous data. A value of γ j t ∈ (0 , 1) interpolates between these two extremes. The process (5) has the property that for any current parameter θ t and γ j t ∈ [0 , 1) , as T → ∞ the distribution of p ( θ T | θ t ) will converge to the prior p ( θ 0 ) . This behaviour depends on interplay between the mean shrinkage and the variance. Other choices of drift variance would result in the variance of p ( θ T | θ t ) either going to 0 or growing to ∞ , harming learning. Thus, the model (5) encourages parameters to move towards the plastic region (initialization). In Appendix C, we discuss this further and other potential choices for the drift model.

3.3 Online estimation of drift

The parameters γ t can themselves be selected within the empirical Bayesian framework by optimizing the predictive likelihood , which quantifies the probability of new data under the current parameters and drift model. We derive the drift estimation procedure in the context of approximate online variational inference [7] with Bayesian Neural Networks (BNN). Let Γ t = ( γ 1 , . . . , γ t ) be the history of observed parameters of the drift model and S t = {B 1 , . . . , B t } be the history of observed data. The objective of approximate online variational inference is to propagate an approximate posterior q t ( θ |S t , Γ t -1 ) over parameters, such that it is constrained to some family Q of probability distributions. In the context of BNNs, it is typical [5] to assume a family Q = { q ( θ ) : q ( θ ) ∼ ∏ D j =1 N ( θ j ; µ j , σ j 2 ); θ = ( θ 1 , . . . , θ D ) } of Gaussian mean-field distributions over parameters θ ∈ R D (separate Gaussian per parameter). Let q t ( θ ) ≜ q t ( θ |S t , Γ t -1 ) ∈ Q be the Gaussian approximate posterior at time t with mean µ t and variance σ 2 t for every parameter. The new approximate posterior q t +1 ( θ ) ∈ Q is found by

where the prior term is the approximate predictive look-ahead prior given by

that has parameters ˜ µ = (˜ µ 1 , . . . , ˜ µ D ) and ˜ σ 2 = (˜ σ 12 , . . . , ˜ σ D 2 ) such that ˜ µ j t ( γ t ) = γ j t µ j t +(1 -γ j t ) µ j 0 , ˜ σ j t 2 ( γ t ) = γ j t 2 σ j t 2 +(1 -γ j t 2 ) σ j 0 2 , see Appendix E.1 for derivation. The form of this prior q t ( θ | γ t ) comes from the non i.i.d. assumption (see Figure 1b) and the form of the drift model (5). For new batch of data B t +1 at time t +1 , the approximate predictive log-likelihood equals to

The log-likelihood (8) allows us to quantify predictions on batch of data B t +1 given our current distribution q t ( θ ) and the drift model from (5). We find such γ ⋆ t that

Using γ ⋆ t in (5) modifies the prior distribution (7) to fit the most recent observations the best by putting more mass on the region where the new parameter could be found (see Figure 1,right).

Gradient-based optimization for γ t . The approximate predictive prior in (7) is Gaussian which allows us to use the so-called reparameterisation trick to optimize (8) via gradient descent. Starting from an initial value of drift parameter γ 0 t at time t , we perform K updates with learning rate η γ

The integral is evaluated by Monte-Carlo (MC) using M samples ϵ i ∼ N ( ϵ ; 0 , I ) , i = 1 , . . . , M

Inductive bias in the drift model is captured by γ 0 t , where γ t, 0 = 1 encourages stationarity, while γ t, 0 = γ t -1 ,K promotes temporal smoothness. In practice, we found γ t, 0 = 1 was the most effective.

Structure in the drift model. The drift model can be defined to be shared across different subsets of parameters which reduces the expressivity of the drift model but also provides regularization to (10). We consider γ t to be either defined for each parameter or for each layer . See Section 5 for details as well as corresponding results in Appendix K.

Interpretation of γ t . By linearising log p ( B t +1 | θ ) around µ t and denoting g t +1 = ∇ ̂ L t +1 ( µ t ) , we compute (8) in a closed form and get the following loss for γ t (see Appendix F) optimizing (9)

Adding the ℓ 2 penalty 1 2 λ || γ t -γ 0 t || 2 encoding the starting point γ 0 t , gives us the closed form for γ t

where we assumed that γ t is shared for a subset of K parameters (see paragraph about structure in drift model) indexed by J K = ( j 1 , . . . , j K ) and ¯ x = ( x j 1 , . . . , x j K ) denotes a vector defined on this subset. We also clip parameters γ t to [0 , 1] . The expression (13) gives us the geometric interpretation for γ t . The value of γ t depends on the angle between (¯ µ t -¯ µ 0 ) and ¯ g t +1 When these vectors are aligned, γ t is high and is low otherwise. When these vectors are orthogonal or the gradient ¯ g t +1 ≈ 0 , the value of γ t is heavily influenced by γ 0 t . Moreover, when ¯ g t +1 ≈ 0 , we can interpret it as being close to a local minimum, i.e., stationary, which means that we want γ t ≈ 1 , therefore adding the ℓ 2 penalty is important. Also, when the norm of the gradients ¯ g t +1 is high, the value of γ t is encouraged to decrease, introducing the drift. This means that using γ t in the parameter update (see Section 3.5) encourages the norm of the gradient to stay small. In practice, we found that update (13) was unstable suggesting that linearization of the log-likelihood might not be a good approximation for learning γ t .

3.4 Approximate Bayesian update of posterior q t ( θ ) with BNNs

The optimization problem (6) for the per-parameter Gaussian q ( θ ) = ∏ j N ( θ j ; µ j , σ j 2 ) with a prior q t ( θ ) = ∏ j N ( θ j ; µ j t , σ j t 2 ) can be written (see Appendix E.1) to minimize the following loss

where λ j t > 0 are per-parameter temperature coefficients. The use of a small temperature parameter λ > 0 (shared for all NN parameters) was shown to improve the empirical performance of Bayesian Neural Networks [54]. Given that in (14), the variance ˜ σ j t 2 ( γ t ) can be small, in order to control the strength of the regularization, we propose to use the temperature per parameter λ j t = λ [ σ j t ] 2 , where λ > 0 is a global constant. This leads to the following objective

where the quantity r j t = [ σ j t ] 2 / [ σ j t ( γ t )] 2 is a relative change in the posterior variance due to the drift. The ratio r j t = 1 when γ j t = 1 . For γ j t < 1 since typically σ j t 2 < σ j 0 2 , the ratio is r j t < 1 .

Algorithm 1 Soft-Reset algoritm

Neural Network (NN) initializing distribution p init ( θ ) and specific initialization θ 0 ∼ p init ( θ Learning rate α t for parameters and η γ for drift parameters

Number of gradient updates K γ on drift parameter γ t

NN initial standard deviation (STD) scaling ν ≤ 1 (see (24)) and ratio s defining σ t νσ 0 .

Receive batch of data B t +1

Initialize drift parameters γ t, 0 = 1

for step

k

= 0

,

1

,

2

, . . . , K

Sample

θ

′

0

∼

p

init

(

θ

)

Stochastic update (22) on drift parameter using specific initialization (26)

end for

Get ˜ θ t ( γ t,K ) with (18) and ˜ α t ( γ t,K ) with (19)

Update parameters

+1

end for

Thus, as long as the non stationarity is detected ( γ j t < 1 ), the objective (15) favors the data term E ϵ ∼N (0; I ) [ ̂ L t +1 ( µ + ϵ ◦ σ ) ] allowing the optimization to respond faster to changes in the data distribution. Let µ t +1 , 0 = ˜ µ t ( γ t ) and σ t +1 , 0 = ˜ σ t ( γ t ) , and perform updates K on (15)

where α µ and α σ are learning rates for the mean and for the standard deviation correspondingly. All derivations are provided in Appendix E.1. The full procedure is described in Algorithm 2.

3.5 Fast MAP update of posterior q t ( θ )

As a faster alternative to propagating the posterior (6), we do MAP updates with the prior p 0 ( θ ) = ∏ j N ( θ j ; µ j 0 ; σ j 0 2 ) and the approximate posterior q t ( θ ) = ∏ j N ( θ j ; θ j t ; σ j t 2 = s 2 σ j 0 2 ) , where s ≤ 1 is a hyperparameter controlling the variance of q t ( θ ) . Since a fixed s may not capture the true parameters variance, using a Bayesian method (see Section 3.4) is preferred but comes at a high computational cost (see Appendix I for discussion). The MAP update is given by (see Appendix E.2 for derivations) finding a minimum of the following proximal objective

where the regularization target for the parameter dimension i is given by

and the per-parameter learning rate is given as (assuming that α t the base SGD learning rate)

Linearising ̂ L t +1 ( θ ) around ˜ θ t ( γ t ) and optimizing (17) for θ leads to (see Appendix E.2)

For γ j t = 1 , we recover the ordinary SGD update, while the values γ j t < 1 move the starting point of the modified SGD closer to the initialization as well as increase the learning rate. In Appendix D we describe additional practical choices made for the Soft Resets algorithm. Algorithm 1 describes the full procedure. Similarly to the Bayesian approach (16), we can do multiple updates on (17). This Soft Resets Proximal algorithm is given in Appendix E.2. Algorithm 3 describes the full procedure.

θ

t

= ˜

θ

t

(

γ

t,K

)

-

˜

α

t

(

γ

t,K

γ

do

)

◦ ∇

θ

L

t

̂

+1

( ˜

θ

t

(

γ

t,K

))

)

4 Related Work

Plasticity loss in Neural Networks. Our model shares similarities with reset-based approaches such as Shrink & Perturb (S&P) [2] and L2-Init [33]; however, whereas we learn drift parameters from data, these methods do not, leaving them vulnerable to mismatch between assumed non-stationarity and the actual realized non-stationarity in the data. Continual Backprop [13] or ReDO [47] apply resets in a data-dependent fashion, e.g. either based on utility or whether units are dead. But they use hard resets, and cannot amortize the cost of removing entire features. Interpretation (13) of γ t connects to the notion of parameters utility from [14], but this quantity is used to prevent catastrophic forgetting by decreasing learning rate for high γ t . Our method increases the learning rate for low γ t to maximize adaptability, and is not designed to prevent catastrophic forgetting.

Non-stationarity. Non-stationarity arises naturally in a variety of contexts, the most obvious being continual and reinforcement learning. The structure of non-stationarity may vary from problem to problem. At one extreme, we have a piece-wise stationary setting, for example a change in the location of a camera generating a stream of images, or a hard update to the learner's target network in value-based deep RL algorithms. This setting has been studied extensively due to its propensity to induce catastrophic forgetting [e.g. 31, 45, 51, 10] and plasticity loss [13, 39, 38, 34]. At the other extreme, we can consider more gradual changes, for example due to improvements in the policy of an RL agent [40, 46, 42, 13] or shifts in the data generating process [36, 55, 20, 53]. Further, these scenarios might be combined, for example in continual reinforcement learning [31, 1, 13] where the reward function or transition dynamics could change over time.

Non-stationary online convex optimization. Non-stationary prediction has a long history in online convex optimization, where several algorithms have been developed to adapt to changing data [see, e.g., 25, 8, 22, 17, 21, 18, 29]. Our approach takes an inspiration from these works by employing a drift model as, e.g., [25, 21] and by changing learning rate as [29, 52]. Further, our OU drift model bears many similarities to the implicit drift model introduced in the update rule of [25] (see also [8, 17]), where the predictive distribution is mixed with a uniform distribution to ensure the prediction could change quickly enough if the data changes significantly, where in our case p 0 plays the same role as the uniform distribution.

Bayesian approaches to non-stationary learning. A standard approach is Variational Continual Learning [41], which focuses on preventing catastrophic forgetting and is an online version of 'Bayes By Backprop' [5]. This method does not incorporate dynamical parameter drift components. In [35], the authors applied variational inference (VI) on non-stationary data, using the OU-process and Bayesian forgetting, but unlike in our approach, their drift parameter is not learned. Further, in [49], the authors considered an OU parameter drift model similar to ours, with an adaptable drift scalar γ and analytic Kalman filter updates, but is applied over the final layer weights only, while the remaining weights of the network were estimated by online SGD. In [28], the authors propose to deal with non-stationarity by assuming that each parameter is a finite sum of random variables following different OU process. They derive VI updates on the posterior of these variables. Compared to this work, we learn drift parameters for every NN parameter rather than assuming a finite set of drift parameters. A different line of research assumes that the drift model is known and use different techniques to estimate the hidden state (the parameters) from the data: in [9], the authors use Extended Kalman Filter to estimate state and in [3], they propagate the MAP estimate of the hidden state distribution with K gradient updates on a proximal objective similar to (47), whereas in Bayesian Online Natural Gradient (BONG) [27], the authors use natural gradient for the variational parameters.

5 Experiments

Soft reset methods. There are multiple variations of our method. We call the method implemented by Algorithm 1 with 1 gradient update on the drift parameter Soft Reset , while other versions show different parameter choices: Soft Reset ( K γ = 10 ) is a version with 10 updates on the drift parameter, while Soft Reset ( K γ = 10 , K θ = 10 ) is the method of Algorithm 3 in Appendix E.2 with 10 updates on drift parameter, followed by 10 updates on NN parameters. Bayesian Soft Reset ( K γ = 10 , K θ = 10 ) is a method implemented by Algorithm 2 with 10 updates on drift parameter followed by 10 updates on the mean µ t and the variance σ 2 t (uncertainty) for each NN parameter. Bayesian method performed the best overall but required higher computational complexity (see Appendix I). Unless specified, γ t is shared for all the parameters in each layer (separately for weight and biases).

0

50

100

150

200

Task id

0.65

0.70

0.75

0.80

0.85

Average Task Accuracy

Permuted MNIST

0

50

100

150

200

Task id

0.2

0.4

0.6

0.8

Average Task Accuracy

Random Label MNIST -- Data Efficient

Online SGD

L2 init

Hard Reset

Shrink and Perturb

K

=10,

K

=10,

per parameter

0

50

100

150

200

Task id

0.2

0.4

0.6

0.8

1.0

Average Task Accuracy

Random Label CIFAR-10 -- Memorization

0

50

100

150

200

Task id

0.65

0.70

0.75

0.80

0.85

Average Task Accuracy

Permuted MNIST

0

50

100

150

200

Task id

0.2

0.4

0.6

0.8

Average Task Accuracy

Random Label MNIST -- Data Efficient

K

=10

Soft Reset Proximal

K

=10,

K

=10

K

=10,

K

=10

K

=10,

K

=10,

per parameter

0

50

100

150

200

Task id

0.2

0.4

0.6

0.8

1.0

Average Task Accuracy

Random Label CIFAR-10 -- Memorization

Loss of plasticity. We analyze the performance of our method on plasticity benchmarks [34, 39, 38]. Here, we have a sequence of tasks, where each task consists of a fixed (for all tasks) subset of 10000 images images from either CIFAR-10 [32] or MNIST, where either pixels are permuted or the label for each image is randomly chosen. Several papers [34, 39, 38] study a memorization random-label setting where SGD can perfectly learn each task from scratch. To highlight the data-efficiency of our approach, we study the data-efficient setting where SGD achieves only 50% accuracy on each task when trained from scratch. Here, we expect that algorithms taking into account similarity in the data, to perform better. To study the impact of the non-stationarity of the input data, we consider permuted MNIST where pixels are randomly permuted within each task (the same task as considered by 34). As baselines, we use Online SGD and Hard Reset at task boundaries. We also consider L2 init [34], which adds L 2 penalty || θ -θ 0 || 2 to the fixed initialization θ 0 as well as Shrink&Perturb [2], which multiplies each parameter by a scalar λ ≤ 1 and adds random Gaussian noise with fixed variance σ . See Appendix H.1 for all details. As metrics, we use average per-task online accuracy (52), which is

where a t i are the online accuracies collected on the task t via N timesteps, corresponding to the duration of the task. In Figure 5, we also use average accuracy over all T tasks, i.e.

The results are provided in Figure 2. We observe that Soft Reset is always better than Hard Reset and most baselines despite the lack of knowledge of task boundaries. The gap is larger in the data efficient regime. Moreover, we see that L2 Init only performs well in the memorization regime, and achieves comparable performance to Hard Reset in the data efficient one. The method L2 Init could be viewed as an instantiation of our Soft Reset Proximal method optimizing (17) with γ t = 0 at every step, which is sub-optimal when there is similarity in the data. Bayesian Soft Reset demonstrates significantly better performance overall, see also discussion below.

In Figure 3, we compare different variants of Soft Reset . We observe that adding more compute for estimating γ t (thus, estimating non-stationarity, K γ = 10 ) as well as doing more updates on NN

1

2

3

4

5

MLP Layer

0.65

0.70

0.75

0.80

0.85

0.90

Minimum

(a) Minimum

t

encountered

MNIST

CIFAR-10

0

20000 40000 60000 80000100000

t

0.5

0.6

0.7

0.8

0.9

1.0

(b) First layer

t

for Random Label MNIST

0

100000

200000

300000

400000

500000

600000

t

0.5

0.6

0.7

0.8

0.9

1.0

(c) First layer

t

for Random Label CIFAR-10

parameters (thus, more accurately adapting to non-staionarity, K θ = 10 ) leads to better performance. All variants of Soft Reset γ t parameters are shared for each NN layer, except for the Bayesian method. This variant is able to take advantage of a more complex per-parameter drift model, while other variants performed considerably worse, see Appendix K.4. We hypothesize this is due to the NN parameters uncertainty estimates σ t which Bayesian method provide, while others do not, which leads to a more accurate drift model estimation, since uncertainty is used in this update (10). But, this approach comes at a higher computational cost, see Appendix I. In Appendix K, we provide ablations of the structure of the drift model, as well as of the impact of learning the drift parameter.

Qualitative behavior of Soft Resets . For Soft Reset , we track the values of γ t for the first MLP layer when trained on random-label tasks studied above (only 20 tasks), as well as the minimum encountered value of γ t for each layer, which highlights the maximum amount of resets. Figure 4b,c shows γ t as a function of t , and suggests that γ t aggressively decreases at task boundaries (red dashed lines). The range of values of γ t depends on the task and on the layer, see Figure 4a. Overall, γ t changes more aggressively for long duration (memorization) random-label CIFAR-10 and less for shorter (data-efficient) random-label MNIST. See Appendix K.2 for more detailed results.

To study the behavior of Soft Reset under input distribution non-stationarity, we consider a variant of Permuted MNIST where each image is partitioned into patches of a given size. The non-stationarity is controlled by permuting the patches (not pixels). Figure 5a shows the minimum encountered γ t for each layer for different patch sizes. As the patch size increases and the problem becomes more stationary, the range of values for γ t is less aggressive. See Appendix K.3 for more detailed results.

Impact of non-stationarity. We consider a variant of random-label MNIST where for each task, an image has either a random or a true label. The label assignment is kept fixed throughout the task and is changed at task boundaries. We consider cases of 20% , 40% and 60% of random labels and we control the duration of each task (number of epochs). In total, the stream contains 200 tasks. In Figure 5b, we show performance of Online SGD , Hard Reset and in Figure 5c, the one of Soft Reset and of Bayesian Soft Reset . See Appendix H.2 for more details. The results suggest that for the shortest duration of the tasks, the performance of all the methods is similar. As we increase the duration of each of the task (moving along the x-axis), we see that both Soft Resets variants perform better than SGD and the gap widens as the duration increases. This implies that Soft Resets is more effective with infrequent data distribution changes. We also observe that Bayesian method performs better in all the cases, highlighting the importance of estimating uncertainty for NN parameters.

5.1 Reinforcement learning

Reinforcement learning experiments. We conduct Reinforcement Learning (RL) experiments in the highly off-policy regime, similarly to [43], since in this setting loss of plasticity was observed. We ran SAC [19] agent with default parameters from Brax [15] on the Hopper -v5 and Humanoid -v4 GYM [6] environments (from Brax [15]). To reproduce the setting from [43], we control the off-policyness of the agent by setting the off-policy ratio M such that for every 128 environment steps, we do 128 M gradient steps with batch size of 256 on the replay buffer. As baselines we consider ordinary SAC , hard-coded Hard Reset where we reset all the parameters K = 5 times throughout training (every 200000 steps), while keeping the replay buffer fixed (similarly to [43]). We employ our Soft Reset method as follows. After we have collected fresh data from the environment, we do one gradient update on γ t (shared for all the parameters within each layer) with batch size of 128 on this new chunk of data and the previously collected one, i.e., two chunks of data in total. Then we initialize ˜ θ t ( γ t ) and we employ the update rule (47) where the regularization ˜ θ t ( γ t ) is kept constant for all the off-policy gradient updates on the replay buffer. See Appendix H.3 for more details.

conv2_d_1

conv2_d_2

conv2_d_3

conv2_d_4

linear

linear_1

0.825

0.850

0.875

0.900

0.925

0.950

0.975

Minimum

t

encountered

(a) Permuted Patch MNIST

Permuted patch size

1

2

4

7

14

30

100

400

Number of epochs per task

0.2

0.4

0.6

0.8

1.0

Average accuracy over all tasks

(b) Online SGD and Hard Reset performance

Methods

Online SGD

Soft Resets

K

=10,

K

=10,

per parameter)

Hard Reset

Random labels %

20

40

60

30

100

400

Number of epochs per task

0.2

0.4

0.6

0.8

1.0

Average accuracy over all tasks

(c) Soft Resets (Bayeisan / non-Bayesian) performance

Methods

Online SGD

Soft Resets

K

=10,

K

=10,

per parameter)

Hard Reset

0

2000

4000

6000

0.0

0.2

0.4

0.6

0.8

1.0

1e6

0

2000

4000

6000

Baseline

Hard Reset

0.0

0.2

0.4

0.6

0.8

1.0

1e6

0.0

0.2

0.4

0.6

0.8

1.0

1e6

The results are given in Figure 6. As the off-policy ratio increases, Soft Reset becomes more efficient than the baselines. This is consistent with our finding in Figure 5b,c, where we showed that the performance of Soft Reset is better when the data distribution is not changing fast. Figure 8 in Appendix H.3 shows the value of learned γ t . It shows γ t mostly change for the value function and not for the policy indicating that the main source of non-stationarity comes from the value function.

6 Conclusion

Learning efficiently on non-stationary distributions is critical to a number of applications of deep neural networks, most prominently in reinforcement learning. In this paper, we have proposed a new method, Soft Resets , which improves the robustness of stochastic gradient descent to nonstationarities in the data-generating distribution by modeling the drift in Neural Network (NN) parameters. The proposed drift model implements soft reset mechanism where the amount of reset is controlled by the drift parameter γ t . We showed that we could learn this drift parameter from the data and therefore we could learn when and how far to reset each Neural Network parameter. We incorporate the drift model in the learning algorithm which improves learning in scenarios with plasticity loss. The variant of our method which models uncertainty in the parameters achieves the best performance on plasticity benchmarks so far, highlighting the promise of the Bayesian approach. Furthermore, we found that our approach is particularly effective either on data distributions with a lot of similarity or on slowly changing distributions. Our findings open the door to a variety of exciting directions for future work, such as investigating the connection to continual learning and deepening our theoretical analysis of the proposed approach.

References

Zaheer Abbas, Rosie Zhao, Joseph Modayil, Adam White, and Marlos C. Machado. Loss of plasticity in continual deep reinforcement learning, 2023.

Jordan T. Ash and Ryan P. Adams. On warm-starting neural network training, 2020.

Gianluca M. Bencomo, Jake C. Snell, and Thomas L. Griffiths. Implicit maximum a posteriori filtering via adaptive optimization, 2023.

Tudor Berariu, Wojciech Czarnecki, Soham De, Jorg Bornschein, Samuel Smith, Razvan Pascanu, and Claudia Clopath. A study on the plasticity of neural networks. arXiv preprint arXiv:2106.00042 , 2021.

Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In Francis Bach and David Blei, editors, Proceedings of the 32nd International Conference on Machine Learning , volume 37 of Proceedings of Machine Learning Research , pages 1613-1622, Lille, France, 07-09 Jul 2015. PMLR.

Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.

Tamara Broderick, Nicholas Boyd, Andre Wibisono, Ashia C Wilson, and Michael I Jordan. Streaming variational bayes. In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Weinberger, editors, Advances in Neural Information Processing Systems , volume 26. Curran Associates, Inc., 2013.

N. Cesa-Bianchi and G. Lugosi. Prediction, Learning, and Games . Cambridge University Press, Cambridge, 2006.

Peter G. Chang, Gerardo Durán-Martín, Alexander Y Shestopaloff, Matt Jones, and Kevin Murphy. Low-rank extended kalman filtering for online learning of neural networks from streaming data, 2023.

Zhiyuan Chen and Bing Liu. Lifelong machine learning. Synthesis Lectures on Artificial Intelligence and Machine Learning , 12(3):1-207, 2018.

Hugh Dance and Brooks Paige. Fast and scalable spike and slab variable selection in highdimensional gaussian processes, 2022.

Shibhansh Dohare, J. Fernando Hernandez-Garcia, Qingfeng Lan, Parash Rahman, Ashique Rupam Mahmood, and Richard S. Sutton. Loss of plasticity in deep continual learning. Nature , 632:768 - 774, 2024.

Shibhansh Dohare, Richard S. Sutton, and A. Rupam Mahmood. Continual backprop: Stochastic gradient descent with persistent randomness, 2022.

Mohamed Elsayed and A. Rupam Mahmood. Addressing loss of plasticity and catastrophic forgetting in continual learning, 2024.

C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax - a differentiable physics engine for large scale rigid body simulation, 2021. http://github.com/google/brax .

Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Yee Whye Teh and Mike Titterington, editors, Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics , volume 9 of Proceedings of Machine Learning Research , pages 249-256, Chia Laguna Resort, Sardinia, Italy, 13-15 May 2010. PMLR.

A. György, T. Linder, and G. Lugosi. Efficient tracking of large classes of experts. IEEE Transactions on Information Theory , IT-58(11):6709-6725, Nov. 2012.

András György and Csaba Szepesvári. Shifting regret, mirror descent, and matrices. In Proceedings of The 33rd International Conference on Machine Learning , pages 2943-2951, 2016.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor, 2018.

Raia Hadsell, Dushyant Rao, Andrei A. Rusu, and Razvan Pascanu. Embracing change: Continual learning in deep neural networks. Trends in Cognitive Sciences , 24(12):1028-1040, 2020.

Eric Hall and Rebecca Willett. Dynamical models and tracking regret in online convex programming. In International Conference on Machine Learning , pages 579-587. PMLR, 2013.

E. Hazan and C. Seshadhri. Efficient learning algorithms for changing environments. In Proc. 26th Annual International Conference on Machine Learning , pages 393-400. ACM, 2009.

Elad Hazan. Introduction to online convex optimization, 2023.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification, 2015.

M. Herbster and M. K. Warmuth. Tracking the best expert. Machine Learning , 32(2):151-178, 1998.

Hemant Ishwaran and J. Sunil Rao. Spike and slab variable selection: Frequentist and bayesian strategies. The Annals of Statistics , 33(2), April 2005.

Matt Jones, Peter Chang, and Kevin Murphy. Bayesian online natural gradient (bong), 2024.

Matt Jones, Tyler R. Scott, and Michael Curtis Mozer. Human-like learning in temporally structured environments. In AAAI Spring Symposia , 2024.

Mikhail Khodak, Maria-Florina Balcan, and Ameet Talwalkar. Adaptive gradient-based metalearning methods, 2019.

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.

James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A. Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, Demis Hassabis, Claudia Clopath, Dharshan Kumaran, and Raia Hadsell. Overcoming catastrophic forgetting in neural networks. Proceedings of the National Academy of Sciences , 114(13):3521-3526, March 2017.

Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009. https://www.cs.toronto.edu/~kriz/ learning-features-2009-TR.pdf .

Aviral Kumar, Rishabh Agarwal, Dibya Ghosh, and Sergey Levine. Implicit underparameterization inhibits data-efficient deep reinforcement learning, 2021.

Saurabh Kumar, Henrik Marklund, and Benjamin Van Roy. Maintaining plasticity in continual learning via regenerative regularization, 2023.

Richard Kurle, Botond Cseke, Alexej Klushyn, Patrick van der Smagt, and Stephan Günnemann. Continual learning with bayesian neural networks for non-stationary data. In International Conference on Learning Representations , 2020.

Zhiqiu Lin, Jia Shi, Deepak Pathak, and Deva Ramanan. The clear benchmark: Continual learning on real-world imagery. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track , 2021.

Clare Lyle, Mark Rowland, and Will Dabney. Understanding and preventing capacity loss in reinforcement learning, 2022.

Clare Lyle, Zeyu Zheng, Khimya Khetarpal, Hado van Hasselt, Razvan Pascanu, James Martens, and Will Dabney. Disentangling the causes of plasticity loss in neural networks, 2024.

Clare Lyle, Zeyu Zheng, Evgenii Nikishin, Bernardo Avila Pires, Razvan Pascanu, and Will Dabney. Understanding plasticity in neural networks, 2023.

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning, 2013.

Cuong V. Nguyen, Yingzhen Li, Thang D. Bui, and Richard E. Turner. Variational continual learning. In International Conference on Learning Representations , 2018.

Evgenii Nikishin, Junhyuk Oh, Georg Ostrovski, Clare Lyle, Razvan Pascanu, Will Dabney, and André Barreto. Deep reinforcement learning with plasticity injection, 2023.

Evgenii Nikishin, Max Schwarzer, Pierluca D'Oro, Pierre-Luc Bacon, and Aaron Courville. The primacy bias in deep reinforcement learning, 2022.

Neal Parikh and Stephen Boyd. Proximal algorithms. Found. Trends Optim. , 1(3):127-239, jan 2014.

German I. Parisi, Ronald Kemker, Jose L. Part, Christopher Kanan, and Stefan Wermter. Continual lifelong learning with neural networks: A review. Neural Networks , 113:54-71, 2019.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.

Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evci. The dormant neuron phenomenon in deep reinforcement learning. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 32145-32168. PMLR, 23-29 Jul 2023.

Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evci. The dormant neuron phenomenon in deep reinforcement learning, 2023.

Michalis K. Titsias, Alexandre Galashov, Amal Rannen-Triki, Razvan Pascanu, Yee Whye Teh, and Jorg Bornschein. Kalman filter for online classification of non-stationary data, 2023.

G. E. Uhlenbeck and L. S. Ornstein. On the theory of the brownian motion. Phys. Rev. , 36:823-841, Sep 1930.

Gido M van de Ven and Andreas S Tolias. Three scenarios for continual learning. arXiv preprint arXiv:1904.07734 , 2019.

Tim van Erven, Wouter M. Koolen, and Dirk van der Hoeven. Metagrad: Adaptation using multiple learning rates in online learning. Journal of Machine Learning Research , 22(161):1-61, 2021.

Eli Verwimp, Rahaf Aljundi, Shai Ben-David, Matthias Bethge, Andrea Cossu, Alexander Gepperth, Tyler L. Hayes, Eyke Hüllermeier, Christopher Kanan, Dhireesha Kudithipudi, Christoph H. Lampert, Martin Mundt, Razvan Pascanu, Adrian Popescu, Andreas S. Tolias, Joost van de Weijer, Bing Liu, Vincenzo Lomonaco, Tinne Tuytelaars, and Gido M. van de Ven. Continual learning: Applications and the road forward, 2024.

Florian Wenzel, Kevin Roth, Bastiaan S. Veeling, Jakub ´ Swi ˛ atkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good is the bayes posterior in deep neural networks really?, 2020.

Runtian Zhai, Stefan Schroedl, Aram Galstyan, Anoop Kumar, Greg Ver Steeg, and Pradeep Natarajan. Online continual learning for progressive distribution shift (OCL-PDS): A practitioner's perspective, 2023.

Martin Zinkevich. Online convex programming and generalized infinitesimal gradient ascent. In Proceedings of the Twentieth International Conference on International Conference on Machine Learning , ICML'03, page 928-935. AAAI Press, 2003.

A Organization of the appendix

The appendix is organized as follows. In Appendix B we describe the Ornstein-Uhlenbeck process and connect it to the drift model (5), we use in the paper. In Appendix C, we discuss other potential choices for a drift model. In Appendix D, we describe additional practical tricks used for learning the drift model parameters. In Appendix E, we describe full Bayesian and maximum a-posteriori (MAP) algorithms to learn both drift parameters and NN parameters. In Appendix F, we provide the proof for the closed form solution derivation (12). In Appendix G, we briefly review the proximal SGD. In Appendix H, we provide experimental details. In Appendix I, we describe computational complexity for all the methods. In Appendix J, we provide sensitivity analysis for different parameters. In Appendix K, we provide additional ablations and experiments showing qualitative behavior of soft resets. In Appendix L, we show a toy example for the failure of SGD in the non-stationary regime. Finally, in Appendix M, we show how the proximal objective can be generalized to use any Gaussian drift model.

B Ornstein-Uhlenbeck process

We make use that the Ornstein-Uhlenbeck process [50] defines a SDE that can be solved explicitly and written as a time-continuous Gaussian Markov process with transition density between on R D

for any pair of times t > s . Based on this as a drift model for the parameters θ t (so θ t is the state x t ) we use the conditional density

where γ j t = e -δ j t and δ j t ≥ 0 corresponds to the learnable discretization time step. In other words, by learning γ j t online we equivalently learn the amount of a continuous 'time shift' δ j t between two consecutive states in the OU process. This essentially models parameter drift since e.g. if γ j t = 1 , then δ j t = 0 and there is no 'time shift' which means that the next state/parameter remains the same as the previous one, i.e. θ j t +1 = θ j t .

C Other choices of drift model

In this section, we discuss alternative choices of a drift model instead of (5).

Independent mean and variance of the drift. We consider the drift model where the mean and the variance are not connected, i.e.,

where parameters γ j t ∈ [0 , 1] control the mean of the distribution and β j t > 0 control the variance. When parameters β j t are fixed to a constant, this would be similar to our experiment in Figure 16 where we assume known task boundaries and we do not estimate the drift parameters but assume it as a hyperparameter. Figure 16, left corresponds to the case when β j t are fixed parameters independent from γ j t whereas Figure 16, right corresponds to the case when β j t = √ 1 -γ j t 2 σ j 0 , i.e., when we use the drift model (5). We see from the results, using drift model (5) leads to a better performance. In case when β j t are learned, estimating the parameters of this model will likely overfit to the noise since there is a lot of degrees of freedom.

Shrink & Perturb [2]. When we do not use the mean of the initialization, we can use the following drift model

for λ j t ∈ [0 , 1] and β j t > 0 . Similarly to the case of (21), estimating both parameters λ j t and β j t from the data will likely overfit to the noise.

Arbitrary linear model. We can use the arbitrary linear model of the form

but estimating the parameters A t and B t has too many degrees of freedom and will certainly overfit.

Gaussian Spike & Slab We consider a Gaussian [11] approximation to Spike & Slab [26] prior

with binary variables γ j t ∈ { 0 , 1 } , which is a product of mixtures of Gaussian distributions p ( θ j t +1 | θ j t ) = N ( θ j t +1 ; θ j t , σ j 2 ) centred around the previous parameter θ j t and initializing distributions p 0 ( θ j t +1 ) = N ( θ j t +1 ; µ j 0 , σ j 0 2 ) . This model, however, implements the mechanism of Hard reset as opposed to the soft ones. Moreover, estimating such a model and incorporating it into a learning update is more challenging since the mixture of Gaussian is not conjugate with respect to a Gaussian which will make the KL term (34) to be computed only approximately via Monte Carlo updates.

D Practical implementations of the drift model estimation

Stochastic approximation for drift parameters estimation In practice, we use M = 1 , which leads to the stochastic approximation

Using NN initializing distribution. In the drift model (5), we assume that the initial distribution over parameters is given by p 0 ( θ ) = ∏ j N ( θ j ; µ j 0 ; σ j 0 2 ) . In practice, we have access to the NN initializer p init ( θ ) = ∏ j N ( θ j ; 0; σ j 0 2 ) where µ j 0 = 0 (for most of the NNs). This means that we can replace ϵ from (10) by 1 σ 0 θ ′ 0 where θ ′ 0 ∼ p init ( θ ) . This means that the term in (22) can be replaced by

where we used the fact that σ 2 t ( γ t ) = γ 2 t σ 2 t +(1 -γ 2 t ) σ 2 0 . Note that in (23), we only need to know the ratio σ 2 σ 2 0 rather than both of these. We will see that in Section 3.5, only this ratio is used for the underlying algorithm. Finally, in practice, we can tie p 0 ( θ ) to the specific initialization θ 0 ∼ p init ( θ ) . It was observed empricially [34] that using a specific initialization in gradient updates led to better performance than using samples from the initial distribution. This would imply that

with ˜ σ j 0 2 = ν 2 σ j 0 2 . The parameter ν ≤ 1 accounts for the fact that the distribution p 0 ( θ ) should have lower than p init ( θ ) variance since it uses the specific initializaiton from the prior. This modification would imply the following modification on the drift model term (23)

If we denote s = σ t νσ 0 , then we can re-write (25) to be

In practice, we can treat s to be a constant hyperparameter.

E Learning parameters with estimated drift models

In this section, we provide a Bayesian Neural Network algorithm to learn the distributions of NN parameters when there is a drift in the data distribution. Moreover, we provide a Maximum aPosteriori (MAP) like inference algorithm which does not require to learn the distributions over parameters, but simply propagates the MAP estimate over these.

E.1 Bayesian Neural Networks algorithm

In this section, we describe an algorithm for parameters update based on Bayesian Neural Networks (BNN). It is based on the online variational Bayes setting described below.

Let the family of distributions over parameters be

which is the family of Gaussian mean-field distributions over parameters θ ∈ R D . Let Γ t = ( γ 1 , . . . , γ t ) be the history of observed parameters of the drift model and S t = {B 1 , . . . , B t } be the history of observed data. We denote by q t ( θ ) ≜ q t ( θ |S t , Γ t -1 ) ∈ Q the Gaussian approximate posterior at time t with mean µ t and variance σ 2 t for every parameter. The approximate predictive look-ahead prior is given by

that has parameters

To see this, we will use the law of total expectation and the law total variance. For two random variables X and Y defined on the same space, law of total expectation says

and the law of total variance says

In our case, from the drift model (5), we have the conditional distribution

From (30), we have and the variance is given by

From here, we have that the mean is given by

Now, we note that q t ( θ | γ t ) is a Gaussian and its parameters are given by E [ θ j ] = γ j t µ j t +(1 -γ j t ) µ j 0 from (31) and by V [ θ j ] = (1 -γ j t 2 ) σ j 0 2 + γ j t 2 θ j t 2 from (32). Then, for new data B t +1 at time t +1 , the approximate predictive log-likelihood equals to

We are looking for a new approximate posterior q t +1 ( θ ) such that

The optimization problem (33) can be written as minimization of the following loss

Using the fact that we are looking for a member q ∈ Q from (27), we can write the objective (34) as

where we used the reparameterisation trick for the loss term. We now expand the regularization term to get

Since the posterior variance of NN parameters may become small, the optimization of (35) may become numerically unstable due to division by ˜ σ j t 2 ( γ t ) . It was shown [54] that using small temperature on the prior led to better empirical results when using Bayesian Neural Networks, a phenomenon known as cold posterior. Here, we define a temperature per-parameter, i.e., λ j t > 0 for every time-step t , such that the objective above becomes

As said above, it is common to use the same temperature λ j t = λ for all the parameters. In this work, we propose the specific choice of the temperature to be

where λ > 0 is some globally chosen temperature parameter. This leads to the following objective

where the quantity r j t is defined as

which represents the relative change in the posterior variance due to the drift. In the exact stationary case, when γ j t = 1 , this ratio is r j t = 1 while for γ j t < 1 , since typically σ j t 2 < σ j 0 2 , we have r j t < 1 . This means that in the non-stationary case, the strength of the regularization in (38) in favor of the data term E ϵ ∼N (0; I ) [ L t +1 ( µ + ϵ ◦ σ )] , allowing the optimization to respond faster to the change in the data distribution. In practice, this data term is approximated via Monte-Carlo, i.e.

To find new parameters, µ t +1 and σ t +1 , we let µ t +1 , 0 = ˜ µ t ( γ t ) and σ t +1 , 0 = ˜ σ t ( γ t ) where corresponding quantities are defined in (29) and perform multiple K updates on (38)

where α µ and α σ are corresponding learning rates.

Practical considerations. As a prior distribution, we use

with some σ j 0 > 0 . In practice, we often have access to the NN library with pre-defined ˜ σ j 0 . We allow more flexbility and define

with some constant ν ∈ (0 , 1] , allowing to scale the NN library variance.

When using Bayesian Neural Networks (BNN), one question is to how to initialize the variational posterior q 0 ( θ ) . Wefound that using specific initialization θ 0 ∼ p init ( θ ) as the mean of the variational posterior q 0 ( θ ) , where p init ( θ ) is the NN initializing distribution, worked significantly better than using 0 (which is the mean of p 0 ( θ ) ). Therefore, we propose to initialize q 0 ( θ ) as where θ 0 ∼ p init ( θ ) and σ j 2 is defined as

and σ j 0 is the variance of the prior distribution p 0 ( θ ) , see (41). Here, π ∈ (0 , 1] is a constant allowing us to scale down the initial variance of q 0 ( θ ) with respect to the prior variance σ j 0 .

The full algorithm of learning the drift parameters γ t as well as learning the Bayesian Neural Network parameters using the procedure above is given in Algorithm 2.

E.2 Modified SGD with drift model

Instead of propagating the posterior (6), we do MAP updates on (4) with the prior p 0 ( θ ) = ∏ j N ( θ j ; µ j 0 ; σ j 0 2 ) and the posterior q t ( θ ) = ∏ j N ( θ j ; θ j t ; s 2 σ j 0 2 ) , where s ≤ 1 is hyperparameter controlling the variance σ 2 t of the posterior q t ( θ ) . Since fixed s may not capture the true parameters variance, using Bayesian method (see Appendix E.1) is preferred but comes at a high computational cost. Instead of Bayesian update (33), we consider maximum a-posteriori (MAP) update

with q t ( θ | γ t ) given by (28). Using the definition of q t ( θ | γ t ) , we get the following problem

where similarly to (36), we use a per-parameter temperature λ j t ≥ 0 . We choose temperature to be equal to

where λ is some constant. Such choice of temperature is motivated by the same logic as in (37) - it is a constant multiplied by the posterior variance σ j t 2 = s 2 σ j 0 2 . With such choice of temperature, maximizing (44) is equivalent to minimizing

where the regularization target for the dimension i is

We can perform K gradient updates (45) with a learning rate α t starting from θ 0 t +1 = ˜ θ t ( γ t ) ,

where the vector-valued learning rate α t ( γ t ) is given by

with α t the base learning rate. Note that doing one update is equivalent to modified SGD method (20). Doing multiple updates on (47) allows us to perform multiple computations on the same data. The corresponding algorithm is given in Algorithm 3.

and the constant r j t ( γ ) is given by

Algorithm 2 Bayesian Soft-Reset algorithm

Global temperature parameter λ ≥ 0 .

Learning rate for the mean α µ and for the standard deviation α σ

Number of gradient updates K θ to be applied on µ and σ

Number of Monte-Carlo samples M θ for estimating µ and σ in (40)

Number of gradient updates K γ on drift parameter γ t in (10)

Number of Monte-Carlo samples M γ to estimate γ t in (11)

Learning rate η γ for drift parameters

Initial drift parameters γ 0 = 1 for every iteration.

Initialization parameters :

Neural Network initial variance for every parameter ˜ σ j 0 2 coming from standard NN library Initial prior variance rescaling ν ∈ (0 , 1] .

Prior variance σ j 0 2 = ν 2 ˜ σ j 0 2 given by (42)

Initial NN parameters θ 0 ∼ p init ( θ )

Initial posterior variance rescaling π ∈ (0 , 1] .

Initial posterior variance σ j init 2 = π 2 σ j 0 2 given by (43)

Initialize zero-mean prior distribution (41) as p 0 ( θ ) = ∏ j N ( θ j ; 0 , σ j 0 2 )

Initialization :

Receive batch B t +1

Current posterior is given by q t = N ( θ ; µ t , σ 2 t )

Estimating the drift

Initialize drift parameter γ 0 t = γ 0 .

Compute drifted prior mean ˜ µ j t ( γ t ) = γ j t µ j t assuming zero-mean prior µ j 0 = 0

end for

Updating variational posterior

Using the definition of ˆ F t from (38), do

end for end for

F Proof of linearisation

Interpretation of γ t . By linearising log p ( B t +1 | θ ) around µ t , we can simplify (8) to get

where ◦ denotes elementwise product, g t = ∇ ̂ L t +1 ( µ t ) is the negative gradient of the loss (1) evaluated at µ t and we added the ℓ 2 -penalty 1 2 λ || γ t -γ t, 0 || 2 to take into account the initialization.

Proof . We assume that the following linearisation is correct

where

Algorithm 3 Proximal Soft-Reset algoritm

Neural Network (NN) initializing distribution p init ( θ ) and specific initialization θ 0 ∼ p init ( θ )

Learning rate α t for parameters and η γ for drift parameters

Number of gradient updates K γ on drift parameter γ t

Number of gradient updates K θ on NN parameters

Global temperature parameter λ ≥ 0

NN initial standard deviation (STD) scaling

Receive batch B t +1

Initialize drift parameter

t

γ

for step

k

= 0

,

1

,

2

Sample

′

0

0

, . . . , K

init

θ

p

(

θ

)

∼

Stochastic update (22) on drift parameter using specific initialization (25)

end for

Initialize

t

+1

t K

t

t

with (46) and use

θ

=

θ

(

γ

)

α

0

end for end for

Then, we have

(

i

t

2

)

γ

ν

1

= 1

do

(

γ

(see (24)) and ratio t K

) =

α

t

(

γ

i

t

)

2

Let's write the integral from (8)

Consider only the exp term inside the integral:

Let's focus on this term

+

s

1

=

-

(

γ

s

2

σ

t

νσ

0

)

.

with (48)

where

Therefore, the integral could be written as

Now, we only keep the terms depending on γ t

Now, we add an l2-penalty λ 2 || γ t -γ t, 0 || 2 and we get

Let's take the gradient wrt γ j t , we get that the j -the component of the gradient is given by

We have then,

In case when γ t are shared for a subset of K parameters, indexes by the index set J K = ( j 1 , . . . , j K ) , then, we can sum the gradients (49) and we get

this gives us

where ¯ a is vector defined on the index set J K , i.e., ¯ a = ( a j 1 , . . . , a j K ) .

G Proximal SGD

Each step of online SGD can be seen in terms of a regularized minimization problem referred to as the proximal form [44]:

In general, we cannot solve (50) directly, so we consider a Taylor expansion of ̂ L t +1 around θ t , giving

Here we see the role of α t > 0 as both enforcing that the Taylor expansion around θ t is accurate, and regularising θ t +1 towards the old parameters θ t (hence ensuring that the learning from past data is not forgotten). Solving (51) naturally leads to the well known SGD update:

where α t can now also be interpreted as the learning rate.

H Experimental details

H.1 Plasticity experiments

Tasks In this section we provide experimental details. As plasticity tasks, we use a randomly selected subset of size 10000 from CIFAR-10 [32] and from MNIST. This subset is fixed for all the tasks. Within each task, we randomly permute labels for every image; we call such problems random-label classification problems. We study two regimes data efficient , where we do 400 epochs on a task with a batch size of 128 , and memorization , a regime where we do only 70 epochs with a batch size of 128 . As the main backbone architecture, we use MLP with 4 hidden layers each having a hidden dimension of 256 hidden units. We use ReLU activation function and do not use any batch or layer normalization. For the incoming data, we apply random crop, for MNIST to produce images of size 24 × 24 and for CIFAR-10 to produce images of size 28 × 28 . We normalize images to be within [0 , 1] range by dividing by 255 . On top of that, we consider permuted MNIST task with a similar training ragime as in [34] - we consider a subset of 10000 images, with batch size 16 and each task is one epoch. As a backbone, we still use MLP with ReLU activation and 4 hidden layers. Moreover, we considered permuted Patch MNIST , where we permute patches, not individual pixels. In this case, we used a simple 4 layer convolutional neural network with 2 fully connected layers at the end.

Metrics We use online accuracy as first metric with results reported in Appendix K. Moreover we use per-task Average Online Accuracy which is

where a t i are the online accuracies collected on the task t via N timesteps.

Baselines First baseline is Online SGD which sequentially learns over the sequence of task, with a fixed learning rate. Hard Reset is the Online SGD which resets all the parameters at task boundaries. L2 init [34] adds a regularizer λ || θ -θ 0 || 2 term to each Online SGD update where the regularization strength λ is a hyperparameter. Shrink & Perturb applies the transformation λθ t + σϵ, ϵ ∼ N ( ϵ ; 0 , I ) to each parameter before the gradient update. The hyperparameters are λ and σ .

Soft Reset corresponds to one update (10) starting from 1 using 1 Monte Carlo estimate. We always use 1 Monte Carlo estimate for updating γ t as we found that it worked well in practice on these tasks. The hyperparameters of the method σ 2 0 initial variance of the prior, which we set to be equal to ν 2 1 N where N is the width of the hidden layer and ν is a constant (hyperparameter). It always equals to ν = 0 . 1 . On top of that the second hyperparameter is s , such that σ t = sσ 0 , which controls the relative decrease of the constant posterior variance. This is the hyperparameter over which we sweep

over. Another hyperparameter is the learning rate for learning γ t . For Soft Reset Proximal , we also have a proximal coefficient regularization constant λ . Besides that, we also sweep over the learning rate for the parameter. For the Bayesian Soft Reset , we just add an additional learning rate for the variance α σ and we do 1 Monte Carlo sample for each ELBO update.

Hyper parameters selection and evaluation For all the experiments, we run a sweep over the hyperparameters. We select the best hyperparameters based on the smallest cumulative error (sum of all 1 -a t i throughout the training). We then report the mean and the standard deviation across 3 seeds in all the plots.

Hyperparameter ranges . Learning rate α which is used to update parameters, for all the methods, is selected from { 1 e -4 , 5 e -4 , 1 e -3 , 5 e -3 , 1 e -2 , 5 e -2 , 1 e -1 , 5 e -1 , 1 . 0 } . The λ init parameter in L2 Init , is selected from { 10 . 0 , 1 . 0 , 0 . 0 , 1 e -1 , 5 e -1 , 1 e -2 , 5 e -2 , 1 e -3 , 5 e -3 , 1 e -4 , 5 e -4 , 1 e -5 , 5 e -5 , 1 e -6 , 5 e -6 , 1 e -7 , 5 e -7 , 1 e -8 , 5 e -8 , 1 e -9 , 5 e -9 , 1 e -10 , } . For S&P, the shrink parameter λ is selected from { 1 . 0 , 0 . 99999 , 0 . 9999 , 0 . 999 , 0 . 99 , 0 . 9 , 0 . 8 , 0 . 7 , 0 . 5 , 0 . 3 , 0 . 2 , 0 . 1 } , and the perturbation parameter σ is from { 1 e -1 , 1 e -2 , 1 e -3 , 1 e -4 , 1 e -5 , 1 e -6 } . As noise distribution, we use the Neural Network initial distribution. For Soft Resets , the learning rate for γ t is selected from { 0 . 5 , 0 . 1 , 0 . 05 , 0 . 01 , 0 . 005 , 0 . 001 , 0 . 0005 , 0 . 0001 } , the constant s is selected from { 1 . 0 , 0 . 95 , 0 . 9 , 0 . 8 , 0 . 7 , 0 . 6 , 0 . 5 , 0 . 3 , 0 . 1 } , the temperature λ in (45) is selected from { 1 . 0 , 0 . 1 , 0 . 01 } , the same is true for the temperature in the Bayesian method (38). Initial prior std rescaling ν = 0 . 05 . On top of that for the Bayesian method, we always use ν (see Algorithm 2) equal to ν = 0 . 05 and π = 0 . 9 , i.e. the posterior is always slightly smaller than the prior. Finally for the Bayesian method we had to learn the variance with learning rate from { 0 . 01 , 0 . 1 , 1 , 10 } range.

In practice, we found that there is one learning rate of 0 . 1 , which was always the best in practice for most of the methods and only proximal Soft Resets on memorization CIFAR-10 required smaller learning rate 0 . 01 . This allowed us to significantly reduce the hyperparameter sweep.

H.2 Impact of non-stationarity experiments

In this experiment, we consider a subset of 10000 images from MNIST (fixed throughtout all experiment) and a sequence of tasks. Each task is constructed by assigning either a true or a random label to each image from MNIST, where the probability of assignment is controlled by the experiment. The duration of each is controlled by the number of epochs with batch size of 128 . As backbone we use MLP with 4 hidden layers and 256 hidden units and ReLU activation. For all the methods, the learning rate is 0 . 1 . For Soft Resets , we use s = 0 . 9 and ν = 1 and η γ = 0 . 01 . Bayesian method uses temperature λ = 0 . 01 . Detailed results are given in Figure 7.

H.3 Reinforcement learning experiments

We conduct experiments in the RL environments. We take the canonical implementation of Soft-Actor Critic(SAC) from Brax [15] repo in github, which uses 2 layer MLPs for both policy and Q-function. It employs ReLU activation functions for both. On top of that, it uses 2 MLPnetworks to parameterize Q -function (see Brax [15]) for more details. To employ Soft Reset , we do the following. After we have collected a chunk of data ( 128 ) time-steps, we do one update (10) on γ t starting from 1 at every update of γ t , where γ t is shared for all the parameters within each layer of a Neural Network, separately for weights and biases. On top of that, since we have policy and value function networks, we have separate γ t for each of these. After the update on γ t , we compute θ t ( γ t ) and α t ( γ t ) , see Section 3.5. After that, we employ the proximal objective (45) with a fixed regularization target θ t ( γ t ) . Concretely, we use the update rule (47) where for each update the gradient is estimate on the batch of data from the replay buffer. This is not exactly the same as what we did with plasticity benchmarks since there the update was applied to the same batch of data, multiple times. Nevertheless, we found this strategy effective and easy to implement on top of a SAC algorithm. In practice, we swept over the parameter s (similar for both, policy and the value function) which controls the relative learning rate increase in (19). Moreover, we swept over the temperature ˜ λ from eqn. (45), which was different for the policy and for the value function. In practice, we found that using temperature of 0 for the policy led to the best empirical results. The range for the temperatures ˜ λ was { 0 . 1 , 0 . 01 , 0 . 001 } and for s was { 0 . 8 , 0 . 9 , 0 . 95 , 0 . 97 , 1 . 0 } . We used ν = 1 for all the experiments. For each experiment, we used a 3 hours of the A 100 GPU with 40 Gb of memory.

0

25

50

75

100

125

150

175

200

0.0

0.2

0.4

0.6

0.8

1.0

20% of random labels | 30 epochs

0

25

50

75

100

125

150

175

200

20% of random labels | 100 epochs

0

25

50

75

100

125

150

175

200

20% of random labels | 400 epochs

0

25

50

75

100

125

150

175

200

0.0

0.2

0.4

0.6

0.8

1.0

40% of random labels | 30 epochs

0

25

50

75

100

125

150

175

200

40% of random labels | 100 epochs

0

25

50

75

100

125

150

175

200

40% of random labels | 400 epochs

0

25

50

75

100

125

150

175

200

Task id

0.0

0.2

0.4

0.6

0.8

1.0

Task accuracy

60% of random labels | 30 epochs

0

25

50

75

100

125

150

175

200

Task id

Task accuracy

60% of random labels | 100 epochs

0

25

50

75

100

125

150

175

200

Task id

Task accuracy

60% of random labels | 400 epochs

Methods

Online SGD

Soft Resets

K

=10,

K

=10,

per parameter)

Hard Reset

Methods

Online SGD

Soft Resets

K

=10,

K

=10,

per parameter)

Hard Reset

0.0

0.2

0.4

0.6

0.8

1.0

1e6

7

8

9

1e

5+9.999e

1

for policy

0.0

0.2

0.4

0.6

0.8

1.0

1e6

0.2

0.4

0.6

0.8

1.0

for policy

0.0

0.2

0.4

0.6

0.8

1.0

1e6

0.4

0.6

0.8

1.0

for policy

Layer # 0 : Weights

Layer # 1 : Biases

Layer # 1 : Weights

Layer # 1 : Biases

Layer # 2 : Weights

0.0

0.2

0.4

0.6

0.8

1.0

1e6

0.96

0.97

0.98

0.99

1.00

for Q1

0.0

0.2

0.4

0.6

0.8

1.0

1e6

0.7

0.8

0.9

1.0

for Q1

0.0

0.2

0.4

0.6

0.8

1.0

1e6

0.4

0.6

0.8

1.0

for Q1

Layer # 0 : Weights

Layer # 1 : Biases

Layer # 1 : Weights

Layer # 1 : Biases

Layer # 2 : Weights

0.0

0.2

0.4

0.6

0.8

1.0

1e6

0.96

0.97

0.98

0.99

1.00

for Q2

0.0

0.2

0.4

0.6

0.8

1.0

1e6

0.7

0.8

0.9

1.0

for Q2

0.0

0.2

0.4

0.6

0.8

1.0

1e6

0.4

0.6

0.8

1.0

for Q2

Layer # 0 : Weights

Layer # 1 : Biases

Layer # 1 : Weights

Layer # 1 : Biases

Layer # 2 : Weights

I Computational complexity

We provide the study of computational cost for all the proposed methods. Notations:

P be the number of parameters in the Neural Network

L is the number of layers

O ( S ) is the cost of SGD backward pass.

M γ - number of Monte Carlo samples for the drift model

M θ - number of Monte Carlo samples for the parameter updates (Bayesian Method).

K γ - number of updates for the drift parameter

K θ - number of NN parameter updates.

The general theoretical cost of all the proposed approaches is given in Table 1. In practice, for all the experiments, we assume that M γ = 1 and M θ = 1 . Moreover, we used K γ = 1 and K θ = 1 for Soft Reset , K γ = 10 and K θ = 1 for Soft Reset with more computation. On top of that, for Soft Reset proximal and all Bayesian methods, we used K γ = 10 and K θ = 10 . Table 2, quantifying the complexity of all the methods from Figure 2.

The complexity O (2 S ) of Soft Resets comes from one update on drift parameter and one updat eon NN parameters. The memory complexity requires storing O ( L ) parameters gamma (one for each layer), parameters θ t with O ( P ) and sampled parameters for drift model update which requires O ( P ) .

Note that as Figure 9 suggests, it is beneficial to spend more computational cost on optimizing gamma and on doing multiple updates on parameters. However, even the cheapest version of our method Soft Resets still leads to a good performance as indicated in Figure 2.

The complexity of soft resets in reinforcement learning setting requires only one gradient update on γ after each new chunk of fresh data from the environment. In SAC, we do G gradient updates on parameters for every new chunk of data. Assuming that complexity of one gradient update in SAC is O ( S ) , soft reset only requires doing one additional gradient update to fit γ parameter.

The computation complexity of Soft Reset in Reinforcement Learning is marginally higher than SAC but leads to better empirical performance in a highly off-policy regime, see Appendix H.3.

K

=10

Soft Reset Proximal

K

=10,

K

=10

K

=10,

K

=10

K

=10,

K

=10,

per parameter

0.81

0.82

0.83

0.84

0.85

0.86

0.87

Permutted MNIST

K

=10

Soft Reset Proximal

K

=10,

K

=10

K

=10,

K

=10

K

=10,

K

=10,

per parameter

0.55

0.60

0.65

0.70

0.75

0.80

0.85

Random Label MNIST -- Data Efficient

K

=10

Soft Reset Proximal

K

=10,

K

=10

K

=10,

K

=10

K

=10,

K

=10,

per parameter

0.80

0.85

0.90

0.95

Random Label CIFAR-10 -- memorization

J Sensitivity analysis

We study the sensitivity of Soft Resets where γ is defined per layer when trained on random-label MNIST (data efficient). We fix the learning rate to α = 0 . 1 . We study the sensitivity of learning rate for the drift parameter, η γ , as well as ν - initial prior standard deviation rescaling, and s - posterior standard deviation rescaling parameter.

On top of that, we conduct the sensitivity analysis of L2 Init [34] and Shrink&Perturb [2] methods. The x-axis of each plot denotes one of the studied hyperparameters, whereas y-axis is the average performance across all the tasks (see Experiments section for tasks definition). The standard deviation is reported over 3 random seeds. A color indicates a second hyperparameter which is studied, if available. In the title of each plot, we write hyperparameters which are fixed. The analysis is provided in Figure 10 for Soft Resets and in Figure 11 for the baselines.

The most important parameter is the learning rate of the drift model η γ . For each method, there exists a good value of this parameter and performance is sensitive to it. This makes sense since this parameter directly impacts how we learn the drift model.

The performance of Soft Resets is robust with respect to the posterior standard deviation scaling s parameter as long as it is s ≥ 0 . 5 . For s < 0 . 5 , the performance degrades. This parameter is defined from σ t = sσ 0 and affects relative increase in learning rate given by 1 γ 2 +(1 -γ 2 ) /s 2 ) which could be ill-behaved for small s .

We also study the sensitivity of the baseline methods. We find that L2 Init [34] is very sensitive to the parameter λ , which is a penalty term for λ || θ -θ 0 || 2 . In fact, Figure 11, left shows that there is only one good value of this parameter which works. Shrink&Perturb [2] is very sensitive to the shrink parameter λ . Similar to L2 Init, there is only one value which works, 0.9999 while values 0.999 and values 0.99999 lead to bad performance. This method however, is not very sensitive to the perturb parameter σ provided that σ ≤ 0 . 001 .

Compared to the baselines, our method is more robust to the hyperparameters choice. Below, we also add sensitivity analysis for other method variants. Figure 12 shows sensitivity of Soft Resets , K γ = 10 , Figure 13 shows sensitivity of Soft Resets , K γ = 10 , K θ = 10 , Figure 14 shows sensitivity of Bayesian Soft Resets , K γ = 10 , K θ = 10 with γ t per layer, Figure 15 shows sensitivity of Bayesian Soft Resets , K γ = 10 , K θ = 10 with γ t per parameter.

0.2

0.4

0.6

0.8

1.0

Posterior standard deviation scaling `

s

`

0.1

0.2

0.3

0.4

0.5

0.6

Average accuracy

p

=0.05

Learning rate

0.0001

0.0005

0.001

0.005

0.01

0.05

0.1

0.5

0.2

0.4

0.6

0.8

1.0

Posterior standard deviation scaling `

s

`

0.1

0.2

0.3

0.4

0.5

0.6

Average accuracy

=0.1

Prior std scaling

p

0.01

0.05

0.1

10

9

10

7

10

5

10

3

10

1

10 1

L2 term regularization cost `

s

`

0.0

0.2

0.4

0.6

0.8

Average accuracy

Sensitivity L2 Init (MNIST). Learning rate

=0.1

10

6

10

5

10

4

10

3

10

2

10

1

Perturb parameter `

`

0.10

0.15

0.20

0.25

0.30

0.35

0.40

0.45

Average accuracy

Sensitivity Shrink and Perturb (MNIST). Learning rate

=0.1

Shrink parameter

0.9

0.99

0.999

0.9999

0.99999

1.0

0.2

0.4

0.6

0.8

1.0

Posterior standard deviation scaling `

s

`

0.1

0.2

0.3

0.4

0.5

0.6

Average accuracy

K

=10 (MNIST). Prior standard deviation scaling

p

=0.05

Learning rate

0.0001

0.0005

0.001

0.005

0.01

0.05

0.1

0.5

0.2

0.4

0.6

0.8

1.0

Posterior standard deviation scaling `

s

`

0.1

0.2

0.3

0.4

0.5

0.6

Average accuracy

K

=10 (MNIST). Drift learning rate

=0.01

Prior std scaling

p

0.01

0.05

0.1

0.2

0.4

0.6

0.8

1.0

Posterior standard deviation scaling `

s

`

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Average accuracy

K

=10,

K

heta

=10 (MNIST). Prior standard deviation scaling

p

=0.05

Learning rate

0.0001

0.0005

0.001

0.005

0.01

0.2

0.4

0.6

0.8

1.0

Posterior standard deviation scaling `

s

`

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Average accuracy

K

=10,

K

heta

=10 (MNIST). Drift learning rate

=0.0001

Prior std scaling

p

0.01

0.05

0.1

0.5

1.0

10

3

10

2

10

1

10 0

Prior standard deviation scaling `

p

`

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Average accuracy

per layer (MNIST). KL cost

=0.01

Learning rate

0.0001

0.0005

0.001

0.005

0.01

0.1

10

3

10

2

10

1

10 0

KL cost,

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Average accuracy

per layer (MNIST). Initial prior rescaling

p

=0.05

Learning rate

0.0001

0.0005

0.001

0.005

0.01

0.1

10

3

10

2

10

1

10 0

Prior standard deviation scaling `

p

`

0.2

0.4

0.6

0.8

1.0

Average accuracy

per parameter (MNIST). KL cost

=0.01

Learning rate

0.0001

0.0005

0.001

0.005

0.01

10

3

10

2

10

1

10 0

KL cost,

0.0

0.2

0.4

0.6

0.8

1.0

Average accuracy

per parameter (MNIST). Initial prior rescaling

p

=0.05

Learning rate

0.0001

0.0005

0.001

0.005

0.01

K Qualitative behavior of soft resets and additional results on Plasticity benchmarks

K.1 Perfect Soft Resets

0

50

100

150

200

Task id

0.0

0.2

0.4

0.6

0.8

Average Task Accuracy

0

50

100

150

200

Task id

Average Task Accuracy

Online SGD

Hard Reset

Hard Reset (only last)

=0.2)

=0.4)

=0.6)

=0.8)

0

200000

400000

600000

Step

0.5

0.6

0.7

0.8

0.9

1.0

CIFAR-10

for linear

0

200000

400000

600000

Step

for linear_1

0

200000

400000

600000

Step

for linear_2

0

200000

400000

600000

Step

for linear_3

0

200000

400000

600000

Step

for linear_4

boundaries, and Hard Reset (only last) which resets only the last layer. We use Soft Reset method (20) where γ t = 1 all the time and becomes γ t = ˆ γ t (with manually chosen ˆ γ t ) at task boundaries. We consider constant learning rate α t ( γ t ) and increasing learning rate (19) at task boundary for Soft Reset . On top of that, we run Soft Reset method unaware of task boundaries which learns γ t . We report Average training task accuracy metric in Figure 16. See Appendix H.1 for details. The results suggest that with the appropriate choice of ˆ γ t , Soft Reset is much more efficient than Hard Reset and the effect becomes stronger if the learning rate α t ( γ t ) increases. We also see that Soft Reset could learn an appropriate γ t without the knowledge of task boundary.

K.2 Qualitative Behaviour on Soft Resets on random-label tasks.

We observe what values of γ t we get as we train Soft Reset method on random-label MNIST (dataefficient) and CIFAR-10 (memorization). The results are given in Figure 17 for MNIST and in Figure 18 for CIFAR-10. We report these for the first 20 tasks.

K.3 Qualitative Behaviour on Soft Resets on permuted patches of MNIST.

We consider a version of permuted MNIST where instead of permuting all the pixels, we permute patches of pixels with a patch size varying from 1 to 14 . The patch size of 1 corresponds to permututed MNIST and therefore the most non-stationary case, while patch size of 14 corresponds to least nonstationary case. We use a convolutional Neural Network in this case. In Figure 19, we report the behavior of γ for different convolutional and fully connected layers on first few tasks.

0

200000

400000

600000

Step

0.5

0.6

0.7

0.8

0.9

1.0

CIFAR-10

for linear

0

200000

400000

600000

Step

for linear_1

0

200000

400000

600000

Step

for linear_2

0

200000

400000

600000

Step

for linear_3

0

200000

400000

600000

Step

for linear_4

0

10000

20000

30000

40000

Step

0.80

0.85

0.90

0.95

1.00

for conv2_d_1

Permuted patch size

1

2

4

7

14

0

10000

20000

30000

40000

Step

0.80

0.85

0.90

0.95

1.00

for conv2_d_2

0

10000

20000

30000

40000

Step

0.80

0.85

0.90

0.95

1.00

for conv2_d_3

0

10000

20000

30000

40000

Step

0.80

0.85

0.90

0.95

1.00

for conv2_d_4

0

10000

20000

30000

40000

Step

0.80

0.85

0.90

0.95

1.00

for linear

0

10000

20000

30000

40000

Step

0.80

0.85

0.90

0.95

1.00

for linear_1

K.4 Bayesian method is better than non-Bayesian

As discussed in Section 5, we found that in practice Soft Reset and Soft Reset Proximal where γ is learned per-parameter, did not perform well on the plasticity benchmarks. However, the Bayesian variant described in Section E.1, actually benefited from specifying γ for every parameter in Neural Network. We report these additional results in Figure 20. We see that the non Bayesian variants where γ t is specified per parameter, do not perform well. The fact that the Bayesian method performs better here suggests that it is important to have a good uncertainty estimate σ 2 t for the update (10) on γ t . When, however, we regularize γ t to be shared across all parameters within each layer, this introduces useful inductive bias which mitigates the lack of uncertainty estimation in the parameters. This is because for non-Bayesian methods, we assume that the uncertainty is fixed, given by a hyperparameter - assumption which would not always hold in practice.

0

50

100

150

200

Task id

0.3

0.4

0.5

0.6

0.7

0.8

Average Task Accuracy

Random Label MNIST -- Data Efficient

Hard Reset

Soft Reset Proximal

Soft Reset Proximal (

per parameter)

per parameter)

per parameter)

0

50

100

150

200

Task id

0.4

0.6

0.8

1.0

Average Task Accuracy

Random Label CIFAR-10 -- Memorization

K.5 Qualitative behavior of soft resets

In this section, we zoom-in in the data-efficient experiment on random-label MNIST. We use Soft Reset Proximal ( γ per layer) method with separate γ for layer (different for each weight and for each bias) and run it for 20 tasks on random-label MNIST. In Figure 21 we show the online accuracy as we learn over this sequence of tasks. In Figure 22, we visualize the dynamics of parameters γ for each layer. First of all, we see that γ t seems to accurately capture the task boundaries. Second, we see that the amount by which each γ t changes depends on the parameter type - weights versus biases, and it depends on the layer. The architecture in this setting starts form linear and goes up to linear 4 , which represent the 4 MLP hidden layers with a last layer linear 4 .

K.6 Impact of specific initialization

In this section, we study the impact of using specific initialization θ 0 ∼ p init ( θ ) in p 0 ( θ ) as discussed in Appendix D. Using the specific initialization in Soft Resets leads to fixing the mean of the p 0 ( θ ) to be θ 0 , see (24). This, in turn, leads to the predictive distribution (25). In case when we are not using specific initialization θ 0 , the mean of p 0 ( θ ) is 0 and the predictive distribution is given by (23). To understand the impact of this design decision, we conduct an experiment on random label MNIST with Soft Reset , where we either use the specific initialization or not. For each of the variants, we

0

20000

40000

60000

80000

100000

0.0

0.2

0.4

0.6

0.8

1.0

Time-step

0

20000

40000

60000

80000

100000

0.00014

0.00012

0.00010

0.00008

0.00006

0.00004

0.00002

0.00000

+1

Per Layer linear b

0

20000

40000

60000

80000

100000

0.00006

0.00004

0.00002

0.00000

0.00002

0.00004

0.00006

0.00008

0.00010

+9.999e

1

Per Layer linear_1 b

0

20000

40000

60000

80000

100000

0.9992

0.9994

0.9996

0.9998

1.0000

Per Layer linear_2 b

0

20000

40000

60000

80000

100000

0.99982

0.99984

0.99986

0.99988

0.99990

0.99992

0.99994

0.99996

Per Layer linear_3 b

0

20000

40000

60000

80000

100000

4.0

4.5

5.0

5.5

6.0

1e

5+9.999e

1

Per Layer linear_4 b

0

20000

40000

60000

80000

100000

0.9970

0.9975

0.9980

0.9985

0.9990

0.9995

1.0000

Per Layer linear w

0

20000

40000

60000

80000

100000

0.9970

0.9975

0.9980

0.9985

0.9990

0.9995

1.0000

Per Layer linear_1 w

0

20000

40000

60000

80000

100000

0.9975

0.9980

0.9985

0.9990

0.9995

1.0000

Per Layer linear_2 w

0

20000

40000

60000

80000

100000

0.9965

0.9970

0.9975

0.9980

0.9985

0.9990

0.9995

1.0000

Per Layer linear_3 w

0

20000

40000

60000

80000

100000

0.9965

0.9970

0.9975

0.9980

0.9985

0.9990

0.9995

1.0000

Per Layer linear_4 w

do a hyperparameters sweep. The results are given in Figure 23. We see that both variants perform similarly.

0

50

100

150

200

0.3

0.4

0.5

0.6

p

0(

)

Soft Reset -- with specific init in

p

0(

)

L Toy illustrative example for SGD underperformance in the non-stationary regime

Illustrative example of SGD on a non-stationary stream. We consider a toy problem of tracking a changing mean value. Let the observations in the stream S t follow y t = µ t + σϵ , where ϵ ∼ N (0 , 1) , σ = 0 . 01 . Every 50 timesteps the mean µ t switches from -2 to 2 . We fit a 3-layer MLP with layer sizes (10 , 5 , 1) and ReLU activations, using SGD with two different choices for the learning rate: α = 0 . 05 and α = 0 . 15 . Moreover, given that we know when a switch of the mean happens, we reset (or not reset) all the parameters at every switch as we run SGD. Only during the reset, we use different learning rate β = 0 . 05 or β = 0 . 15 . Using higher learning rate during reset allows SGD to learn faster from new data. We also ran SGD with α = 0 . 05 and β = 0 . 15 , where the higher learning rate is used during task switch but we do not reset the parameters. We found that it performed the same as SGD with α = 0 . 05 , which highlights the benefit of reset.

0

50

100

150

200

250

300

Number of timesteps

2

1

0

1

2

Predicted mean

True data

SGD (

=0.05)

SGD (

=0.05) + reset (

=0.05)

SGD (

=0.15)

SGD (

=0.05) + reset (

=0.15)

We report the predicted mean ˆ µ t for all SGD variants in Figure 24. We see that after the first switch of the mean, the SGD without reset takes more time to learn the new mean compared to the version with parameters reset. Increasing the learning rate speeds up the adaptation to new data, but it still remains slower during the mean change from 2 to -2 compared to the version that resets parameters. This example highlights that resets could be highly beneficial for improving the performance of SGD which could be slowed down by the implicit regularization towards the previous parameters θ t and the impact of the regularization strength induced by the learning rate.

M Using arbitrary drift models

The approach described in section 3.5 provides a general strategy of incorporating arbitrary Gaussian drift models p ( θ | θ t ; ψ t ) = N ( θ ; f ( θ t ; ψ t ); g 2 ( θ t ; ψ t )) which induces proximal optimization problem

The choice of f ( θ t ; ψ t ) and g ( θ t ; ψ t ) affects the behavior of the estimate θ t +1 from (53) and ultimately depends on the problem in hand. The objective function of the form (53) was studied in context of online convex optimization in [21],[29], where the underlying algorithms estimated the deterministic drift model online. These worked demonstrated improved regret bounds depending on model estimation errors. This approach could also be used together with a Bayesian Neural Network (BNN).

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We outline main contributions of the paper in the introduction and abstract.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss limitations in the experimental section

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We disclose the experimental information in Experimental and Appendix sections.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: Unfortunately, due to IP constrains, we cannot release the code for the paper.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide experimental details in the appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We specify that we report results with 3 random seeds with mean and standard deviation.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide information about compute resources required in the appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Based on our understanding, our work conforms to the every aspect of NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We cite the works which introduced the publicly available datasets

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper focuses on non-stationary learning of neural networks. The paper proposes a method that automatically adjusts a parameter that influences the stochastic gradient descent (SGD) update to account for non-stationarity.\n\nSTRENGTHS:\nThe problem tackled is relevant and the related work is clearly discussed. The method is well-interpreted and relevant for the community. The paper is also very well presented.\n\nWEAKNESSES:\nIn my view, the paper has some weaknesses, especially regarding the experimental validation:\n- The experimental validation does not clearly validate the benefits of the method, possibly due to lack of clarity. In Figure 3, it is not clear to me what method this paper proposes. I believe the paper should propose a method and validate it, instead of comparing all its possible variants. Nevertheless, I can see that most (if not all) of such variants are outperforming the chosen baselines. However, only two of the baselines referred in the related work were chosen. The chosen benchmark problems are also not motivated to me. \n- The reinforcement learning part is not clear to me. Is there an explanation for why, on-policy, the method is outperformed by the simple baseline? Is not on-policy the most non-stationary setting? Why? And why does the proposed method stop earlier than the competitors in the top-right figure? \n- It is not clear to me what the authors intend to show with Figure 2, and the perfect soft-resets regime. It would be nice to have more explanation here, and possibly move this validation to the end of the experimental section, as appearing before the main validation of the method is confusing. \n\nFinally, even though the method is well motivated, it is not theoretically analyzed. We have no theoretical proof that the method will outperform traditional SGD in terms of efficiency, or preventing plasicity loss and catastrophic forgetting, or other dimension. \n\nMinor: in page 6, Equation (15), what is the parameter $\\tilde{\\lambda}$? Was it introduced before?\n\nQUESTIONS:\nI have the following questions:\n- Could the authors make some clarifications with respect to the experimental validation? Specifically, on the choice of baselines, problems, which method or methods the authors in fact propose to the community. \n- Can the authors also clarify my concern regarding the reinforcement learning setting and the performance in on-policy settings?\n- Can the authors clarify what is the intended take-away from the analysis in Figure 2?\n- Can the authors clarify the Monte-Carlo samples in Equation (7)? Specifically, what is sampled, for instance in the reinforcement learning setting, and if the sampling requires access to a free simulator of the environment.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- The authors study a learning algorithm that can handle the non-stationarity of the data distribution.\n- They propose a parameter drift model based on the Ornstein-Uhlenbeck process, which models a form of “soft parameter reset” adaptive to the data stream. The drift model has an adaptive parameter $\\gamma_t$ which a gradient-based optimizer can learn online.\n- They illustrate the update rule of the main parameters of a neural network incorporating the learned drift model. The update rule is first proposed under the Bayesian neural network framework and then later adapted to a non-Bayesian neural network.\n- They numerically corroborate the efficacy of their method on plasticity benchmarks and reinforcement learning experiments.\n\nSTRENGTHS:\n- S1. They propose a novel way to learn online the amount of resetting the parameters.\n- S2. The method is general enough to be applied broadly to continual learning problems and reinforcement learning problems.\n\nWEAKNESSES:\n- W1. Modeling parameter drift is not well-motivated.\n    - The learnable drift model $p(\\theta_{t+1} \\mid \\theta_t, \\gamma_t)$ is the main contribution of this work. In my opinion, however, it is unclear why we should care about the drift of “currently learning” neural network parameters.\n    - According to the beginning of Section 3, the local minima of the objective function may change over time. This is of course acceptable. However, I cannot understand what is relevant between the change of the local optima $\\theta^{\\star}_t$ and the current parameter $\\theta_t$. To be more specific, the two sentences in lines 94-96 are not connected well.\n- W2. The paper is NOT self-contained overall.\n    - The authors claim that they chose the Ornstein-Uhlenbeck (OU) process as a drift model. However, it is unclear how a continuous stochastic differential equation (OU process) can be converted to a discrete Markovian chain defined in Equation (4). The paper does not even introduce the actual form of the OU process. Moreover, it is not very clear why the authors chose the particular model defined in Equation (4) because it is not known (at least in the paper) that the model is a unique option to choose.\n    - There are a lot of equations whose derivation is omitted. Let me list them: Equations (9), (10), (11), (14), (15). In particular, I am suspicious of both the validity and usefulness of Equation (10).\n    - For these reasons, the paper is not so easy to follow.\n- W3. The soft reset method seems computationally heavy.\n    - I am worried about the computational cost. Although it is good to learn the forgetting parameter $\\gamma_t$ online, it increases the computational cost almost twice. It might be more than just twice because learning the $\\gamma_t$ requires Monte-Carlo (MC) sampling.\n    - In addition, there seem too many hyperparameters to tune.\n- W4. Minor comments on typos/mistakes\n    - Do not capitalize the word “neural network”. (e.g., line 35-36)\n    - The term “plasticity loss” may be read as a type of loss function for some people (new to this field). I recommend using the term “loss of plasticity”.\n    - Line 84: “non-stationary” → “non-stationarity”\n    - Line 109: “mportant” → “important”\n    - Line 127: “wrt” → “with respect to”. Do not use abbreviations.\n    - Line 137: “property” → “properties”\n    - Line 154: what is “s” next to a parenthesis?\n    - Line 168: “$\\theta = (\\theta_i, \\ldots, \\theta_D)$” → “$\\theta = (\\theta_1, \\ldots, \\theta_D)$”\n    - Line 169: “$\\mathbb{R}^D$” is a more standard notation than “$\\mathcal{R}^D$”.\n    - Line 222: In “$\\lambda\\_i = \\hat{\\lambda}\\_i \\sigma^2_{t,i}$\", I think $\\hat{\\lambda}$ must not have a subscript.\n    - Line 225: “sinc” → “since”\n    - Equation (14): What is “$\\tilde{\\mathcal{F}}$”? (Partial) gradient of “$\\mathcal{F}$”?\n    - Line 232: What does it mean by “we assume that $\\mu_0 = \\theta_0, \\sigma_0^2.$”?\n    - Line 236: “$\\theta_{t,i}(\\gamma_{t,i}) = \\theta_t$” → “$\\theta_{t,i}(\\gamma_{t,i}) = \\theta_{t,i}$”\n    - Line 240: “linearisng” → “linearizing” or “linearising”\n    - Line 292: duplicate “See”\n    - Line 354: “modelling” → “modeling”\n    - Lines 492-493: I guess “*data efficient*” and “*memorization*” are swapped.\n\nQUESTIONS:\nQ1. As far as I know, the Hard Reset method typically resamples the model parameter every time it resets the parameter. With this in mind, what if we slightly modify the drift model as: $p(\\theta \\mid \\theta_t, \\gamma_t) = \\mathcal{N} (\\theta; \\gamma_t \\theta_t + (1-\\gamma_t) \\theta’_0 ; (1-\\gamma_t^2)\\sigma_0^2)$, where we re-sample $\\theta’_0 \\sim p_0(\\theta_0)$ every time $t$?\n\nLIMITATIONS:\nThe paper discusses its limitations in the experiment section. I think it might be better to mention them explicitly in the conclusion section.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a method to effectively learn the neural network parameters in non-stationary environments. They propose a modified learning algorithm that adapts to non-stationarity through an Ornstein-Uhlenbeck process with an adaptive drift parameter. Drift parameter is used to track the non-stationarity in the data - specifically, adaptive drift pulls the NN parameters towards initial parameters when the data is not well-modelled by current parameters while on the other hand reducing the degree of regularization (towards initial parameters) when the data is relatively stationary. The paper introduces a Bayesian Soft-Reset algorithm in Bayesian Neural Network (BNN) framework and a General Soft-Reset algorithm in deterministic non-Bayesian framework to effectively learn in non-stationary environments. Empirically, authors demonstrate that their method performs effectively in non-stationary supervised and off-policy reinforcement learning settings.\n\nSTRENGTHS:\nTo take into account the non-stationarity in the data - the methodology in the paper modifies the SGD algorithm by modifying the regularization strength and the regularization target of the SGD algorithm - a significant contribution of the work is to do this modification in a principled manner using a drift parameter that is updated online from the stream of data.\n\nOverall, the paper is well written. The need and novelty of the methodology are clearly presented, and the illustration of the use of methodology on realistic instances clearly demonstrates how it works and improves on the existing methodologies for learning in non-stationary data regime.\n\nWEAKNESSES:\nSome parts can be improved - \n1) Section 2 - Simplification of equation (6) to (9) requires some explanation [page -5].\n2) Section 3 - Some intuition behind the objective (11) [page - 6] would be good to understand it better (especially the second term in that objective)\n3) Some discussion on how does the methodology perform relative to the degree of non-stationarity in the data?\n4) Some minor typos in line 109, 154, 190\n\nQUESTIONS:\nHow does the methodology perform relative to the degree of non-stationarity in the data? Specifically, to what extent does this methodology remain effective under varying levels of non-stationarity?\n\nLIMITATIONS:\nAuthors have adequately addressed the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work focuses on the problem of plasticity loss in non-stationary problems. This work proposes a new solution which adaptively drifts the parameters toward the initial distribution. The proposed solution is a form of soft resetting and can be seen as a meta version of L2-init where the degree of drift towards initialization is learned from data.\n\nSTRENGTHS:\nThe proposed solution is novel. Current soft resetting methods reset the parameters of the model at a constant rate, irrespective of the degree of non-stationarity. The proposed solution introduces adaptive soft resetting, where the degree of resetting is calculated based on the degree of non-stationarity. This adaptiveness allows the solution to adapt faster and retain more prior knowledge. The experimental results show the effectiveness of the proposed solution.\n\nWEAKNESSES:\nAlthough the proposed solution is novel and effective, the paper has many weaknesses:\n1. Hyperparameter sensitivity. Given that the proposed solutions introduce many new hyperparameters and the authors say in line 207 that one of the solutions is sensitive to the choice of hyperparameter, the paper needs to contain hyperparameter sensitivity curves for all new hyperparameters.\n2. Computational cost. The proposed solutions seem computationally expensive. It would be good to include the wall of each algorithm beside one of the figures, and Figure 3 might be best.\n3. Limited performance improvement. The primary purpose of high replay-ratio is to improve the sample efficiency of algorithms. However, results in Figure 4 show that the proposed solution does not utilize a high replay-ratio. Its sample efficiency at replay ratio 1 is the same as its sample efficiency at replay ratio 128.\n\nQUESTIONS:\nSee Weaknesses\n\nLIMITATIONS:\nComputational cost and hyperparameter sensitivity need to be discussed in the conclusion of the main paper.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-18"}