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

Authors: alexandre  galashov
michalis k titsias
google  deepmind
andrás  györgy
clare  lyle
razvan  pascanu
yee  whye
teh  google deepmind
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 18:45:41

Abstract: Neural networks are traditionally trained under the assumption that data come from a stationary distribution. However, settings which violate this assumption are becoming more popular;  examples include supervised learning under distributional shifts, reinforcement learning, continual learning and non-stationary contextual bandits. In this work we introduce a novel learning approach that automatically models and adapts to non-stationarity, via an Ornstein-Uhlenbeck process with an adaptive drift parameter. The adaptive drift tends to draw the parameters towards the initialisation distribution, 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.

Main paper: 
 introduction: 
 Neural networks (NNs) are typically trained using algorithms like stochastic gradient descent (SGD), assuming data comes from a stationary distribution. This assumption fails in scenarios such as continual learning, reinforcement learning, non-stationary contextual bandits, and supervised learning with distribution shifts [20, 53] . A phenomenon occurring in non-stationary settings is the loss of plasticity [12, 2, 13] , manifesting either as a failure to generalize to new data despite reduced training loss [4, 2] , or as an inability to reduce training error as the data distribution changes [13, 37, 1, 42, 34] .In [38] , the authors argue for two factors that lead to the loss of plasticity: preactivation distribution shift, leading to dead or dormant neurons [47] , and parameter norm growth causing training instabilities. To address these issues, strategies often involve hard resets based on heuristics like detecting dormant units [47] , assessing neuron utility [13, 12] , or simply after a fixed number of steps [43] . Though effective at increasing plasticity, hard resets can be inefficient as they can discard valuable knowledge captured by the parameters.We propose an algorithm that implements a mechanism of soft parameter resets, in contrast to the hard resets discussed earlier. A soft reset partially moves NN parameters towards the initialization while keeping them close to their previous values. It also increases learning rate of the learning algorithm, allowing new NN parameters to adapt faster to the changing data. The amount by which the parameters move towards the initialization and the amount of learning rate increase are controlled by the drift parameters which are learned online. The exact implementation of soft reset mechanism is based on the use of a drift model in NN parameters update before observing new data. Similar ideas which modify the starting point of SGD as well as increase the learning rate of SGD depending on non-stationarity were explored (see [21, 29] ) in an online convex optimization setting. Specifically, in [21] , the authors assume that the optimal parameter of SGD changes according to some dynamical model out of a finite family of models. They propose an algorithm to identify this model and a way to leverage this model in a modified SGD algorithm. Compared to these works, we operate in a general non-convex setting. Proposed drift model can be thought as a dynamical Bayesian prior over Neural Network parameters, which is adapted online to new data. We make a specific choice of drift model which implements soft resets mechanism.Our contributions can be summarized as follows. First, we propose an explicit model for the drift in NN parameters and describe the procedure to estimate the parameters of this model online from the stream of data. Second, we describe how the estimated drift model is incorporated in the learning algorithm. Third, we empirically demonstrate the effectiveness of this approach in preventing the loss of plasticity as well as in an off-policy reinforcement learning setting.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: We outline main contributions of the paper in the introduction and abstract.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
non-stationary learning with online sgd: 
 In a non-stationary learning setting with changing data distributions p t (x, y), where x ∈ R L , y ∈ R K , we define the loss function for parameters θ ∈ R D asEQUATIONOur goal is to find a parameter sequence Θ = (θ 1 , . . . , θ T ) that minimizes the dynamic regret:EQUATIONwith a reference sequence Θ ⋆ = (θ ⋆ 1 , . . . , θ ⋆ T ), satisfying θ ⋆ t = arg min θ L t (θ). A common approach to the online learning problem is online stochastic gradient descent (SGD) [23] . Starting from initial parameters θ 0 , the method updates these parameters sequentially for each batch of data {(x i t , y i t )} B i=1 s.t. (x i t , y i t ) ∼ p t (x t , y t ). The update rule is:θ t+1 = θ t -α t ∇ θ L t+1 (θ t ),where ∇ θ L t+1 (θ t ) = 1 B B i=1 ∇ θ L t+1 (θ t , x i t , y i t ) and α t is learning rate. See also Appendix G for the connection of SGD to proximal optimization. t+1 (blue dot) online Bayesian estimation can be less data efficient and take time to recover when the change-point occurs. (e) The use of p(θ|θ t , γ t ) and the estimation of γ t 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 p t+1 (θ) (blue dashed circle) can faster track θ * t+1 .non-i.i.d. fashion such that a change in the data distribution is modeled by the drift in the parameters p(θ t+1 |θ t , γ t ) at every time t + 1 before new data is observed. We assume a class of drift models p(θ t+1 |θ t , γ t ) which encourages the parameters to move closer to the initialization. The amount of drift (and level of non-stationarity) is controlled by γ t which are estimated online from the data.As can be seen below, in the context of SGD, this approach adjusts the starting point θ t of the update to a point θt (γ t ), which is closer to the initialization and increases the learning rate proportionally to the drift. In the context of Bayesian inference, this approach shrinks the mean of the estimated posterior towards the prior and increases the variance proportional to γ t . This approach is inspired by prior work in online convex optimization for non-stationary environments [e.g., 25, 21, 8, 18, 29] .Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We discuss limitations in the experimental section Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
toy illustration of the advantage of drift models: 
 Consider online Bayesian inference with 2-D observations y t = θ ⋆ + ϵ t , where θ ⋆ ∈ R 2 are unknown true parameters and ϵ t ∼ N (0; σ 2 I) is Gaussian noise with variance σ 2 . Starting from a Gaussian prior p 0 (θ) = N (θ; µ 0 ; Σ 0 ), the posterior distribution p t+1 (θ) = p(θ|y 1 , . . . , y t ) = N (θ; µ t+1 , Σ t+1 ) is updated using Bayes' rule p t+1 (θ) ∼ p(y t+1 |θ)p t (θ).The posterior update (3) comes from the i.i.d. assumption on the data generation process (Figure 1a ), since p t+1 (θ) ∼ p 0 (θ) t+1 s=1 p(y s |θ). By the Central Limit Theorem (CLT), the posterior mean µ t converges to θ ⋆ and the covariance matrix Σ t shrinks to zero (the radius of red circle in Figure 1c ). Suppose now that the true parameters θ ⋆ t (kept fixed before t) change to new parameters θ ⋆ t+1 at time t + 1. The i.i.d. assumption (Figure 1a ) is violated and the update (3) becomes problematic because the low uncertainty (small radius of red dashed circle in Figure 1d ) in p t (θ) causes the posterior p t+1 (θ) (see blue circle) to adjust slowly towards θ ⋆ t+1 (blue dot) as illustrated in Figure 1d . To address this issue, we assume that before observing new data, the parameters drift according to p(θ t+1 |θ t , γ t ) where the amount of drift is controlled by γ t . The corresponding conditional independence structure is shown in Figure 1b . The posterior update then becomes:EQUATIONFor a suitable choice of drift model p(θ t+1 |θ t , γ t ), this modification allows p t+1 (θ) (blue circle) to adjust more rapidly towards the new θ ⋆ t+1 (blue dot), see Figure 1e . This is because the new prior p(θ|θ ′ t , γ t )p t (θ ′ t )dθ ′ t has larger variance (green circle) than p t (θ) and its mean is closer to the center of the circle. Ideally, the parameter γ t should capture the underlying non-stationarity in the data distribution in order to control the impact of the prior p(θ|θ ′ t , γ t )p t (θ ′ t )dθ ′ t . For example, if at some point the non-stationarity disappears, we want the drift model to exhibit no-drift to recover the posterior update (3) . This highlights the importance of the adaptive nature of the drift model.We motivate the specific choice of the drift model which is useful for maintaining plasticity. We assume that our Neural Network has enough capacity to learn any stationary dataset in a fixed number of iterations starting from a good initialization θ 0 ∼ p 0 (θ) [see, e.g., 24, 16] . Informally, we call the initialization θ 0 plastic and the region around θ 0 a plastic region.Consider now a piecewise stationary datastream that switches between a distribution p a , with a set of local minima M a of the negative likelihood L(θ), to a distribution p b at time t + 1, with a set of local minima M b . If M b is far from M a , then hard reset might be beneficial, but if M b is close to M a , resetting parameters is suboptimal. Furthermore, since θ is high-dimensional, different dimensions might need to be treated differently. We want a drift model that can capture all of these scenarios.Drift model. The drift model p(θ t+1 |θ t , γ t ) which exhibits the above properties is given byEQUATIONwhich is separately defined for every parameter dimension θ i where p 0(θ i 0 ) ∼ N (θ i 0 ; µ i 0 ; σ i 0 2) is the per-parameter prior distribution and γ t = (γ 1 t , . . . , γ D t ). The model is a discretized Ornstein-Uhlenbeck (OU) process [50] (see Appendix A for the derivation).The parameter γ t ∈ [0, 1] is a drift parameter and controls the amount of non-stationarity in each parameter. For γ t = 1, there is no drift and for γ t = 0, the drift model reverts the parameters back to the prior. A value of γ t ∈ (0, 1) interpolates between these two extremities. A remarkable property of ( 5) is that starting from the current parameter θ t , if we simulate a long trajectory, as T → ∞, the distribution of p(θ T |θ t ) will converge to the prior p(θ 0 ). This is only satisfied (for γ t ∈ (0, 1)) due to the variance σ 2 0 (1 -γ 2 t ). Replacing it by an arbitrary variance σ 2 would result in the variance of p(θ T |θ t ) either going to 0 or growing to ∞, harming learning. Thus, the model ( 5) encourages parameters to move towards plastic region (initialization). In Appendix B, we discuss this further and other potential choices for the drift model.The drift model p(θ t+1 |θ t , γ t ) quantifies prior belief about the change in parameters before seeing new data. A suitable choice of an objective to select γ t is predictive likelihood which quantifies the probability of new data under our current parameters and drift model. From Bayesian perspective, it means selecting the prior distribution which explains the future data the best.We derive the drift estimation procedure in the context of approximate online variational inference [7] with Bayesian Neural Networks (BNN). Let Γ t = (γ 1 , . . . , γ t ) be the history of observed parameters of the drift model and S t = {(x 1 , y 1 ), . . . , (x t , y t )} be the history of observed data. The objective of approximate online variational inference is to propagate an approximate posterior q t (θ|S t , Γ t-1 ) over parameters, such that it is constrained to some family Q of probability distributions. In the context of BNNs, it is typical [5] to assume a family Q = {q(θ) : q(θ) ∼ D i=1 N (θ i ; µ i , σ i 2 ); θ = (θ 1 , . . . , θ D )} of Gaussian mean-field distributions over parameters θ ∈ R D (separate Gaussian per parameter). For simplicity of notation, we omit the index i. Let q t (θ) ≜ q t (θ|S t , Γ t-1 ) ∈ Q be the Gaussian approximate posterior at time t with mean µ t and variance σ 2 t for every parameter. The new approximate posterior q t+1 (θ) ∈ Q is found byEQUATIONwhere the prior term is the approximate predictive look-ahead prior given by q t (θ|γ t ) = q t (θ t )p(θ|θ t , γ t )dθ t = N (θ; μt (γ t ), σ2 t (γ t ))that has parameters μt(γ t ) = γ t µ t + (1 -γ t )µ 0 , σ2 t (γ t ) = γ 2 t σ 2 t + (1 -γ 2 t )σ 2 0, see Appendix I.1 for derivation. The form of this prior q t (θ|γ t ) comes from the non i.i.d. assumption (see Figure 1b ) and the form of the drift model (5) . For new batch of data (x t+1 , y t+1 ) at time t + 1, the approximate predictive log-likelihood equals to log q t (y t+1 |x t+1 , γ t ) = log p(y t+1 |x t+1 , θ)q t (θ|γ t )dθ.The log-likelihood (8) allows us to quantify predictions on new data (x t+1 , y t+1 ) given our current distribution q t (θ) and the drift model from (5) . We want to find such γ ⋆ t that γ ⋆ t ≈ arg max γt log q t (y t+1 |x t+1 , γ t )Using γ ⋆ t in (5) modifies the prior distribution (7) to fit the most recent observations the best by putting more mass on the region where the new parameter could be found (see Figure 1 ,right).Gradient-based optimization for γ t . The approximate predictive prior in (7) is Gaussian which allows us to use the so-called reparameterisation trick to optimize (8) via gradient descent. Starting from an initial value of drift parameter γ 0 t at time t, we perform K updates with learning rateEQUATIONThe integral is evaluated by Monte-Carlo (MC) using M samples ϵ i ∼ N (ϵ; 0,EQUATION) Inductive bias in the drift model is captured by γ 0 t , where γ t,0 = 1 encourages stationarity, while γ t,0 = γ t-1,K promotes temporal smoothness. In practice, we found γ t,0 = 1 was the most effective.Structure in the drift model. The drift model can be defined to be shared across different subsets of parameters which reduces the expressivity of the drift model but also provides regularization to (10) . We consider γ t to be either defined for each parameter or for each layer. See Section 5 for details as well as corresponding results in Appendix H.Interpretation of γ t . By linearising log p(y t+1 |x t+1 , θ) around µ t , we can compute (8) in a closed form and get the following loss for γ t (see Appendix J for the proof) optimizing ( 9)F(γ t ) = 0.5(σ 2 t (γ t ) ⊙ g t+1 ) ⊤ g t+1 -(γ t ⊙ µ t + (1 -γ t ) ⊙ µ 0 ) ⊤ g t+1, where g t = ∇L t+1 (µ t ) and L t+1 (θ) = -log p(y t+1 |x t+1 , µ t ), and where ⊙ denotes element-wise product performed only over parameters for which γ t is shared (see paragraph about structure in drift model). The transpose operation is also defined on a subset of parameters for which γ t is shared. Adding the ℓ 2 penalty 1 2 λ||γ t -γ 0 t || 2 encoding the starting point γ 0 t , gives us the closed form for γ tEQUATIONwhere we also clip parameters γ t to [0, 1]. The expression (12) gives us the geometric interpretation for γ t . The value of γ t depends on the angle between (µ t -µ 0 ) and g t+1 When these vectors are aligned, γ t is high and is low otherwise. When these vectors are orthogonal or the gradient g t+1 ≈ 0, the value of γ t is heavily influenced by γ 0 t . Moreover, when g t+1 ≈ 0, we can interpret it as being close to a local minimum, i.e., stationary, which means that we want γ t ≈ 1, therefore adding the ℓ 2 penalty is important. Also, when the norm of the gradients g t+1 is high, the value of γ t is encouraged to decrease, introducing the drift. This means that using γ t in the parameter update (see Section 3.5) encourages the norm of the gradient to stay small. In practice, we found that update (12) was unstable suggesting that linearization of the log-likelihood might not be a good approximation for learning γ t .The optimization problem (6) for the per-parameter Gaussian q(θ) = N (θ; µ, σ 2 ) with Gaussian prior q t (θ) = N (θ; µ t , σ 2 t ), both defined for every parameter of NN, can be written (see Appendix I.1) to minimize the following lossEQUATION)where λ i t > 0 are per-parameter temperature coefficients. The use of small temperature λ > 0 parameter (shared for all NN parameters) was shown to improve empirical performance of Bayesian Neural Networks [54] . Given that in (13) , the variance σ2 t (γ t ) can be small, in order to control the strength of the regularization, we propose to use per parameter temperatureλ i t = λ × σ i t 2, where λ > 0 is a global constant. This leads to the following objectiveEQUATIONwhere the quantityr i t = [σ i t ] 2 /[σ i t (γ t )]2 is a relative change in the posterior variance due to the drift. The ratio r i t = 1 when γ t = 1. For γ t < 1 since typically σ 2 t < σ 2 0 , the ratio is r i t < 1. Thus, as long as there is non-stationarity (γ t < 1), the objective (14) favors the data term E ϵ∼N (0;I) [L t+1 (µ + ϵσ)] γ t,k+1 = γ t,k + η γ ∇ γ log p(y t+1 |x t+1 , γ t θ t + (1 -γ t )θ 0 + θ ′ 0 p 1 -γ 2 t + γ 2 t s 2 ) γt=γ t,kend for Get θt (γ t,K ) with (17) and αt (γ t,K ) with ( 18) Update parameters θ t+1 = θt (γ t,K ) -αt (γ t,K ) • ∇ θ L t+1 ( θt (γ t,K )) end for allowing the optimization to respond faster to changes in the data distribution. To find new parameters, let µ t+1,0 = μt (γ t ) and σ t+1,0 = σt (γ t ), and perform updates K on ( 14)EQUATIONwhere α µ and α σ are learning rates for the mean and for the standard deviation correspondingly. All derivations are provided in Appendix I.1. The full procedure is described in Algorithm 2.As a faster alternative to propagating the posterior (6), we do MAP updates with the prior p 0 (θ) = N (θ; µ 0 ; σ 2 0 ) and the approximate posterior q t (θ) = N (θ; θ t ; σ 2 t = s 2 σ 2 0 ), where s ≤ 1 is a hyperparameter controlling the variance σ 2 t of q t (θ). Since a fixed s may not capture the true parameters variance, using a Bayesian method (see Section 3.4) is preferred but comes at a high computational cost (see Appendix E for discussion). The MAP update is given by (see Appendix I.2 for derivations) finding a minimum of the following proximal objectiveEQUATIONwhere the regularization target for the parameter dimension i is given byEQUATION)and the per-parameter learning rate is given as (assuming that α t the base SGD learning rate)EQUATION)Linearising L t+1 (θ) around θt (γ t ) and optimizing (16) for θ leads to (see Appendix I.2)EQUATION)where • is element-wise multiplication. For γ t = 1, we recover the ordinary SGD update, while the values γ t < 1 move the starting point of the modified SGD closer to the initialization as well as increase the learning rate. Algorithm 1 describes the full procedure. In Appendix C we describe additional practical choices made for the Soft Resets algorithm. Similarly to the Bayesian approach (15), we can do multiple updates on (16) . We describe this Soft Resets Proximal algorithm in Appendix I.2 and full procedure is given in Algorithm 3.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA] Justification: Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
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 (12) of γ t connects to the notion of parameters utility from [14] , but this quantity is used to prevent catastrophic forgetting by decreasing learning rate for high γ t . Our method increases the learning rate for low γ t to maximize adaptability, and is not designed to prevent catastrophic forgetting.Non-stationarity. Non-stationarity arises naturally in a variety of contexts, the most obvious being continual and reinforcement learning. The structure of non-stationarity may vary from problem to problem. At one extreme, we have a piece-wise stationary setting, for example a change in the location of a camera generating a stream of images, or a hard update to the learner's target network in value-based deep RL algorithms. This setting has been studied extensively due to its propensity to induce catastrophic forgetting [e.g. 31, 45, 51, 10] and plasticity loss [13, 39, 38, 34] . At the other extreme, we can consider more gradual changes, for example due to improvements in the policy of an RL agent [40, 46, 42, 13] or shifts in the data generating process [36, 55, 20, 53] . Further, these scenarios might be combined, for example in continual reinforcement learning [31, 1, 13] where the reward function or transition dynamics could change over time.Non-stationary online convex optimization. Non-stationary prediction has a long history in online convex optimization, where several algorithms have been developed to adapt to changing data [see, e.g., 25, 8, 22, 17, 21, 18, 29] . Our approach takes an inspiration from these works by employing a drift model as, e.g., [25, 21] and by changing learning rate as [29, 52] . Further, our OU drift model bears many similarities to the implicit drift model introduced in the update rule of [25] (see also [8, 17] ), where the predictive distribution is mixed with a uniform distribution to ensure the prediction could change quickly enough if the data changes significantly, where in our case p 0 plays the same role as the uniform distribution.Bayesian approaches to non-stationary learning. A standard approach is Variational Continual Learning [41] , which focuses on preventing catastrophic forgetting and is an online version of "Bayes By Backprop" [5] . This method does not incorporate dynamical parameter drift components. In [35] , the authors applied variational inference (VI) on non-stationary data, using the OU-process and Bayesian forgetting, but unlike in our approach, their drift parameter is not learned. Further, in [49] , the authors considered an OU parameter drift model similar to ours, with an adaptable drift scalar γ and analytic Kalman filter updates, but is applied over the final layer weights only, while the remaining weights of the network were estimated by online SGD. In [28] , the authors propose to deal with non-stationarity by assuming that each parameter is a finite sum of random variables following different OU process. They derive VI updates on the posterior of these variables. Compared to this work, we learn drift parameters for every NN parameter rather than assuming a finite set of drift parameters. A different line of research assumes that the drift model is known and use different techniques to estimate the hidden state (the parameters) from the data: in [9] , the authors use Extended Kalman Filter to estimate state and in [3] , they propagate the MAP estimate of the hidden state distribution with K gradient updates on a proximal objective similar to (43) , whereas in Bayesian Online Natural Gradient (BONG) [27] , the authors use natural gradient for the variational parameters.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: We disclose the experimental information in Experimental and Appendix sections.Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
experiments: 
 Soft reset methods. There are multiple variations of our method. We call the method implemented by Algorithm 1 with 1 gradient update on the drift parameter Soft Reset, while other versions show different parameter choices: Soft Reset (K γ = 10) is a version with 10 updates on the drift parameter, while Soft Reset (K γ = 10, K θ = 10) is the method of Algorithm 3 in Appendix I.2 with 10 updates on drift parameter, followed by 10 updates on NN parameters. Bayesian Soft Reset (K γ = 10, K θ = 10) is a method implemented by Algorithm 2 with 10 updates on drift parameter followed by 10 updates on the mean µ t and the variance σ 2 t (uncertainty) for each NN parameter. Bayesian method performed the best overall but required higher computational complexity (see Appendix E). Unless specified, γ t is shared for all the parameters in each layer (separately for weight and biases). performance on random-label MNIST (data efficient). Right: performance on random-label CIFAR-10 (memorization). The x-axis is the task id and the y-axis is the per-task training accuracy (25) .We analyze the performance of our method on plasticity benchmarks [34, 39, 38] . Here, we have a sequence of tasks, where each task consists of a fixed (for all tasks) subset of 10000 images images from either CIFAR-10 [32] or MNIST, where either pixels are permuted or the label for each image is randomly chosen. Several papers [34, 39, 38 ] study a memorization random-label setting where SGD can perfectly learn each task from scratch. To highlight the data-efficiency of our approach, we study the data-efficient setting where SGD achieves only 50% accuracy on each task when trained from scratch. Here, we expect that algorithms taking into account similarity in the data, to perform better. To study the impact of the non-stationarity of the input data, we consider permuted MNIST where pixels are randomly permuted within each task (the same task as considered by 34). As baselines, we use Online SGD and Hard Reset at task boundaries. We also consider L2 init [34] , which adds L2 penalty ||θ -θ 0 || 2 to the fixed initialization θ 0 as well as Shrink&Perturb [2] , which multiplies each parameter by a scalar λ ≤ 1 and adds random Gaussian noise with fixed variance σ. See Appendix D.1 for all details. As metrics, we use average per-task online accuracy (25) , which isA t = 1 N N i=1 a t i, where a t i are the online accuracies collected on the task t via N timesteps, corresponding to the duration of the task. In Figure 5 , we also use average accuracy over all T tasks, i.e.A T = 1 T T t=1 A tThe results are provided in Figure 2 . We observe that Soft Reset is always better than Hard Reset and most baselines despite the lack of knowledge of task boundaries. The gap is larger in the data efficient regime. Moreover, we see that L2 Init only performs well in the memorization regime, and achieves comparable performance to Hard Reset in the data efficient one. The method L2 Init could be viewed as an instantiation of our Soft Reset Proximal method optimizing (16) with γ t = 0 at every step, which is sub-optimal when there is similarity in the data. Bayesian Soft Reset demonstrates significantly better performance overall, see also discussion below.In Figure 3 , we compare different variants of Soft Reset. We observe that adding more compute for estimating γ t (thus, estimating non-stationarity, K γ = 10) as well as doing more updates on NN parameters (thus, more accurately adapting to non-staionarity, K θ = 10) leads to better performance. All variants of Soft Reset γ t parameters are shared for each NN layer, except for the Bayesian method. This variant is able to take advantage of a more complex per-parameter drift model, while other variants performed considerably worse, see Appendix H.4. We hypothesize this is due to the NN parameters uncertainty estimates σ t which Bayesian method provide, while others do not, which leads to a more accurate drift model estimation, since uncertainty is used in this update (10) . But, this approach comes at a higher computational cost, see Appendix E. In Appendix H, we provide ablations of the structure of the drift model, as well as of the impact of learning the drift parameter.Qualitative behavior of Soft Resets. For Soft Reset, we track the values of γ t for the first MLP layer when trained on random-label tasks studied above (only 20 tasks), as well as the minimum encountered value of γ t for each layer, which highlights the maximum amount of resets. Figure 4b, c shows γ t as a function of t, and suggests that γ t aggressively decreases at task boundaries (red dashed lines). The range of values of γ t depends on the task and on the layer, see Figure 4a . Overall, γ t changes more aggressively for long duration (memorization) random-label CIFAR-10 and less for shorter (data-efficient) random-label MNIST. See Appendix H.2 for more detailed results.To study the behavior of Soft Reset under input distribution non-stationarity, we consider a variant of Permuted MNIST where each image is partitioned into patches of a given size. The non-stationarity is controlled by permuting the patches (not pixels). Figure 5a shows the minimum encountered γ t for each layer for different patch sizes. As the patch size increases and the problem becomes more stationary, the range of values for γ t is less aggressive. See Appendix H.3 for more detailed results.Impact of non-stationarity. We consider a variant of random-label MNIST where for each task, an image has either a random or a true label. The label assignment is kept fixed throughout the task and is changed at task boundaries. We consider cases of 20%, 40% and 60% of random labels and we control the duration of each task (number of epochs). In total, the stream contains 200 tasks.In Figure 5b , we show performance of Online SGD, Hard Reset and in Figure 5c , the one of Soft Reset and of Bayesian Soft Reset. See Appendix D.2 for more details. The results suggest that for the shortest duration of the tasks, the performance of all the methods is similar. As we increase the duration of each of the task (moving along the x-axis), we see that both Soft Resets variants perform better than SGD and the gap widens as the duration increases. This implies that Soft Resets is more effective with infrequent data distribution changes. We also observe that Bayesian method performs better in all the cases, highlighting the importance of estimating uncertainty for NN parameters.Reinforcement learning experiments. We conduct Reinforcement Learning (RL) experiments in the highly off-policy regime, similarly to [43] , since in this setting loss of plasticity was observed. We ran SAC [19] agent with default parameters from Brax [15] on the Hopper-v5 and Humanoid-v4 GYM [6] environments (from Brax [15] ). To reproduce the setting from [43] , we control the off-policyness of the agent by setting the off-policy ratio M such that for every 128 environment steps, we do 128M gradient steps with batch size of 256 on the replay buffer. As baselines we consider ordinary SAC, hard-coded Hard Reset where we reset all the parameters K = 5 times throughout training (every 200000 steps), while keeping the replay buffer fixed (similarly to [43] ). We employ our Soft Reset method as follows. After we have collected fresh data from the environment, we do one gradient update on γ t (shared for all the parameters within each layer) with batch size of 128 on this new chunk of data and the previously collected one, i.e., two chunks of data in total. Then we initialize θt (γ t ) and we employ the update rule (43) where the regularization θt (γ t ) is kept constant for all the off-policy gradient updates on the replay buffer. See Appendix D.3 for more details. The results are given in Figure 6 . As the off-policy ratio increases, Soft Reset becomes more efficient than the baselines. This is consistent with our finding in Figure 5b, c , where we showed that the performance of Soft Reset is better when the data distribution is not changing fast. Figure 8 in Appendix D.3 shows the value of learned γ t . It shows γ t mostly change for the value function and not for the policy indicating that the main source of non-stationarity comes from the value function.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [No] Justification: Unfortunately, due to IP constrains, we cannot release the code for the paper.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion: 
 Learning efficiently on non-stationary distributions is critical to a number of applications of deep neural networks, most prominently in reinforcement learning. In this paper, we have proposed a new method, Soft Resets, which improves the robustness of stochastic gradient descent to nonstationarities in the data-generating distribution by modeling the drift in Neural Network (NN) parameters. The proposed drift model implements soft reset mechanism where the amount of reset is controlled by the drift parameter γ t . We showed that we could learn this drift parameter from the data and therefore we could learn when and how far to reset each Neural Network parameter. We incorporate the drift model in the learning algorithm which improves learning in scenarios with plasticity loss. The variant of our method which models uncertainty in the parameters achieves the best performance on plasticity benchmarks so far, highlighting the promise of the Bayesian approach. Furthermore, we found that our approach is particularly effective either on data distributions with a lot of similarity or on slowly changing distributions. Our findings open the door to a variety of exciting directions for future work, such as investigating the connection to continual learning and deepening our theoretical analysis of the proposed approach.NeurIPS Paper ChecklistQuestion: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: We provide experimental details in the appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: We specify that we report results with 3 random seeds with mean and standard deviation.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: We provide information about compute resources required in the appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA]Justification:Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification:Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: We cite the works which introduced the publicly available datasets Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13 • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Justification: Based on our understanding, our work conforms to the every aspect of NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).We make use that the Ornstein-Uhlenbeck process [50] defines a SDE that can be solved explicitly and written as a time-continuous Gaussian Markov process with transition density p(x t |x s ) = N (x s e -(t-s) , (1 -e -2(t-s) )σ 2 0 I), for any pair of times t > s. Based on this as a drift model for the parameters θ t (so θ t is the state x t ) we use the conditional density p(θ t+1 |θ t ) = N (θ t γ t , (1 -γ 2 t )σ 2 0 I), where γ t = e -δt and δ t ≥ 0 corresponds to the learnable discretization time step. In other words, by learning γ t online we equivalently learn the amount of a continuous "time shift" δ t between two consecutive states in the OU process. This essentially models parameter drift since e.g. if γ t = 1, then δ t = 0 and there is no "time shift" which means that the next state/parameter remains the same as the previous one, i.e. θ t+1 = θ t .In this section, we discuss alternative choices of a drift model instead of (5) .Independent mean and variance of the drift. We consider the drift model where the mean and the variance are not connected, i.e., p(θEQUATIONwhere γ t ∈ [0, 1] is the parameters controlling the mean of the distribution and β t is the learned variance. When β is fixed, this would be similar to our experiment in Figure 16 where we assumed 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 β t is a fixed parameter independent from γ t whereas Figure 16 , right corresponds to the case when β t = 1 -γ 2 t σ 0 , i.e., when we use the drift model ( 5). We see from the results, using drift model ( 5) leads to a better performance. In case when β t are learned, estimating the parameters of this model will likely overfit to the noise since there is a lot of degrees of freedom. [2] . When we do not use the mean of the initialization, we can use the following drift model p(θ t+1 |θ t , λ t , β t ) = N (θ t+1 ; λ t θ t ; β 2 t ) Similarly to the case of (20), estimating both parameters λ t and β t from the data will likely overfit to the noise.Arbitrary linear model. We can use the arbitrary linear model of the form p(θ t+1 |θ t , A t , B t ) = N (θ t+1 ; A t θ t ; B t ), but estimating the parameters A t and B t has too many degrees of freedom and will certainly overfit.We consider a Gaussian [11] approximation to Spike & Slab [26] prior p(θ t+1 |θ t , γ t ) = γ t p(θ t+1 |θ t ) + (1 -γ t )p 0 (θ t+1 ), which is a mixture of two distributions -a Gaussian p(θ t+1 |θ t ) = N (θ t+1 ; θ t , σ 2 ) centered around the previous parameter θ t and an initializing distribution p 0 (θ t+1 ) = N (θ t+1 ; µ 0 , σ 2 0 ). This model, however, implements the mechanism of Hard reset as opposed to the soft ones. Moreover, estimating such a model and incorporating it into a learning update is more challenging since the mixture of Gaussian is not conjugate with respect to a Gaussian which will make the KL term (34) to be computed only approximately via Monte Carlo updates. Using NN initializing distribution. In the drift model ( 5), we assume that the initial distribution over parameters is given by p 0 (θ) = N (θ; µ 0 ; σ 2 0 ). In practice, we have access to the NN initializer p init (θ) = N (θ; 0; σ 2 0 ) where µ 0 = 0 (for most of the NNs). This means that we can replace ϵ from (10) by 1 σ0 θ ′ 0 where θ ′ 0 ∼ p init (θ). This means that the term in ( 21) can be replaced byEQUATIONwhere we used the fact that σ 2 t (γ t ) = γ 2 t σ 2 t + (1 -γ 2 t )σ 2 0 . Note that in (22) , we only need to know the ratio σ 2 σ 2 0 rather than both of these. We will see that in Section 3.5, only this ratio is used for the underlying algorithm. Finally, in practice, we can tie p 0 (θ) to the specific initialization θ 0 ∼ p init (θ). It was observed empricially [34] that using a specific initialization in gradient updates led to better performance than using samples from the initial distribution. This would imply thatEQUATIONwith σ2 0 = p 2 σ 2 0 . The parameter p ≤ 1 accounts for the fact that the distribution p 0 (θ) should have lower than p init (θ) variance since it uses the specific initializaiton from the prior. This modification would imply the following modification on the drift model term ( 22)EQUATION)Tasks In this section we provide experimental details. As plasticity tasks, we use a randomly selected subset of size 10000 from CIFAR-10 [32] and from MNIST. This subset is fixed for all the tasks. Within each task, we randomly permute labels for every image; we call such problems random-label classification problems. We study two regimes -data efficient, where we do 400 epochs on a task with a batch size of 128, and memorization, a regime where we do only 70 epochs with a batch size of 128. As the main backbone architecture, we use MLP with 4 hidden layers each having a hidden dimension of 256 hidden units. We use ReLU activation function and do not use any batch or layer normalization. For the incoming data, we apply random crop, for MNIST to produce images of size 24 × 24 and for CIFAR-10 to produce images of size 28 × 28. We normalize images to be within [0, 1] range by dividing by 255. On top of that, we consider permuted MNIST task with a similar training ragime as in [34] -we consider a subset of 10000 images, with batch size 16 and each task is one epoch. As a backbone, we still use MLP with ReLU activation and 4 hidden layers. Moreover, we considered permuted Patch MNIST, where we permute patches, not individual pixels.In this case, we used a simple 4 layer convolutional neural network with 2 fully connected layers at the end.Metrics We use online accuracy as first metric with results reported in Appendix H. Moreover we use per-task Average Online Accuracy which isEQUATIONwhere a t i are the online accuracies collected on the task t via N timesteps.Baselines First baseline is Online SGD which sequentially learns over the sequence of task, with a fixed learning rate. Hard Reset is the Online SGD which resets all the parameters at task boundaries.L2 init [34] adds a regularizer λ||θ -θ 0 || 2 term to each Online SGD update where the regularization strength λ is a hyperparameter. Shrink & Perturb applies the transformation λθ t + σϵ, ϵ ∼ N (ϵ; 0, I) to each parameter before the gradient update. The hyperparameters are λ and σ.Soft Reset corresponds to one update (10) starting from 1 using 1 Monte Carlo estimate. We always use 1 Monte Carlo estimate for updating γ t as we found that it worked well in practice on these tasks. The hyperparameters of the method -σ 2 0 initial variance of the prior, which we set to be equal top 2 1N where N is the width of the hidden layer is a constant (hyperparameter). It always equals to p = 0.1. On top of that the second hyperparameter is s, such that σ t = sσ 0 , which controls the relative decrease of the constant posterior variance. This is the hyperparameter over which we sweep over. Another hyperparameter is the learning rate for learning γ t . For Soft Reset Proximal, we also have a proximal coefficient regularization constant λ. Besides that, we also sweep over the learning rate for the parameter. For the Bayesian Soft Reset, we just add an additional learning rate for the variance α σ and we do 1 Monte Carlo sample for each ELBO update.Hyper parameters selection and evaluation For all the experiments, we run a sweep over the hyperparameters. We select the best hyperparameters based on the smallest cumulative error (sum of all 1 -a t i throughout the training). We then report the mean and the standard deviation across 3 seeds in all the plots.Hyperparameter ranges . Learning rate α which is used to update parameters, for all the methods, is selected from {1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 1e-1, 5e-1, 1.0}. The λ init parameter in L2 Init, is selected from {10.0, 1.0, 0.0, 1e-1, 5e-1, 1e-2, 5e-2, 1e-3, 5e-3, 1e-4, 5e-4, 1e-5, 5e -5, 1e -6, 5e -6, 1e -7, 5e -7, 1e -8, 5e -8, 1e -9, 5e -9, 1e -10, }. For S&P, the shrink parameter λ is selected from {1.0, 0.99999, 0.9999, 0.999, 0.99, 0.9, 0.8, 0.7, 0.5, 0.3, 0.2, 0.1}, and the perturbation parameter σ is from {1e -1, 1e -2, 1e -3, 1e -4, 1e -5, 1e -6}. As noise distribution, we use the Neural Network initial distribution. For Soft Resets, the learning rate for γ t is selected from {0.5, 0.1, 0.05, 0.01, 0.005, 0.001, 0.0005, 0.0001}, the constant s is selected from {1.0, 0.95, 0.9, 0.8, 0.7, 0.6, 0.5, 0.3, 0.1}, the proximal cost λ in ( 41) is selected from {1.0, 0.1, 0.01}, the same is true for the proximal cost in the Bayesian method (38) . On top of that for the Bayesian method, we always use p (see Algorithm 2) equal to p = 0.05 and s = 0.9, i.e. the posterior is always slightly smaller than the prior. Finally for the Bayesian method we had to learn the variance with learning rate from {0.01, 0.1, 1, 10} range.In practice, we found that there is one learning rate of 0.1, which was always the best in practice for most of the methods and only proximal Soft Resets on memorization CIFAR-10 required smaller learning rate 0.01. This allowed us to significantly reduce the hyperparameter sweep.In this experiment, we consider a subset of 10000 images from MNIST (fixed throughtout all experiment) and a sequence of tasks. Each task is constructed by assigning either a true or a random label to each image from MNIST, where the probability of assignment is controlled by the experiment. The duration of each is controlled by the number of epochs with batch size of 128. As backbone we use MLP with 4 hidden layers and 256 hidden units and ReLU activation. For all the methods, the learning rate is 0.1. For Soft Resets, we use s = 0.9 and p = 1 and η γ = 0.01. Bayesian method uses proximal cost λ = 0.01. Detailed results are given in Figure 7 .We conduct experiments in the RL environments. We take the canonical implementation of Soft-Actor Critic(SAC) from Brax [15] repo in github, which uses 2 layer MLPs for both policy and Q-function. It employs ReLU activation functions for both. On top of that, it uses 2 MLP networks to parameterize Q-function (see Brax [15] ) for more details. To employ Soft Reset, we do the following. After we have collected a chunk of data (128) time-steps, we do one update (10) on γ t starting from 1 at every update of γ t , where γ t is shared for all the parameters within each layer of a Neural Network, separately for weights and biases. On top of that, since we have policy and value function networks, we have separate γ t for each of these. After the update on γ t , we compute θ t (γ t ) and α t (γ t ), see Section 3.5. After that, we employ the proximal objective (41) with a fixed regularization target θ t (γ t ). Concretely, we use the update rule (43) where for each update the gradient is estimate on the batch of data from the replay buffer. This is not exactly the same as what we did with plasticity benchmarks since there the update was applied to the same batch of data, multiple times. Nevertheless, we found this strategy effective and easy to implement on top of a SAC algorithm. In practice, we swept over the parameter s (similar for both, policy and the value function) which controls the relative learning rate increase in (18) . Moreover, we swept over the proximal regularization constant λ from eqn. (41) , which was different for the policy and for the value function. In practice, we found that using proximal constant of 0 for the policy led to the best empirical results. The range for the proximal constants λ was {0.1, 0.01, 0.001} and for s was {0.8, 0.9, 0.95, 0.97, 1.0}. We used p = 1 for all the experiments. For each experiment, we used a 3 hours of the A100 GPU with 40 Gb of memory.We provide the study of computational cost for all the proposed methods. Notations:• P be the number of parameters in the Neural Network• L is the number of layers• O(S) is the cost of SGD backward pass.• M γ -number of Monte Carlo samples for the drift model• M θ -number of Monte Carlo samples for the parameter updates (Bayesian Method).• K γ -number of updates for the drift parameter• K θ -number of NN parameter updates. Method Comp. cost Memory SGD O(S) O(P ) Soft resets γ per layer O(K γ M γ S + S) O(L + (M γ + 1)P ) Soft resets γ per param. O(K γ M γ S + S) O(P + (M γ + 1)P ) Soft resets γ per layer + proximal (K θ iters) O(K γ M γ S + K θ S) O(L + (M γ + 1)P ) Soft resets γ per param. + proximal (K θ iters) O(K γ M γ S + K θ S) O(P + (M γ + 1)P ) Bayesian Soft Reset Proximal (K θ iters) γ per layer O(K γ M γ S + 2M θ K θ S) P (L + (M γ + 2)P ) Bayesian Soft Reset Proximal (K θ iters) γ per param. O(K γ M γ S + 2M θ K θ S) P (P + (M γ + 2)P )Table 1 : Comparison of methods, computational cost, and memory requirements The general theoretical cost of all the proposed approaches is given in Table 1 . In practice, for all the experiments, we assume M γ = 1 and M = 1. Moreover, we used K γ = 1 and K θ = 1 for Soft Reset, K γ = 10 and K θ = 1 for Soft Reset with more computation. On top of that, for Soft Reset proximal and all Bayesian methods, we used K γ = 10 and K θ = 10. O(30S) 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 O(L) parameters gamma (one for each layer), parameters θ t with O(P ) and sampled parameters for drift model update which requires O(P ).Note that as Figure 9 suggests, it is beneficial to spend more computational cost on optimizing gamma and on doing multiple updates on parameters. However, even the cheapest version of our method Soft Resets still leads to a good performance as indicated in Figure 2 .The complexity of soft resets in reinforcement learning setting requires only one gradient update on γ after each new chunk of fresh data from the environment. In SAC, we do G gradient updates on parameters for every new chunk of data. Assuming that complexity of one gradient update in SAC is O(S), soft reset only requires doing one additional gradient update to fit γ parameter.The computation complexity of Soft Reset in Reinforcement Learning is marginally higher than SAC but leads to better empirical performance in a highly off-policy regime, see Appendix D.3. 2 for complexity analysis. The y-axis is the average performance on all the tasks across the stream.Comp. cost Memory SAC O(GS) O(P ) Soft resets γ per layer O(S + GS) O(L + 2P ) Table 3 : Comparison of methods, computational cost, and memory requirements for methods in for RL.We study the sensitivity of Soft Resets where γ is defined per layer when trained on random-label MNIST (data efficient). We fix the learning rate to α = 0.1. We study the sensitivity of learning rate for the drift parameter, η γ , as well as p -initial prior standard deviation rescaling, and s -posterior standard deviation rescaling parameter.On top of that, we conduct the sensitivity analysis of L2 Init [34] and Shrink&Perturb [2] methods. The x-axis of each plot denotes one of the studied hyperparameters, whereas y-axis is the average performance across all the tasks (see Experiments section for tasks definition). The standard deviation is reported over 3 random seeds. A color indicates a second hyperparameter which is studied, if available. In the title of each plot, we write hyperparameters which are fixed. The analysis is provided in Figure 10 for Soft Resets and in Figure 11 for the baselines.The most important parameter is the learning rate of the drift model η γ . For each method, there exists a good value of this parameter and performance is sensitive to it. This makes sense since this parameter directly impacts how we learn the drift model.The performance of Soft Resets is robust with respect to the posterior standard deviation scaling s parameter as long as it is s ≥ 0.5. For s < 0.5, the performance degrades. This parameter is defined from σ t = sσ 0 and affects relative increase in learning rate given by 1 γ 2 +(1-γ 2 )/s 2 ) which could be ill-behaved for small s.We also study the sensitivity of the baseline methods. We find that L2 Init [34] is very sensitive to the parameter λ, which is a penalty term for λ||θ -θ 0 || 2 . In fact, Figure 11 , left shows that there is only one good value of this parameter which works. Shrink&Perturb [2] is very sensitive to the shrink parameter λ. Similar to L2 Init, there is only one value which works, 0.9999 while values 0.999 and values 0.99999 lead to bad performance. This method however, is not very sensitive to the perturb parameter σ provided that σ ≤ 0.001.Compared to the baselines, our method is more robust to the hyperparameters choice. Below, we also add sensitivity analysis for other method variants. Figure 12 shows sensitivity of Soft Resets, K γ = 10, Figure 13 shows sensitivity of Soft Resets, K γ = 10, K θ = 10, Figure 14 shows sensitivity of Bayesian Soft Resets, K γ = 10, K θ = 10 with γ t per layer, Figure 15 shows sensitivity of Bayesian Soft Resets, K γ = 10, K θ = 10 with γ t per parameter. (Left) shows sensitivity of L2 Init with respect to the L2 penalty regularization cost λ applied to ||θ -θ 0 || 2 term. We do not use an additional hyperparameter, therefore there is only one color.(Right) shows sensitivity of Shrink&Perturb method where the x-axis is the perturb parameter σ while the color indicates the shrink parameter λ.Each step of online SGD can be seen in terms of a regularized minimization problem referred to as the proximal form [44] :EQUATION)In general, we cannot solve (26) directly, so we consider a Taylor expansion of L t+1 around θ t , givingEQUATION)Here we see the role of α t > 0 as both enforcing that the Taylor expansion around θ t is accurate, and regularising θ t+1 towards the old parameters θ t (hence ensuring that the learning from past data is not forgotten). Solving (27) naturally leads to the well known SGD update:θ t+1 = θ t -α t ∇ θ L t+1 (θ t ),where α t can now also be interpreted as the learning rate. Figure 16 : Perfect soft-resets on data-efficient random-label MNIST. Left, Soft Reset method does not use higher learning rate when γ < 1. Right, Soft Reset increases the learning rate when γ < 1, see (18) . The x-axis represents task id, whereas the y-axis is the average training accuracy on the task.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 boundaries, and Hard Reset (only last) which resets only the last layer. We use Soft Reset method (19) where γ t = 1 all the time and becomes γ t = γt (with manually chosen γt ) at task boundaries. We consider constant learning rate α t (γ t ) and increasing learning rate (18) We observe what values of γ t we get as we train Soft Reset method on random-label MNIST (dataefficient) and CIFAR-10 (memorization). The results are given in Figure 17 for MNIST and in Figure 18 for CIFAR-10. We report these for the first 20 tasks. We consider a version of permuted MNIST where instead of permuting all the pixels, we permute patches of pixels with a patch size varying from 1 to 14. The patch size of 1 corresponds to permututed MNIST and therefore the most non-stationary case, while patch size of 14 corresponds to least nonstationary case. We use a convolutional Neural Network in this case. In Figure 19 , we report the behavior of γ for different convolutional and fully connected layers on first few tasks.As discussed in Section 5, we found that in practice Soft Reset and Soft Reset Proximal where γ is learned per-parameter, did not perform well on the plasticity benchmarks. However, the Bayesian variant described in Section I.1, actually benefited from specifying γ for every parameter in Neural Network. We report these additional results in Figure 20 . We see that the non Bayesian variants where γ t is specified per parameter, do not perform well. The fact that the Bayesian method performs better here suggests that it is important to have a good uncertainty estimate σ 2 t for the update (10) on γ t . When, however, we regularize γ t to be shared across all parameters within each layer, this introduces useful inductive bias which mitigates the lack of uncertainty estimation in the parameters. This is because for non-Bayesian methods, we assume that the uncertainty is fixed, given by a hyperparameter -assumption which would not always hold in practice.In this section, we zoom-in in the data-efficient experiment on random-label MNIST. We use Soft Reset Proximal (γ per layer) method with separate γ for layer (different for each weight and for each bias) and run it for 20 tasks on random-label MNIST. In Figure 21 we show the online accuracy as we learn over this sequence of tasks. In Figure 22 , we visualize the dynamics of parameters γ for each layer. First of all, we see that γ t seems to accurately capture the task boundaries. Second, we see that the amount by which each γ t changes depends on the parameter type -weights versus biases, and it depends on the layer. The architecture in this setting starts form linear and goes up to linear4, which represent the 4 MLP hidden layers with a last layer linear4. In this section, we study the impact of using specific initialization θ 0 ∼ p init (θ) in p 0 (θ) as discussed in Appendix C. Using the specific initialization in Soft Resets leads to fixing the mean of the p 0 (θ) to be θ 0 , see (23) . This, in turn, leads to the predictive distribution (24) . In case when we are not using specific initialization θ 0 , the mean of p 0 (θ) is 0 and the predictive distribution is given by (22) . To understand the impact of this design decision, we conduct an experiment on random label MNIST with Soft Reset, where we either use the specific initialization or not. For each of the variants, we do a hyperparameters sweep. The results are given in Figure 23 . We see that both variants perform similarly.In this section, we provide a Bayesian Neural Network algorithm to learn the distributions of NN parameters when there is a drift in the data distribution. Moreover, we provide a MAP-like inference algorithm which does not require to learn the distributions over parameters, but simply propagates the MAP estimate over these.In this section, we describe an algorithm for parameters update based on Bayesian Neural Networks (BNN). It is based on the online variational Bayes setting described below.Let the family of distributions over parameters beEQUATIONwhich is the family of Gaussian mean-field distributions over parameters θ ∈ R D (separate Gaussian per parameter). For simplicity of notation, we omit the index i. Let Γ t = (γ 1 , . . . , γ t ) be the history of observed parameters of the drift model and S t = {(x 1 , y 1 ), . . . , (x t , y t )} be the history of observed data. We denote by q t (θ) ≜ q t (θ|S t , Γ t-1 ) ∈ Q the Gaussian approximate posterior at time t with mean µ t and variance σ 2 t for every parameter. The approximate predictive look-ahead prior is given by q t (θ|γ t ) = q t (θ t )p(θ|θ t , γ t )dθ t = N (θ; µ t (γ t ), σ 2 t (γ t )),that has parameters µt (γ t ) = γ t µ t + (1 -γ t )µ 0 , σ 2 t (γ t ) = γ 2 t σ 2 t + (1 -γ 2 t )σ 2 0 .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 saysE[Y ] = E[E[Y |X]]and the law of total variance saysV[Y ] = E[V[Y |X]] + V[E[Y |X]]In our case, from the drift model ( 5), we have the conditional distributionEQUATIONFrom (30), we haveE[θ|θ t ] = γ t θ t + (1 -γ t )µ 0 V[θ|θ t ] = (1 -γ 2 t )σ 2 0From here, we have that the mean is given byEQUATIONand the variance is given byEQUATION)Now, we note that q t (θ|γ t ) is a Gaussian and its parameters are given by E[θ] = γ t µ t + (1 -γ t )µ 0 from (31) and by V[θ] = (1 -γ 2 t )σ 2 0 + γ 2 t θ 2 t from (32) . Then, for new data (x t+1 , y t+1 ) at time t + 1, the approximate predictive log-likelihood equals to log q t (y t+1 |x t+1 , γ t ) = log p(y t+1 |x t+1 , θ)q t (θ|γ t )dθ.We are looking for a new approximate posterior q t+1 (θ) such that q t+1 (θ) = arg min q KL [q(θ)||p(y t+1 |x t+1 , θ)q t (θ|γ t )]The optimization problem (33) can be written as minimization of the following loss F t (θ, γ t ) = E q [L t+1 (θ)] + KL [q(θ)||q t (θ|γ t )] ,since L t+1 (θ) = -log p(y t+1 |x t+1 , θ). Using the fact that we are looking for a member q ∈ Q from (28), we can write the objective (34) as F t (µ, σ, γ t ) = E ϵ∼N (0;I) [L t+1 (µ + ϵσ)] + KL [q(θ)||q t (θ|γ t )] ,where we used the reparameterisation trick for the loss term. We now expand the regularization term to getEQUATION)Since the posterior variance of NN parameters may become small, the optimization of ( 35) may become numerically unstable due to division by σ 2 t,i (γ t ). It was shown [54] that using small temperature on the prior led to better empirical results when using Bayesian Neural Networks, a phenomenon known as cold posterior. Here, we define a temperature per-parameter, i.e., λ t,i > 0 for every time-step t, such that the objective above becomes Ft (µ, σ, γ t ; {λ t,i } i ) = E ϵ∼N (0;I) [L t+1 (µ + ϵσ)]EQUATION)As said above, it is common to use the same temperature λ t,i = λ for all the parameters. In this work, we propose the specific choice of the temperature to be λ t,i = λσ 2 t,i ,where λ > 0 is some globally chosen temperature parameter. This leads to the following objective Ft (µ, σ, γ t ; {rt, i} i ) = E ϵ∼N (0;I) [L t+1 (µ + ϵσ)]+ 1 2 i r t,i (µ i -µ t,i (γ t )) 2 + σ 2 i -σ 2 t,i (γ t ) log σ 2 i , (38) where the quantity r t,i is defined asr t,i = σ 2 t,i σ 2 t (γt) = σ 2 t,i γ 2 t σ 2 t,i +(1-γ 2 t )σ 2 0 ,which represents the relative change in the posterior variance due to the drift. In the exact stationary case, when γ t = 1, this ratio is r t,i = 1 while for γ t < 1 , since typically σ 2 t < σ 2 0 , we have r t,i < 1. This means that in the non-stationary case, the strength of the regularization in (38) in favor of the data term E ϵ∼N (0;I) [L t+1 (µ + ϵσ)], allowing the optimization to respond faster to the change in the data distribution. In practice, this data term is approximated via Monte-Carlo, i.e. To find new parameters, µ t+1 and σ t+1 , we let µ 0 t+1 = µ t (γ t ) and σ 0 t+1 = σ t (γ t ) and perform multiple K updates on (38) µ k+1 t+1 = µ k t+1 -α µ Ft (µ k t+1 , σ k t+1 , γ t , {rt, i} i ), σ k+1 t+1 = σ k t+1 -α σ Ft (µ k t+1 , σ k t+1 , γ t , {rt, i} i ), where α µ and α σ are corresponding learning rates. The full algorithm of learning the drift parameters γ t as well as learning the Bayesian Neural Network parameters using the procedure above is given in Algorithm 2. where λ is some constant. Such choice of temperature is motivated by the same logic as in (37) -it is a constant multiplied by the posterior variance σ 2 t,i = s 2 σ 2 0,i . With such choice of temperature, maximizing (40) is equivalent to minimizing G(θ; λ) = L t+1 (θ) + λ 2 D i=1 |θ i -θ i t (γt)| 2 rt,i(γ) (41) where the regularization target for the dimension i isEQUATION)and the constant r t,i (γ) is given byr t,i (γ t ) = (γ i t ) 2 + 1-(γ i t ) 2 s 2We can perform K gradient updates (41) with a learning rate α t starting from θ 0 t+1 = θ t (γ t ),EQUATIONwhere the vector-valued learning rate α t (γ t ) is given byEQUATION)with α t the base learning rate. Note that doing one update is equivalent to modified SGD method (19) . Doing multiple updates on (43) allows us to perform multiple computations on the same data. The corresponding algorithm is given in Algorithm 3.Interpretation of γ t . By linearising log p(y t+1 |x t+1 , θ) around µ t , we can simplify (8) to get F(γ t ) = (γ t ⊙ µ t + (1 -γ t ) ⊙ µ 0 ) T g t+1 -0.5(σ 2 t (γ t ) ⊙ g t+1 ) T g t+1 -λ K i=1 (γ t,i -γ 0 t,i ) 2 , where ⊙ denotes elementwise product, g t = -∇L t+1 (µ t ) is the negative gradient of the loss (1) evaluated at µ t and we added the ℓ 2 -penalty 1 2 λ(γ t,i -γ 0 t,i ) 2 to take into account the initialization.Since Σ t (γ t ) = diag(σ K Toy illustrative example for SGD underperformance in the non-stationary regimeIllustrative example of SGD on a non-stationary stream. We consider a toy problem of tracking a changing mean value. Let the observations in the stream S t follow y t = µ t + σϵ, where ϵ ∼ N (0, 1), σ = 0.01. Every 50 timesteps the mean µ t switches from -2 to 2. We fit a 3-layer MLP with layer sizes (10, 5, 1) and ReLU activations, using SGD with two different choices for the learning rate: α = 0.05 and α = 0.15. Moreover, given that we know when a switch of the mean happens, we reset (or not reset) all the parameters at every switch as we run SGD. Only during the reset, we use different learning rate β = 0.05 or β = 0.15. Using higher learning rate during reset allows SGD to learn faster from new data. We also ran SGD with α = 0.05 and β = 0.15, where the higher learning rate is used during task switch but we do not reset the parameters. We found that it performed the same as SGD with α = 0.05, which highlights the benefit of reset. We report the predicted mean μt for all SGD variants in Figure 24 . We see that after the first switch of the mean, the SGD without reset takes more time to learn the new mean compared to the version with parameters reset. Increasing the learning rate speeds up the adaptation to new data, but it still remains slower during the mean change from 2 to -2 compared to the version that resets parameters. This example highlights that resets could be highly beneficial for improving the performance of SGD which could be slowed down by the implicit regularization towards the previous parameters θ t and the impact of the regularization strength induced by the learning rate. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: Although the proposed solution is novel and effective, the paper has many weaknesses:
1. 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.
2. 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.
3. 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.
questions: questions: See Weaknesses
limitations: limitations: Computational cost and hyperparameter sensitivity need to be discussed in the conclusion of the main paper.
paper_summary: paper_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.
strength_weakness: strength_weakness: Some parts can be improved - 
1) Section 2 - Simplification of equation (6) to (9) requires some explanation [page -5].
2) Section 3 - Some intuition behind the objective (11) [page - 6] would be good to understand it better (especially the second term in that objective)
3) Some discussion on how does the methodology perform relative to the degree of non-stationarity in the data?
4) Some minor typos in line 109, 154, 190
questions: 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: limitations: Authors have adequately addressed the limitations of their work.
paper_summary: paper_summary: - The authors study a learning algorithm that can handle the non-stationarity of the data distribution.
- 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.
- 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.
- They numerically corroborate the efficacy of their method on plasticity benchmarks and reinforcement learning experiments.
strength_weakness: strength_weakness: - W1. Modeling parameter drift is not well-motivated.
    - 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.
    - 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.
- W2. The paper is NOT self-contained overall.
    - 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.
    - 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).
    - For these reasons, the paper is not so easy to follow.
- W3. The soft reset method seems computationally heavy.
    - 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.
    - In addition, there seem too many hyperparameters to tune.
- W4. Minor comments on typos/mistakes
    - Do not capitalize the word “neural network”. (e.g., line 35-36)
    - 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”.
    - Line 84: “non-stationary” → “non-stationarity”
    - Line 109: “mportant” → “important”
    - Line 127: “wrt” → “with respect to”. Do not use abbreviations.
    - Line 137: “property” → “properties”
    - Line 154: what is “s” next to a parenthesis?
    - Line 168: “$\theta = (\theta_i, \ldots, \theta_D)$” → “$\theta = (\theta_1, \ldots, \theta_D)$”
    - Line 169: “$\mathbb{R}^D$” is a more standard notation than “$\mathcal{R}^D$”.
    - Line 222: In “$\lambda\_i = \hat{\lambda}\_i \sigma^2_{t,i}$", I think $\hat{\lambda}$ must not have a subscript.
    - Line 225: “sinc” → “since”
    - Equation (14): What is “$\tilde{\mathcal{F}}$”? (Partial) gradient of “$\mathcal{F}$”?
    - Line 232: What does it mean by “we assume that $\mu_0 = \theta_0, \sigma_0^2.$”?
    - Line 236: “$\theta_{t,i}(\gamma_{t,i}) = \theta_t$” → “$\theta_{t,i}(\gamma_{t,i}) = \theta_{t,i}$”
    - Line 240: “linearisng” → “linearizing” or “linearising”
    - Line 292: duplicate “See”
    - Line 354: “modelling” → “modeling”
    - Lines 492-493: I guess “*data efficient*” and “*memorization*” are swapped.
questions: 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: limitations: The paper discusses its limitations in the experiment section. I think it might be better to mention them explicitly in the conclusion section.
paper_summary: paper_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.
strength_weakness: strength_weakness: In my view, the paper has some weaknesses, especially regarding the experimental validation:
- 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. 
- 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? 
- 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. 

Finally, 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. 

Minor: in page 6, Equation (15), what is the parameter $\tilde{\lambda}$? Was it introduced before?
questions: questions: I have the following questions:
- 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. 
- Can the authors also clarify my concern regarding the reinforcement learning setting and the performance in on-policy settings?
- Can the authors clarify what is the intended take-away from the analysis in Figure 2?
- 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: limitations: Yes.

Comments: title: Response
Part 2
Response to a response: Part 1
Response to a response
Response to authors' rebuttal
The response is not satisfactory yet
comment: We thank the Reviewer 4koq for their response and for sparking interesting discussion! Please see our answer below.

> I am concerned about that, in the author's explanation, the local optimum $\theta^{\ast}_t$ seems a unique deterministic value/vector. However, in the era of overparameterized deep learning and because of its nonconvex nature, the $\theta^{\ast}_t$ may not be unique; rather, there might be a set of (non-unique) local optima.

Please note that when we discuss local optimas $\theta^*_{t}$ and $\theta^*_{t+1}$, we do not assume their uniqueness. Rather, the only assumption we make is that our algorithm may converge to some local optima $\theta^*_{t}$ (which is a member of a set of local optima), and after non-stationarity, a new set of local optima will appear, where one of the member is $\theta^*_{t+1}$. The reason why we use two specific points $\theta^*_{t}$ and $\theta^*_{t+1}$ in Figure 1, is to simplify the understanding and communicate the intuition behind the non-stationary behaviour and the impact of a drift model.

In the paper, in Section 3.1, we mainly refer to a region of new local minima, please see, for example, Lines 110-112: "If the change is such that the region of new good local minima is far away from $\theta^*$, SGD may take considerably higher amount of time to move its parameters towards this region."

Hope this clarifies the concerns.

> About linearization: I am already aware of the derivation explained for reviewer 6GY1. All I wanted is just a kind derivation in your next revised paper. I wanted to take part in improving your paper.

Thanks to the useful discussion with you and the concerns you have raised, we think that including the full derivation of the linearization in the appendix of the paper, will greatly reduce confusion shared by you and other reviewers. We thank you for your suggestions, we will add more context in the paper.
I thank the authors for a couple of more detailed replies.

1. About original drift model: Thank you for resolving most of my confusion. Now I think I understand the concept of drift model. Still, I am concerned about that, in the author's explanation, the local optimum $\theta^{\ast}\_t$ seems a unique deterministic value/vector. However, in the era of overparameterized deep learning and because of its nonconvex nature, the $\theta^{\ast}\_t$ may not be unique; rather, there might be a set of (non-unique) local optima. Because of this, I thought the drift model is a kind of mechanism of choosing a proper/desired $\theta^{\ast}\_t$ among all local optima. Although I admit that $\theta^{\ast}\_t$ not a random object (but still it can be because of randomness in data sampling) and $\theta_t$ is an estimate of it (thus it is definitely a random object), considering the non-uniqueness of local optima, it seems vague and non-rigorous for me to say "getting closer to $\theta^{\ast}\_t$" or "concentrating aroung $\theta^{\ast}\_t$". 

2. About linearization: I am already aware of the derivation explained for reviewer 6GY1. All I wanted is just a kind derivation in your next revised paper. I wanted to take part in improving your paper.

3. About variations of drift model: thank you for clearly answering my question. Personally, I was excited for suggesting a new point of view of your methodology, but it turns out it was somewhat meaningless (which is sad). Anyway, I am happy to understand the significance of the original formulation of the drift model.

Thank you again for the detailed response. Nonetheless, in my opinion, it would be better to judge a completely rewritten manuscript with another review process. Therefore, I maintain my score.
> Eq.9, linearisation

As we stated above in the rebuttal, the response to Reviewer 6GY1 (due to space constraints) contains the partial derivation. Linearisation is $\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)$, where $g_{t+1}=\nabla_{\theta} -\log p(y_{t+1}|x_{t+1},\theta=\mu_t)$. Then, $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}$. Since, $q_t^{t+1}(\theta | S_{t}, \Gamma_{t-1}, \gamma_t)$ is Gaussian, we compute the integral eq.6 in a closed form and keep only the terms depending on $\gamma_t$. We provide derivation for a scalar case here and full derivation in the appendix. In eq.6, we have
$\log\int p(y_{t+1}|x_{t+1},\theta)\exp^{-\frac{(\theta -\mu_t(\gamma_t))^2}{2\sigma^2(\gamma_t)}}d\theta =\log \frac{1}{\sqrt{2\pi\sigma^2(\gamma_t)}}\int p(y_{t+1}|x_{t+1},\mu_t)\exp^{-g_{t+1}(\theta-\mu_t)} \exp^{-\frac{(\theta-\mu_t(\gamma_t))^2}{2\sigma^2(\gamma_t)}}d\theta$

which becomes

$\log p(y_{t+1}|x_{t+1},\mu_t)+\log \int \frac{1}{\sqrt{2\pi\sigma^2(\gamma_t)}}\exp^{-g_{t+1}(\theta-\mu_t)}\exp^{-\frac{(\theta-\mu_t(\gamma_t))^2}{2\sigma^2(\gamma_t)}}d\theta$ 

Consider only the exponent term in the integral

$\frac{-1}{2\sigma^2(\gamma_t)}(2\sigma^2(\gamma_t)g_{t+1}(\theta-\mu_t)+\theta^2-2\theta \mu_t(\gamma_t)+\mu_t(\gamma_t)^2)$

then

$\frac{-1}{2\sigma^2(\gamma_t)}\left(\theta^2-2\theta \left[\mu_t(\gamma_t)-\sigma^2(\gamma_t)g_{t+1}\right]+\mu_t(\gamma_t)^2-2\sigma^2(\gamma_t)g_{t+1}\mu_t\right)$

then

$\frac{-1}{2\sigma^2(\gamma_t)}\left[ \left(\theta-(\mu_t(\gamma_t)-\sigma^2(\gamma_t)g_{t+1}) \right)^2+2\mu_t(\gamma_t) \sigma^2(\gamma_t)g_{t+1}-\sigma^4(\gamma_t)g_{t+1}^2-2\sigma^2(\gamma_t)g_{t+1}\mu_t)\right]$

then

$\frac{-1}{2\sigma^2(\gamma_t)} \left(\theta-(\mu_t(\gamma_t)-\sigma^2(\gamma_t)g_{t+1})\right)^2-\mu_t(\gamma_t)g_{t+1}+0.5\sigma^2(\gamma_t)g^2_{t+1}+g_{t+1}\mu_t$

Now, get back to the integral

$\log\int\frac{1}{\sqrt(2\pi\sigma^2(\gamma_t)}\exp^{\frac{-1}{2\sigma^2(\gamma_t)}\left(\theta-(\mu_t(\gamma_t)-\sigma^2(\gamma_t)g_{t+1}) \right)^2-\mu_t(\gamma_t)g_{t+1}+0.5\sigma^2(\gamma_t)g^2_{t+1}+g_{t+1}\mu_t}d\theta$ 

Which equals to

$1-\mu_t(\gamma_t)g_{t+1}+0.5 \sigma^2(\gamma_t)g^2_{t+1}+g_{t+1}\mu_t$

We keep only the terms which depend on $\gamma_t$ and get

$G(\gamma_t)=-\mu_t(\gamma_t)g_{t+1}+0.5 \sigma^2(\gamma_t)g^2_{t+1} $

We want to maximize $G(\gamma_t)$ or minimize $F(\gamma_t)=-G(\gamma_t)$ which is exactly the loss we defined in eq.9.

We will add the full derivation in the appendix.

> Different drift model

Assuming that $\theta'\sim\mathcal{N}(\theta;\mu_0;\sigma^2_0)$, the model you wrote will have the mean $\gamma_t\theta_t+(1-\gamma_t)\mu_0$ and the variance $\gamma^2_t \sigma^2_t+(1-\gamma^2_t)\sigma^2_0$, which are the same as for our OU model, see eq.4 and Line 173 after eq.5, therefore they are mathematically equivalent.

In order to discuss it in detail, let $\mu_0$ is the prior mean for eq.4 and $\mu'_0$ is the prior mean in your variant, i.e., $\theta'\sim\mathcal{N}(\theta;\mu'_0;\sigma^2_0)$. When we incorporate the Gaussian model drift in the update, we are only concerned with its mean and variance since these affect SGD updates (see eq.16).

We have the following three cases:

* $\mu_0=\mu_0'=0$, the models have the same mean equal to $0$. They are mathematically equivalent to the model $\theta_{t+1}=\gamma_t\theta_t+\sqrt{1-\gamma^2_t}\theta_0$, where $\theta_0\sim p_0(\theta)=\mathcal{N}(\theta;0; \sigma^2)$ -- the models are the same. The have the mean $\gamma_t\mu_t$ and the variance $\gamma^2_t\sigma^2_t+(1-\gamma^2_t)\sigma^2_0$.

* $\mu_0=\mu'_0=\theta_0\neq0$, the models have the same mean which are the one specific initialization of NN parameters. They are mathematically equivalent with the mean $\gamma_t\mu_t+(1-\gamma_t)\theta_0$ and the same variance (as above)

* $\mu_0=\theta_0$ and $\mu'_0 \neq \theta_0$. These two models have the same variance, $\gamma_t^2 \sigma^2_t + (1-\gamma^2_t)\sigma^2_0$, but have different means: $\gamma_t\theta_t+(1-\gamma_t)\theta_0$ versus $\gamma_t\theta_t+(1-\gamma_t)\mu'_0$.

To implement variant 3, the mean $\mu'_0$ is fixed at the beginning of the training, because otherwise we would be under-counting the variance in the corresponding drift model. It is unclear how to interpret this model nor how it will affect performance.

Thus, in order to answer your question about the impact of the particular initialization, we believe variant 1 is the right answer. We conducted an experiment of using Soft Reset variant 1 and 2 on random label MNIST (data efficient), where we swept over hyperparameters. We found that these variants had similar performance on this benchmark. We will add this experiment to the paper in the Appendix.

We hope your answer clarifies the questions and concerns you have raised, and we thank you for raising interesting points for discussion!
Dear 4koq, thank you for your feedback and please find our answer to your concerns below.

> I don’t think this answer really responds to my question in W1. Here, I am not asking about the motivation of the drift model (although I did it in W2). My question is: why do we care about the conditional distribution of 
Θt+1 , a consequence of a single step of training from  θt , even though we want to estimate or utilize the dynamics of the moving local optima  θt∗↦θt+1∗?
After pondering this issue and staring at Figure 1, I suddenly realized that the drift model is actually modeling the dynamics of  θt+1∗  given  θt.. So in my understanding, the  θt+1  in “p(θt+1∣θt,γt)” actually means  θt+1∗ , the local optima at time  t+1, rather than the parameter we will have by updating the parameter from θt  using a single step of learning algorithm. (Or, at least, we want to find a local optimum θt+1∗ close to the initialization & θt .) Is my understanding correct? If it is, I think this should have been explained in the paper in more detail; also, the term “parameter drift” is a bit misleading in this sense.

We apologize for the confusion, but under the Bayesian perspective depicted in Figure 1 $\theta^*_{t}$ and $\theta^*_{t+1}$ denote **fixed/deterministic** values and therefore are not random variables, while the corresponding random variables which are assigned probability distributions are $\theta_{t}$ and $\theta_{t+1}$. More specifically, we learn the dynamical model  distribution $p(\theta_{t+1}|\theta_{t},\gamma_t)$, by learning $\gamma_t$, so that the specific value $\theta^*_{t+1}$ can become more likely or “explainable” under the distribution $p(\theta_{t+1}|\theta_{t},\gamma_t)$. In other words we hope that $p(\theta_{t+1}|\theta_{t},\gamma_t)$ will place a considerable probability mass around the fixed value $\theta^*_{t+1}$.   We are going to update the paper  to make the above clear and remove the confusion.   

To further explain here this in Figure 1, consider the stationary case depicted in Figure,1a. There the posterior concentrates around an optimal value (which is fixed value $\theta^*$), in other words $q_t(\theta)$ will gradually converge to a delta mass around the value $\theta^*$.  In a non-stationary case, the optimal value can suddenly change over time and e.g. from the value $\theta^*_{t}$ at time $t$ can change to $\theta^*_{t+1}$ . Without a dynamical model  (Figure 1,b)  the new posterior $q_{t+1}(\theta)$ after observing the new data at time $t+1$ has a small radius/variance (blue dashed circle) and it cannot concentrate fast enough towards the new optimum.  The use of  the dynamical or drift model (Figure 1,c) introduces an “intermediate step” that constructs an “intermediate prior distribution”  $p_t(\theta) = \int q_t(\theta) p(\theta | \theta_t, \gamma_t) d \theta_t$ which can have increased variance and shifted mean  (green dashed circle). Then  once we fully incorporate the  new data point at time $t+1$ the updated posterior $q_{t+1}(\theta) \propto p(y_{t+1}|\theta) p_t(\theta)$ can better concentrate around $\theta_{t+1}^*$. We will update the paper to clarify the above.

> citation [2-3] is a typo

Thank you, it should be [1-2].

> Q2

Yes we will add this to the paper (we forgot to explicitly write it).

> Q3

Thank you.

> Eq.11&12 -- usual BNN objective?

Yes, except for the temperature defined per-parameter. In BNN, temperature $\lambda$ is either 1, or fixed to a constant for all the parameters.

> Line 232: What does it mean by “we assume that $\mu_0=\theta_0,\sigma^2_0.$?

It means that the mean of the prior $p_0(\theta)=\mathcal{N}(\theta;\mu_0,\sigma^2)$ is fixed to $\theta_0$, i.e. the initialization and $\sigma^2_0$ is a constant. Normally NNs are initialized from $p_0(\theta)=\mathcal{N}(\theta;0,\sigma^2_0)$, but we make the prior to be initialization-dependent. We will rewrite this sentence to explicitly highlight this.

We hope your answer clarifies the questions and concerns you have raised, and we thank you for raising interesting points for discussion!
Dear Reviewer DfC5, thank you for your response.

> Regarding Figure 3, the authors replied that their method is Soft Reset, but I see it being outperformed by several of its variants. Then why do the authors choose as their proposed method Soft Reset and not one of the variants that outperform it?

It is true that it is outperformed by other variants. Soft Reset is the cheapest method among all the other variants -- it requires only one update on $\gamma_t$ and one update on parameters $\theta_{t+1}$. All the other methods -- Soft Reset more compute, Soft Reset proximal, Bayesian Soft Reset proximal, require significantly more compute. Figure 3 suggests that we can in fact leverage more compute to achieve better empirical performance.

This is why we chose to present Soft Reset as the main method because it is relatively cheap and already achieves good performance compared to external baselines. However, when more compute is available, other methods could also be used to further improve performance.

Hope this clarifies the confusion.
I thank the authors for the response.

I had some points clarified. However, I am inclined to maintain my score, since the experimental validation is somehow confusing. Regarding the RL results, I can not see clear benefits from the approach. Regarding Figure 3, the authors replied that their method is Soft Reset, but I see it being outperformed by several of its variants. Then why do the authors choose as their proposed method Soft Reset and not one of the variants that outperform it?

Thank you.
I have read the author’s rebuttal, but I still have remaining concerns and questions. Let me leave my comments/questions about the author’s responses one at a time.

1. I don’t think this answer really responds to my question in W1. Here, I am not asking about the motivation of the drift model (although I did it in W2). My question is: why do we care about the conditional distribution of $\theta_{t+1}$, a consequence of a single step of training from $\theta_t$, even though we want to estimate or utilize the dynamics of the moving local optima $\theta^{\ast}\_t \mapsto \theta^{\ast}\_{t+1}$?
    - After pondering this issue and staring at Figure 1, I suddenly realized that the drift model is actually modeling the dynamics of $\theta^{\ast}\_{t+1}$ given $\theta_t$. So in my understanding, the $\theta_{t+1}$ in “$p(\theta_{t+1} \mid \theta_t, \gamma_t)$” actually means $\theta^{\ast}_{t+1}$, the local optima at time $t+1$, rather than the parameter we will have by updating the parameter from $\theta_t$ using a single step of learning algorithm. (Or, at least, we want to find a local optimum $\theta^{\ast}\_{t+1}$ close to the initialization & $\theta_t$.) Is my understanding correct? If it is, I think this should have been explained in the paper in more detail; also, the term “parameter drift” is a bit misleading in this sense.
    - By the way, I guess the citation ([2-3]) has a typo, right? It seems that it should include [1].
2. Thank you for the explanation of the relationship between the OU process, the Gaussian Markov process, and the choice of the drift model. Will the authors add this explanation to their paper? Or do they just ignore it? 
3. Thank you for the discussion on the other options for drift models and for considering putting the discussion to the paper.
4. Thank you for derivations of the equations in more detail. However, I think the derivations could be more kind than the current explanation.
    - Eq 9: Please state to which parameter the linearization is taken (e.g., By linearizing log… around $\theta=\mu_t$ (cf. Line 198))
    - Eq. 11 & 12: Is these are just the usual objective functions in BNN training?
5. Thank you for the interesting response based on the drift model and the variance. But what if we change the **variance scheduling**: what if we suitably change the drift model into $p(\theta \mid \theta_t, \gamma_t) = \mathcal{N} (\theta; \gamma_t \theta_t + (1-\gamma_t) \theta’_0(t); 2\gamma_t(1-\gamma_t)\sigma_0^2)$ where $\theta’_0(t) \sim p_0(\theta_0)$? I guess this model does not suffer from the same problem of variance inflation: If $\sigma_t = \sigma_0$, then ${\rm Var}(\theta) = \gamma_t^2 \sigma_t^2 + (1-\gamma_t)^2 \sigma_0^2 + 2\gamma_t (1-\gamma_t) \sigma_0^2 = \\{\gamma_t^2 + (1-\gamma_t)^2 + 2\gamma_t (1-\gamma_t)\\}\sigma_0^2 = \sigma_0^2$...! The motivation behind this is the concern about the dependency on the particular initialization “$\theta_0$” which seems to be fixed throughout the training: “What if the initialization distribution is only the matter?” Please let me know if there are other problems in this drift model, or I would be very happy if you could test this drift model empirically (but I understand if it is impossible due to the time limit)!

Also, there is a question that is not answered yet. Let me bring it here:

- Line 232: What does it mean by “we assume that $\mu_0 = \theta_0, \sigma_0^2.$”?

Although I appreciate the time and effort invested in the rebuttal and am happy with the author's general response, overall, I feel like the author’s response is not really satisfying yet. Even though the authors requested to reconsider my assessment, I cannot do so before I get a satisfying further response. If the further response is not satisfactory as well, sadly and unfortunately, I am **ready** to decrease my score to 3 (but **not yet**) because, in my view, there is still a huge room for improvement in writing and presentation.