Optimal Protocols for Continual Learning via Statistical Physics and Control Theory

We thank the reviewer for their valuable time reviewing our work and for their positive comments.

The theoretical framework assumes a "multi-headed" approach, where each task has its own output layer. However, most replay methods that interleave past tasks typically use a shared output layer. When using a shared head, the best approach is generally to mix all tasks in a batch, as tasks without associated samples get heavily penalized, pushing their output neurons towards inactivity -- neurons become dead. While the multi-headed setup make theory tractably—and is reasonable given the paper’s theoretical focus—this choice limits the immediate applicability of the findings to common practical implementations.

We thank the reviewer for this insightful comment. The idea of increasing the number of classes while using a shared output layer and task-specific classes, all updated at each iteration, is indeed interesting and relevant. At present, the theoretical description of the dynamics, based on statistical physics, has not been extended to multi-label classification. However, we believe this extension is feasible. Once this preliminary step is accomplished, integrating the control theory component would follow naturally. We agree that this represents an interesting direction for future research and have now mentioned it in the revised manuscript. Thank you for highlighting this point.

Questions

While the Fashion-MNIST interpolation setting enables direct control of task similarity and consistent with the theoretical framework, the paper would be stronger if it demonstrated results in settings that deviate somewhat from the theoretical conditions. For example, using MNIST with controlled pixel permutations would provide an alternative way to study task similarity. Such experiments would showcase the generalizability of the theoretical insights to different types of task relationships and data distributions. This additional experiment isn’t necessary, but some discussion in Section 3.2 on why this specific scenario was chosen would be helpful.

We have welcomed the reviewer’s suggestion and run additional experiments on CIFAR10 without introducing any interpolation between tasks. We create two classification tasks taking as task 1 the classes with odd labels and the others as task 2. The results are presented in Appendix C.2 and confirm our observations on pseudo-optimal strategy striking the balance between no-replay and interleaved.

The term 'epoch' in this paper may be misleading. While the authors use it to refer to training time scaled by input dimensions, 'epoch' conventionally means a single pass over the dataset in machine learning. This could cause confusion, especially in Section 3.2, where theoretical and practical results are compared. I suggest using a term like 'normalized training steps' to improve clarity.

We thank the reviewer for the suggestion, that we welcome to improve the clarity of our presentation. We have replaced the “training epoch” with “training step” everywhere in the manuscript.

We hope our response addresses your concerns and welcome further suggestions to improve the clarity of our presentation. If you find our response convincing/useful, please consider increasing the score.

We thank the reviewer for their valuable time reviewing our work and for their positive comments.

Please try to use shorter sentences. The sentences tend to be overly long and complex.

We thank the reviewer for this suggestion. We have improved the clarity of our presentation in the revised version of the manuscript.

The mathematics and equations are dense and could limit accessibility for readers unfamiliar with optimal control or statistical physics.

We have revised Sec. 2 expanding the explanations and introducing additional pointers to previous works. We are happy to further improve our presentation if the reviewer has more specific suggestions in mind.

The methodology might benefit from additional real-world datasets or further exploration in complex architectures for broader generalizability. Additional datasets with higher complexity (e.g., CIFAR-10/100 or domain-specific continual learning datasets) would better demonstrate the generalizability of the proposed protocols.

We have welcomed the reviewer’s suggestion and run additional experiments on CIFAR10, presented in Appendix C.2. These new results confirm our initial observations of Sec. 3.2. In particular, the performance of the pseudo-optimal learning strategy appears to interpolate between no-replay and interleaved depending on the dataset size.

Experiments primarily focus on two-task scenarios. Including more tasks would test the robustness and scalability of the approach, offering more insight into how it handles diverse, real-world continual learning settings.

Our theoretical framework is directly applicable to more than two classes. We have now added a paragraph about the case T=3 in Appendix C.1, our results are shown in the additional Fig. 10. Also in this case, we find that the optimal strategy outperforms the heuristic benchmarks. The structure of the optimal strategies for T=3 is more complicated, but remarkably it still shares some similarities with the pseudo-optimal strategy described in the main text. Specifically, task 1 is replayed only when its loss is comparable to the losses on the other two tasks. Notably, as the number of tasks grows, the number of possible heuristic strategies expands significantly (we consider 6), making it difficult to identify effective solutions through intuition alone. This highlights the importance of a framework for systematically determining the optimal strategy. While the additional paragraph in Appendix C is mainly a proof of concept, leveraging insights from our theoretical framework to deal with multi-class continual learning in real-world settings is an exciting and feasible application of our theory that we leave for future work.

Although the paper suggests a posteriori interpretations, actionable guidelines on why specific phases are structured as they are (e.g., conditions where certain strategies excel) would improve applicability and user understanding.

We thank the reviewer for this insightful suggestion. In response, we have expanded our explanation in Section 3.2. Specifically, we find that the size of the downstream task plays a crucial role. For small datasets, the no-replay strategy performs well, as the loss reduction on the new task outweighs the forgetting of previous tasks. Conversely, for large datasets, the interleaved strategy is more effective, as it mitigates forgetting while maintaining reasonable performance on the new task. The transition between these two phases depends on the specific dataset, as evidenced by the differences between Fig.6 (MNIST) and Fig.11 (CIFAR10). Identifying this transition point can be challenging in practice. However, the pseudo-optimal strategy we propose effectively interpolates between these two extremes, achieving near-optimal performance across both regimes.

We thank the reviewer for the precious feedback, which helped us to improve the communication of our findings.

The first problem is that, continual learning is motivated from the costly nature of retraining. With small/linear models, the costs (and challenges) of retraining the model are diminished. There does not seem to be an obvious extension for this theory to larger models, and so the applicability of this approach seems largely limited to linear/small models.

First, we would like to clarify that our model is not linear. Both teachers have nonlinearities and a nonlinearity is also applied to the pre-activations of the student network. Our theory holds for arbitrary nonlinear activation functions. In the results section, we explicitly evaluate it using error functions for both teacher and student networks. Furthermore, we do not rely on linearization or any other approximations of the dynamics to obtain our analytical results. Crucially, our dynamical equations are exact in the high-dimensional limit. Regarding the applicability of our approach to larger models, we emphasize that it may not be essential to extend the full analytic derivation to state-of-the-art architectures and tasks in order to derive meaningful insights into optimal strategies. Instead, an effective approach would be to use interpretable strategies inferred from simpler models as a basis for training more complex, realistic models. The rationale is that certain hyperparameter schedules can be broadly effective and are not strongly dependent on model-specific details, much like widely used heuristics such as learning rate annealing. Furthermore, key extensions of our theory that would bring it closer to real-world data and models are within reach and are the focus of ongoing work. One promising example is the use of structured data models, which would allow us to incorporate some structure of real tasks directly into the theoretical framework. Structured data models have been extensively studied using statistical physics methods in recent years, and these recent advancements offer a promising foundation for applying our optimal control framework.

References:

• Mézard, Marc. Indian Journal of Physics (2023): 1-12.
• Cagnetta, Francesco, et al. Physical Review X 14.3 (2024): 031001.
• Goldt, Sebastian, et al. Physical Review X 10.4 (2020): 041044.
• Loureiro, Bruno, et al. Advances in Neural Information Processing Systems 34 (2021): 18137-18151.
• Wakhloo, Albert J., Tamara J. Sussman, and SueYeon Chung. Physical Review Letters 131.2 (2023): 027301.

The second is the "optimallity" of the proposed protocol is ultimately about speed of minimizing training error, and not the generalization quality of the model. On this front, the proposed approach could be generalized. Indeed, the results in Section 6 seem to report test loss, whereas most other plots report train loss.

We would like to clarify that in the online learning setting there is no distinction between train and test loss, as each example is used only once. Consequently, there is no concept of a “loss landscape”, that is instead induced by the presence of a training dataset. For this reason, the online setting does not permit the study of the generalization gap, but this limitation is inherent to online learning itself and not specific to our application. Nonetheless, the online learning regime remains highly relevant to understand learning dynamics and continues to be widely studied in machine learning theory, see example references below. In the discussion section, we highlight extending our framework to batch learning as an exciting future direction, which would enable us to study the role of memorization, among other aspects.

References:

• Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. The Journal of Machine Learning Research, 22(1):4788–4838, 2021.
• Goldt, Sebastian, et al. Advances in neural information processing systems 32 (2019).
• Song Mei, Andrea Montanari, and Phan-Minh Nguyen. Proceedings of the National Academy of Sciences, 115(33):E7665–E7671, 2018.
• Lenaic Chizat and Francis Bach. Advances in neural information processing systems, 31, 2018.
• Grant Rotskoff and Eric Vanden-Eijnden. Communications on Pure and Applied Mathematics, 75(9):1889–1935, 2022.

We address the reviewer's questions in a separate comment. We hope our response addresses your concerns and welcome further suggestions to improve the clarity of our presentation. If you find our response convincing/useful, please consider increasing the score.

I thank the authors for their detailed reply, you have cleared up many questions. I want to like this paper, the optimal control approach is very interesting. I think my main issue is still the first point, on which I would appreciate further discussion:

• Context: Continual learning is primarily about dealing with the trade-off of stability and plasticity. Catastrophic forgetting often occurs, meaning that performance on an old task can suffer while training on a new task. This is indeed a problem for which many have devised methods to solve. In particular, these approaches try to avoid logging experience from past tasks. However, if the experience from the previous task is logged then training can be interleaved, as your paper uses as a baseline.

• The problem being studied here is, for the most part, task selection. My primary concern is that it is not clear if the task selection problem is an existing problem that is of interest. The authors themselves note in the related works: "to the best of our knowledge, the problem of optimal task selection has not been explored yet." I can see certain problems for which this can matter, such as in curriculum learning, particularly in reinforcement learning. But this literature may have differences in the specific problem formulation, as well as existing strong baselines.

• Reprieve: I understand that your theoretical contributions are novel and interesting, but I think the underlying problem addressed by theoretical considerations should also be interesting. Step-size adaptation can be potentially much more influential in learning, but this would require comparison with the many adaptive and hyperparameter free optimisers.

(Edit: I want to be clear that I am not looking for additional results here. Rather, I am looking for more discussion and justification on the choice of problem being studied.)

I also have some further clarifications regarding some of your answers to my questions:

• On your high-dimensional limit vs NTK: To clarify, your "overlaps" involve a limit of the input dimension and the NTK involves a limit in the width. Is this the primary difference? This is worth highlighting, and discussing in relation to NTK so as to familiarize readers.

• On (non)linearity of student and teacher: yes, I understand that there is a nonlinear activation function (even for the teacher), but this is basically a generalized linear family which seems necessarily limited compared to deeper hierarchical models. Is there a universal approximation theorem in the limit of large input dimension on a generalized linear function? (Not that universal approximation of the teacher-family is required for your problem of study to be interesting, but it would be a limitation.)

• On test vs train in online setting: I did see that you consider the online setting, but some of your figures reference "epoch" which would involve training on the same dataset more than once. You should clarify what epoch means in this case, and perhaps be more explicit.

• Your references in the rebuttal (pt1): listing papers without any context does little to help connect these paper with any particular claims you are making in the rebuttal.

• About task similarity and linearity: Correlation and cosine similarity are inherently linear relationships. That is, there exist nonlinear relationships which can exhibit 0 correlation. My question was directed at whether this linearity assumption (of task similarity) is a limitation.

Reprieve: I understand that your theoretical contributions are novel and interesting, but I think the underlying problem addressed by theoretical considerations should also be interesting. Step-size adaptation can be potentially much more influential in learning, but this would require comparison with the many adaptive and hyperparameter free optimisers.

We hope to have clarified that our theory addresses the core problem of determining the optimality of replay-based approaches to mitigate catastrophic forgetting. While replay is a well-established technique for tackling the broader challenge of continual learning, current methods are still predominantly heuristic. Our work moves beyond these heuristics to provide a principled framework.

The paragraph on “Optimal learning rate schedules” (Sec. 3.1) is precisely addressing the question: What is the optimal adaptive learning rate for this problem? What is interesting is that: (1) we can find the optimal learning rate dynamics without resorting to heuristic methods; (2) this is achieved in conjunction with the optimization of the replay schedule. If the reviewer is suggesting that adaptive step-size methods alone could resolve catastrophic forgetting, we respectfully disagree. This phenomenon persists even in applications where such methods are widely used. Anyways, our theory already provides the optimal learning rate schedule, and removing the optimization over replay can only degrade performance.

We thank the reviewer for the precious feedback, which helped us to improve the communication of our findings.

The learning rate experiments need more experimental support to be convincing. The learned schedule is being compared to a constant learning rate schedule. It should also be compared to more standard learning rate schedules.

We have run additional experiments on the effect of learning rate annealing, shown in the updated Fig. 9 of Appendix C.1. In particular, we consider exponential and power law learning rate schedules, in combination with interleaved replay protocol. For each value of task similarity, we optimize over the schedule parameters via grid search. We still find a performance gap with respect to the optimal strategy, which highlights the relevance of the joint optimization of training protocols.

It’s unclear how much these results can transfer to harder/“real world” settings. See question 1 below. On the one “real world” task, there does not seem to be any noticeable difference between the interleaved and pseudo-optimal strategy (Figure 7).

The optimal strategy depends on the number of available samples for the downstream task. For small downstream tasks, no-replay outperforms interleaved (second and third panels). The interleaved strategy gets closer to optimal as the size of the downstream dataset increases. The pseudo-optimal strategy, because of its structure, results in an interpolation between no-replay and interleaved strategies that automatically adapts to the dataset size, leading to the best performance in both regimes.

We have added a new experiment on CIFAR10 that confirms this observation. This result is reported and explained in Appendix C.2.

This is precisely one of the messages that Fig. 7 aims to convey. We thank the reviewer for reporting their concern, this led us to include additional explanations improving clarity in the main text. We hope that now it is clear that, from left to right in the plots, the performance curves of the interleaved and no-replay strategies gradually shift relative to each other, with the first becoming closer to optimal as the other moves further away.

Section 2 explaining the machinery behind finding the optimal control parameters can be explained better. It is a bit difficult to follow for someone not already familiar with the topic, which many people who are attempting to use the results of this paper (CL researchers) would likely not be.

We have revised Sec. 2 expanding the explanations and introducing additional pointers to previous works. We are happy to further improve our presentation if the reviewer has more specific suggestions in mind.

Due to space limitations, we address your question in a following comment.

We hope our response addresses your concerns and welcome further suggestions to improve the clarity of our presentation. If you find our response convincing/useful, please consider increasing the score.