We thank Reviewer DfC5 for their feedback. Please find our response below.

Figure 3, hard to read

Thanks to your feedback, we will modify the way we present the results, using the extended page limit for the camera-ready version. First, we will present Soft Reset method and compare it to baselines. Second, we will present other variants of the method and compare it to the main method.

baselines and benchmarks motivation

The benchmarks are motivated by Plasticity Loss literature [1-4]. In this setting, Neural Networks lose ability to learn (lose plasticity) when exposed to certain forms of non-stationarity. Random label MNIST and CIFAR10 benchmarks are most common ones. Moreover, in this setting, resetting is highly beneficial (see [1]).

We used 3 baselines, Shrink&Perturb [5], L2 Init [1] and Hard Resets. The Shrink&Perturb[5] is used to prevent plasticity loss by perturbing parameters according to $\theta_{t+1}=\lambda\theta_t+\sigma \epsilon$ where $\epsilon\sim\mathcal{N}(0;I)$. L2 Init approach adds $||\theta - \theta_0||^2$ to the loss and was shown to outperform many other approaches in preventing plasticity loss. Both baselines are highly related to our drift model and discussed in Related Work. We hope this clarifies the motivations.

RL is not clear...

Please see our response to Reviewer Hmj6.

Why in on-policy, simple method performs better than ours, given that it is more non-stationary setting?

In RL, the non-stationary rises from changing input data distribution (since we change the policy) as well as from changing learning targets (when learning value functions). In the off-policy RL, we expect a high impact of changing learning target non-stationarity. It was observed [4,6-7], that in this setting we can exhibit the loss of plasticity. The effect becomes more present as we increase the replay ratio (increase the number of updates per collected data). Under high replay ratio since we do many updates on replay buffer, we could overfit to the replay buffer and when we start collecting new data, the local minima of NN can switch drastically. This is the regime which motivated our method (see Figure 1 and beginning of Section 3). In this setting, hard resets [4] are effective but our method shows to be even more effective.

When we operate in on-policy setting, the situation described above is less likely to happen. If data distribution changes quickly, then we are less likely to overfit to it due to the noise in the update. Using the language from Figure 1, the uncertainty of the Bayesian posterior will not be shrinking if the data changes quickly. This, however, depends on the environment. The Hopper environment (Figure 4, bottom) is a relatively simple task to solve, meaning that the agent will progress fast and the corresponding input data distribution will change fast. The Humanoid environment (Figure 4, top) is a much more challenging environment, and the simple agent might struggle initially. This implies that even on-policy algorithm can see a lot of similar data and start to overfit to it, putting it closer to the off-policy setting. This explains why our method is more effective in Humanoid rather than Hopper environment when relpay ratio is 1.

why stops earlier?

The experiment was not fully finished at the time of submission. See attached pdf for final version of the Figure.

Figure 2 ?

In Figure 2, we demonstrate that drift model eq.4 coupled with update rule eq.16 is a good strategy of resetting parameters when task boundaries are known. For appropriately chosen $\gamma \neq 0$, we can achieve significantly better performance than hard reset. In Figure 2, left, we use constant l.r. $\alpha_t$ whereas in Figure 2,right, we use $\alpha_t(\gamma_t)$ from eq.17, which is more beneficial. We will add more clarifications in the text about the nature of this experiment.

Theoretical analysis

Unfortunately, the phenomenon of plasticity loss in Neural Networks is not theoretically analyzed (there is no model for that), it is an empirical phenomenon [1-4]. We would require this theoretical framework for plasticity loss to be developed in order to analyze our method. Our future plans involve analyzing the method in the context of online convex optimization.

MC samples in eq.7?

MC samples are not coming from a simulator but from a Gaussian distribution eq.5 induced by the drift model eq.4. MC samples are needed to approximate integral in eq.7. The MC samples correspond to NN parameters.

Overall, we hope our answer clarifies the points which you raised during your review. If we have addressed your questions, we would be grateful if you would consider increasing your score. We would be happy to answer any further questions you might have.

References:

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

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

[3] Disentangling the Causes of Plasticity Loss in Neural Networks, Clare Lyle, Zeyu Zheng, Khimya Khetarpal, Hado van Hasselt, Razvan Pascanu, James Martens, Will Dabney, 2024

[4] Understanding and Preventing Capacity Loss in Reinforcement Learning, Clare Lyle, Mark Rowland, Will Dabney, 2022

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

[6] The Primacy Bias in Deep Reinforcement Learning, Evgenii Nikishin, Max Schwarzer, Pierluca D'Oro, Pierre-Luc Bacon, Aaron Courville, 2022

[7] The Dormant Neuron Phenomenon in Deep Reinforcement Learning, Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, Utku Evci, 2023

We thank the reviewer 6GY1 for their positive feedback. Please find our answer below.

Section 2 - Simplification of equation (6) to (9) requires some explanation [page -5]

The equation 9 is obtained from 6 as follows. First, we linearise the log-likelihood function $\log p(y_{t+1} | x_{t+1}, \theta)$ around $\mu_t$ which equals to $\log p(y_{t+1} | x_{t+1}, \theta) \sim log p(y_{t+1} | x_{t+1}, \mu_t) + g_{t+1}^T (\theta - \mu_t)$. Then, we notice that $p(y_{t+1} | x_{t+1}, \theta) = \exp^{\log p(y_{t+1} | x_{t+1}, \theta)}$ which means that $p(y_{t+1} | x_{t+1}, \theta) \sim p(y_{t+1} | x_{t+1}, \mu_t) \exp^{-g_{t+1}^T \mu_t} \exp^{g_{t+1}^T \theta}$. Then, since $q_{t}^{t+1}(\theta|S_t,\Gamma_{t-1},\gamma_t)$ is a Gaussian, we can write compute the integral in a closed form. After that, we keep only the terms which depend on $\gamma_{t}$, since we only are interested in the optimization of $\gamma_t$.

We will add a full derivation of this expression in the Appendix.

Section 3 - Some intuition behind the objective (11) [page - 6] would be good to understand it better (especially the second term in that objective)

The equation in objective 11 is negative Evidence Lower Bound (ELBO) on approximate predictive log likelihood eq 6, if all $\lambda_i = 1$. It could be derived by considering the integral in the right-hand side of eq. 6, dividing and multiplying it by $q(\theta)$ and applying Jensen Inequality. The fact that we introduce $\lambda_i \neq 1$ means that we introduce a temperature parameter (see [1]) on the prior for each dimension $i$. It was shown empirically [1] that using a temperature in the prior leads to better empirical results. We will add more explicit explanations of how to derive the objective 11 and will add full derivation in the Appendix.

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?

Our experiments in Figure 3 and Figure 4 partially answer this question. The difference between Data Efficient and Memorization regimes in Figure 3 is the amount of epochs given to a task – 70 epochs for Data Efficient and 400 for Memorization. There is more non-stationarity in the Data-Efficient setting since it changes more frequently and our methodology remains more effective than baselines.

In general, we expect our method to be very helpful in scenarios when there are relatively long stationary phases which are succeeded by large changes of the optimal parameters. Moreover, given that we can define $\gamma_t$ per parameter or per layer, we can have a combination of these scenarios -- some parameters stay stationary whereas other change significantly. In general, such scenario would occur when we have stationary segments which are followed by non-stationary ones. In case, when the non-stationarity is very strong and the data distribution changes at every step, we think that it is less likely that our method will bring an additional benefit over SGD since in this setting, SGD will be constantly refreshed by the noise from the data. SGD, however, will struggle when there is a large change after a long stationary phase. Compared to reset-type algorithms, our method is much more adaptive, and relearns the parameters only when it "has to" (i.e., when the data changes sufficiently).

We see a partial support for this claim in our RL experiments in Figure 4, where as we become more off-policy, we see much more benefit of our method over standard learning approach. Moreover, we also see a benefit of our method over baseline in a more on-policy regime in Figure 4, left, top on Humanoid environment. Since this environment is hard to solve, the baseline agent will see a lot of similar data initially and therefore the corresponding RL algorithm could overfit. Constantly refreshing the parameters is beneficial here. On the other hand, if we study Figure 4, left, bottom which shows performance in Hopper environment, we can see that Soft Resets is less effective. In this setting, since the problem is easy, the baseline method sees a lot of different data constantly and is less prone to overfit.

Some minor typos in line 109, 154, 190

Thank you for the catch, we will make the adjustments.

References:

[1] How Good is the Bayes Posterior in Deep Neural Networks Really?, Florian Wenzel, Kevin Roth, Bastiaan S. Veeling, Jakub Świątkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, Sebastian Nowozin, 2020

We thank reviewer 4koq for their response. Please find our detailed answer below.

why drift model...

In Section 3 and in Appendix D, we presented the reasons to use a drift model together with learning NN parameters. Figure 1 illustrates the high-level intuition in case of online Bayesian estimation. Using SGD language, assuming that the data is stationary up to the time $T$, SGD estimate $\theta_{T}$ would tend towards a local optima $\theta^*$. If data stays stationary at time $T+1$, $\theta_{T+1}$ will move closer to the $\theta^*$. However, if at time $T+1$, the data distribution changes and so is the set of local optima, SGD might struggle to move towards this set starting from $\theta_T$. Drift model allows the learning algorithm to make larger moves towards this new set of local optima. Such idea was also used in the context of online convex optimization, see [2-3]. The form in eq.4 encourages parameters to return back towards the initialization over time by shrinking the mean and increasing the variance (l.r. in SGD), allowing to make bigger steps towards new local minima.

We will change the wording of the beginning of Section 3 to better reflect this explanation.

form of OU process

OU process defines a SDE that can be solved explicitly and written as a time-continuous Gaussian Markov process with transition density $p(x_t|x_s)=\mathcal{N}(x_s e^{-(t-s)},(1-e^{-2(t-s)})\sigma_0^2I)$ for any pair of times $t>s$. Based on this, as a drift model for the parameters $\theta_t$ (so $\theta_t$ is the state $x_t$) we use the conditional density $p(\theta_{t+1}|\theta_t)=\mathcal{N}(\theta_t\gamma_t,(1-\gamma_t^2)\sigma_0^2I)$ where $\gamma_t=e^{-\delta_t}$ and $\delta_t\geq0$ corresponds to the learnable discretization time step. In other words, by learning $\gamma_t$ online we equivalently learn the amount of a continuous “time shift” $\delta_t$ between two consecutive states in the OU process. This essentially models parameter drift since e.g. if $\gamma_t=1$, then $\delta_t=0$ and there is no “time shift" implying $\theta_{t+1}=\theta_t$.

...why eq.4

In Section 3, we argued for choosing eq.4 because it pushes the learning towards the initialization as well as towards previous parameters, allows to use gradient-based methods to estimate drift parameters and keeps positive finite variance (assuming we learn over the infinite amount of time) which avoids degenerate cases of $0$ or $\infty$ variance. Moreover, it couples the mean and the variance via $\gamma$, making it easier to learn $\gamma$ and less likely to overfit. Other potential choices for the drift models:

• $\theta_{t+1}=\gamma\theta_t+(1-\gamma)\mu_0+\beta\epsilon$, with $\epsilon\sim\mathcal{N}(0,I)$

  • Fixed $\beta$ was explored in our experiment in Figure 2, left, where we used constant l.r. $\alpha_t$ for parameters update. We found that it performed worse than using rescaled l.r. $\alpha(\gamma_t)$ from eq.17 which is derived using our model, see Figure 2, right. For $\mu_0=0$, we recover Shrink&Perturb [4]
  • Learning $\gamma$ and $\beta$ will likely overfit

• Mixture $p(\theta_{t+1}|\theta_t,\gamma_t)=\gamma_tp(\theta_{t+1}|\theta_t)+(1-\gamma_t)p_0(\theta_{t+1})$, where $p(\theta_{t+1}|\theta_t)=\mathcal{N}(\theta_t;\sigma^2)(\theta_{t+1})$ and $p_0(\theta_{t+1})=\mathcal{N}(\mu_0;\sigma_0^2)(\theta_{t+1})$, which is a Gaussian version of Spike&Slab [5]. This encourages hard resets instead of soft ones. Moreover, using mixtures is problematic since the KL in eq.11 cannot be computed exactly and needs to be approximated.

We will add the discussion of drift model choices in the appendix.

derivations...

Please see our reply to Reviewer 6GY1 about the eq.9 and eq.11. Eq.10 is obtained from eq.9 by finding fixed points. In eq.10, we have a typo, the term $\lambda\gamma_t^0$ in the denominator should be replaced by $\lambda$. Eq.14 is a gradient descent rule applied to eq.11(eq.12) where $\tilde{F}$ is actually a derivative of $F$ wrt $\mu_t$ and $\sigma_t$. Eq.15 is Maximum a-posteriori (MAP) update for $\theta$ from the posterior $p(y_{t+1}|x_{t+1},\theta)q^{t+1}{t}(\theta |\gamma_t)$ and where we use temperature in the prior (see [3]).

We will add all the derivations and explanations in the appendix.

$\theta'$ every time

While a possible modification for the drift model, it could be problematic to use, since resampling $\theta_{0}$ has higher variance than OU process. The variance of this model is $\gamma_t^2\sigma^2_t+(1-\gamma_t)^2\sigma^2_0+(1-\gamma^2_t)\sigma^2_0$, which equals to $2\sigma^2_0$ for $\gamma_t=0$ which is 2x larger than the variance of the initialization. This implies that such a model might inflate the variance over time since e.g. if $\gamma_t=0.5$ for all time steps $t$, then $\sigma_{t+1}^2=\gamma_t^2 \sigma^2_t+2(1-\gamma_t)\sigma^2_0=0.5^2\sigma^2_t+\sigma^2_0$ which grows with the iterations.

Thus, we believe that our current drift model more accurately implements the mechanism of parameter resets since it always brings us back to the exact initialization distribution.

We hope we have been able to address your concerns and we are happy to answer further questions on the above subjects. If we were able to address your concerns, we would be grateful if you would consider increasing your review score.

References:

[1] Dynamical Models and Tracking Regret in Online Convex Programming, Eric C. Hall, Rebecca M. Willett, 2013

[2] Adaptive Gradient-Based Meta-Learning Methods, Mikhail Khodak, Maria-Florina Balcan, Ameet Talwalkar, 2019

[3] How Good is the Bayes Posterior in Deep Neural Networks Really?, Florian Wenzel, Kevin Roth, Bastiaan S. Veeling, Jakub Świątkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, Sebastian Nowozin, 2020

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

[5] Spike and slab variable selection: Frequentist and Bayesian strategies, Hemant Ishwaran, J. Sunil Rao, 2005

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.

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!

We thank the reviewer Hmj6 for their feedback. Please find our detailed answer below.

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.

We have provided the answer about hyperparameters sensitivity above in the common answer for all the reviewers.

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.

We have provided the answer about computational complexity above in the common answer for all the reviewers.

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.

Our choice of Reinforcement Learning experiment is motivated by the Primacy Bias [1] paper, where a significant plasticity loss occurs and where the benefits of parameter resets were observed. In this setting, it was shown that as replay ratio increases, the performance of the algorithm (Soft Actor Critic) could significantly degrade and hard-resetting parameters once in a while, can significantly improve its performance. With this in mind, our Figure 4 demonstrates that Soft Reset can in fact be an even more effective strategy when off-policy ratio increases – it can learn when to reset parameters and by what amount, compared to Hard Resets which require setting these parameters manually.

We are happy to take further questions on any of the above matters, and we hope that we have been able to address your concerns. If we have done so, we would be grateful if you would consider increasing your review score.

References:

[1] The Primacy Bias in Deep Reinforcement Learning, Evgenii Nikishin, Max Schwarzer, Pierluca D'Oro, Pierre-Luc Bacon, Aaron Courville, 2022