Non-Stationary Learning of Neural Networks with Automatic Soft Parameter Reset
Authors: Alexandre Galashov
Michalis Titsias
András György
Clare Lyle
Razvan Pascanu
Yee Whye Teh
Maneesh Sahani
Venue: NeurIPS 2024 poster (2024)
Keywords: Non-stationarity
plasticity loss
online learning
deep learning
Primary area: optimization_for_deep_networks
TL;DR: Learning to reset Neural Networks parameters using Ornstein-Uhlenbeck to learn on non-stationary data

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
Alexandre Galashoy* Gatsby Unit, UCL Google DeepMind agalashhv@googlee.coom
Michalis K. Titsias
András György Google DeepMind agyorgyggoogle.cm com
Google DeepMind
mtitsiaas@gogle..om
Clare Lyle Google Deep clarelyle@google..o com
Razvan Pascanu Google DeepMind razp@google. com
Yee Whye Teh Google DeepMind University of Oxford yutehhgoogle..om com
Maneesh Sahani Gatsby Unit, UCL maneesh@gataby. acllac.uk
Abstract
Neural networks are most often trained under the assumption that data come from 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 Ornstein- Uhlenbeck process with an adaptive drift parameter. The adaptive drift draws the parameters towards the distribution used initialisation, be understood as a form of soft parameter reset. We show empirically that our approach performs well in non-stationary supervised, and off off-policy reinforcement learning settings.
the approach
1 Introduction
Neural networks (NNs) are typically trained using algorithms such as stochastic gradient descent (SGD), which implicitly assume that the data come from stationary distribution. This assumption is incorrect for scenarios such as continual learning, reinforcement learning, non-stationary contextual
be updated online as new data are encountered, this approach often leads to loss of plasticity [3]. manifesting either as failure generalise to new data despite reduced training loss inability to reduce training error as the data distribution changes 13
bandits, and supervised learning with distribution shifts [20] 53) Although the parameters
2
as an
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
often addressed using hard resets based on heuristics such as detecting dormant units [47]. assessing unit utility 11 [2), or simply after a fixed number of steps has elapsed[43] (3) Although effective at increasing plasticity, hard resets be inefficient they can discard valuable knowledge captured by the parameters.
*Corresponding author
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

We propose an algorithm that instead implements form of soft parameter reset, avoiding the pitfalls associated with hard resets, A soft reset moves the parameters the direction of the initial-value distribution, while maintaining dependence on their previous values. It also increases the learning rate applied to gradient- baase updates, allowing new parameters adapt faster to the changing data. The magnitude of the soft reset is governed by an in time, with the scale of the drift itself chosen adaptively. In effect, the model approximates the action of dynamical Bayesian prior over the NN parameters, which is adapted online to new data.
Caateinn-Uhleebckk drift process linking the parameters
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 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 avoids the challenge of plasticity loss, as well as an off-policy reinforcement learning setting.
2 Non-stationary learning with Online SGD
Consider a non-stationary supervised learning setting with changing data distributions (t(2, y), where IRL and y RKK (y may be restricted "one-hot" categorical indicator vectors). Given single-point loss L(0, x. y), we define the time-dependent expected loss function for parameters 0 E IR be
(+) = ~p [(,, 24, 21,3t).
Our goal is to find parameter sequence θ = (01,... ...., 0r) that reduces the dynamic regret
(r(0,07) = 1 - 1= (6, - (,))))
(2)
relative to an oracular reference sequence O (0) 07.) satisfying 0Ú E arg ming (0). One
common approach to online learning employs online stochastic gradient descent (SGD) 23 An
initial parameter value 0g is updated sequentially for each batch of data B. = {(1) yi H= s.t. (xt, vt) ~ (1 m1). The update rule is:
0: = 8 04 (OL.())
where = = Limi L(0, vi) is the empirical loss the batch, and ag is learning rate. See also Appendix G for the connection of SGD to proximal optimization.
Convex settings. For convex losses, online SGD with fixed learning rate can track non- stationarity [56]. By selecting a
potentill 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 a,, while in [2I] and in 29). the authors suggest modifying the starting point of SGD from 0g to an adjusted 62 proportional to the level of non-stationarity.
Non-convex settings. Non-stationary learning with NNs is more complex, since now there is 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 may be too aggressive if the data distributions are similar.
34]. Parameter resets [13 48 [2] partially mitigate the problem, but
3 Online non-stationary learning with learned soft parameter resets
Notation. We denote by A0 μ1, o²) a Gaussian distribution on 0 with mean and variance a². We denote by 0³ the j-th component of the vector = (0) ......... PD). We assume that the NN outputs define a distribution over targets y, so that the single- loss is the negative log likelihood C(0, x, 3) - log p (,,0), 8). As above, B, {( }
batch of data sampled from pt y)
with fixed batch size B (which could be 1) and (0) = log (BBB)) (8) = 1 DD L(B, x;, yf) is the average negative log likelihood on batch B,. Element-wise product denoted by
Our proposed soft resel scheme tracks changes in the distribution of non-stationary data using an explicit drift model for the parameters p(0 (, 7t). Specifically, the drift model we adopt assumes

---

Graphical model
Toy Bayesian Inference example
AMO
Pelo)
DELLS exsurption
DO RissniLd massaption
WETIOUT demariccal mode
10 Moraiatooerr CUDE
Figure 1: Left: graphical model for data generating process in the (a) stationary case and (b) non- stationary case with drift model p(0 10,. yt). Right: (c) In stationary online learning regime, the Bayesian posterior (red dashed circles) the long run will concentrate around 6* (red dot). (d) In a non-stationary regime where the optimal parameters suddenly change from current value 87 to new value 8 (blue dot) online Bayesian estimation can be less data efficient and take time to recover when the mhange-point occurs. (e) The use of p(0 0,, Vt)) and the estimation of % allows to increase the uncertainty, by soft resetting the posterior to make it closer to the prior (green dashed circle), so that the updated Bayesian posterior pr+1 (0) (blue dashed circle) can faster track 0/41-
that the parameters change in the direction of the initialization distribution at each time f. The amount of change, and thus the strength of reset, is controlled by non-stationary hyperparameters " These hyperparameters are themselves estimated online from the data.
From a probabilistic standpoint, the drift model implies an empirical Bayesian prior on parameters 8.+1, which is displaced from the posterior on Og towards the prior distribution, and has increased variance (controlled by 71). 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 to point 6, (rt) between the previous estimate 0g and the mean of initializing distribution, with 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
29
3.1 Toy illustration of the advantage of drift models
Consider online Bayesian inference with 2-D observations = 0" + where 0" E IR2 are
unknown true parameters and & (N( (2; is Gaussian noise with variance 0², Starting from a Gaussian prior (o) N0;; Po: Eo), the posterior distribution P++1(0) (0)= p(0 [10/ly...... ,9411) =
(O; pc+1, 1141) updated using Bayes' rule
(+) x p(y++110)p
The posterior update 3 arises from assumption that data are i.i.d. (Figure D.) since then Pe+1 (f) x (o) I 10). By the Central Limit Theorem, the posterior mean converges to 0" and the covariance matrix E, shrinks to zero (the radius of red circle in Figure[o).
Suppose now that the true parameters Of (which were fixed up to time t) change to new parameters 0i+1 time f + 1. The i.i.d. assumption is thus violated and the update B becomes problematic because the low uncertainty (small radius of red dashed circle in Figure Id) in pt (0) causes the posterior Pt+1 (0) (see blue circle) to adjust slowly towards 6/ (blue dot) as illustrated in Figure TH.
The issue addressed by allowing explicitly for the possibility that, before observing new data, the parameters drift according 10,, vx). The corresponding conditional independence structure is shown in Figure To. The posterior update then becomes:
(0) x p()++110) f p(010% ())p((0))))))
For suitable choice of drift model p(, this modification allows P++ (0) (blue circle) to adjust more rapidly towards the new 6:41 (blue dot), see Figure Dp. This is because the new prior (0000', 7pp, (8))dd) has larger variance (green circle) than (0( and its mean is closer to the center of the circle. Ideally, the parameter " should capture the underlying non-stationarity in the data distribution in order to control the impact of the prior p(0 02 7x 7x) (ppx(00)) (O))d'' For example, if at some point the non- non-stationarity disappears, we want the drift model to adaptively eliminate changes and recover the stationary posterior update 33. This highlights the importance of the adaptive nature of the drift model.
3

---

The parameters > E [0, 1] control the degree to which each parameter is reset. When ?i = 1, the model has 0/41 and 0l equal, and the data in B are used to refine the estimate of the common parameter. When 7/ = 0, 0/41 is drawn from the prior po independently of 0j, and so Bi-1 is used independently of previous data. A value of of E (0, 1) interpolates between these two extremes. The process (5) has the property that for any current parameter 0r and of E [0, 1). as T + co the distribution of (0+7(0,) will converge to the prior p p(O)). This behaviour depends on interplay between the mean shrinkage and the variance. Other choices of drift variance would result in the variance of ((6))1.)) either going to 0 or growing to cc, harming learning. Thus, the model (5) we discuss 3.2 Ornstein-Uhlenbeed 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 fixed number of iterations starting from a good initialization 8g po [see, e.g., [6] Informally, we refer to the region of high probability under po as plastic region.
In non-stationary data setting, learning for the batch B++1 is initialized based on the learnt parameter at the previous time step, 0z. Empirically, if this parameter far from the initial plastic region, the NN may suffer from loss of plasticity with SGD and similar online methods failing to adapt the new data. In this case, a hard reset of the parameters to do (or an alternative draw from po) may be helpful. However, if changes in the data distribution are gradual then performance may be improved by retaining information about 8,. Furthermore, since @ is high- nigh-dimensional, each dimension might benefit from different updates.
Drift model. A drift model p(0 00.. 7)) which captures these desiderata is given by p(0,0,, 7x) = I1, (N):00 of + (1 - o8 P6 (1 =)))=(
(5) defined independently for each parameter dimension 05, where Pa (0)) N (O0; pd: o is the per- parameter prior distribution and " (7/..... P) controls drift in each parameter separately. The model is a discretized Ornstein- -hleenbec (OU) process (50) (see Appendix B
for the derivation).
encourages parameters to move towards the plastic region (initialization). In Appendix C 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 opti- mizing 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 approxi- mate online variational inference 7 with Bayesian Neural Networks (BNN). Let I2 (?i..... 7)) be the history of observed parameters of the drift model and S = [BI..... B, be the history of observed data. The objective of approximate online variational inference is to propagate an approximate posterior qe (0) r, over parameters, such that it is constrained to some fam- ily of probability distributions. In the context of BNNs, it is typical 55 to assume a family (('.....,")) of Gaussian mean field distributions over parameters 0 Rº (separate Gaussian per parameter). Let qo (0) & (,,(0S,, 1'--i) Q be the Gaussian approximate posterior at time f with mean Me and variance of for every parameter. The new approximate posterior qt+1 (f) Q is found by
Q {q(0) q()) Π p10 ,
(++(() (0) = arg min, KIL
(6)
where the prior term the approximate predictive look-ahead prior given by
q: (0()))) f
dz) (
that has parameters == ( ...,, (iP) and /² (212 ...... .2) such that ind (1x) = 14 + (1 o juy. (T)) = ?/? + (1 20 see Appendix E. for derivation. The form of this prior
(fl))) comes from the non i.i.d. assumption (see Figure [p) and the form of the drift model (5)
For new batch of data B.+1 at time t + 1, the approximate predictive lokelikdl equals to
4. = log f 1 00 (0)x))d0.
The log-likelihood (8) allows us to quantify predictions on batch of data Bi+1 given our current distribution qe (0) and the drift model from (5). We find such 72 that
7/ R arg max,, log qe (B:+1 (nr))
(9)

---

Using Yi 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 rright).
Gradient-based optimization for "/t. The approximate predictive prior in 7 is Gaussian which allows us use the so-called reparameterisation trick to optimize (8) via gradient descent. Starting from an initial value of drift parameter of at time t, we perform K updates with learning rate T/y
7,,+1 Yrkk + 77, V log p(B ([BB++llA((n.. (T..)) + co (Hr,))) I)de.
(10)
The integral is evaluated by Monte-Carlo (MC) using M samples ~ ([ee 7), 1= M
J p(B co (o.,)) I)de
P(B+ (71,)) + c (((()))) (11)
Inductive bias in the drift model is captured by Tp. where Yr.0 = encourages stationarity, while 7t.0 , promotes temporal smoothness. In practice, we found 7r.0 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 (TO). We consider 7 to be either defined for each parameter or for each layer. See Section (3 details as well as corresponding results in Appendix
Interpretation of 7/- By linearising log p(B++1100) around μ4 and denoting 1 V E (4z). we
compute (8) in a closed form and get the following loss for Yt (see Appendix optimizing (9
F(1x) 0.5(o) 9(+1) ( ge+1 - (74 b4 + (1 7+) (40 g++1,
(12)
Adding the le penalty 1 - 7p 122 encoding the starting point 7p gives us the closed form for 7
3 (i, K
(13)
(50 - 1)7 +K+
where we assumed that n/2 is shared for subset of K parameters (see paragraph about structure in drift model) indexed by JK- = (fi,... ........ and ā = (z),.... Fix denotes a vector defined on this subset. We also clip parameters "/ to [0, 1]. The expression (33 gives us the geometric interpretation for "/- The value of " depends on the angle between (īe (o) and Øc+1 When these vectors are aligned, "e is high and is low otherwise. When these vectors are orthogonal or the gradient Ør+ the value of 7/ is heavily influenced by 7p. Moreover, when Ø++1 RJ 0, we can interpret it as being close to a local minimum, i.e., stationary, which means that we want 7t 1, therefore adding the & penalty is important. Also, when the norm of the gradients Ge+1 is high, the value of 7/. is encouraged to decrease, introducing the drift. This means that using "i 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 good approximation for learning 7t-
3.4 Approximate Bayesian update of posterior q: (0) with BNNs
The optimization problem (6) for the per-parameter Gaussian (0) N ,0,00 o02) with a prior (0) = 11, (00); pl/ of²) can be written (see Appendix E.1 to minimize the following loss
Fe(,.0,7) 7c) BonN(o:))) (E.1lu+ ΣE D
X
log [od]?]
(xx))
(14)
where X > 0 are per-parameter temperature coefficients. The use of a small temperature parameter (shared for all NN parameters) was shown to improve the empirical performance of Bayesian Neural Networks
(54) Given that in [4. the variance (xx(7)) can be small, in order to control the
strength of the regularization, we propose to use the temperature per parameter X > [o?] where
is a global constant. This leads to the following objective
Fe(p,a, m.) = Ee-N((:)) (EmilM+oo + Σ, r! (ws (0-(((0)))) + [ojp2
|o] (1)) log(osj3
(15)
where the quantity rf |o,1 /o/ (1:)) is a relative change in the posterior variance due to the drift. The ratio of 1 when 7/ 1. For of < 1 since typically of² < of² the ratio is rf < 1.
5

---

Algorithm Soft-Reset algoritm
Input:
Data-stream ST = {B.)7-)
Neural Network (NN) initializing distribution Pinat (f) and specific initialization 80 Panit (0)
Learning rate a, for parameters and 77, for drift parameters
Number of gradient updates Ky on drift parameter *t
NN initial standard deviation (STD) scaling v S (see (24) and ratio defining uoro
for step t = 0,
T do Receive batch of data 1 Initialize drift parameters 7r.0 for step k 2,... K, do
Sample 0 Pimit
Stochastic update 22 on drift parameter using specific initialization (26)
7r,k+1 -t.k 7/, log + 7000 + 0b uvv - 7/ + 82 end for
Get (,(%,,) with 18 and (, with 19
Update parameters 0.+1 (.(Tk)) k((T.,) VoLe+1(0((n.x) end for
Thus, as long as the non stationarity detected 6o < 1), the objective 15 favors the data term
(ow)((0:1) I e coa] allowing the optimization respond faster to changes in the data distribution. Let Mc+1,0 (7)) and 0+100 5z(T)), and perform updates K 15
M141641 M4+1 - (1 (HLA)) 04411+1 01114 - (M4+1k (
111 7e),
(16)
where an and ag learning rates for the mean and for the standard deviation correspondingly. All derivations are provided Appendix E.T The full procedure is described Algorithm
3.5 Fast MAP update of posterior (0)
As faster alternative to propagating the posterior 6 we do MAP updates with the prior po())
II, NO00; ( o and the approximate posterior 4: (0) = II N00000/;0 = =²0g where 8 S 1 is hyperparameter controlling the variance of (0). Since fixed may not capture the true parameters variance, using Bayesian method (see Section 3.4. preferred but comes at a high computational cost (see Appendix forr discussion). The MAP update is given by (see Appendix E.2 for derivations) finding a minimum of the following proximal objective
(0) = 21+))))))000)) + + E [0-00
(17)
(T.)
where the regularization target for the parameter dimension given by
(1x) 7f 02 (1 )/)
(18)
and the per-parameter learning rate given (assuming that ag the base SGD learning rate)
((()) (++))++-(22))
(19)
Linearising (0) around (7x) and optimizing 17 for 0 leads to (see Appendix E.2
O = 0,(,) ā.(x) V To
(20)
For ? = 1, we recover the ordinary SGD update, while the values of < 1 move the starting point of 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 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.

---

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 (33) of % connects to the notion of parameters utility from [14], but this quantity is used to prevent catastrophic forgetting by decreasing learning rate for high "/- 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 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 piece-wise stationary setting, for example change in the location of camera generating a stream of images, or a hard update to the learner's '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 5 [0] and plasticity loss [[3] 38 34]. At the other extreme,
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 13 where the reward function or transition dynamics could change over time.
Non-stationary online convex optimization. Non-stationary prediction has long history in online convex optimization, where several algorithms have been developed to adapt to changing data [see, 8 [2]
e.g., 25
18 29]. Our approach takes an inspiration from these works by employing 17
21
drift model as, c.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 [7) D. 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 Po plays the same role as the uniform distribution.
Bayesian approaches to non-stationary learning. A standard approach is Variational Continual Learning [1], which focuses on preventing catastrophic forgetting and is an online version of "Bayes By Backpropp (5). This method does not incorporate dynamical parameter drift components. In (5) the authors applied variational inference (VI) on non-stationary data, using the OU-process and Bayesian forgetting, but unlike in 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 y 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 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 | with gradient update on the drift parameter Soft Reset, while other versions show different parameter choices: Soft Reset (K, = 10) is version with 10 updates on the drift parameter, while Soft Reset (K, = 10, K₂ 10) the method of Algorithm B in Appendix E.2 with 10 updates on drift parameter, followed by 10 updates on NN parameters. Bayesian Soft Reset (K, = 10, Kg = 10) a method implemented by Algorithm 2 with 10 updates on drift parameter followed by 10 updates on the mean μr and the variance of (uncertainty) for each NN parameter. Bayesian method performed the best overall but required higher computational complexity (see Appendix [). Unless specified, 2t is shared for all the parameters in each layer (separately for weight and biases).
7

---

Furmited ST
Random Lobol MINIST Data Eifciant
1.0
Random Labal CHAR-110 Memotaaa
Aectig Accursey
(eerne Tokk
0.75
Aeeruge Talk
0.6
M/
0.6
PA
2.4
0.55
160 Tissid
150
203
50
100 Task 11
355
250
100 Tascid
150
200
Online SGD
Hard Roset Strins and Perturo
Sult Reses
1n
Bayes on Soft Reser - 10, K2 10, yorr parameter
Figure 2: Plasticity benchmarks. Left: performance on permuted MNIST. Center: performance random-labe MNIST (data efficient). Right: performance on random-lade labbel CIFAR-10 (memoriza- tion). The x-axis is the task id and the y-axis is the per-task training accuracy
52
Formuted SINST
Random Labul MNIST Data Blicient
1.0
handom Labul 11AAR-10- Memotasl on
Nrioje Took Tox. Moone
88
0.6
0.75
4
2.4
0.2
0.59
50
100
150
203
50
100 Task ad
150
256
100
150
260
Tascid
Tascid
Soft Ruset
Soft Sweet Prosimel Ky" 10, K,
Dayesan Seft Reset Ky 10.K, Ky 10, yor paramater
Soft R,
Dayecaa Soft Rasat 6, 10, Ky-K
Figure 3: Different variants of Soft Resets. Left: performance on permuted MNIST: Center: performance on random-label MNIST (data efficient). Right: performance on random-late CIFAR- 10 (memorization). The x-axis is the task id and the y-axis is the per-task training accuracy 52)
Loss of plasticity. We analyze the performance of 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] MNIST, where either pixels are permuted or the label for each image is randomly chosen. Several papers [34] 33 (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- 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 L.2 init [3]]. which adds L.2 penalty ||0 loll2 to the fixed initialization &n as well as &hrink&Perturb (2). which multiplies each parameter by scalar À < I and adds random Gaussian noise with fixed variance See Appendix H.T for all details. As metrics, we use average per-task online accuracy (52). which is
A.
at
where of are the online accuracies collected on the task 1 via N timesteps, corresponding to the duration of the task. In Figure we also use average accuracy over all T tasks, i.e.
Ar = -
The results are provided in Figure 2 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 L. Init could be viewed an instantiation of our Soft Reset Proximal method optimizing (17) with 7t = 0 at every step, which sub-optimal when there is similarity in the data. Bayesian Soft Reset demonstrates significantly better performance overall, see also discussion below.
We observe that Soft Reset is always better than Hard Reset
In Figure 3 we compare different variants of Soft Reset. We observe that adding more compute for estimating Yt (thus, estimating non- non-stationarity, K, 10) as well as doing more updates on NN
8

---

(24 Miaraan anccoarrreedd
Fapt sgal Aaxion
10 lppr Aandon Label 20
2.00
0.4
Layer
X0801 12-A-10
Figure 4: Left: the minimum encountered Yt for each layer on random-label MNIST and CIFAR-10. Center: the dynamics of y/ the first 20 tasks MNIST. Right: the same CIFAR-10. 10.
parameters (thus, more accurately adapting to non-staionarity, K₂- 10) leads to better performance. All variants of Soft Reset "/ 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.44 We hypothesize this is due to the NN parameters uncertainty estimates or which Bayesian method provide, while others do not, which leads to more accurate drift model estimation, since uncertainty is used in this update (TO). But, this approach comes at a higher computational cost, see Appendix [1] In Appendix K we provide ablations of the structure of the drift model, 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 %/ for each layer, which highlights the maximum amount of resets. Figure4p.. shows 7/2 as function of f, and suggests that 7/t aggressively decreases at task boundaries (red dashed lines). The range of values of 7/2 depends on the task and on the layer, see Figure 4p. Overall, " changes more aggressively for long duration (memorization) random-label CIFAR-TO 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 given size. The con-stationarity is controlled by permuting the patches (not pixels). Figure shows the minimum encountered Y for each layer for different patch sizes. As the patch size increases and the problem becomes more stationary, the range of values for y is less aggressive. See Appendix K.3 for more detailed results.
Impact of non-stationarity. We consider 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 [5., 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 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 [55 on the Hopper-v5 and Humanoid-v4 GYM 6 environments (from Brax (15). To reproduce the setting from [43] we control the off policycesss of the agent by setting the off-policy ratio M such that for every 128 environment steps, we do 128M gradient steps with batch size of 256 on the replay buffer. As basclines we consider ordinary SAC, hard-coded Hard Reset where we reset all the parameters K times throughout training (every 200000 steps), while keeping the replay buffer fixed (similarly to 43) We employ our Soft Reser method as follows. After we have collected fresh data from the environment, we do one gradient update on % (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 0, (7x) and we employ the update rule (47) where the regularization 0, (7t) kept constant for all the off-policy gradient updates on the replay buffer. See Appendix H.3 for more details.

---

(40 Permatte Path MesT
exeo
csso
Figure 5: (a) the x-axis denotes the layer, the y-axis denotes the minimum encountered " for each convolutional and colly-connettdd layer when trained on permuted Patches MNIST, color is the patch size. The impact of -on-stationarity on performance on random-label MNIST of Online SGD and Hard Reset shown in (b) while the one of Soft Resets is shown in (c). The x-axis denotes the number of epochs each task lasts, while the marker and line styles denote the percentage of random labels within each task, circle (solid) represents 20%, rectangle(dashed) 40%, while rhombus (dashed and dot) 60%. The y-axis denotes the average performance (over seeds) on the stream of 200 tasks.
Humanoid. Replay Rabio - 1.0
Humancld, Rep by Rot 32.0
Humaadds Replay Raclo 128.0
6003
2003
Ruplay
Happor Rotio
32.0
Rup 128.0
6202
4003
0.2
3.4
0.6
0.4
3.0
0.0
0.1
0..
1.0
100
165
Reset
Figure 6: RL results. First row is humanoid, second is hopper. Each column corresponds to different replay ratio. The x-axis is the number of total timesteps, the y-axis the average reward. The shaded area denotes the standard deviation across 3 random seeds and the solid line indicates the mean.
The results are given in Figure As the off- off-policy ratio increases, Soft Reser becomes more efficient than the baselines. This is consistent with our finding in Figure Sp.c, where we showed that the performance of Soft Reset better when the data distribution not changing fast. Figure g in Appendix shows the value of learned " It shows % 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 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 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 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 lot of similarity or on slowly changing distributions. Our findings open the door to variety of exciting directions for future work, such as investigating the connection to continual learning and deepening our theoretical analysis of the proposed approach.
10

---

References
[1] Zaheer Abbas, Rosie Zhao, Joseph Modayil, Adam White, and Marlos C. Machado. Loss of plasticity in continual deep reinforcement learning, 2023.
[2] Jordan T. Ash and Ryan P. Adams. On warm-starting neural network training, 2020.
[3] Gianluca M. Bencomo, Jake C. Snell, and Thomas L.. Griffiths. Implicit maximum posteriori filtering via adaptive optimization, 2023.
[4] 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.0004 000444 2021.
[5] 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 Inter- national Conference Machine Learning. volume 37 of Proceedings of Machine Learning Research, pages 1613-1622, Lille, France, 07-09 Jul 2015. PMLR.
[6] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.
[7] 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.
[8] N. Cesa-Bianchi and G. Lugosi. Prediction, Learning, and Games. Cambridge University Press, Cambridge, 2006.
[9] Peter G. Chang, Gerardo Martnn, Alexander Y Shestopaloff, Matt Jones, and Kevin Murphy. Low-rank extended kalman filtering for online learning of neural networks from streaming data, 2023.
[10] Zhiyuan Chen and Bing Liu. Lifelong machine learning. Synthesis Lectures on Artificial Intelligence and Machine Learning, 12(3):1-207, 2018.
[11] Hugh Dance and Brooks Paige Fast and scalable spike and slab variable selection in high- dimensional gaussian processes, 2022.
[12] Shibhansh Dohare, J. Fernando Bernandez-Garche Qingfeng Lan, Parash Rahman, Ashique Ru- pam Mahmood, and Richard S. Sutton. Loss of plasticity in deep continual learning. Nature, 632:768 774, 2024.
[13] Shibhansh Dohare, Richard S. Sutton, and A. Rupam Mahmood. Continual backprop: Stochastic gradient descent with persistent randomness, 2022.
[14] Mohamed Elsayed and A. Rupam Mahmood. Addressing loss of plasticity and catastrophic forgetting in continual learning, 2024.
[15] 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 e/brax
[16] 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.
[17] A. György, Linder, and G. Lugosi. Efficient tracking of large classes of experts. IFEE Transactions Information Theory, IT50()))77.11 Nov. 2012.

---

Barminng-features-2000 [18] 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.
[19] Tuomas Haamoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- Off-policy maximum entropy deep reinforcement learning with a stochastic actor, 2018.
[20] Raia Hadsell, Dushyant Rao, Andrei A. Rusu, and Razvan Pascanu. Embracing change: Continual learning deep neural networks. Trends in Cognitive Sciences, 24(12): 2020.
[21] Eric Hall and Rebecca Willett. Dynamical models and tracking regret in online convex pro- gramming. In International Conference on Machine Learning, pages 579-587. PMLR, 2013.
[22] 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.
[23] Elad Hazan. Introduction to online convex optimization, 2023.
[24] Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification, 2015.
[25] M. Herbster and M. K. Warmuth. Tracking the best expert. Machine Learning, 32(2):151-178,
1998.
[26] Hemant Ishwaran and J. Sunil Rao. Spike and slab variable selection: Frequentist and bayesian strategies. The Annals of Statistics, 33(2), April 2005.
[27] Matt Jones, Peter Chang, and Kevin Murphy. Bayesian online natural gradient (bong), 2024.
[28] Matt Jones, Tyler R. Scott, and Michael Curtis Mozer. Human-like learning in temporally structured environments. In AAAI Spring Symposia, 2024.
[29] Mikhail Khodak, Maria-Florina Balcan, and Ameet Talwalkar. Adaptive gradient-based meta- learning methods, 2019.
[30] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
[31] James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A. Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Overcoming catas- Demis Hassabis, Claudia Clopath, Dharshan Kumaran, and Raia Hadsell. trophic forgetting in neural networks. Proceedings of the National Academy of Sciences, 114(13:3521-3526, March 2017.
[32] Alex Krizhevsky. Learning multiple layers of features from tiny images. Tech- nical report, University of Toronto, 2009. https://www.co. ca. wes.tooontoo...m www.oooo.ooooooooooooooooo edu/ "kriz/
[33] Aviral Kumar, Rishabh Agarwal, Dibya Ghosh, and Sergey Levine. Implicit under- parameterization inhibits data-efficient deep reinforcement learning, 2021.
[34] Saurabh Kumar, Henrik Marklund, and Benjamin Van Roy. Maintaining plasticity in continual learning via regenerative regularization, 2023.
[35] 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 Learning Representations, 2020.
[36] Zhiqiu Lin, Jia Shi, Deepak Pathak, and Deva Ramanan. The clear benchmark: Continual learning on real-world imagery. In Thirty-fithh Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2021.
[37] Clare Lyle, Mark Rowland, and Will Dabney. Understanding and preventing capacity loss in reinforcement learning, 2022.
12

---

[38] 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.
[39] Clare Lyle, Zeyu Zheng, Evgenii Nikishin, Bernardo Avila Pires, Razvan Pascanu, and Will Dabney. Understanding plasticity in neural networks, 2023.
[40] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning, 2013.
[41] Cuong V. Nguyen, Yingzhen Li, Thang D. Bui, and Richard E. Turner. Variational continual learning. In International Conference on Learning Representations, 2018.
[42] Exgenii Nikishin, Junhyuk Oh, Georg Ostrovski, Clare Lyle, Razvan Pascanu, Will Dabney, and André Barreto. Deep reinforcement learning with plasticity injection, 2023.
[43] Evgenii Nikishin, Max Schwarzer, Pierluca D'Oro, Pierre-Luc Bacon, and Aaron Courville. The primacy bias in deep reinforcement learning, 2022.
[44] Neal Parikh and Stephen Boyd. Proximal algorithms. Found. Trends Optim., 1(3):127-339, jan 2014.
[45] German I. Parisi, Ronald Kemker, Jose L. Part, Christopher Kanan, and Stefan Wermter. Continual lifelong learning with neural networks: A review. Neural Newworks, 113:54-71, 2019.
[46] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.
[47] 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 Engelhaardt, 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-32166 PMLR, 23-29 Jul 2023.
[48] Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evei. The dormant neuron phenomenon in deep reinforcement learning, 2023.
[49] Michalis K. Titsias, Alexandre Galashov, Amal Rannen- Trinen-Triki, Razvan Pascanu, Yee Whye Teh, and Jorg Bornschein. Kalman filter for online classification of non-stationary data, 2023.
[50] G. E. Uhlenbeck and L. S. Ornstein. On the theory of the brownian motion. Phys. Rev., 36:823-808 841, Sep 1930.
[51] Gido M van de Ven and Andreas S Tolias. Three scenarios for continual learning. arXiv preprint 1rXiv:1904.07734 2019.
[52] 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.
[53] 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 Popeseu, 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.
[54] Florian Wenzel, Kevin Roth, Bastiaan Veeling, Jakub Swiętkowski, 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.
[55] Runtian Zhai, Stefan Schroedl, Aram Galstyan, Anoop Kumar, Greg Ver Steeg, and Pradeep Natarajan. Online continual learning for progressive distribution shift (OCL-PDS): A practi- tioner's perspective, 2023.
[56] 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, '03 page 928-935 AAAI Press, 2003.
13

---

A Organization of the appendix
The appendix is organized as follows. In Appendix we describe the Ornstein- Uhiiebee- 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. Appendix D we describe additional practical tricks used for learning the drift model parameters. In Appendix D we describe full Bayesian and maximum a-posteriori (MAP) algorithms to learn both drift parameters and NN parameters. In Appendix we provide the proof for the closed form solution derivation (T2). In Appendix G we briefly review the proximal SGD. In Appendix H we provide experimental details. Appendix we describe computational complexity for all the methods. In Appendix D provide sensitivity analysis for different parameters. In Appendix K we provide additional ablations and experiments showing qualitative behavior of soft resets. In Appendix[[ [ we show 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-Uhlenbed process
We make use that the and written as a continuus Gaussian Markov process with transition density between on R
process [50] defines SDE that can be solved explicitly
p(=) = (2 age (t-a), (1 - 2(t 20
for any pair of times s. Based on this as a drift model for the parameters 0₁ (so 0g is the state m4) we use the conditional density
p(0++1 71) NN (1 )o)))
where 7/ = e and 8f 2 0 corresponds to the learnable discretization time step. In other words, by learning oo online we equivalently learn the amount of a continuous "time shift" of between two consecutive states in the OU process. This essentially models parameter drift since e.g. if if = 1, then 62 = 0 and there is no "time shift" which means that the next state/parameter remains the same the previous one, i.e. 0f 02.
C Other choices of drift model
In this section, we discuss alternative choices of drift model instead of (5)
Independent mean and variance the drift. We consider the drift model where the mean and the variance are not connected, i.e..
p(0r+1 1r. 7e, B,) W1 of ((1+++1100 (1- 211-))m66 33²3)
(21)
where parameters rd E [o, 1] control the mean of the distribution and sy O control the variance. When parameters Bg 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 Py are fixed parameters independent
from 7/ whereas Figure 16 right corresponds to the case when B - / 7/² o i.e., when
the drift model ( We from the results, using drift model (5) leads better performance. In case when 18) 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 the mean of the initialization, we the following drift model
p(0++1 0,, Ar, 3c) = II AN X/0,, (8(2)
for X E [0, 1] and 88 > Similarly to the case of (2T), estimating both parameters X and B] from the data will likely overfit to the noise.
14

---

Arbitrary linear model.
We can use the arbitrary linear model of the form
p(0++1|01, A4, B,) (0(++)1; A,01, Bt),
but estimating the parameters A, and B, has too many degrees of freedom and will certainly overfit.
Gaussian Spike & Slab
We consider Gaussian approximation Spike & Slab 26 prior
p(0++1 10, II, ( + (1 (66((())))))
with binary variables 7h E {0,1}, which is a product of mixtures of Gaussian distributions p(0(+110!) N (N 2 centred around the previous parameter 02 and initializing distri- butions po(Ol...)
(1 =6)") This model, however, implements the mechanism of Hard reset as opposed to the soft ones. Moreover, estimating such model and incorporating it into learning update is more challenging since mixture 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
J ((Be+tl((e(E)) + + T)de (B (1(()t t (cco((tt)
(22)
Using NN initializing distribution. In the drift model (5) we assume that the initial distribution over parameters is given by Po( = II, NN m6: oy²). In practice, we have access to the NN initializer Pinit (0) = II, N 0; o where H6 (for most the NNs). This means that we can replace from 10 by do % where 0/ Panac (0). This means that the term in 22 be replaced by
(Bet1 +coo = ( (Be4)) + 0 / 7 + 2
(23)
where we used the fact that of (1)) 2/0? + 7? )o? Note that in 23 only need know the ratio oz 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 (0) to the specific initialization Og Pinit(0). 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
Po NO :0. o3.0)
(24)
with 0j²² =2=== The parameter v S accounts for the fact that the distribution (f) should have lower than Pinit (0) variance since uses the specific initializaiton from the prior. This modification would imply the following modification on the drift model term (3
p (B +o00.(7))) = = p (1 0 o 7
of 78
(25)
we denote s = -ي_ then we can re-write (25) to be
(Ber (ma()) + -eoo
Ban +0 u/1 - -7 7
(26)
In practice, treat to be constant hyperparameter.
15

---

E
Learning parameters with estimated drift models
In this section, we provide a Bayesian Neural Network algorithm learn the distributions of NN parameters when there is a drift in the data distribution. Moreover, we provide a Maximum a- Posteriori (MAP) like inference algorithm which does not require learn the distributions over parameters, but simply propagates the MAP estimate over these.
E.1 Bayesian Neural Networks algorithm
In this section, we describe algorithm for parameters update based Bayesian Neural Networks (BNN). It is based on the online variational Bayes setting described below.
Let the family of distributions over parameters be
{q(0) (()) I[ N(0:
(0) (0',...,0)))
(27)
which is the family of Gaussian mean trean-field distributions over parameters E RDD Let T
(?..... 7h) be the history of observed parameters of the drift model and S, = the history of observed data. We denote by q:(0) _ (e S,, I't-)) E Q the Gaussian approximate
B. be
posterior at time with mean ft and variance of for every parameter. The approximate predictive
look-ahead prior is given by
@e (0 = 4:((0))((100 (00pp(000,,00000 = (())
(28)
that has parameters
pl(m)) = A/ + (1 -1-74) 72) = = %/ + 20
(29)
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
E Y| [X]]
and the law of total variance says
V[y] = [[==[[V(((]] + V[E[v|]]
In our case, from the drift model 5 we have the conditional distribution
0 0'10 - + (1 7o uub e² (1 ((0;1)
(30)
From 30 we have
E 011 o/ + - -
V10 [00] (1 yo02
From here, we have that the mean is given by
e = 0 [ = 7/ pr + (1 - -
(31)
and the variance is given by
V(0) = [[V[0 +
V10] = (1 (1-7) ( 203 7002
(32)
Now, note that (0)x)) is Gaussian and its parameters are given by E[P'] = + (1 - 72 3m46
from (3T) and by Viei] = (1 )/ + 7/ ?² from (32p Then, for new data B441 at time the approximate predictive log-likelilhn likelihood equals to
log (B = log f p(B (()()))))
We are looking for a new approximate posterior (0) such that (0) = arg min KLL (B00)))(((((+==1)))0
(33)
16

---

The optimization problem 33 can be written as minimization of the following loss
n0) = By [[++1(0] + IKL [g(0) (0)%)),
(34)
Using the fact that we are looking for a member E ② from 27 we can write the objective (34) as
Felm, a, 7) = E.-,((o.)) (u+1lu + KLL (0()) Ige (0))))
where we used the reparametcrisation trick for the loss term. We now expand the regularization term to get
Σ (wl (T)))² o02
1 log or²²
ax = -B-
(35)
26?2(Tx)
2
Since the posterior variance of NN parameters may become small, the optimization of (35) may become numerically unstable due division by (x4) was shown [54] that using small temperature on the prior led to better empirical results when using Bayesian Neural Networks, phenomenon known as cold posterior. Here, we define a temperature per-parameter, i.e., X > 0 for every time-step t, such that the objective above becomes
Felp, no (X1,) E.-Aioon)
Σ
X (n² - ((ee)) + 0²2 2072(nn)
1 log o52
(36)
As said above, common to use the same temperature X X for all the parameters. In this work, we propose the specific choice of the temperature to be
X Xo??
(37)
where À > 0 is some globally chosen temperature parameter. This leads to the following objective
+ 2 Felu, o, ni (r/b)) = Ee-N(O.)) [co1)
Σ rl [ive - pol (x1)) 12 - o²²² (7t) log g2)
(38)
where the quantity rf is defined as
(39)
8 (T))
72 o +(1- =0²
which represents the relative change in the posterior variance due to the drift. In the exact stationary case, when od 1, this ratio is rf while for of 1, since typically of? oj² we have rf 1. This means that in the non-stationary case, the strength of the regularization in (38) in favor of the
data term (o:,) ( + J. 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.
Be-N((::1) [Zetilu + co7]
Σ
Ze (u + Evo)
(40)
M
To find new parameters, 44+1 and d+t1, we let M4+1.0 (le and 0+11.0 5x(%) where corresponding quantities are defined in (29) and perform multiple K updates on (38)
Mr.1841 = Pr+1k ag F 01 To (r(b),
0441641 0 a fa Fe 041, n, (r1))),
where Or and 0r are corresponding learning rates.
Practical considerations.
As prior distribution, we use
(0) II _V(000,00)))
(41)
17

---

with some of > O. In practice, we often have access to the NN library with pre-defined o.- We allow more flexbility and define
og vā Us with some constant v E (0, 1]. allowing to scale the NN library variance.
(42)
When using Bayesian Neural Networks (BNN), one question is to how to initialize the variational posterior go (0). We found that using specific initialization 0g Pinat(() as the mean of the variational
posterior (0), where Pinit (0) is the NN initializing distribution, worked significantly better than using 0 (which is the mean of po (0)). Therefore, propose initialize (o)
(0) = II N(O 0g ommtt" ?),
where 8g Pinit (0) and of 2 is defined as
inst
o 2 = n²= =²o²²
(43)
and of is the variance of the prior distribution po( see 41 Here, (0, 1] is constant allowing us scale down the initial variance of (o( with respect to the prior variance oy-
The full algorithm 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 with the prior po(O)
NV p/; og and the posterior (0) = 11, NV 6/: s 2 where S 1 is hyperparam- eter controlling the variance of of the posterior (e(0). Since fixed 8 may not capture the true parameters variance, using Bayesian method (see Appendix E.I is preferred but comes at a high computational cost. Instead of Bayesian update (33). we consider maximum a-posteriori (MAP) update
max log (Br+ + log qe (fl:a).
with qe (6) given by (28) Using the definition of qe (0) )Y)) we get the following problem
max ( - Σ (Eu (x)))
(44)
20 (1))
where similarly to (6). we use per- per-parameter temperature X 2 0. We choose temperature to be equal
= oo?
where À is some constant. Such choice of temperature is motivated by the same logic in (37 is constant multiplied by the posterior variance of²² soof² With such choice of temperature, maximizing (44) is equivalent to minimizing
Get1 (0:1(0;A) (0) +
Eラ=1
10 10-00
(45)
(7())
where the regularization target for the dimension is
(1()x) = 0% (1 )nb
(46)
and the constant (7) is given by
(1())) ((7/)2 + 1-()))
We can perform K gradient updates
(45) with a learning rate a starting from 0441 = =c(1)),
gh+11 = o. a)(%)) o VoGr. ()) A),
(47)
where the vector-valued learning rate (7))))) is given by
ox(7.) = a. = or (62))² 1 1- 1-
(48)
s2
with ay the base learning rate. Note that doing one update is equivalent to modified SGD method
200
Doing multiple updates on (47) allows us to perform multiple computations on the same data. The
corresponding algorithm is given in Algorithm
18

---

Algorithm 2 Bayesian Soft-Resel algorithm
Input:
Data-stream Sy = {Br }?-
Global temperature parameter À 0.
Learning rate for the mean dn and for the standard deviation &a
Number of gradient updates Kg to be applied l and
Number of Monte-Carlo samples Mg for estimating lr and o in (40)
Number of gradient updates Ky drift parameter " in (10)
Number of Monte-Carlo samples M, to estimate ? in Learning rate 7, for drift parameters
Initial drift parameters "Yo for every iteration. Initialization parameters:
Neural Network initial variance for every parameter o6² coming from standard NN library
Initial prior variance rescaling // E (0, 1].
Prior variance of v²002 given by
42
Initial NN parameters 0g ~ Pinse (0)
Initial posterior variance rescaling 10
Initial posterior variance o'sst" 2 T o given by 43
Initialization:
Initialize zero-mean prior distribution
41 po( - -1,,000,,0 Initialize posterior qo (0) 11, (10); 00 of omst?? ?)
for step - 0, ,, do
Receive batch B++1
Current posterior is given by gi N(O; Me. Estimating the drift
Initialize drift parameter y8 70-
Compute drifted prior mean jud (7x) 7d it assuming zero-mean prior ng 0
Compute drifted prior variance 2/2 (74) = 7 22 + - 200 for & = O..... Ky do Sample Gaussian noise E; N(e;; 0, 1). l.... Mr Mr
.++1 r 77, V. log 1 M,
Σ P(B (xx) - c, o. (nt))
end for Updating variational posterior
Let p441,0 71 1// 0441,0 2 + +1 2 o2,
29
of² 120:"+(1-)) 2 ooz
Let o o?(Te)
using 39
for k k= 0,... Kg do
Using the definition of Fr from 388 do
Mc-1,++1 Me+1A a Fi (1t l,k 0 Tt, (r();,))
0441k41 = 0+11 ag F ( 0441. no, (r))::)) end for end for
F Proof of linearisation
Interpretation of Yr- By linearising log p(B(+110) around μrs we can simplify 8 to get
F(x:) = (1+ o pe + (1 - 7e) o uo)" gt+1 - 0.5(a?((x) 91+1)² g++1
ト (nd 7,002
where denotes elementwise product, 9r V is the negative gradient of the loss evaluated at d4 and we added the Ez-penalty 1 - 7 to take into account the initialization.
Proof. We assume that the following linearisation correct
log (Bg + g (0 pc).
where
9++1 = - Vologp(B.+1le mi) voCrz1(4)
19

---

Algorithm 3 Proximal Soft-Resel algoritm
Input:
Data-stream Sy = {B
Neural Network (NN) initializing distribution Pinat (f) and specific initialization 80 Panit (0) Learning rate a, for parameters and 77, for drift parameters
Number of gradient updates K, on drift parameter 7
Number of gradient updates Kg on NN parameters
Global temperature parameter À 2 0
___ NN initial standard deviation (STD) scaling v S 1 (see (24) and ratio
T for step / = 0, 1, ,....,T do
Receive batch B Initialize drift parameter 7: =1 for step 0, 1...,.., 2, ................................ K, do
Sample 8/ Pinse
Stochastic update 22 on drift parameter using specific initialization (25)
s+1 = 77 V
log pB (1 7000 + 8 vvvi== 11=2 + 7/ 82
7=
end for
Initialize 40+ = ((( K) with 44 and use an = a ((68)² + --(221 with 48
for step 0, 1, 2, : K do
05+ f++1 av(%) (0
end for end for
Then, we have
p(B 10) ~ p(Be 114)) exp=)i(00mm)
Let's write the integral from
log
p(B 10) exp
do
((2T) Ex(:o)
log / p(B,+ 11x) 1(4) expeñ 10 Uo) exp 100 ()))) ((:00- d8
V(2T) [2((1))
log (Bpp((((lu)))) + log / expet1 ni) exp 1: -00-0 d0
1
(2T))1I,(()))
Consider only the exp term inside the integral:
1 -0-
o mi)
1 - (n)) + 2. (m))H((7))
1
0 0 g 14
2 uel (x)) (70(0))))))
ET 0 + 8T (ET (nd) + 90+1) - s.
2 2 (1)) TE, (m)mm))
Hu(1))
1 (0"2)-0 20 (m))e())) + 9c+1))
g =u((1)) (x))p(())
2
Let's focus on this term
1
10T 102-10-20000 (1) + =
2
2 (0"x, 10
28 ,
1 (0 ),)0
20 + ) ((7)) b(x)) b(1x))
2
(0'22 0 - -0 b(7) 6( 3 x, b(1)) + 2 (0)))) Tx, (6)))))))))))))))))))))
1 (0 - b(x))) (0 b(7)) + _ (,))),, (
2
20

---

where
b( = Ex(m)) [ (mo)met))) + 9411]
Therefore, the integral could written as
log / expoč11 (0-w4) exp =(0--((m))) (xx00000"""""=""0x0000000=0000 (9n)(0-ul1r)) dB
1
((2T) [Ea))
[ b(x)) ",,((),))))))))))))))))))))))" + log / exp 2 ( d0
1
F+1 "1((11) Ex, (76)4e(7))=
) PEE(()))
2 b(t)) ( " 924114 1 m =x (x)m()
Now, only keep the terms depending %
b(%)) 2, (1:)6())
(
1
Σ (1i)a(())) m+12l* En(7)) [Σm (3) (1)) + +c++1]
Me (e()6) (1))m()))
2
1
[u(t)) (x(x))22 (T))u()))) + (t4))(()))
g 21 (7))9444
2 uu (71) (7)))((()))
2
( + 2 E
1 (70))441
Since Ex(r)) diaglo? o or? + (1 72) og), we recover
0 o Me + (1 11) a no) + 3044 (lo? o r2 + (1 (1--ni 72) oî) o g++1)
Now, we add an 12-penalty all - 7r,ol? and we get
1 20 (lo o or (1- og) g4+1)
A
F(T.) = (1-
2 - 70,0l
Let's take the gradient wrt y²' we get that the j-the component of the gradient given by
(VF()))) 1441 (ut vb) + 96+1 (oo o002 (1)11144))) 1(TH 7(,0)
(Hd pb) + Avt.0 gt X> + g4zz (lod" of?) o (
(49)
We have then,
b. s (nr) mm)) + X
Ø4. 200² o
In case when ?t are shared for a subset of K parameters, indexes by the index set JK = (f1...., then, sum the gradients 449 and get
Th mb) + Avc. v [A + 9641 (108² - o,2) 00++1)+ 0, JEJk
this gives us
TT (āg (1o)= + KNy0
7t (ō++1 (ōo² 2 0. 2)) 1)T g T) where ā is vector defined on the index set JK, i.c., ā (a (aj...... ajx):
KX
21

---

G Proximal SGD
Each step of online SGD can be seen terms of regularized minimization problem referred as the proximal form 44
0+11 arg ming (0) +
1 2n; 10 00
(50)
In general, we cannot solve 50 directly, so we consider a Taylor expansion of Êe+1 around t,. giving
0 arg ming VoLe+1 0,)
2n, 110 - 0|1
(51)
Here we see the role of ar >> O as both enforcing that the Taylor expansion around Og is accurate, and regularising 0.+1 towards the old parameters Og (hence ensuring that the learning from past data is not forgotten). Solving (5T) naturally leads to the well known SGD update:
0c+1
a₁ Vo (0+);
where at now also be interpreted 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 x 24 and for CIFAR-10 to produce images of size 28 x 28. We normalize images to be within [0, 1] range by dividing by 255. On top of that, we consider permuted MNIST task with similar training ragime as in [34] we consider a subset of 10000 images, with batch size 16 and each task is one epoch. As backbone,
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, used simple 4 layer convolutional neural network with 2 fully connected layers the end.
Metries We use online accuracy first metric with results reported in Appendix K Moreover use per-task Average Online Accuracy which is
A
Σ d
(52)
N i=1
where af are the online accuracies collected on the task & 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 Alle- full2 term to each Online SGD update where the regularization strength is hyperparameter. Shrink & Perturb applies the transformation X0, oe, e (e( to each parameter before the gradient update. The hyperparameters are À and a.
Soft Reset corresponds to one update (0)) starting from using Monte Carlo estimate. We always use 1 Monte Carlo estimate for updating Yt as we found that it worked well in practice on these tasks. The hyperparameters of the method - og initial variance of the prior, which we set to be equal to 22 N where N is the width of the hidden layer and v is constant (hyperparameter), It always equals to v = 0.1. On top of that the second hyperparameter 8, such that a; 800, which controls the relative decrease of the constant posterior variance. This is the hyperparameter over which we sweep
22

---

over. Another hyperparameter is the learning rate for learning "t. For Soft Reset Proximal, we also have a proximal coefficient regularization constant X. 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 0g and we do Monte Carlo sample for each ELBO update.
Hyper parameters selection and evaluation For all the experiments, we run sweep over the hyperparameters. We select the best hyperparameters based on the smallest cumulative error (sum of all 1 af throughout the training). We then report the mean and the standard deviation across 3 seeds in all the plots.
Hyperparameter ranges Learning rate a which is used to update parameters, for all the methods,
1.0). The Xinit parameter in is selected from (1e-4, 5 1e-3, 3, 1e- 2, 1.2 Init, is selected from
-e, 1,-3, 5e.5e-3, -3, le 1,-4, 5e-5e-4, le- }. For S&P, the shrink
5, 5e - 5, le 6, 5e 6, le
7, 7. 7.
1,1e--8,5c. 8,5c 155--,, 8, 1e- -,
parameter À is selected from {1.0, 0.99999, 0.9999, 0.999, 0.99, the perturbation parameter is from {le - 1e 2, 3, le 4. 5, le 63. As noise distribution, we use the Neural Network initial distribution. For Soft Resets, the learning rate for y is selected from (0.5, 0.1, 0.05, 0.01, 0.005, 0.001, 0.0005, 0.0001), ) the constant 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.00.. 0.05. On top of that for the Bayesian method, we always use v (see Algorithm 2 equal to v = 0.05 and 0.9, i.e. the posterior is always slightly smaller than the prior. Finally for the Bayesian method
0.3,0.0..2... 0.2, 0.13. and
the same is true for the temperature in the Bayesian method (38) Initial prior std rescaling
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 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 0.9 and temperature À = 0.01. Detailed results are given in Figure 7
and n7 0.01. Bayesian method uses
H.3 Reinforcement learning experiments
We conduct experiments in the RL enviromments. We take the canonical implementation of Soft-Actor Critic(SAC) from Brax [13] 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 MLP networks to parameterize Q-function (see Brax 15) for more details. To employ Soft Reset, we do the following. After have collected a chunk of data (128) time-steps, we do one update (T)) on 7 starting from 1 at every update of nt. where Y is shared for all the parameters within each layer of 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 Yts we compute (c(7)e) and (x(T)), see Section 3.5 After that, we employ the proximal objective (45) with a fixed regularization target 0, (7t). Concretely, we use the update rule (47) where for each update the gradient 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 0.9, 0.95, 0.97, 1.0). We used v for all the experiments. For each experiment, we used 3 hours of the A100 GPU with 40 Gb of memory.
23

---

10
ADC-fonteo 01011
Figure 7: Non-stationarity impact. The x-axis denotes task id, each column denotes the duration, whereas row denotes the amount of label noise. Each color denotes the method studied. The y-axis denotes average over 3 seeds online accuracy.
Roplay Radio 1.0. for poicy
Peplay
Deplay 128.0
2.0
1.0
0.5
0.6
Layer D: D-Weights
Layer Stoee
Layer W-Weghhs
Layer Dram
Layer W:uiigti Weights
0.0
0.5
08
0.4
0.5
3.0
0.2
1.0
for
Rostay
Replay
0.32
2.0
0.8
0.05
Layer # Weights
Brams
Weishts
0.7
0.4
0.93
Layer
0.4
0.2
Keplay Rat
1.5 for 0F
140
Andlay Rabo
32.0, 02
Tel
Anp RADS 126. yor
1.00
0.98
0.8
0.5
0.97
Layer
Blaess
Layer Weights
0.2
Figure 8: Visualization the Yr dynamics for the run on Humanoid environment. Each column corresponds to the replay ratio studied. First row denotes the Y for the policy æ. The second and the third rows denote the " for the two Q-functions.
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.
My number of Monte Carlo samples for the drift model

---

Mg 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.
Method SGD
Comp. cost
Memory
(
O(P)
Soft resets y per layer
O(K,M,S-
O (M, + 1)P)
Soft resets 7 per param.
O(K,M,5
O (M, + I)P)
Soft resets 7 per layer proximal (Kg iters)
MKK,M,5KKKp
O +M,, 1.,+700 1]P)
Soft resets per param. + proximal (Kg iters)
M,K,M 2MgKo5) OKK,M, 2MgKas
O(P (M, [M,+10 I)P)
Bayesian Soft Reset Proximal (Ke iters) y per layer
P + (M, +2)P)
Bayesian Soft Reset Proximal (Kg iters) y per param.
P(P + (M, +2)
Table 1: Comparison of methods, computational cost, and memory requirements
The general theoretical cost of all the proposed approaches is given in Table [ In practice. for all the experiments, we assume that M,, for Soft Reset, K-- = 10 and Kg = 1 for Soft Reser with more computation. On top of that, for Soft Reset proximal and Bayesian methods, we used K, 10 and K₃ 10. Table quantifying the complexity of all the methods from Figure
and Mo Moreover, we used K = and K₂ =
Method SGD
Comp. cost
Memory
O(S)
O[P)
Soft resets " per layer
0(2S)
O 2P)
Soft resets per param.
0(25)
0(3P)
Soft resets y per layer proximal (Ko = 10 iters)
0(205)
O(I + 2P)
Soft resets y per param. + proximal (Ko iters)
0(205)
0(3P)
Bayesian Soft Reset Proximal (Kg iters) 7 per layer
O(305)
P(L 3P)
Bayesian Soft Reset Proximal (Kg iters) per param.
((305)
P(4)
Table 2: Comparison of methods, computational cost, and memory requirements for methods in
Figure 2
The complexity O(2S) of Soft Resets comes from one update on drift parameter and one updat eon NN parameters. The memory complexity requires storing (L) parameters gamma (one for each layer), parameters 0g with O(P) and sampled parameters for drift model update which requires (PP).
Note that as Figure suggests, is beneficial to spend more computational cost on optimizing gamma and doing multiple updates on parameters. However, the cheapest version of our method Soft Resets still leads to good performance as indicated in Figure
The complexity of soft resets in reinforcement learning setting requires only gradient update y after each new chunk of fresh data from the environment. SAC, we do G gradient updates on parameters for every new chunk of data. Assuming that complexity of one gradient update in SAC is OSS), soft reset only requires doing one additional gradient update to fit y parameter.
Method SAC
Comp. cost O(GS)
Memory O(P)
Soft resets y per layer
O(S + GS)
OLL ML+22)) 2P)
Table 3: Comparison of methods, computational cost, and memory requirements for methods in for RI.
The computation complexity of Soft Reset in Reinforcement Learning is marginally higher than SAC but Icads to better empirical performance in highly off- off-policy regime, see Appendix H.3
25

---

Nammlted RIMST
Random Lebel MRIST Duta efident
Ransem Label CPAR-10 mameraaiionn
0.95
0.86
0.82
0.75
0.92
0.83
0.83
0.82
0.50
0:90
Sont nese N, -10
Dayosan Sont fleeel K, Ko-30
Son Sarne X 110
mayonae San tloee &y K0-00
1111 ner prommeer
Son maat Prosmaa -10 -10
Banisaa Saft Reseet 10 & 10 so sarandor
Soft
K
Sott Rraan Prosmmel,
Exyeeae flese K,
Daypater Sat faaeeedy
tradeoff.. The x-axis indicates the method going from the cheapest (left) to the most expensive (right). See Table 2 for complexity analysis. The y-axis the average performance on all the tasks across the stream.
Figure 9:
J Sensitivity analysis
We study the sensitivity of Soft Resets where y 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 initial prior standard deviation rescaling, and s posterior
for the drift parameter, 7hy, as well standard deviation rescaling parameter.
On top of that, conduct the sensitivity analysis of 1.2 Init 34 and Shrink&Perturb 2 methods. The x-axis each plot denotes one of the studied hyperparameters, whereas y-axis 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 second hyperparameter which is studied, if available. In the title each plot, we write hyperparameters which are fixed. The analysis 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 17,, For each method, there exists a good value of this parameter and performance is sensitive to This makes sense since this parameter directly impacts how learn the drift model.
The performance of Soft Resets is robust with respect to the posterior standard deviation scaling s parameter long it is 0.5. For 0.5, the performance degrades. This parameter is defined from o; = soo and affects relative increase in learning rate given by -+((->7)/57 which could be ill-behaved for small s.
We also study the sensitivity of the baseline methods. We find that 1.2 Init parameter X, which is a penalty term for 0 - 0ol12. In fact, Figure left shows that there is only 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. 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 a provided that o < 0.001.
very sensitive the
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 Soft Resets, K, = 10, Figure shows sensitivity of Soft Resets, K,- = 10, Kg 10, Figure 14 shows sensitivity of Bayesian Soft Resets, K2 10, Ko
with 7t per layer, Figure 15 shows sensitivity of Bayesian Soft Resets. Ky, 10. Ko 10 with % per parameter.
26

---

Senristy Sww.errrrrrerr staddrr cotegannGd
Pastrii strdee Barionssiinng
2.
honworrasde deveter scuiir
Figure 10: Soft Reset, sensitivity analysis of performance with respect to the hyperparameters on data-efficient random-label MNIST. The x-axis denotes the studied hyperparameter, whereas the y-axis denotes the average performance across the tasks. The standard deviation is computed over 3 random seeds. The color indicates additional studied hyperparameter. (Left) shows sensitivity analysis where the x-axis is the posterior standard deviation scaling and the color indicates the drift model learning rate 7 (Right) shows sensitivity of Soft Reset where the x-axis the posterior standard deviation scaling and the color indicates initial prior standard deviation scaling v.
Sem by OPNEST. Lowrirereteeeeo
Senst and
Learine
0.45
0.22
0,90
03009
0.00
36-5
35t5
20-7
304
Figure 11: L2 Init and Shrink&Perturb sensitivity analysis of performance with respect to the hyperparameters on data-efficient random-label MNIST. The x-axis denotes the studied hyperparame- ter, whereas the y-axis denotes the average performance across the tasks. The standard deviation computed over 3 random seeds. The color optionally indicates additional studied hyperparameter. (Left) shows sensitivity of L2 Init with respect to the L2 penalty regularization cost applied to ||0 - la term. We do not use an additional hyperparameter, therefore there is only one color. method where the x-axis is the perturb parameter a
(Right) shows sensitivity of Shrink& while the color indicates the shrink parameter X.
Now ahboainn
Postorte stodded Aediiossnnnn
Figure 12: Soft Reset, K, 10, sensitivity analysis of performance with respect to the hyperparam- eters on data-efficient uundom-labe MNIST. The x-axis denotes the studied hyperparameter, whereas the y-axis denotes the average performance across the tasks. The standard deviation is computed over random seeds. The color indicates additional studied hyperparameter. (Left) shows sensitivity analysis where the x-axis is the posterior standard deviation scaling and the color indicates the drift model learning rate 7yy- (Right) shows sensitivity analysis where the x-axis the posterior standard deviation scaling s and the color indicates initial prior standard deviation scaling v.
27

---

Hor edacting
Patari Patwiiiarrded iarded
21
honworratdis cevisrracuir
0.
Figure 13: Soft Reset, K, 10, Kg 10, sensitivity analysis of performance with respect to the hyperparameters on data-efficient random-label MNIST. The x-axis denotes the studied hyperparameter, whereas the y-axis denotes the average performance across the tasks. The standard deviation is computed over 3 random seeds. The color indicates additional studied hyperparameter (Left) shows sensitivity analysis where the x-axis is the posterior standard deviation scaling and the color indicates the drift model learning rate 77, (Right) shows sensitivity analysis where the x-axis is the posterior standard deviation scaling and the color indicates initial prior standard deviation scaling v.
ypur
Sexulivily Eoyeslla ScB-hasst you layar SHRST) Intta pt rusndd og 0.85
0.501
5,3005
0355
0.2
22 234d 10 217
Figure 14: Bayesian Soft Reset, K,, = 10, Kg 10 with % per layer, sensitivity analysis of performance with respect to the hyperparameters on data-efficient random-label MNIST. The x-axis denotes the studied hyperparameter, whereas the y-axis denotes the average performance across the tasks. The standard deviation is computed over 3 random seeds. The color indicates additional studied hyperparameter. (Left) shows sensitivity analysis where the x-axis is the prior standard deviation initial scaling y and the color indicates the drift model learning rate 7/,, (Right) shows sensitivity analysis where the x-axis is the KL divergence coefficient. À while the color indicates the learning rate 177--
ypur peramiter
3.5
6. Leariin
Leareeg
03505
0332
031
0.2
TeB Mendard devictce
Figure 15: Bayesian Soft Reset, K, = 10, = 10 with Y per parameter, sensitivity analysis of performance with respect to the hyperparameters on data-efficient random-label MNIST. The x-axis denotes the studied hyperparameter, whereas the y-axis denotes the average performance across the tasks. The standard deviation is computed over 3 random seeds. The color indicates additional studied hyperparameter. (Left) shows sensitivity analysis where the x-axis is the prior standard deviation initial scaling v and the color indicates the drift model learning rate 77,- (Right) shows sensitivity analysis where the x-axis the KL divergence coefficient À while the color indicates the learning rate 1/--
K Qualitative behavior of soft resets and additional results on Plasticity benchmarks
K.1
Perfect Soft Resets
28 To understand the impact of drift model (5). we study the data efficient random-label MNIST setting where task boundaries are known. We run Online SGD, Hard Reset which resets all parameters at task

---

Puffect Soft with conslart.
Porfect Solteleset was nigher i.
0.6
Tock Accursty
Online 3DD
lakk
Resat Jonty last
Sont
Soft Resed ty
02-
Soft Reset Ty
Son Reset (y 0.00
Sont Seset 1y 0.00
150
150
250
50
1.50
200
Teet V1
Tock 15
Figure 16: Perfect soft-resets on data-efficient random-label MNIST. Left, Soft Reset method does not use higher learning rate when y < 1. Right, Soft Reset increases the learning rate when 7 < 1, see (19). The x-axis represents task id, whereas the y-axis is the average training accuracy on the task.
Figure 17: Behaviour of 7/ for different layers on random-label MNIST (data efficient) for the first
20 tasks.
boundaries, and Hard Reset (only last) which resets only the last layer. We use Soft Reset method (20)
where ") 1 all the time and becomes Yt = "e (with manually chosen Dt) at task boundaries. We
consider constant learning rate (x{(1)) and increasing learning rate (19) at task boundary for Soft
Reset. On top of that, we run Soft Reser method unaware of task boundaries which learns %/- We report Average training task accuracy metric in Figure [6] See Appendix for details. The results suggest that with the appropriate choice of 2r Soft Reset is much more efficient than Hard Reset and the effect becomes stronger if the learning rate ar (7e()) increases. We also see that Soft Reset could learn an appropriate 7/ without the knowledge of task boundary.
K.2 Qualitative Behaviour on Soft Resets on random-labe tasks.
We observe what values of "t we get as we train Soft Reset method on random-label MNIST (data- efficient) 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 Soft Resets on permuted patches of MNIST.
We consider a version permuted MNIST where instead of permuting all the pixels, we permute patches of pixels with a patch size varying from to 14. The patch size of corresponds to permututed MNIST and therefore the most non non-stationary case, while patch size 14 corresponds to least non- stationary case. We convolutional Neural Network in this case. In Figure [9 report the behavior of y for different convolutional and fully connected layers on first few tasks.
2350
Figure 18: Behaviour of 7/ for different layers on random-label CIFAR-10 (memorization) for the first 20 tasks.
29

---

Figure 19: Behaviour of 7/e for different layers on permuted MNIST
K.4
Bayesian method is better than non-Bayesian
As discussed in Section 5 found that in practice Soft Reset and Soft Reset Proximal where is learned per paramereeetr, did not perform well on the plasticity benchmarks. However, the Bayesian variant described in Section EI_ 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 "e specified per parameter, do not perform well. The fact that the Bayesian method performs better here suggests that it is important to have good uncertainty estimate of for the update (TO) on %e. When, however, we regularize Yr 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 -ay-Baaen methods, assume that the uncertainty fixed, given by hyperparameter assumption which would not always hold in practice.
Random Label MNIST Data Efficient
Random Label CIFAR-10 10 Memarization
Average Task Accuracy
0.8
Aveerage Task Accuracy
0.7
Hard Reset
0.6
Soft Reset
0.5
Soft Reset Proximal (y per parameteri
Soft Reset (y per parameter)
Beyesian
Proximal
0.3
Bayeslan Soft Reset Proximal (y parameter)
0
50
100
150
200
50
100
150
200
Task id
Task
Figure 20: Performance of per-parameter methods. By default, we use per layer and if specified, we use 7 per parameter.
K.5
Qualitative behavior of soft resets
In this section, zoom-in in in the data-efficient experiment random-label MNIST. We Soft Reset Proximal (7 per layer) method with separate 7 for layer (different for each weight and for each bias) and for 20 tasks on random-label MNIST. In Figure we show the online accuracy learn over this sequence of tasks. In Figure 22 we visualize the dynamics of parameters y for each layer. First of all, we that 7t seems accurately capture the task boundaries. Second, see that the amount by which each 7/2 changes depends on the parameter type weights versus biases, and depends on the layer. The architecture this setting starts form lênear and goes up linear4, which represent the 4 MLP hidden layers with last layer linear4.
K.6 Impact of specific initialization
In this section, we study the impact of using specific initialization Ag (P) in Pu()) as discussed in Appendix D Using the specific initialization in Soft Resets leads to fixing the mean of the Pu(O) to
be 0g. see 24 This, turn, leads to the predictive distribution (25). In case when we are not using specific initialization Og the mean of Po()) is O and the predictive distribution is given by (23). To understand the impact this design decision, we conduct an experiment on random label MNIST with Soft Reset, where either use the specific initialization not. For each of the variants,
30

---

0.2
23630
45630
65620
B5630
100000
Time-step
Figure 21: Visualization of accuracy when trained data efficient random-label MNIST task. The dashed red lines correspond to a task boundary.
Figure 22: Visualization of y and task boundaries data-efficient Random-label MNIST.
do a hyperparameters sweep. The results are given in Figure 23 We see that both variants perform similarly.
0.4
0.3
Soft Res no specific init Pş(0) Soft Reset with specific init in Pule)
50
100
150
200
Figure 23: Impact of specific initialization 6g mean of pa (0) in Soft Resets. The x-axis represents task id. The y-axis represents the average task accuracy with standard deviation computed over random seeds. The task is random label MNIST data efficient.
L Toy illustrative example for SGD underperformance in the non-stationary regime
Illustrative example of SGD on non-stationary stream. We consider toy problem of tracking changing mean value. Let the observations in the stream S, follow ye = + ae, where ~ N NNO 1). 0.01. Every 50 timesteps the mean μ4 switches from -2 2. We fit 3-layer MLP with layer sizes (10, 5. 1) and ReLU activations, using SGD with two different choices for the learning rate: 0.15. Moreover, given that we know when switch of the mean happens, we reset (or not reset) all the parameters at every switch as run SGD. Only during the reset, we use different learning rate 3 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 3 = 0.15, where the higher learning rate used during task switch but we do not reset the parameters. We found that it performed the same SGD with 0.05, which highlights the benefit of reset.
0.05 and
31

---

Thue data
SGD (a 0.05) reset (8 0.05 SGD (a= 0.15)
Prodicted
0.05) reset (8 0.15)
100
150
Number timesteps
Figure 24: Non-stationary mean tracking with SGD.
We report the predicted mean Pa for all SGD variants in Figure 24 We see that after the first switch of the mean, the SGD without reset takes more time 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 compared 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 and the impact of the regularization strength induced by the learning rate.
M Using arbitrary drift models
The approach described section 3.5 provides general strategy of incorporating arbitrary Gaussian drift models p(0 wh) (0; (0t whe): (0,; v/e) which induces proximal optimization problem
0.11 arg min E + 2g(0,; v.)
1
||0 e/)11
(53)
The choice of f6;; a and ( (/)) affects the behavior of the estimate 0s+1 from 553 and ulti- mately depends the problem in hand. The objective function of the form (53) was studied in context of online convex optimization in [21], (2)]. where the underlying algorithms estimated the deterministic drift model online. These worked demonstrated improved regret bounds depending model estimation errors. This approach could also be used together with a Bayesian Neural Network (BNN).
32

---

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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long 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 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might 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 faimess.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 complete (and correct) proof?
Answer: [NA]
33

---

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 cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide 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 supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
Answer: [Yes]
Justification:
We disclose the experimental information Experimental and Appendix
sections.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer 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 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 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 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 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 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 submis- sions 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 reproduce that algorithm. (b) If the contribution is primarily
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 way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe 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., 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 instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
34

---

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.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.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 guwdes/ddeeeeee Cooeeeoodeesssseeeeee //nips.com htte://alllllgiiiee///oeeeuaaass ssiniiiiccy 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 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 supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
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 the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, 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: deviation.
We specify that report results with 3 random seeds with mean and standard
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, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall 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).
35

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 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 range (e.g. negative enror rates).
If error bars reported in tables 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
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide information about compute resources required 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 compute required for each of the individual experimental runs well 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 failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics http:p////eerrr....co//////////a
Answer: [Yes]
Justification: Based on our understanding, our work conforms to the every aspect of NeurlPS 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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration 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 no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact 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.
36

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Lieenses for existing assets
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?
Question:
Answer: [Yes]
Justification: We cite the works which introduced the publicly available datasets
Guidelines:
The answer NA means that the paper does 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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
37

---

If this information is not available online, the authors are encouraged to reach out to the asset's '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/moodl 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 details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, 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 involve crowdsourcing 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 NeurlPS 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.
38


Reviews:
[{"review_id": "CCOxpekU3c", "reviewer": "Reviewer_Hmj6", "summary": "This 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.", "questions": "See Weaknesses", "limitations": "Computational cost and hyperparameter sensitivity need to be discussed in the conclusion of the main paper.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The 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.", "weaknesses": "Although 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer Hmj6 for their feedback. Please find our detailed answer below.\n\n> 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.\n\nWe have provided the answer about hyperparameters sensitivity above in the common answer for all the reviewers.\n\n> 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.\n\nWe have provided the answer about computational complexity above in the common answer for all the reviewers.\n\n> 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\nOur choice of Reinforcement Learning experiment is motivated by the Primacy Bias [1] paper, where a significant plasticity loss occurs and where the benefits of parameter resets were observed. In this setting, it was shown that as replay ratio increases, the performance of the algorithm (Soft Actor Critic) could significantly degrade and hard-resetting parameters once in a while, can significantly improve its performance. With this in mind, our Figure 4 demonstrates that Soft Reset can in fact be an even more effective strategy when off-policy ratio increases – it can learn when to reset parameters and by what amount, compared to Hard Resets which require setting these parameters manually.\n\nWe are happy to take further questions on any of the above matters, and we hope that we have been able to address your concerns. If we have done so, we would be grateful if you would consider increasing your review score.\n\n**References**:\n\n[1] The Primacy Bias in Deep Reinforcement Learning, Evgenii Nikishin, Max Schwarzer, Pierluca D'Oro, Pierre-Luc Bacon, Aaron Courville, 2022", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1Egd7ixmN1", "reviewer": "Reviewer_6GY1", "summary": "The 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.", "questions": "How 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?", "limitations": "Authors have adequately addressed the limitations of their work.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "To 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.", "weaknesses": "Some 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer 6GY1 for their positive feedback. Please find our answer below.\n\n> Section 2 - Simplification of equation (6) to (9) requires some explanation [page -5]\n\nThe equation 9 is obtained from 6 as follows. First, we linearise the log-likelihood function $\\log p(y_{t+1} | x_{t+1}, \\theta)$ around $\\mu_t$ which equals to $\\log p(y_{t+1} | x_{t+1}, \\theta) \\sim log p(y_{t+1} | x_{t+1}, \\mu_t) + g_{t+1}^T (\\theta - \\mu_t)$. Then, we notice that $p(y_{t+1} | x_{t+1}, \\theta) = \\exp^{\\log p(y_{t+1} | x_{t+1}, \\theta)}$ which means that $p(y_{t+1} | x_{t+1}, \\theta) \\sim p(y_{t+1} | x_{t+1}, \\mu_t) \\exp^{-g_{t+1}^T \\mu_t} \\exp^{g_{t+1}^T \\theta}$. Then, since $q_{t}^{t+1}(\\theta|S_t,\\Gamma_{t-1},\\gamma_t)$ is a Gaussian, we can write compute the integral in a closed form. After that, we keep only the terms which depend on $\\gamma_{t}$, since we only are interested in the optimization of $\\gamma_t$.\n\nWe will add a full derivation of this expression in the Appendix.\n\n> Section 3 - Some intuition behind the objective (11) [page - 6] would be good to understand it better (especially the second term in that objective)\n\nThe equation in objective 11 is negative Evidence Lower Bound (ELBO) on approximate predictive log likelihood eq 6, if all $\\lambda_i = 1$. It could be derived by considering the integral in the right-hand side of eq. 6, dividing and multiplying it by $q(\\theta)$ and applying Jensen Inequality. The fact that we introduce $\\lambda_i \\neq 1$ means that we introduce a temperature parameter (see [1]) on the prior for each dimension $i$. It was shown empirically [1] that using a temperature in the prior leads to better empirical results. We will add more explicit explanations of how to derive the objective 11 and will add full derivation in the Appendix.\n\n\n> How 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\nOur experiments in Figure 3 and Figure 4 partially answer this question. The difference between Data Efficient and Memorization regimes in Figure 3 is the amount of epochs given to a task – 70 epochs for Data Efficient and 400 for Memorization. There is more non-stationarity in the Data-Efficient setting since it changes more frequently and our methodology remains more effective than baselines.\n\nIn general, we expect our method to be very helpful in scenarios when there are relatively long stationary phases which are succeeded by large changes of the optimal parameters. Moreover, given that we can define $\\gamma_t$ per parameter or per layer, we can have a combination of these scenarios -- some parameters stay stationary whereas other change significantly. In general, such scenario would occur when we have stationary segments which are followed by non-stationary ones. In case, when the non-stationarity is very strong and the data distribution changes at every step, we think that it is less likely that our method will bring an additional benefit over SGD since in this setting, SGD will be constantly refreshed by the noise from the data. SGD, however, will struggle when there is a large change after a long stationary phase. Compared to reset-type algorithms, our method is much more adaptive, and relearns the parameters only when it \"has to\" (i.e., when the data changes sufficiently).\n\nWe see a partial support for this claim in our RL experiments in Figure 4, where as we become more off-policy, we see much more benefit of our method over standard learning approach. Moreover, we also see a benefit of our method over baseline in a more on-policy regime in Figure 4, left, top on Humanoid environment. Since this environment is hard to solve, the baseline agent will see a lot of similar data initially and therefore the corresponding RL algorithm could overfit. Constantly refreshing the parameters is beneficial here. On the other hand, if we study Figure 4, left, bottom which shows performance in Hopper environment, we can see that Soft Resets is less effective. In this setting, since the problem is easy, the baseline method sees a lot of different data constantly and is less prone to overfit.\n\n> Some minor typos in line 109, 154, 190\n\nThank you for the catch, we will make the adjustments.\n\n**References**:\n\n[1] How Good is the Bayes Posterior in Deep Neural Networks Really?, Florian Wenzel, Kevin Roth, Bastiaan S. Veeling, Jakub Świątkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, Sebastian Nowozin, 2020", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MjwueBLtbd", "reviewer": "Reviewer_4koq", "summary": "- 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.", "questions": "Q1. 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$?", "limitations": "The paper discusses its limitations in the experiment section. I think it might be better to mention them explicitly in the conclusion section.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer 4koq for their response. Please find our detailed answer below.\n\n> why drift model...\n\nIn Section 3 and in Appendix D, we presented the reasons to use a drift model together with learning NN parameters. Figure 1 illustrates the high-level intuition in case of online Bayesian estimation. Using SGD language, assuming that the data is stationary up to the time $T$, SGD estimate $\\theta_{T}$ would tend towards a local optima $\\theta^*$. If data stays stationary at time $T+1$, $\\theta_{T+1}$ will move closer to the $\\theta^*$. However, if at time $T+1$, the data distribution changes and so is the set of local optima, SGD might struggle to move towards this set starting from $\\theta_T$. Drift model allows the learning algorithm to make larger moves towards this new set of local optima. Such idea was also used in the context of online convex optimization, see [2-3]. The form in eq.4 encourages parameters to return back towards the initialization over time by shrinking the mean and increasing the variance (l.r. in SGD), allowing to make bigger steps towards new local minima.\n\nWe will change the wording of the beginning of Section 3 to better reflect this explanation.\n\n> form of OU process\n\nOU process defines a SDE that can be solved explicitly and written as a time-continuous Gaussian Markov process with transition density $p(x_t|x_s)=\\mathcal{N}(x_s e^{-(t-s)},(1-e^{-2(t-s)})\\sigma_0^2I)$ for any pair of times  $t>s$. Based on this, as a drift model for the parameters $\\theta_t$ (so $\\theta_t$ is the state $x_t$) we use the conditional density $p(\\theta_{t+1}|\\theta_t)=\\mathcal{N}(\\theta_t\\gamma_t,(1-\\gamma_t^2)\\sigma_0^2I)$ where $\\gamma_t=e^{-\\delta_t}$ and $\\delta_t\\geq0$ corresponds to the learnable discretization time step. In other words, by learning $\\gamma_t$ online we equivalently learn the amount of a continuous “time shift” $\\delta_t$ between two consecutive states in the OU process. This essentially models parameter drift since e.g. if $\\gamma_t=1$, then $\\delta_t=0$ and there is no “time shift\" implying $\\theta_{t+1}=\\theta_t$.\n\n> ...why eq.4\n\nIn Section 3, we argued for choosing eq.4 because it pushes the learning towards the initialization as well as towards previous parameters, allows to use gradient-based methods to estimate drift parameters and keeps positive finite variance (assuming we learn over the infinite amount of time) which avoids degenerate cases of $0$ or $\\infty$ variance. Moreover, it couples the mean and the variance via $\\gamma$, making it easier to learn $\\gamma$ and less likely to overfit. Other potential choices for the drift models:\n\n* $\\theta_{t+1}=\\gamma\\theta_t+(1-\\gamma)\\mu_0+\\beta\\epsilon$, with $\\epsilon\\sim\\mathcal{N}(0,I)$\n\n  * Fixed $\\beta$ was explored in our experiment in Figure 2, left, where we used constant l.r. $\\alpha_t$ for parameters update. We found that it performed worse than using rescaled l.r. $\\alpha(\\gamma_t)$ from eq.17 which is derived using our model, see Figure 2, right. For $\\mu_0=0$, we recover Shrink&Perturb [4]\n  * Learning $\\gamma$ and $\\beta$ will likely overfit\n\n* Mixture $p(\\theta_{t+1}|\\theta_t,\\gamma_t)=\\gamma_tp(\\theta_{t+1}|\\theta_t)+(1-\\gamma_t)p_0(\\theta_{t+1})$, where $p(\\theta_{t+1}|\\theta_t)=\\mathcal{N}(\\theta_t;\\sigma^2)(\\theta_{t+1})$ and $p_0(\\theta_{t+1})=\\mathcal{N}(\\mu_0;\\sigma_0^2)(\\theta_{t+1})$, which is a Gaussian version of Spike&Slab [5]. This encourages **hard** resets instead of **soft** ones. Moreover, using mixtures is problematic since the KL in eq.11 cannot be computed exactly and needs to be approximated.\n\nWe will add the discussion of drift model choices in the appendix.\n\n> derivations...\n\nPlease see our reply to Reviewer 6GY1 about the eq.9 and eq.11. Eq.10 is obtained from eq.9 by finding fixed points. In eq.10, we have a typo, the term $\\lambda\\gamma_t^0$ in the denominator should be replaced by $\\lambda$. Eq.14 is a gradient descent rule applied to eq.11(eq.12) where $\\tilde{F}$ is actually a derivative of $F$ wrt $\\mu_t$ and $\\sigma_t$. Eq.15 is Maximum a-posteriori (MAP) update for $\\theta$ from the posterior $p(y_{t+1}|x_{t+1},\\theta)q^{t+1}{t}(\\theta |\\gamma_t)$ and where we use temperature in the prior (see [3]).\n\nWe will add all the derivations and explanations in the appendix.\n\n> $\\theta'$ every time\n\nWhile a possible modification for the drift model, it could be problematic to use, since resampling $\\theta_{0}$ has higher variance than OU process. The variance of this model is $\\gamma_t^2\\sigma^2_t+(1-\\gamma_t)^2\\sigma^2_0+(1-\\gamma^2_t)\\sigma^2_0$, which equals to $2\\sigma^2_0$ for $\\gamma_t=0$ which is 2x larger than the variance of the initialization. This implies that such a model might inflate the variance over time since e.g. if $\\gamma_t=0.5$ for all time steps $t$, then $\\sigma_{t+1}^2=\\gamma_t^2 \\sigma^2_t+2(1-\\gamma_t)\\sigma^2_0=0.5^2\\sigma^2_t+\\sigma^2_0$ which grows with the iterations.\n\nThus, we believe that our current drift model more accurately implements the mechanism of parameter resets since it always  brings us back to the exact initialization distribution.\n\nWe hope we have been able to address your concerns and we are happy to answer further questions on the above subjects. If we were able to address your concerns, we would be grateful if you would consider increasing your review score.\n\n**References**:\n\n[1] Dynamical Models and Tracking Regret in Online Convex Programming, Eric C. Hall, Rebecca M. Willett, 2013\n\n[2] Adaptive Gradient-Based Meta-Learning Methods, Mikhail Khodak, Maria-Florina Balcan, Ameet Talwalkar, 2019\n\n[3] How Good is the Bayes Posterior in Deep Neural Networks Really?, Florian Wenzel, Kevin Roth, Bastiaan S. Veeling, Jakub Świątkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, Sebastian Nowozin, 2020\n\n[4] On Warm-Starting Neural Network Training, Jordan T. Ash, Ryan P. Adams, 2020\n\n[5] Spike and slab variable selection: Frequentist and Bayesian strategies, Hemant Ishwaran, J. Sunil Rao, 2005", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rZzHWbvaEP", "reviewer": "Reviewer_DfC5", "summary": "The 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.", "questions": "I 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.", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The 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.", "weaknesses": "In 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer DfC5 for their feedback. Please find our response below.\n\n> Figure 3, hard to read\n\nThanks to your feedback, we will modify the way we present the results, using the extended page limit for the camera-ready version. First, we will present **Soft Reset** method and compare it to baselines. Second, we will present other variants of the method and compare it to the main method.\n\n> baselines and benchmarks motivation\n\nThe benchmarks are motivated by Plasticity Loss literature [1-4]. In this setting, Neural Networks lose ability to learn (lose plasticity) when exposed to certain forms of non-stationarity. Random label MNIST and CIFAR10 benchmarks are most common ones. Moreover, in this setting, resetting is highly beneficial (see [1]).\n\nWe used 3 baselines, Shrink&Perturb [5], L2 Init [1] and Hard Resets. The Shrink&Perturb[5] is used to prevent plasticity loss by perturbing parameters according to $\\theta_{t+1}=\\lambda\\theta_t+\\sigma \\epsilon$ where $\\epsilon\\sim\\mathcal{N}(0;I)$. L2 Init approach adds $||\\theta - \\theta_0||^2$ to the loss and was shown to outperform many other approaches in preventing plasticity loss. Both baselines are highly related to our drift model and discussed in Related Work. We hope this clarifies the motivations.\n\n> RL is not clear...\n\nPlease see our response to Reviewer Hmj6.\n\n> Why in on-policy, simple method performs better than ours, given that it is more non-stationary setting?\n\nIn RL, the non-stationary rises from changing input data distribution (since we change the policy) as well as from changing learning targets (when learning value functions). In the off-policy RL, we expect a high impact of changing learning target non-stationarity. It was observed [4,6-7], that in this setting we can exhibit the loss of plasticity. The effect becomes more present as we increase the replay ratio (increase the number of updates per collected data). Under high replay ratio since we do many updates on replay buffer, we could overfit to the replay buffer and when we start collecting new data, the local minima of NN can switch drastically. This is the regime which motivated our method (see Figure 1 and beginning of Section 3). In this setting, hard resets [4] are effective but our method shows to be even more effective.\n\nWhen we operate in on-policy setting, the situation described above is less likely to happen. If data distribution changes quickly, then we are less likely to overfit to it due to the noise in the update. Using the language from Figure 1, the uncertainty of the Bayesian posterior will not be shrinking if the data changes quickly. This, however, depends on the environment. The Hopper environment (Figure 4, bottom) is a relatively simple task to solve, meaning that the agent will progress fast and the corresponding input data distribution will change fast. The Humanoid environment (Figure 4, top) is a much more challenging environment, and the simple agent might struggle initially. This implies that even on-policy algorithm can see a lot of similar data and start to overfit to it, putting it closer to the off-policy setting. This explains why our method is more effective in Humanoid rather than Hopper environment when relpay ratio is 1.\n\n> why stops earlier?\n\nThe experiment was not fully finished at the time of submission. See attached pdf for final version of the Figure.\n\n> Figure 2 ? \n\nIn Figure 2, we demonstrate that drift model eq.4 coupled with update rule eq.16 is a good strategy of resetting parameters when task boundaries are known. For appropriately chosen $\\gamma \\neq 0$, we can achieve significantly better performance than hard reset. In Figure 2, left, we use constant l.r. $\\alpha_t$ whereas in Figure 2,right, we use $\\alpha_t(\\gamma_t)$ from eq.17, which is more beneficial. We will add more clarifications in the text about the nature of this experiment.\n\n> Theoretical analysis \n\nUnfortunately, the phenomenon of plasticity loss in Neural Networks is not theoretically analyzed (there is no model for that), it is an empirical phenomenon [1-4]. We would require this theoretical framework for plasticity loss to be developed in order to analyze our method. Our future plans involve analyzing the method in the context of online convex optimization.\n\n> MC samples in eq.7?\n\nMC samples are not coming from a simulator but from a Gaussian distribution eq.5 induced by the drift model eq.4. MC samples are needed to approximate integral in eq.7. The MC samples correspond to NN parameters.\n\nOverall, we hope our answer clarifies the points which you raised during your review. If we have addressed your questions, we would be  grateful if you would consider increasing your score. We would be happy to answer any further questions you might have.\n\n**References**:\n\n[1] Maintaining Plasticity in Continual Learning via Regenerative Regularization, Saurabh Kumar, Henrik Marklund, Benjamin Van Roy, 2023\n\n[2] Understanding plasticity in neural networks, Clare Lyle, Zeyu Zheng, Evgenii Nikishin, Bernardo Avila Pires, Razvan Pascanu, Will Dabney, 2023\n\n[3] Disentangling the Causes of Plasticity Loss in Neural Networks, Clare Lyle, Zeyu Zheng, Khimya Khetarpal, Hado van Hasselt, Razvan Pascanu, James Martens, Will Dabney, 2024\n\n[4] Understanding and Preventing Capacity Loss in Reinforcement Learning, Clare Lyle, Mark Rowland, Will Dabney, 2022\n\n[5] On Warm-Starting Neural Network Training, Jordan T. Ash, Ryan P. Adams, 2020\n\n[6] The Primacy Bias in Deep Reinforcement Learning, Evgenii Nikishin, Max Schwarzer, Pierluca D'Oro, Pierre-Luc Bacon, Aaron Courville, 2022\n\n[7] The Dormant Neuron Phenomenon in Deep Reinforcement Learning, Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, Utku Evci, 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear reviewers, thank you for your feedback. In this section we provide answers to recurrent points from some of you.

# Computational Complexity

The following tables will be added to the Appendix together with a reference to this appendix at the end of Section 3.

**Notations**:

* P - number of parameters of Neural Network (NN)
* L - number of layers.
* $O(S)$ - cost of backwards pass of SGD
* $K$ - number of Monte Carlo (MC) samples in eq.7
* $J$ - number of iterations in eq. 7
* $I$ - number of parameter updates for proximal methods in eq.14 and eq.18
* $M$ - number of MC samples for Bayesian method in eq. 12


| Method                                      | Comp. cost.  | Memory           |
|----------------------------------------------|--------------|------------------|
| SGD                                          | O(S)         | O(P)             |
| Soft resets $\gamma$ p. layer                  | O(JKS + S)    | O(L+(K+1)P)     |
| Soft resets $\gamma$ p. param.                 | O(JKS + S)    | O(P+(K+1)P)     |
| Soft resets $\gamma$ p. layer + proximal (I iters) | O(JKS + I S)  | O(L+(K+1)P)     |
| Soft resets $\gamma$ p. param. + proximal (I iters) | O(JKS + I S)  | O(P+(K+1)P)     |
| Bayesian Soft Reset Proximal (I Iter) $\gamma$ p.layer | O(JKS + 2M I S)| P(L+(K+2)P)     |
| Bayesian Soft Reset Proximal (I Iter) $\gamma$ p.param | O(JKS + 2M I S)| P(P+(K+2)P)     |
|

The table above denotes the theoretical cost of each of the methods. In practice, we use $J=1$ for Soft Resets and $J=10$ for proximal methods. The table below quantifies the exact cost for methods from Figure 3.

| Method                                      | Comp. cost.  | Memory           |
|----------------------------------------------|--------------|------------------|
| SGD                                          | O(S)         | O(P)             |
| Soft resets                                  | O(2S)         | O(L+2P)          |
| Soft resets more compute                     | O(11S)        | O(L+2P)          |
| Soft resets proximal                         | O(10S + 10S) | O(L+2P)          |
| Bayesian Soft Reset Proximal                 | O(10S + 20S) | P(L+3P)          |
| Bayesian Soft Reset Proximal $\gamma$ p.param | O(10S + 20S) | P(P+3P)          |
|

From Figure 3, we see that it is beneficial to spend more compute on optimizing $\gamma$ and NN parameters. However, even the cheapest Soft Resets leads to a good performance.

In Reinforcement Learning (RL) experiment, we do **one** update on $\gamma$ after each new chunk of fresh data. We do $G$ updates on NN parameters with a cost $O(S)$ each. The table below denotes the complexities.

| Method        | Complexity  | Memory       |
|---------------|-------------|--------------|
| SAC           | O(G S)      | O(P)         |
| Soft Reset    | O(S + G S)  | O(L + 2 * P) |
|

Soft Reset is marginally more expensive than SAC in RL but leads to a better performance (see Figure 4) in a highly off-policy regime.

# Hyperparameters sensitivity

We study the sensitivity of **Soft Resets** where $\gamma$ is defined per layer.

**Fixed parameters:**

* Number of MC samples $K=1$ and $M=1$
* Learning rate for parameters was tuned prior to that and equals to $\alpha=0.1$.

**Sensitivity parameters** (see Algorithm 1 for precise definitions):

* $\eta_{\gamma}$ - learning rate for the drift model
* $f$ - initial prior standard deviation rescaling, i.e., $\sigma^l_0 = f \frac{1}{\sqrt{H}}$ where $H$ is the width of the layer $l$
* $s$ - posterior scaling, i.e., $\sigma^l_t = s * \sigma^l_0$.

We provide the plots for the sensitivity analysis of Soft Reset on MNIST (data efficient) in the attached pdf. On top of that, we conduct the sensitivity analysis of L2 Init [1] 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.

**Takeaways**

The most important parameter is the learning rate $\eta_{\gamma}$ 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\geq0.5$. For $s<0.5$, the performance degrades. This parameter is defined from $\sigma_{posterior} = s \sigma_{prior}$ and effects relative increase in learning rate (see eq. 13 and eq. 17). This increase is given by $1/(\gamma^2 + (1-\gamma^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 is very sensitive to the parameter $\lambda$ which is a penalty on $||\theta-\theta_0||^2$ term. In fact, Figure 2, left shows that there is only one good value of this parameter which works. Shrink\&Perturb is very sensitive to the shrink parameter ($\lambda$). Similar to L2 Init, there is only one value which works, $0.9999$ while values $0.999$ and $0.99999$ lead to bad performance. This method however, is not very sensitive to the perturb parameter $\sigma$ provided that $\sigma \leq 0.001$.

Compared to the baselines, our method is more robust to the hyperparameters choice.

We also conduct sensitivity analysis for other variants of the method, but due to space constraints, will only include the results in the camera ready version. The take-aways are similar.

**References:**

[1] Maintaining Plasticity in Continual Learning via Regenerative Regularization, Saurabh Kumar, Henrik Marklund, Benjamin Van Roy, 2023

[2] On Warm-Starting Neural Network Training, Jordan T. Ash, Ryan P. Adams, 2020

Meta-review: This paper proposes to adapt to non-stationarity in neural network training using an Ornstein-Uhlenbeck process with an adaptive drift parameter. The proposed method estimates the drift parameter online and allows for automatic soft resets of parameters when needed. The approach is evaluated on plasticity benchmarks and off-policy reinforcement learning tasks. The reviewers all agree that the main idea is interesting and novel. 

However, the significance of the contributions are not clearly demonstrated due to the lack of clarity in the submission. In its current form, even the most engaged readers have some difficulty grasping the core technical ideas, e.g., see the discussion thread with Reviewer 4koq.  In the rebuttal process, the authors commit to resolving several of these issues—I will not repeat them here—which will be critical to the paper having broad impact. In particular, I strongly recommend adding a candid discussion on the computational costs of the proposed approach.
Decision: Accept (poster)