Optimal Protocols for Continual Learning via Statistical Physics and Control Theory

Authors: Francesco Mori, Stefano Sarao Mannelli, Francesca Mignacco
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: machine learning theory, statistical physics, online learning, continual learning, optimal control theory

Abstract:
Artificial neural networks often struggle with _catastrophic forgetting_ when learning multiple tasks sequentially, as training on new tasks degrades the performance on previously learned tasks. Recent theoretical work has addressed this issue by analysing learning curves in synthetic frameworks under predefined training protocols. However, these protocols relied on heuristics and lacked a solid theoretical foundation assessing their optimality. In this paper, we fill this gap by combining exact equations for training dynamics, derived using statistical physics techniques, with optimal control methods. We apply this approach to teacher-student models for continual learning and multi-task problems, obtaining a theory for task-selection protocols maximising performance while minimising forgetting. Our theoretical analysis offers non-trivial yet interpretable strategies for mitigating catastrophic forgetting, shedding light on how optimal learning protocols modulate established effects, such as the influence of task similarity on forgetting. Finally, we validate our theoretical findings with experiments on real-world data.

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 8

Reviews and discussion:
<thread>
creation date:
1730577018698

thread_type:
official_review

<thread_object>
creator: reviewer_4ows

<content>
summary:
The authors present a theoretical framework based on a teacher-student setup to define optimal strategies for task switching and learning rate adjustments in continual learning settings. Their approach starts with a dimensionality reduction, which makes the problem more tractable by focusing on three types of overlaps between student and teacher instead of tracking each network weight. In the limit of infinite input dimension, this reduction effectively captures the high-dimensional stochastic dynamics of SGD through ODEs. The next step is to frame the problem of optimal task selection and learning rate scheduling as an optimal control problem, aiming to minimize the final generalization error across all tasks. They address this problem through a variational approach that iterates between forward dynamics (tracking SGD), backward dynamics (using conjugate parameters), and updating the control strategy (task selection/learning rate) until reaching an optimal protocol. Using synthetic experiments, they show that their optimal protocol has a structured design: an initial focus phase on the new task, followed by a revision phase with interleaved training, which performs best at intermediate task similarities. They validate their theoretical insights on Fashion-MNIST by constructing controlled task pairs through linear interpolation.

strengths:
- The paper is well-written and organized, with a comprehensive discussion of related work.
- The topic addresses a significant gap in the theoretical understanding of continual learning. The insights presented are valuable for the continual learning community and likely of interest to a broader machine learning audience.
- The experiments are detailed, especially for a theoretical work.
- While the optimal theory relies on access to a 'teacher' model, which is not realistic in practical applications, the theoretical derivations remain insightful. Moreover, the authors introduce a 'pseudo-optimal' version that operates without direct access to this information, requiring only an estimate of generalization error on the tasks—something feasible in real-world scenarios.

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

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.
- 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730652331750

thread_type:
official_review

<thread_object>
creator: reviewer_papx

<content>
summary:
The paper addresses the problem of catastrophic forgetting in neural networks during continual learning. The paper aims to find optimal task-selection protocols to minimize forgetting and maximize performance, moving beyond heuristic approaches. This is done through combining statistical physics with optimal control theory to derive closed-form formulas for training dynamics.

**Contributions:**
- Derive optimal task-selection protocols and learning rate schedules as functions of task similarity and problem parameters.
- Proposes a "pseudo-optimal" replay strategy with distinct focus and revision phases to minimize forgetting.

strengths:
- Figures are clear, well explained, and easy to understand.
- The paper uses an original and refreshing approach with strong theoretical backing.
- Demonstrates robust alignment between theoretical predictions and experimental results.
- Complex mathematical formulations are well-supported.

weaknesses:
- Please try to use shorter sentences. The sentences tend to be overly long and complex.
- The mathematics and equations are dense and could limit accessibility for readers unfamiliar with optimal control or statistical physics.
- 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.
- 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.
- 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.
- Theoretical findings rely on idealized assumptions (e.g., i.i.d. Gaussian inputs, simplified two-layer networks) which may not hold in complex, structured data environments. Expanding to other data models would strengthen claims of broader applicability.
- The model does not address relative task difficulty or task order impact, which are critical in real-world scenarios. Future extensions could explore adaptive protocols that consider task-specific attributes dynamically.

questions:
- In many results, the pseudo-optimal method is comparable in performance to the naive interleaved. Could you better explain the advantages of using this pseudo-optimal method instead of a naive interleave?
- The "pseudo-optimal" strategy provides practical insights, but the underlying mechanics of the optimal task-selection structure could be further explained to make the protocol more interpretable.
- (see weaknesses for more suggestions)
</content>
</thread_object>

</thread>
<thread>
creation date:
1730703311136

thread_type:
official_review

<thread_object>
creator: reviewer_pchb

<content>
summary:
This paper provides a theoretical framework for optimal task selection and step-size optimization in continual learning.
The authors use the student-teacher formalism where a task is determined by a linear teacher, and the student uses a neural network with 1 hidden layer.
Using linearized dynamics for the student's network, and ideas from optimal control, they identify an "optimal protocol" for sampling data from tasks, as well as for setting the step-size.
The ideas are thoroughly evaluated in a toy setting with simulated data, with some additional experiments on Fashion MNIST and MLPs.

strengths:
- The optimal control perspective provided by the problem formulation and theoretical analysis is interesting. Using this optimal control perspective to inform continual learning seems novel. While I have questions and suggestions below, I found this section to be clear overall.

- The empirical analysis of the proposed method on the toy problem to be thorough and well-motivated.

weaknesses:
The problem formulation, while interesting, seems ultimately misguided for most settings in which continual learning is desirable.

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

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

questions:
- Can you clarify how the multi-headed approach deals with training on previous tasks? Are the readout weights from task 1 updated with interleaved training on task 1?

- I understand that relevant details can be found in the appendix, but the presentation of the results in Section 2 can be improved. In particular, the overlaps that are introduced should be briefly described in how they dictate the optimization dynamics.

- The "dimensional reduction techiques from statistical physics" approach you describe, and the "overlaps" seem to be related to more familiar concepts in the deep learning theory literature. In particular, they seem analogous to the linearized dynamics described by the neural tangent kernel. Can you comment on any similarities or differences?

- "we consider K [hidden dim] = T [number of tasks] to guarantee that the student network has enough capacity to learn all tasks perfectly." I do not think this is necessarily the case. Is the additional assumption that the data is a scalar linear regression problem? If this is the case, this seems like a very limited setting.

- Section 3 and Figure 2: Why are the number of epochs in task 2 comparatively low? Is it the case that interleaved training can achieve perfect error on task 1 and task 2 asymptotically?

- Section 3.1 (agreement with theory): Can you comment on whether there are any differences between the assumptions made in the problem formulation and the experimental results in this section?

- Section 3 (impact of task similarity): Task similarity seems to be measured in terms of the absolute value of the cosine similarity, which makes sense for the linearity assumption. I wonder what can be said when this assumption does not hold, can any results be generalized to nonlinear data?

- Section 3 (Optimal learning rate schedules): It seems like your results were using vanilla SGD< how does the "optimal learning rate" compare to adaptive optimisers (e.g., momentum, gradient normalization, or hyperparameter-free optimisers)

- Section 3.2 (experiments on real data): You mention that training was on a squared loss, but this is a classification problem. How were the targets represented?


** Minor Comments
- I do not think the title accurately describes the contributions. First, the term optimal protocol is not common in the continual learning literature, and only defined towards the end of Section 1. Also, the connection to statistical physics is relegated to the Appendix.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730926704079

thread_type:
official_review

<thread_object>
creator: reviewer_vgxp

<content>
summary:
The paper proposes an ODE relating network parameters, training control parameters (such as learning rate and which task to learn on), and the final performance of the model. Using this, it derives optimal protocols for the training control parameters (i.e. at each step, which task to train on and what learning rate to use) during continual learning. They notice that the task protocol can be approximated by a heuristic where the new task is trained on by the network until its loss matches the loss of the previous task, and then data from the two tasks is interleaved. They also derive an optimal learning rate schedule. They test their protocols on a synthetic task of matching output generated by two randomly initialized neural networks and a version of Fashion MNIST where each task is composed of picking 2 classes from the dataset. On the synthetic tasks, their protocols significantly outperform previously proposed protocols, while the improvements are more modest on Fashion MNIST.

strengths:
- The idea itself is very interesting, and I liked the fact that the theoretically optimal protocol can be approximated by a simple heuristic. If the results transfer, the “pseudo-optimal” protocol could be very useful.
- At least on the synthetic tasks, the replay protocol seems to outperform the standard interleaved protocol across essentially all settings.

weaknesses:
- 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.
- 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).
- 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.

questions:
- It would be nice to see if some of these results transfer to harder continual learning datasets. Specifically, these results are all done on two class classification tasks. Do they hold with more classes? Also do they hold for non MLP networks? Do they hold for datasets such as split Imagenet or split CIFAR? Do they hold across more than 2 tasks?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732031701743

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to reviewer feedback

comment:
We thank the reviewers for their valuable time and constructive comments. 

All reviewers recognize the value of our work, highlighting that our theoretical framework is “interesting” (vgxp, PcHB, 4ows), “original and refreshing” (pAPx), “useful” (vgxp), and “insightful” (4ows). Additionally, all reviewers acknowledge that our theoretical results are well-supported and thoroughly validated by numerical simulations.

In response to the reviewers' suggestions for improving clarity and expanding experiments, we have made the following updates:

- Enhanced the clarity of our technical presentation, particularly in Sec. 2, and better explained our experimental findings in Sec. 3.2.
- Conducted additional simulations on learning rate annealing, benchmarking exponential and power-law schedules (Appendix C.1).
- Extended our analysis to a continual learning problem with T=3 tasks, comparing the theory-derived optimal schedule with multiple benchmarks (Appendix C.1).
- Tested the pseudo-optimal strategy with new experiments on CIFAR10 (Appendix C.2).

Detailed responses to specific reviewer comments will be provided in individual replies.

In summary, our work offers a novel and promising theoretical framework with potential practical implications, validated by extensive simulations and interpretable results. If your concerns have been addressed, we kindly ask you to consider increasing your score.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732483419664

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Discussion period ending soon

comment:
Dear Area Chair and Reviewers,

We have submitted detailed responses to all the reviewers' concerns, along with a general response summarizing the changes made. With only one day remaining in the discussion period, we kindly ask for confirmation on whether our replies have addressed the reviewers' concerns, and, if so, to consider adjusting your score accordingly.

Your feedback is crucial to improving the quality of the paper, and we would greatly appreciate your engagement before the deadline.

Thank you for your time and consideration.

Best regards,
The Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1732804191827

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers and AC,

We believe we have addressed the concerns raised in your reviews in the updated version of the paper. While we can no longer update the submission, we are happy to continue engaging and answer any questions that may arise during the reviewers-AC discussion phase.

Kind regards,
The Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1733076431070

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Follow-up on updated responses before discussion ends

comment:
Dear Area Chair and Reviewers,

Thank you for your work. We believe that our revised version reflects your recommendations. In particular: We clarified the technical presentation, expanded on experimental findings, and conducted additional analyses and experiments based on your suggestions.

While Reviewers pAPx and vgxp don't seem to need further clarifications, Reviewers 4ows and PcHB expressed some concerns that we addressed in our revision. 

In particular, we addressed Reviewer 4ows's concerns about the generality of our results by adding a new experiment on CIFAR10, as requested. For Reviewer PcHB, who raised questions about the problem setup and theoretical framework, we clarified that our context is well-established in the field and highlighted the originality of our contribution.

We believe our responses address the reviewers' concerns, but we are happy to discuss further if needed. If our revisions address your concerns, we kindly ask you to reflect this in your evaluation before the discussion phase ends tomorrow.

We greatly appreciate your time and insights.

Best regards,

The Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper theoretically approaches continual learning using ODEs (for learning curves of online SGD), applying it to teacher-student continual learning methods, which is novel. They derive an optimal learning rate schedule and study the influence of task similarity on forgetting, which is useful / significant in continual learning. There are also experiments are on Fashion MNIST (and after rebuttal, CIFAR10), which is nice to see and improves the impact and reach of the paper. Overall, the novelty and significance are key strengths of the paper, and therefore I recommend accepting (with spotlight). I think the community will benefit from reading and discussing this paper and its theory techniques for continual learning.