Optimal Protocols for Continual Learning via Statistical Physics and Control Theory

## Abstract

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 ones. 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 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 can modulate established effects, such as the influence of task similarity on forgetting. Finally, we validate our theoretical findings on real-world data.

## 1 Introduction

Mastering a diverse range of problems is crucial for both artificial and biological systems. In the context of training a neural network on a series of tasks—a.k.a. *multi-task learning* (Caruana, ; ; ; )—the ability to learn new tasks can be improved by leveraging knowledge from previous ones (Suddarth & Kergosien, ). However, this process can lead to *catastrophic forgetting*, where learning new tasks degrades performance on older ones. This phenomenon has been observed in both theoretical neuroscience (McCloskey & Cohen, ; Ratcliff, ) and machine learning (Srivastava et al., ; Goodfellow et al., ), and occurs when the network parameters encoding older tasks are overwritten while training on a new task. Several mitigation strategies have been proposed (French, ; Kemker et al., ), including semi-distributed representations (French, ; ), regularisation methods (Kirkpatrick et al., ; Zenke et al., ; Li & Hoiem, ), dynamical architectures (Zhou et al., ; Rusu et al., ), and others (see e.g. Parisi et al. (); De Lange et al. () for thorough reviews). A common strategy, known as *replay*, is to present the network with examples from the old tasks while training on the new one to minimise forgetting (Shin et al., ; Draelos et al., ; Rolnick et al., ).

On the theoretical side, Baxter () pioneered the research on continual learning by deriving PAC bounds. More recently, further performance bounds have been obtained in the context of multi-task learning, few-shot learning, domain adaptation, and hypothesis transfer learning (Wang et al., ; Zhang & Yang, ; Zhang & Gao, ).
However, these results focused on worst-case analysis, offering bounds that may not reflect the typical performance of algorithms. In contrast, Dhifallah & Lu () began investigating the typical-case scenario, providing a precise characterisation of transfer learning in simple neural network models. Gerace et al. (); Ingrosso et al. () extended this analysis to more complex architectures and generative models, allowing for a better description of the relation between tasks.
Finally, Lee et al. (; ) proposed a theoretical framework for the study of the dynamics of continual learning with a focus on catastrophic forgetting. Their work provided a theoretical explanation for the surprising empirical results of Ramasesh et al. (), which revealed a non-monotonic relation between forgetting and task similarity, where maximal forgetting occurs at intermediate task similarity. Analogously, Shan et al. () studied a Gibbs formulation of continual learning in deep linear networks, and demonstrated how the interplay between task similarity and network architecture influences forgetting and knowledge transfer.

Despite the significant interest in transfer learning and catastrophic forgetting, mitigation strategies considered thus far were pre-defined heuristics, offering no guarantees of optimality. In contrast, here we address the problem of identifying the optimal protocol to minimise forgetting. Specifically, we focus on *replay* as a prototypical mitigation strategy and employ optimal control theory to determine the optimal training protocol that maximises performance across different tasks.

##### Our contribution.

In this work, we combine dimensionality-reduction techniques from statistical physics (Saad & Solla, ; ; Biehl & Schwarze, ) and Pontryagin’s maximum principle from control theory (Feldbaum, ; Pontryagin, ; Kopp, ) to derive optimal task-selection protocols for the training dynamics of a neural network in a continual learning setting. Pontryagin’s principle works efficiently in low-dimensional deterministic systems, hence requiring the statistical physics approach to neural networks (Engel, ), where the evolution of high-dimensional stochastic systems are condensed to a few key order parameters governed by ordinary differential equations (ODEs) (Saad & Solla, ; ; Biehl & Schwarze, ).
Specifically, we consider the teacher-student framework of Lee et al. ()—a prototype continual learning setting amenable to analytic characterisation. Our main contributions are:

- •
We leverage the ODEs for the learning curves of online SGD to derive closed-form formulae for the optimal training protocols. In particular, we provide equations for the optimal task-selection protocol and the optimal learning rate schedule, as a function of the task similarity $\gamma$ and the problem parameters. Our framework is broadly applicable beyond the specific context of continual learning, and we outline several potential extensions.
- •
We evaluate our equations for a range of problem parameters and find highly structured protocols. Interestingly, we are nonetheless able to interpret these strategies a posteriori, formulating a criterion for “pseudo-optimal” task-selection. This strategy consists of an initial *focus* phase, where only the new task is presented to the network, followed by a *revision* phase, where old tasks are replayed.
- •
We clarify the impact of task similarity on catastrophic forgetting. At variance with what previously observed (Ramasesh et al., ; Lee et al., ; ), catastrophic forgetting is minimal at intermediate task similarity when learning with optimal task selection. We provide a mechanistic explanation of this phenomenon by disentangling dynamical effects at the level of first-layer and readout weights.
- •
We demonstrate that the insights from our optimal strategies in synthetic settings transfer to real datasets. Specifically, we show the efficacy of our pseudo-optimal strategy on a continual learning task using the Fashion-MNIST dataset, where it optimally interpolates between simple heuristic strategies based on the problem’s parameters.

##### Further related works.

In recent years, several theoretical works on online learning dynamics in one-hidden-layer neural networks have addressed a range of machine learning problems, including over-parameterisation (Goldt et al., ), algorithmic analysis (Refinetti et al., ; Srinivasan et al., ), and learning strategies (Lee et al., ; Saglietti et al., ; Sarao Mannelli et al., ). However, these studies have not explored the problem from an optimal control perspective.

Early works addressed the optimality of hyperparameters in high-dimensional online learning for committee machines via control theory. These studies focused on optimising the learning rate (Saad & Rattray, ; Rattray & Saad, ; Schlösser et al., ), the regularisation (Saad & Rattray, ), and the learning rule (Rattray & Saad, ). However, to the best of our knowledge, the problem of optimal task selection has not been explored yet.
Carrasco-Davis et al. () and Li et al. () applied optimal control to the dynamics of connectivist models of behaviour, but their analysis was limited to low-dimensional settings. Urbani () extended the Bellman equation to high-dimensional mean-field dynamical systems, though without considering learning processes.

Several other works have combined ideas from machine learning and optimal control. Notably, Han et al. () interpreted deep learning as an optimal control problem on a dynamical system, where the control variables correspond to the network parameters. Chen & Hazan () formulated meta-optimization as an optimal control problem, but their analysis did not involve dimensionality reduction techniques nor did it address task selection.

## 2 Model-based theoretical framework

We adopt a model-based approach to investigate the supervised learning of multiple tasks. Following Lee et al. (; ), we consider a teacher-student framework (Gardner & Derrida, ) where a “student" neural network is trained on synthetic inputs $\bm{x}\in\mathbb{R}^{N}$, drawn i.i.d. from a standard Gaussian distribution, $x_{i}\sim\mathcal{N}(0,1)$. The labels for each task $t=1,\ldots,T$ are generated by single-layer “teacher" networks: $y^{(t)}=g_{*}(\bm{x}\cdot\bm{w}_{*}^{(t)}/\sqrt{N})$, where $\bm{W}_{*}=(\bm{w}_{*}^{(1)},\ldots,\bm{w}_{*}^{(T)})^{\top}\in\mathbb{R}^{T
\times N}$ denote the corresponding teacher vectors, and $g_{*}$ the activation function. The student is a two-layer neural network with $K$ hidden units, first-layer weights $\bm{W}=\left(\bm{w}_{1},\ldots,\bm{w}_{K}\right)^{\top}\in\mathbb{R}^{K\times N}$, activation function $g$, and second-layer weights $\bm{v}\in\mathbb{R}^{K}$, that outputs the prediction:

$$ $\displaystyle\hat{y}\left(\bm{x};\bm{W},\bm{v}\right)=\sum_{k=1}^{K}v_{k}\,g \left(\frac{\bm{x}\cdot\bm{w}_{k}}{\sqrt{N}}\right)\;.$ (1) $$

Following a standard *multi-headed* approach to continual learning (Zenke et al., ; Farquhar & Gal, ), we allow for task-dependent readout weights: $\bm{V}=(\bm{v}^{(1)},\ldots,\bm{v}^{(T)})^{\top}\in\mathbb{R}^{T\times K}$. While the readout is switched during training according to the task under consideration, the first-layer weights are shared across tasks. A pictorial representation of this model is displayed in Fig. 1. Training is performed via Stochastic Gradient Descent (SGD) on the squared loss of $y^{(t)}$ and $\hat{y}^{(t)}=\hat{y}(\bm{x};\bm{W},\bm{v}^{(t)})$, in the *online* regime, where at each training step the algorithmic update is computed using a new sample $(\bm{x},y^{(t)})$.

Figure: Figure 1: Pictorial representation of the continual learning task in the teacher-student setting. A “student” network is trained on i.i.d. inputs from two teacher networks, defining two different tasks (panel a). The student has sufficient capacity to learn both tasks. However, sequential training results in catastrophic forgetting, where the performance on a previously learned task significantly deteriorates when a new task is introduced (panel b). Parameters: $K=T=2$.
Refer to caption: x1.png

The generalisation error of the student on task $t$ is given by

$$ $\displaystyle\varepsilon_{t}\left(\bm{W},\bm{V},\bm{W}_{*}\right)\coloneqq \frac{1}{2}\Big{<}\Big{(}y^{(t)}-\hat{y}^{(t)}\Big{)}^{2}\Big{>}=\frac{1}{2} \mathbb{E}_{\bm{x}}\left[\left(g^{*}\left(\frac{\bm{w}_{*}^{(t)}\cdot\bm{x}}{ \sqrt{N}}\right)-\hat{y}(\bm{x};\bm{W},\bm{v}^{(t)})\right)^{2}\right]\,,$ (2) $$

where we use the angular brackets $\langle\cdot\rangle$ to denote the expectation over the input distribution for a given set of teacher and student weights. As shown in Lee et al. (; ), we can derive a set of dynamical equations for the generalisation error across training in the high-dimensional limit. We leverage this low-dimensional description and optimal control theory to derive *optimal training protocols* for multi-task learning. In particular, we optimise over task selection and learning rate.

##### Forward training dynamics.

First, we consider the set of equations describing the training dynamics for a given task-selection protocol.
As further discussed in Appendix A, in the limit of large input dimension $N\to\infty$ with $K,T\sim\mathcal{O}_{N}(1)$, the dynamics of the generalisation error is entirely captured by the evolution of the student readouts $\bm{V}$ and a set of low dimensional matrices—called *overlaps* in the statistical physics literature:

$$ $\displaystyle M_{kt}\coloneqq\frac{\bm{w}_{k}\cdot\bm{w}_{*}^{(t)}}{N}\,,$ $\displaystyle Q_{kh}\coloneqq\frac{\bm{w}_{k}\cdot\bm{w}_{h}}{N}\,,$ $\displaystyle S_{tt^{\prime}}\coloneqq\frac{\bm{w}_{*}^{(t)}\cdot\bm{w}_{*}^{( t^{\prime})}}{N}\,,$ (3) $$

for all $k,h=1,\ldots,K$ and $t=1,\ldots,T$. For the remainder of the paper we consider $K=T$, to guarantee that the student network has enough capacity to learn all tasks perfectly. Teacher vectors are normalised, while the task similarity is tuned by a parameter $\gamma$, so that $S_{tt^{\prime}}=\delta_{t,t^{\prime}}+\gamma(1-\delta_{t,t^{\prime}})$. For simplicity, it is useful to encode all the dynamical degrees of freedom of interest—the overlaps and the readout weights—in the same vector. We use the shorthand notation $\mathbb{Q}=({\rm vec}(\bm{Q}),{\rm vec}(\bm{M}),{\rm vec}(\bm{V}))^{\top}\in
\mathbb{R}^{K^{2}+2KT}$. Following Lee et al. (), we write a set of ODEs

$$ $\displaystyle\frac{{\rm d}{\mathbb{Q}}(\alpha)}{{\rm d}\alpha}=f_{\mathbb{Q}} \left(\mathbb{Q}(\alpha),\bm{u}(\alpha)\right)\quad\quad\text{with }\alpha\in( 0,\alpha_{F}]\;,$ (4) $$

$\alpha$ denoting the effective training *time*—i.e., the ratio between training epochs and input dimension $N$, as detailed in Appendix A—and $\bm{u}$ the dynamical variables that we want to control optimally. In particular, we study the optimal schedules for task-selection $t_{c}(\alpha)$ and learning rate $\eta(\alpha)$. Here, $t_{c}(\alpha)\in\{1\,,\ldots\,,T\}$ indicates on which task the student is trained at time $\alpha$. The specific form of the functions $f_{\mathbb{Q}}$ is derived in Appendix A. The initial condition $\mathbb{Q}(0)$ matches the initialisation of the SGD algorithm. In particular, the initial first-layer weights and readout weights are drawn i.i.d. from a normal distribution with variances of $10^{-3}$ and $10^{-2}$, respectively. Notably, the trajectory appears to be largely independent of the specific initialisation of the first-layer weights. For instance, in Fig. 2, simulations and theory correspond to different initialisations, yet the curves show excellent agreement. Let us stress that Eq. 4 is a low-dimensional deterministic equation that fully captures the high-dimensional stochastic dynamics of SGD as $N\to\infty$. This dimensionality reduction is crucial to apply the optimal control techniques presented in the next section.

##### Optimal control framework and backward conjugate dynamics.

Our first main contribution is to derive training strategies that are “optimal” with respect to the generalisation performance *at the end of training* and on *all tasks*. In practice, the goal of the optimisation process is to minimise a linear combination of the generalisation errors on the different tasks at the final training time $\alpha_{F}$:

$$ $\displaystyle h(\mathbb{Q}(\alpha_{F}))=\sum_{t=1}^{T}c_{t}\,\varepsilon_{t}( \mathbb{Q}(\alpha_{F}))\;\quad\quad\text{with}\quad c_{t}\geq 0\;\text{and}\; \sum_{t=1}^{T}c_{t}=1\;,$ (5) $$

where the coefficients $c_{t}$ identify the relative importance of different tasks and $\varepsilon_{t}$ denotes the infinite-dimensional limit of the average generalisation error on task $t$, as defined in Eq. 2. Crucially, we have an analytic expression for $\varepsilon_{t}$, derived in Appendix A. In the remainder of the paper, we assume equally important tasks $c_{t}=1/T$. As customary in optimal control theory (Pontryagin, ), we adopt a variational approach to solve the problem. We define the cost function

$$ $\displaystyle\mathcal{F}[\mathbb{Q},\mathbb{\hat{Q}},{\bm{u}}]=h\left(\mathbb{ Q}(\alpha_{F})\right)+\int_{0}^{\alpha_{F}}{\rm d}\alpha\;\mathbb{\hat{Q}}( \alpha)^{\top}\left[-\frac{{\rm d}\mathbb{Q}(\alpha)}{{\rm d}\alpha}+f_{ \mathbb{Q}}\left(\mathbb{Q}(\alpha),{\bm{u}}(\alpha)\right)\right]\;,$ (6) $$

where the *conjugate order parameters*
$\hat{\mathbb{Q}}=({\rm vec}(\hat{\bm{Q}}),{\rm vec}(\hat{\bm{M}}),{\rm vec}(
\hat{\bm{V}}))^{\top}$ enforce the training dynamics in the training interval $\alpha\in[0,\alpha_{F}]$. Finding the optimal protocol amounts to minimising the cost function $\mathcal{F}$ with respect to $\mathbb{Q}$, $\hat{\mathbb{Q}}$, and $\bm{u}$. We defer the details of this variational procedure to Appendix A, and present only the main steps here. The minimisation with respect to $\mathbb{Q}$ provides a set of equations for the *backward* dynamics of the conjugate parameters

$$ $\displaystyle-\frac{{\rm d}\hat{\mathbb{Q}}(\alpha)^{\top}}{{\rm d}\alpha}= \hat{\mathbb{Q}}(\alpha)^{\top}\nabla_{\mathbb{Q}}f_{\mathbb{Q}}(\mathbb{Q}( \alpha),\bm{u}(\alpha))\quad\quad\text{with }\alpha\in[0,\alpha_{F}).$ (7) $$

The final condition for the dynamics is given by

$$ $\hat{\mathbb{Q}}(\alpha_{F})=\nabla_{\mathbb{Q}}h(\mathbb{Q}_{F})=\sum_{t=1}^{ T}c_{t}\nabla_{\mathbb{Q}}\varepsilon_{t}(\mathbb{Q}(\alpha_{F}))\,.$ (8) $$

The optimal control curve $\bm{u}^{*}(\alpha)$ is obtained as the solution of the minimisation:

$$ $\displaystyle\bm{u}^{*}(\alpha)=\underset{\bm{u}\in\mathcal{U}}{\operatorname{ argmin}}\left\{\hat{\mathbb{Q}}(\alpha)^{\top}f_{\mathbb{Q}}(\mathbb{Q}(\alpha ),\bm{u})\right\}\;,$ (9) $$

where $\mathcal{U}$ is the set of allowed controls, that can be either continuous or discrete. For instance, for task selection we take ${u}(\alpha)=t_{c}(\alpha)$ and $\mathcal{U}=\{1\,,2\,\ldots\,,T\}$, where we use the notation $t_{c}(\alpha)$ to indicate the current task, i.e., the task on which the student is trained at time $\alpha$.
When optimising over both task selection and learning rate schedule we take $\bm{u}=(t_{c},\eta)$ and $\mathcal{U}=\{1\,,2\,\ldots\,,T\}\times\mathbb{R}^{+}$. Crucially, the optimal control equations 4, 7, and 9 must be iterated until convergence, starting from an initial guess on $\bm{u}$. Let us stress that the space $\mathcal{U}$ of possible controls is high-dimensional and hence it is not feasible to explore it via greedy search strategies.

## 3 Results and applications

Figure: Figure 2: How to learn a new task without forgetting the old one? The student is trained on task $1$ until convergence during the first phase ($\alpha\in[0,1000]$), then task $2$ is introduced. During the second phase ($\alpha\in(1000,1025]$), task $1$ may be replayed to prevent forgetting. For better visibility, we only display the regions $\alpha\in[0,20]\cup[990,1025]$. We compare three strategies: a) no replay, b) interleaved replay, i.e., alternating between the two tasks, c) the optimal strategy derived in Sec. 2. Crosses mark numerical simulations of a single trajectory at $N=20000$, lines mark the solution of Eq. 4. Colour bars represent the protocol $t_{c}$. Parameters: $\gamma=0.3$, $K=T=2$, and $\eta=1$.
Refer to caption: x2.png

The theoretical framework of Sec. 2 is extremely flexible and can be applied to a variety of settings. We will present a technical paper focused on the method in a specialised venue, while in this section we focus on a particular setting and investigate in detail the impact of optimal training protocols on both synthetic and real tasks. First, we consider the synthetic teacher-student framework. We compute optimal task-selection protocols, investigating their structure and interplay with task similarity and learning rate schedule. Then, we transfer the insights gained from the interpretation of the optimal strategies in synthetic settings to applications on real datasets.

### 3.1 Experiments on synthetic data

We formulate the problem of continual learning as follows. During a first training phase, the student learns perfectly task $t=1$. Then, the goal is to learn a new task $t=2$ without forgetting the previous one. A given time window of duration $\alpha_{F}$ is assigned for the second training phase. In particular, we investigate the role of *replay*—i.e., presenting samples from task $1$ during the second training phase—and the structure of the optimal replay strategy.

We use the equations derived in Sec. 2 to study optimal replay during the second phase of training. To this end, we take the task-selection variable as our control $u(\alpha)\!=\!t_{c}(\alpha)\!\in\!\{1,2\}$, while we set $t_{c}\!=\!1$ during the first training phase. The result of the optimisation in Eq. 9 strikes the balance between training on the new task and replaying the old task. We do not enforce any constraints on the number of samples from task $1$ to use in the second phase, therefore our method provides both the optimal *fraction* of replayed samples and the optimal task *ordering*, depending on the time window $\alpha_{F}$. Fig. 2 compares the learning dynamics of three different strategies, depicting the loss on task $1$ (full orange line), task $2$ (dashed green line), and their average (dotted black line) as a function of the training time $\alpha$. Numerical simulations—marked by cross symbols—are in excellent agreement with our theory, with deviations smaller than $1/\sqrt{N}$ and hence compatible with finite-size effects.

The student is trained exclusively on task $1$ until $\alpha=1000$, when the task is perfectly learned with loss $\sim 10^{-6}$. Then, the student is trained on a combination of new and old tasks for a training time of duration $\alpha_{F}=25$. A colour bar above each plot illustrates the associated task-selection strategy $t_{c}(\alpha)$. Panel a) shows training without replay, where only task $2$ is presented in the second phase. We observe catastrophic forgetting of task $1$. Panel b) shows a heuristic “interleaved” strategy, where training alternates one sample from the new task to one sample from the old one. As observed in Lee et al. (), the interleaved strategy already provides a performance gain, demonstrating the relevance of replay to mitigate catastrophic forgetting. Panel c) of Fig. 2 shows the loss dynamics for the optimal replay strategy. Notably, the optimal strategy exhibits a complex structure and displays a clear performance improvement over the other two strategies. In particular, we find that the optimal task-selection strategy always presents an initial phase where training is performed only on the new task. This behaviour is observed consistently across a range of task similarities.

Figure: Figure 3: The impact of task similarity on continual learning. a) Average loss on both tasks at the end of the second training phase as a function of the task similarity $\gamma$ under the replay setting from Fig. 2. Different lines correspond to different strategies: no replay (purple crosses), optimal replay (red dots), interleaved (blue squares), pseudo-optimal replay (cyan dashed line). b-e) Optimal replay strategies for different values of $\gamma=0.1,0.3,0.6,0.9$. Colour bars represent the protocol $t_{c}(\alpha)$.
Refer to caption: x3.png

##### The impact of task similarity.

To understand the structure of the optimal strategy, we examine its performance in relation to task similarity $\gamma$. Fig. 3a) depicts the average loss at the end of training as a function of $\gamma$. For the no-replay strategy, we reproduce the findings from previous works (Lee et al., ; ): the highest error occurs at intermediate task similarity. Lee et al. () explained this non-monotonicity as a trade-off between node re-use and node activation. Specifically, for small $\gamma$, there is minimal interference between tasks, and one hidden neuron predominantly aligns with the new task, while the other neuron retains the knowledge of the old task, leading to task *specialisation*. At large $\gamma$, features from task $1$ are reused for task $2$, avoiding forgetting. However, at intermediate $\gamma$, interference is maximal, and both neurons quickly align with task $2$, while task $1$ is forgotten. Remarkably, Fig. 3a) shows that replay reverses this trend: the minimal error occurs at intermediate $\gamma$. To explain this nontrivial behaviour, we must first understand the optimal replay protocol.

##### Interpretation of the optimal replay structure.

The optimal replay dynamics is illustrated in panels b-e) of Fig. 3 for different values of $\gamma$, displaying a highly structured protocol. We can interpret this structure a posteriori: an initial *focus phase* without replay is followed by a *revision phase* involving interleaved replay. The transition between these two phases corresponds approximately to the point at which the loss on the new task matches the loss on the old one. To investigate the significance of this structure, we also test an interleaved strategy, plotted in Fig. 3a), where the task ordering in the second training phase is fully randomised while maintaining the same overall replay fraction of the optimal strategy. This protocol has a performance gap compared to the optimal one, showing the importance of a properly structured replay scheme. Additionally, we test a “pseudo-optimal” variant, where the *focus phase* is retained, but the *revision phase* is randomised. This variant performs comparably to the optimal strategy, suggesting that while the specific order of the revision phase is largely unimportant, it is key to precede it with a training phase on the new task.

We can now attempt to understand the inverted non-monotonic behaviour of the average loss as a function of $\gamma$ under the optimal protocol. First, as shown in Fig. 8 of Appendix C, the optimal protocol achieves a good level of node specialisation across all values of $\gamma$. Thus, replay prevents the task interference that typically causes performance deterioration at intermediate $\gamma$. The non-monotonic behaviour of the optimal curve in Fig. 3a) arises from a different origin, involving two opposing effects related to the first-layer weights and the readout. The initial decrease of the loss with $\gamma$ is quite intuitive, as only minimal knowledge can be transferred from task $1$ to task $2$ when $\gamma$ is small. Consequently, the focus phase on task $2$ must be longer for smaller $\gamma$, leaving less time to revise task $1$, thereby reducing performance. On the other hand, the performance decrease observed in Fig. 3a) for $\gamma>0.3$ is more subtle and is related to the readout layer. Once the two hidden neurons have specialised—each aligning with one of the teacher vectors—we expect the readout weights corresponding to the incorrect teacher to be suppressed. Specifically, if $\bm{w}_{1}=\bm{w}_{*}^{(1)}$ and $\bm{w}_{2}=\bm{w}_{*}^{(2)}$, the learning dynamics should drive the readout weights $\bm{v}^{(1)}=(v^{(1)}_{1},v^{(1)}_{2})^{\top}$ and $\bm{v}^{(2)}=(v^{(2)}_{1},v^{(2)}_{2})^{\top}$ towards $\bm{v}^{(1)}=(1,0)^{\top}$ and $\bm{v}^{(2)}=(0,1)^{\top}$, representing full recovery of the teacher network. As shown in Fig. 8 of Appendix C, the time required to suppress the off-diagonal weights $v^{(1)}_{2}$ and $v^{(2)}_{1}$ increases as $\gamma\to 1$. This is intuitive, as higher task similarity $\gamma$ reduces the distinction between tasks, slowing the suppression of the off-diagonal weights. In Appendix A, we derive analytically the convergence timescale $\alpha_{\rm conv}$ of the readout layer, showing that

$$ $\alpha_{\rm conv}=\frac{3\pi}{\eta(\pi-6\operatorname{arcsin}\left(\gamma/2 \right))}\,,$ (10) $$

where $\eta$ is the learning rate.
This timescale is a monotonically increasing function of $\gamma$ and diverges as $\gamma\!\to\!1$ with $\alpha_{\rm conv}\!\approx\!\sqrt{3}\pi/(2\eta(1-\gamma))$. This result explains the performance decrease of the optimal strategy as $\gamma\to 1$. In summary, the performance decrease for $\gamma\to 0$ is due to the first-layer weights, while for $\gamma\to 1$ it is related to the readout weights.

Figure: Figure 4: Joint optimisation of task-selection and learning-rate. Optimal learning rate as a function of training time $\alpha$ for the same parameters as Fig. 2. There is a single optimal learning rate curve, but for visibility purposes we show it as a solid orange line when training on task $1$ and a dashed green line on task $2$. The task-selection protocol $t_{c}(\alpha)$ is illustrated in the colour bar.
Refer to caption: x4.png

##### Optimal learning rate schedules.

While the above results were obtained with constant learning rate, annealing schedules are widely used in machine learning. Thus, it is relevant to study the optimal interplay between replay and learning rate scheduling. Optimal learning rate dynamics have been studied with a similar approach in Saad & Rattray (). We jointly optimise the task-selection protocol together with the learning rate to investigate its impact on continual learning.
Fig. 4 shows the optimal learning rate schedule for task similarity $\gamma\!=\!0.3$ in the second training phase of duration $\alpha_{F}\!=\!25$. Similarly as for constant learning rate (panel c) of Fig. 3), optimal task-selection is characterised by an initial focus phase. Notably, this phase coincides with a strong annealing of the learning rate to achieve optimal performance. Intuitively, when learning the new task, the learning rate starts high and is gradually decreased over time. Interestingly, while entering the revision phase, the optimal learning rate schedule exhibits a highly nontrivial structure (see Fig. 4). Indeed, despite the optimal learning rate curve is unique, we find that effectively it can be seen as two different curves, associated to the respective tasks. In practice, the optimal learning rate curve “jumps” between these two curves according to the task selected at a given training time. Overall, the joint optimisation over task selection and learning rate provides a significant improvement in performance. Indeed, the average loss obtained by this strategy—shown in Fig. 9 of Appendix C—exhibits a large gap compared to all the other protocols.

##### Multi-task learning from scratch.

We also consider a multi-task setting (see Fig. 5) where both tasks must be learned from scratch within a fixed number of epochs corresponding to a total training time $\alpha_{F}$. We first consider sequential learning, i.e., training only on task $1$ for $\alpha<\alpha_{F}/2$, then only on task $2$, or vice versa. As shown in Fig. 5a) sequential learning leads to catastrophic forgetting, with the worst performance observed at intermediate task similarity. In contrast, using a randomly interleaved strategy, where examples from both tasks are presented in equal proportion but in random order, shows significant improvement. This approach can exploit task similarity, leading to a monotonic decrease in average loss as $\gamma$ increases. The optimal strategy, displayed in Fig. 5b-e) for various values of $\gamma$, follows a structured interleaved protocol that further enhances performance. Contrary to the continual learning framework, the optimal structure gives only marginal gain over the plain interleaved stratetegy and does not offer a clear a posteriori interpretation.

### 3.2 Experiments on real data

Figure: Figure 5: The impact of task similarity on multi-task learning. a) Average loss as a function of task similarity $\gamma$ at the end of training ($\alpha_{F}=10$). Different lines correspond to different strategies: sequential (purple crosses), optimal (red dots), and randomly interleaved with $50\%$ samples from each task (blue squares). b-e) Optimal replay strategy for different values of $\gamma=0.1,0.3,0.6,0.9$. Parameters: $K\!=\!T\!=\!2$ and $\eta\!=\!10$.
Refer to caption: x5.png

We consider the experimental framework established in Ramasesh et al. (); Lee et al. () for the study of task similarity in relation to catastrophic forgetting.
We use the Fashion-MNIST dataset (Xiao et al., ) to generate upstream and downstream tasks: The upstream dataset—$\mathcal{D}_{1}=\{\boldsymbol{x}_{i}^{(1)},y_{i}^{(1)}\}_{i}$—consists in a pair of classes from the standard dataset, while the downstream dataset is generated by a linear interpolation of the upstream dataset with a second auxiliary dataset—$\tilde{\mathcal{D}}=\{\tilde{\boldsymbol{x}}_{i},\tilde{y}_{i}\}_{i}$—containing a new pair of classes,

$$ $\mathcal{D}_{2}=\{\boldsymbol{x}_{i}^{(2)},y_{i}^{(2)}\}_{i}=\{\gamma \boldsymbol{x}_{i}^{(1)}+(1-\gamma)\tilde{\boldsymbol{x}}_{i},\gamma y_{i}^{(1 )}+(1-\gamma)\tilde{y}_{i}\}_{i}$ (11) $$

where the parameter $\gamma$ controls the task similarity.
We then train a standard two-layer feedforward ReLU neural network on the two datasets using online SGD on a squared error loss. We consider a dynamical architecture (Zhou et al., ; Rusu et al., ) where the readout weights are changed switching from one task to another, but the hidden layer is shared. During training, we apply the three strategies discussed in the previous sections: a no-replay strategy, a strategy with interleaved replay, and a “pseudo-optimal” strategy. Recall that the latter is inspired by the optimal protocol derived in the previous section. It consists of an initial phase of training exclusively on the new task until performance on both tasks becomes comparable, followed by a phase of interleaved replay. Crucially, this protocol can be easily implemented in practice, as it only requires an estimate of the generalization error on the two tasks, which can be obtained in real-world settings.

Figure: Figure 6: Training dynamics. Training curves on the modified fashion MNIST task at similarity $\gamma=0.5$. The network is trained for $10.000$ epochs on the first task before switching to the second task and being trained for additional $10.000$ epochs. The results are obtained from $100$ realisations of the problem. The first three panels show the test loss on task $1$ (solid orange), task $2$ (dashed green), and their average (dotted black) for three training strategies, from left to right: no-replay, interleaved, and pseudo-optimal. The rightmost panel shows the average loss over the entire training.
Refer to caption: x6.png

Fig. 6 shows the training loss under the different training protocols for $\gamma=0.5$. While the no-replay strategy appears to be successful for small downstream datasets (i.e., a few epochs in the online framework) in the longer run it leads to strong forgetting and high average loss. The interleaved is beneficial in the long run but largely slows down learning of the new task. Overall, the pseudo-optimal protocol identified in Sec. 3.1 shows a better performance over the entire trajectory.

Figure: Figure 7: Average loss comparison. The figure focuses on the average loss and shows the final loss achieved by the three strategies as we increase the size of task $2$ (from left to right: $500$, $1.000$, $2.000$, $4.000$, and $8.000$ samples) while task $1$ has always $10.000$ samples. Individual panels show the performance of the three strategies as we span the value of $\gamma$ form $0.05$ to $0.95$.
Refer to caption: x7.png

This result is not limited to the specific value of $\gamma$. In Fig. 7 we show snapshots of the average loss for different downstream task sizes while spanning over a range of $\gamma$s. This figure provides additional support to the observation reported previously that the no-replay strategy is optimal for small downstream tasks, the interleaved strategy is convenient for large downstream tasks, and the pseudo-optimal one combines the benefits of the two leading to the best performance overall. In summary, the pseudo-optimal strategy derived for the synthetic model performs well on real-world data. Notably, despite the differences between the synthetic and real settings—such as data structure—the pseudo-optimal strategy remains effective, demonstrating its robustness across problems.

## 4 Discussion

##### Conclusion.

In this work, we introduce a systematic approach for identifying and interpreting optimal task-selection strategies in synthetic learning settings. We consider a teacher-student scenario as a prototypical continual learning problem to achieve analytic understanding of supervised multi-task learning. We incorporate prior results on exact ODEs for high-dimensional online SGD dynamics into a control-theory framework that allows us to derive exact equations for the optimal protocols. Our theory reveals that optimal task-selection protocols are typically highly structured—alternating between focused learning and interleaved replay phases—and display a nontrivial interplay with task similarity. We also identify highly structured optimal learning rate schedules that synchronise with optimal task-selection to enhance overall performance. Finally, leveraging insights from the synthetic setting, we extract a pseudo-optimal strategy applicable to real tasks.

##### Limitations and Perspectives.

This work marks an initial step toward a systematic theoretical understanding of optimal training protocols in synthetic learning settings. Here, we discuss limitations and highlight promising future directions. First, Pontryagin’s maximum principle provides a necessary condition for optimality but does not guarantee a global optimum. Nevertheless, the strategies derived from this approach performed significantly better than previously proposed heuristics. Additionally, Pontryagin’s principle does not easily extend to stochastic problems. This limitation is overcome in the high-dimensional limit where concentration results provide deterministic dynamical equations. For simplicity, we focus on i.i.d. Gaussian inputs, but our analysis can be extended to more structured data models (Goldt et al., ; Loureiro et al., ; Adomaityte et al., ) to study how input distribution affects task selection. In particular, we do not model the relative task difficulty—an important extension that naturally connects to the theory of curriculum learning (Weinshall & Amir, ; Saglietti et al., ; Cornacchia & Mossel, ; Abbe et al., ). Furthermore, it would be interesting to go beyond the study of online dynamics to understand the impact of memorisation in batch learning settings (Sagawa et al., ). Existing results in the spurious correlations (Ye et al., ) and fairness (Ganesh et al., ) literature suggest a strong dependence of the classifier’s bias on the presentation order in batch learning. Our method can be applied to mean-field models—like (Mannelli et al., ; Jain et al., )—to theoretically investigate this phenomenon. An interesting extension of our work involves applying recently-developed statistical physics methods to the study of deeper networks (Bordelon & Pehlevan, ; ) and more complex learning architectures (Rende et al., ; Tiberi et al., ). Another interesting direction concerns finding optimal protocols for shaping, where task order significantly impacts both animal learning and neural networks (Skinner, ; Tong et al., ; Lee et al., ).

#### Acknowledgments

We thank Rodrigo Carrasco-Davis, Sebastian Goldt, and Andrew Saxe for useful feedback.
This work was supported by a Leverhulme Trust International Professorship grant [number LIP-2020-014],
the Wallenberg AI, Autonomous Systems, and Software Program (WASP),
and by the Simons Foundation (Award Number: 1141576).

## Appendix A Details on the theoretical derivations

In this appendix, we provide detailed derivations of the equations in Sec. 2 of the main text. In the interest of completeness, we also report the derivation of the ODEs describing online SGD dynamics and the generalisation error as a function of the order parameters, first derived in (Lee et al., ). We remind that inputs are $N-$dimensional vectors $\bm{x}\in\mathbb{R}^{N}$ with independent identically distributed (i.i.d.) standard Gaussian entries ${x}_{i}\sim\mathcal{N}(0,1)$, while the labels are generated by single-layer teacher networks: $y^{(t)}=g_{*}(\bm{x}\cdot\bm{w}_{*}^{(t)}/\sqrt{N})$, $t=1,\ldots,T$, with a different teacher for each task. The student is a one-hidden layer network that outputs the prediction:

$$ $\displaystyle\hat{y}^{(t)}=\sum_{k=1}^{K}v_{k}^{(t)}g\left(\frac{\bm{x}\cdot \bm{w}_{k}}{\sqrt{N}}\right)\;.$ (12) $$

We focus on the *online* (on *one-pass*) setting, so that at each training step the student
network is presented with a fresh example $\bm{x}^{\mu}$, $\mu=1,\ldots,P$, and $P/N\sim\mathcal{O}_{N}(1)$. The weights of the student are updated through gradient descent on $\frac{1}{2}(\hat{y}^{(t)}-y^{(t)})^{2}$ following the task-selection protocol $t_{c}$:

$$ $\displaystyle\begin{split}\bm{w}_{k}^{\mu+1}&=\bm{w}_{k}^{\mu}-\eta^{\mu} \Delta^{(t_{c})\mu}\,v_{k}^{(t_{c})\mu}\,g^{\prime}\left(\lambda^{\mu}_{k} \right)\frac{\bm{x}^{\mu}}{\sqrt{N}}\;,\\ v_{k}^{(t)\mu+1}&=v_{k}^{(t)\mu}-\frac{\eta^{\mu}}{N}\Delta^{(t)\mu}\,g( \lambda^{\mu}_{k})\delta_{t,t_{c}}\;,\\ \Delta^{(t)\mu}&\coloneqq\hat{y}^{(t)\mu}-y^{(t)\mu}=\sum_{k=1}^{K}v_{k}^{(t)} g(\lambda^{\mu}_{k})-g_{*}(\lambda_{*}^{(t)\mu})\;,\end{split}$ (13) $$

where $\eta^{\mu}$ denotes the (possibly time-dependent) learning rate and we have rescaled it by $N$ in the dynamics of the readout weights for future convenience. We have defined the preactivations, a.k.a. *local fields*,

$$ $\displaystyle\lambda^{\mu}_{k}\coloneqq\frac{\bm{x}^{\mu}\cdot\bm{w}^{\mu}_{k} }{\sqrt{N}}\;,$ $\displaystyle\lambda_{*}^{(t)\mu}\coloneqq\frac{\bm{x}^{\mu}\cdot\bm{w}^{(t)}_ {*}}{\sqrt{N}}\;.$ (14) $$

Notice that, due to the online-learning setup, at each training epoch the input $\bm{x}$ is independent of the weights. Therefore, due to
the Gaussianity of the inputs, the local fields are also jointly Gaussian with zero mean and second
moments given by the *overlaps*:

$$ $\displaystyle\begin{split}M_{kt}&\coloneqq\mathbb{E}_{\bm{x}}\left[\lambda_{k} \lambda_{*}^{(t)}\right]=\frac{\bm{w}_{k}\cdot\bm{w}_{*}^{(t)}}{N}\;,\\ Q_{kh}&\coloneqq\mathbb{E}_{\bm{x}}\left[\lambda_{k}\lambda_{h}\right]=\frac{ \bm{w}_{k}\cdot\bm{w}_{h}}{N}\;,\\ S_{tt^{\prime}}&\coloneqq\mathbb{E}_{\bm{x}}\left[\lambda_{*}^{(t)}\lambda_{*} ^{(t^{\prime})}\right]=\frac{\bm{w}_{*}^{(t^{\prime})}\cdot\bm{w}_{*}^{(t)}}{N }\;.\end{split}$ (15) $$

### A.1 Generalisation error as a function of the order parameters

We can write the generalisation error (Eq. 2 of the main text) as an average over the local fields:

$$ $\displaystyle\begin{split}\varepsilon_{t}\left(\bm{W},\bm{V},\bm{W}_{*}\right) &=\frac{1}{2}\sum_{k,h}v_{k}^{(t)}v_{h}^{(t)}\mathbb{E}_{\bm{\lambda},\bm{ \lambda}_{*}}\left[g(\lambda_{k})g(\lambda_{h})\right]+\frac{1}{2}\mathbb{E}_{ \bm{\lambda},\bm{\lambda}_{*}}\left[g_{*}(\lambda^{(t)}_{*})^{2}\right]\\ &\quad-\sum_{k}v_{k}^{(t)}\mathbb{E}_{\bm{\lambda},\bm{\lambda}_{*}}\left[g( \lambda_{k})g_{*}(\lambda_{*}^{(t)})\right]\;.\end{split}$ (16) $$

where the expectation is computed over the multivariate Gaussian distribution

$$ $\displaystyle\begin{split}P(\bm{\lambda},\bm{\lambda}_{*})&=\frac{1}{\sqrt{(2 \pi)^{K+T}|\bm{C}|}}\exp\left(-\frac{1}{2}(\bm{\lambda},\bm{\lambda}_{*})^{ \top}\bm{C}^{-1}(\bm{\lambda},\bm{\lambda}_{*})\right)\;,\\ \bm{C}&=\left(\begin{array}[]{cc}\bm{Q}&\bm{M}\\ \bm{M}^{\top}&\bm{S}\end{array}\right)\;.\end{split}$ (17) $$

From now on, we adopt the unified notation

$$ $\displaystyle I_{2}(\beta,\rho)\coloneqq\mathbb{E}_{\bm{\lambda},\bm{\lambda}_ {*}}\left[g_{\beta}(\lambda_{\beta})g_{\rho}(\lambda_{\rho})\right]\;,$ (18) $$

where $\beta,\rho$ can refer both to the indices of the student weights $k,h$ or the tasks $t,t^{\prime}$. We can then rewrite the generalisation error as

$$ $\displaystyle\begin{split}\varepsilon_{t}\left(\bm{W},\bm{V},\bm{W}_{*}\right) &=\frac{1}{2}\sum_{k,h}v_{k}^{(t)}v_{h}^{(t)}I_{2}(k,h)+\frac{1}{2}I_{2}(t,t)- \sum_{k}v_{k}^{(t)}I_{2}(k,t)\;.\end{split}$ (19) $$

In all the results presented in Sec. 2, we consider $g(z)=g_{*}(z)={\rm erf}\left(z/\sqrt{2}\right)$. In this case, there is an analytic expression for the integral $I_{2}$ (Saad & Solla, ):

$$ $\displaystyle I_{2}(\beta,\rho)=\frac{2}{\pi}\arcsin\frac{q_{\beta\rho}}{\sqrt {1+q_{\beta\beta}}\sqrt{1+q_{\rho\rho}}}\;,$ (20) $$

and we use the symbol $q$ to denote generically an overlap from Eq. 15, according to the choice of indices $\beta,\rho$, e.g., $q_{kh}=Q_{kh}$, $q_{kt}=M_{kt}$, and $q_{tt_{c}}=S_{tt_{c}}$. In this special case, the generalisation error can be written explicitly as a function of the overlaps

$$ $\displaystyle\begin{split}\varepsilon_{t}\left(\bm{W},\bm{V},\bm{W}_{*}\right) &=\frac{1}{2\pi}\sum_{k,h}v_{k}^{(t)}v_{h}^{(t)}\arcsin\frac{Q_{kh}}{\sqrt{1+Q _{kk}}\sqrt{1+Q_{hh}}}+\frac{1}{2\pi}\arcsin\frac{S_{tt}}{1+S_{tt}}\\ &\quad-\frac{1}{\pi}\sum_{k}v_{k}^{(t)}\arcsin\frac{M_{kt}}{\sqrt{1+Q_{kk}} \sqrt{1+S_{tt}}}\;.\end{split}$ (21) $$

### A.2 Ordinary differential equations for the forward training dynamics

Given that the generalisation error depends only on the overlaps, in order to characterise the learning curves we need to compute the equations of motion for the overlaps from the SGD dynamics of the weights given in Eq. 13. The order parameter $S_{tt^{\prime}}$ associated to the teachers is constant in time. We obtain an ODE for $M_{kt}$ by multiplying both sides of the first of Eq. 13 by $\bm{w}_{*}^{(t)}$ and dividing by $N$:

$$ $\displaystyle\frac{\bm{w}_{k}^{\mu+1}\cdot\bm{w}_{*}^{(t)}}{N}-\frac{\bm{w}_{k }^{\mu}\cdot\bm{w}_{*}^{(t)}}{N}=-\frac{\eta^{\mu}}{N}\Delta^{(t_{c})\mu}v_{k} ^{(t_{c})\mu}g^{\prime}(\lambda_{k}^{\mu})\lambda_{*}^{(t)\mu}\;,$ (22) $$

where we stress the difference between $t_{c}$, the task selected for training at epoch $\mu$, and $t$, the task for which we compute the overlap. We define a “training time” $\alpha=\mu/N$ and take the infinite-dimensional limit $N\rightarrow\infty$. The parameter $\alpha$ becomes continuous and $M_{kt}$ concentrates to the solution of the
following ODE:

$$ $\displaystyle\frac{{\rm d}M_{kt}}{{\rm d}\alpha}=-\eta v_{k}^{(t_{c})}\, \mathbb{E}_{\bm{\lambda},\bm{\lambda}_{*}}\left[\Delta^{(t_{c})}g^{\prime}( \lambda_{k})\lambda_{*}^{(t)}\right]\coloneqq f_{\bm{M},kt}\;,$ (23) $$

where the expectation is computed over the distribution in Eq. 17. The ODE for $Q_{kh}$ is obtained similarly from Eq. 13:

$$ $\displaystyle\begin{split}\frac{\bm{w}_{k}^{\mu+1}\cdot\bm{w}_{h}^{\mu+1}}{N}- \frac{\bm{w}_{k}^{\mu}\cdot\bm{w}_{h}^{\mu}}{N}&=-\frac{\eta^{\mu}}{N}\Delta^{ (t_{c})\mu}v_{k}^{(t_{c})\mu}g^{\prime}(\lambda^{\mu}_{k})\lambda^{\mu}_{h}- \frac{\eta^{\mu}}{N}\Delta^{(t_{c})\mu}v_{h}^{(t_{c})\mu}g^{\prime}(\lambda^{ \mu}_{h})\lambda^{\mu}_{k}\\ &\quad+(\eta^{\mu})^{2}\left(\Delta^{(t_{c})\mu}\right)^{2}v_{k}^{(t_{c})\mu}v _{h}^{(t_{c})\mu}g^{\prime}(\lambda^{\mu}_{k})g^{\prime}(\lambda^{\mu}_{h}) \frac{\bm{x}\cdot\bm{x}}{N}\;.\end{split}$ (24) $$

In the infinite-dimensional limit, we find

$$ $\displaystyle\frac{{\rm d}Q_{kh}}{{\rm d}\alpha}$ $\displaystyle=-\eta v_{k}^{(t_{c})}\mathbb{E}_{\bm{\lambda},\bm{\lambda}_{*}} \left[\Delta^{(t_{c})}g^{\prime}(\lambda^{\mu}_{k})\lambda^{\mu}_{h}\right]- \eta v_{h}^{(t_{c})}\mathbb{E}_{\bm{\lambda},\bm{\lambda}_{*}}\left[\Delta^{(t _{c})}g^{\prime}(\lambda^{\mu}_{h})\lambda^{\mu}_{k}\right]$ (25) $\displaystyle\quad+\eta^{2}v_{k}^{(t_{c})}v_{h}^{(t_{c})}\mathbb{E}_{\bm{ \lambda},\bm{\lambda}_{*}}\left[\left(\Delta^{(t_{c})}\right)^{2}g^{\prime}( \lambda_{k})g^{\prime}(\lambda_{h})\right]\coloneqq f_{\bm{Q},kh}\;.$ (26) $$

Finally, taking the infinite dimensional limit of the second Eq. 13, we find the ODE for the readout:

$$ $\displaystyle\frac{{\rm d}v_{k}^{(t)}}{{\rm d}\alpha}$ $\displaystyle=-\eta\,\mathbb{E}_{\bm{\lambda},\bm{\lambda}_{*}}\left[\Delta^{( t)}g(\lambda_{k})\right]\delta_{t,t_{c}}\coloneqq f_{\bm{V},tk}\;.$ (27) $$

It is useful to write this system of ODEs in a more compact form. With the shorthand notation $\mathbb{Q}=\left({\rm vec}(\bm{Q}),{\rm vec}(\bm{M}),{\rm vec}(\bm{V})\right)^
{\top}$, $f_{\mathbb{Q}}=\left({\rm vec}(f_{\bm{Q}}),{\rm vec}(f_{\bm{M}}),{\rm vec}(f_{
\bm{V}})\right)^{\top}$, we can write

$$ $\displaystyle\frac{{\rm d}{\mathbb{Q}}(\alpha)}{{\rm d}\alpha}=f_{\mathbb{Q}} \left(\mathbb{Q}(\alpha),\bm{u}(\alpha)\right)\;,$ $\displaystyle\alpha\in(0,\alpha_{F}]\;.$ (28) $$

The initial condition for $\mathbb{Q}(0)$ is chosen to reproduce the random initialisation of the SGD algorithm. In particular, the initial first-layer weights and readout weights are drawn i.i.d. from a normal distribution with variances of $10^{-3}$ and $10^{-2}$, respectively.

It is useful to write explicit expressions for the integrals involved in $f_{\mathbb{Q}}$ (Lee et al., ). First, expanding the terms in $\Delta^{(t)}$, we can write

$$ $\displaystyle\begin{split}f_{\bm{Q},kh}&=-\eta v_{k}^{(t_{c})}\left[\sum_{n=1} ^{K}v_{n}^{(t_{c})}I_{3}(n,k,h)-I_{3}(t_{c},k,h)\right]\\ &\quad-\eta v_{h}^{(t_{c})}\left[\sum_{n=1}^{K}v_{n}^{(t_{c})}I_{3}(n,h,k)-I_{ 3}(t_{c},h,k)\right]\\ &\quad+\eta^{2}v_{k}^{(t_{c})}v_{h}^{(t_{c})}\left[\sum_{n,m=1}^{K}v_{n}^{(t_{ c})}v_{m}^{(t_{c})}I_{4}(n,m,k,h)+I_{4}(t_{c},t_{c},k,h)\right.\\ &\left.\qquad-2\sum_{n=1}^{K}v_{n}^{(t_{c})}I_{4}(n,t_{c},k,h)\right]\end{ split}\;,$ (29) $$

$$ $\displaystyle\begin{split}f_{\bm{M},kt}=-\eta v_{k}^{(t_{c})}\sum_{n=1}^{K}v_{ n}^{(t_{c})}I_{3}(n,k,t)+\eta v_{k}^{(t_{c})}I_{3}(t_{c},k,t)\;,\end{split}$ (30) $$

$$ $\displaystyle f_{\bm{V},tk}=\eta\left[-\sum_{n=1}^{K}v_{n}^{(t_{c})}I_{2}(k,n) +I_{2}(k,t_{c})\right]\delta_{t,t_{c}}\;.$ (31) $$

Similarly as in Eq. 18, we adopt the unified notation for the integrals

$$ $\displaystyle\begin{split}I_{3}(\beta,\rho,\zeta)&\coloneqq\mathbb{E}_{\bm{ \lambda},\bm{\lambda}_{*}}\left[\lambda_{\beta}g^{\prime}_{\rho}(\lambda_{\rho })g(\lambda_{\zeta})\right]\;,\\ I_{4}(\beta,\rho,\zeta,\tau)&\coloneqq\mathbb{E}_{\bm{\lambda},\bm{\lambda}_{* }}\left[g_{\beta}(\lambda_{\beta})g_{\rho}(\lambda_{\rho})g^{\prime}_{\zeta}( \lambda_{\zeta})g^{\prime}_{\tau}(\lambda_{\tau})\right]\;,\end{split}$ (32) $$

where $\beta,\rho,\zeta,\tau$ can refer both to the indices of the student weights $k,h,n,m$ or the tasks $t,t_{c}$. In the special case $g(z)=g_{*}(z)={\rm erf}(z/\sqrt{2})$, the integrals have explicit expressions as a function of the overlaps

$$ $\displaystyle\begin{split}I_{3}(\beta,\rho,\zeta)&=\frac{2q_{\rho\zeta}(1+q_{ \beta\beta})-2q_{\beta\rho}q_{\beta\zeta}}{\pi\sqrt{\Lambda_{3}}(1+q_{\beta \beta})}\;,\,\\ I_{4}(\beta,\rho,\zeta,\tau)&=\frac{4}{\pi^{2}\sqrt{\Lambda_{4}}}\arcsin\frac{ \Lambda_{0}}{\sqrt{\Lambda_{1}\Lambda_{2}}}\;,\end{split}$ (33) $$

the symbol $q$ denotes generically an overlap from Eq. 15, and

$$ $\displaystyle\begin{split}\Lambda_{0}&=\Lambda_{4}\,q_{\beta\rho}-q_{\beta\tau }\,q_{\rho\tau}\,(1+q_{\zeta\zeta})-q_{\beta\zeta}\,q_{\rho\zeta}\,(1+q_{\tau \tau})+q_{\zeta\tau}\,q_{\beta\zeta}\,q_{\rho\tau}+q_{\zeta\tau}\,q_{\rho\zeta }\,q_{\beta\tau}\;,\\ \Lambda_{1}&=\Lambda_{4}\,(1+q_{\beta\beta})-q_{\beta\tau}^{2}\,(1+q_{\zeta \zeta})-q_{\beta\zeta}^{2}\,(1+q_{\tau\tau})+2q_{\zeta\tau}q_{\beta\zeta}\,q_{ \beta\tau}\;,\\ \Lambda_{2}&=\Lambda_{4}\,(1+q_{\rho\rho})-q_{\rho\tau}^{2}\,(1+q_{\zeta\zeta} )-q_{\rho\zeta}^{2}\,(1+q_{\tau\tau})+2q_{\zeta\tau}q_{\rho\zeta}q_{\rho\tau} \;,\\ \Lambda_{3}&=(1+q_{\beta\beta})(1+q_{\rho\rho})-q_{\beta\rho}^{2}\;,\\ \Lambda_{4}&=(1+q_{\zeta\zeta})\,(1+q_{\tau\tau})-q_{\zeta\tau}^{2}\;.\end{split}$ (34) $$

### A.3 Informal derivation of Pontryagin maximum principle

Let us consider the augmented cost function

$$ $\displaystyle\mathcal{F}[\mathbb{Q},\mathbb{\hat{Q}},{\bm{u}}]=h\left(\mathbb{ Q}(\alpha_{F})\right)+\int_{0}^{\alpha_{F}}{\rm d}\alpha\;\mathbb{\hat{Q}}( \alpha)^{\top}\left[-\frac{{\rm d}\mathbb{Q}(\alpha)}{{\rm d}\alpha}+f_{ \mathbb{Q}}\left(\mathbb{Q}(\alpha),{\bm{u}}(\alpha)\right)\right]\;,$ (35) $$

where the conjugate variables $\mathbb{\hat{Q}}(\alpha)$ act as Lagrange multipliers, enforcing the dynamics at time $\alpha$. Setting to zero variations with respect to $\mathbb{\hat{Q}}(\alpha)$ results in the forward dynamics

$$ $\frac{\delta\mathcal{F}[\mathbb{Q},\mathbb{\hat{Q}},{\bm{u}}]}{\delta\mathbb{ \hat{Q}}(\alpha)}=0\Rightarrow\frac{{\rm d}\mathbb{Q}(\alpha)}{{\rm d}\alpha}= f_{\mathbb{Q}}\left(\mathbb{Q}(\alpha),{\bm{u}}(\alpha)\right)\,.$ (36) $$

Integrating by parts, we find

$$ $\displaystyle\mathcal{F}[\mathbb{Q},\mathbb{\hat{Q}},{\bm{u}}]$ $\displaystyle=h\left(\mathbb{Q}(\alpha_{F})\right)+\int_{0}^{\alpha_{F}}{\rm d }\alpha\;\mathbb{\hat{Q}}(\alpha)^{\top}f_{\mathbb{Q}}\left(\mathbb{Q}(\alpha) ,{\bm{u}}(\alpha)\right)+\int_{0}^{\alpha_{F}}{\rm d}\alpha\frac{{\rm d} \mathbb{\hat{Q}}(\alpha)}{{\rm d}\alpha}^{\top}\mathbb{Q}(\alpha)$ (37) $\displaystyle-\mathbb{\hat{Q}}(\alpha_{F})\mathbb{Q}(\alpha_{F})+\mathbb{\hat{ Q}}(0)\mathbb{Q}(0)\;.$ $$

Setting to zero variations with respect to $\mathbb{Q}(\alpha)$ for $0<\alpha<\alpha_{F}$, we find the backward dynamics

$$ $\displaystyle-\frac{{\rm d}\hat{\mathbb{Q}}(\alpha)^{\top}}{{\rm d}\alpha}= \hat{\mathbb{Q}}(\alpha)^{\top}\nabla_{\mathbb{Q}}f_{\mathbb{Q}}\left(\mathbb{ Q}(\alpha),\bm{u}(\alpha)\right)\,,$ (38) $$

while for $\alpha=\alpha_{F}$ we get the final condition

$$ $\hat{\mathbb{Q}}(\alpha_{F})=\nabla_{\mathbb{Q}}h(\mathbb{Q}(\alpha_{F}))\,.$ (39) $$

Note that we do not consider variations with respect to $\mathbb{Q}(0)$ as this quantity is fixed by the initial condition $\mathbb{Q}(0)=\mathbb{Q}_{0}$. Finally, minimizing the cost function with respect to the control ${\bm{u}}$, we get the optimality condition in Eq. 9 of the main text.

### A.4 Optimal control framework

To determine the optimal control, we iterate Eqs. 4, 7, and 9 of the main text until convergence (Bechhoefer, ). Let us consider first the case where the control is the current task $t_{c}(\alpha)$, such that $t_{c}(\alpha)=t$ if the network is trained on task $t\in\{1\,,\ldots\,,T\}$ at training time $\alpha$. For simplicity, we focus on the case $T=2$, but the following discussion is easily generalised to any $T$. In particular, since here ${u}(\alpha)=t_{c}(\alpha)$ the evolution equation 4 can be written as

$$ $\displaystyle\frac{{\rm d}{\mathbb{Q}}(\alpha)}{{\rm d}\alpha}=f_{\mathbb{Q}} \left(\mathbb{Q}(\alpha),t_{c}(\alpha)\right)\,,\quad\mathbb{Q}(0)=\mathbb{Q}_ {0}\,.$ (40) $$

Similarly, the backward dynamics reads

$$ $\displaystyle-\frac{{\rm d}\hat{\mathbb{Q}}(\alpha)^{\top}}{{\rm d}\alpha}= \hat{\mathbb{Q}}(\alpha)^{\top}\nabla_{\mathbb{Q}}f_{\mathbb{Q}}\left(\mathbb{ Q}(\alpha),t_{c}(\alpha)\right)\,,$ (41) $$

with final condition

$$ $\hat{\mathbb{Q}}(\alpha_{F})=\frac{1}{2}\nabla_{\mathbb{Q}}\varepsilon_{1}( \mathbb{Q}(\alpha_{F}))+\frac{1}{2}\nabla_{\mathbb{Q}}\varepsilon_{2}(\mathbb{ Q}(\alpha_{F}))\,.$ (42) $$

The optimality equation 9 yields

$$ $\displaystyle t_{c}^{*}(\alpha)=\underset{t_{c}\in\{1,2\}}{\operatorname{ argmin}}\left\{\hat{\mathbb{Q}}(\alpha)^{\top}f_{\mathbb{Q}}\left(\mathbb{Q}( \alpha),t_{c}(\alpha)=t_{c}\right)\right\}\;.$ (43) $$

Therefore, we find the explicit formula for the optimal task protocol

$$ $t_{c}^{*}(\alpha)=\begin{cases}1\,\quad&\text{ if }\hat{\mathbb{Q}}(\alpha)^{ \top}\left[f_{\mathbb{Q}}\left(\mathbb{Q}(\alpha),t_{c}(\alpha)=2\right)-f_{ \mathbb{Q}}\left(\mathbb{Q}(\alpha),t_{c}(\alpha)=1\right)\right]>0\\ 2\,\quad&\text{ otherwise. }\end{cases}$ (44) $$

Then, we start from a guess for the control variable $t_{c}(\alpha)$. We integrate Eq. 40 forward, obtaining the trajectory $\mathbb{Q}(\alpha)$ for $\alpha\in(0,\alpha_{F})$. Then, we integrate the backward equation 41, starting from the final condition 42, obtaining the trajectory $\hat{\mathbb{Q}}(\alpha)$ for $\alpha\in(0,\alpha_{F})$. Then, the control variable can be updated using Eq. 44 and used in the next iteration of the algorithm. These equations 40, 41, and 44 are iterated until convergence.

We next consider the joint optimisation of the learning rate schedule $\eta(\alpha)$ and the task protocol $t_{c}(\alpha)$. The optimality condition 9 can be written as

$$ $\displaystyle(t_{c}^{*}(\alpha),\eta(\alpha))=\underset{t_{c}\in\{1,2\},\eta \in\mathbb{R}^{+}}{\operatorname{argmin}}\left\{\hat{\mathbb{Q}}(\alpha)^{\top }f_{\mathbb{Q}}\left(\mathbb{Q}(\alpha),(t_{c}(\alpha),\eta(\alpha))=(t_{c}, \eta)\right)\right\}\;.$ (45) $$

Crucially, the function $\hat{\mathbb{Q}}^{\top}f_{\mathbb{Q}}(\mathbb{Q},(t_{c},\eta))$ turns out to be quadratic in $\eta$. Explicitly,

$$ $\hat{\mathbb{Q}}^{\top}f_{\mathbb{Q}}(\mathbb{Q},(t_{c},\eta))=a\eta^{2}+b\eta\,,$ (46) $$

where

$$ $\displaystyle a=$ $\displaystyle\sum_{k,h=1}^{K}\hat{Q}_{kh}v_{k}^{(t_{c})}v_{h}^{(t_{c})}\left[ \sum_{n,m=1}^{K}v_{n}^{(t_{c})}v_{m}^{(t_{c})}I_{4}(n,m,k,h)+I_{4}(t_{c},t_{c} ,k,h)\right.$ (47) $\displaystyle\left.\qquad-2\sum_{n=1}^{K}v_{n}^{(t_{c})}I_{4}(n,t_{c},k,h) \right]\,,$ $$

and

$$ $\displaystyle b$ $\displaystyle=-\sum_{k,h=1}^{K}\hat{Q}_{kh}\left\{v_{k}^{(t_{c})}\left[\sum_{n =1}^{K}v_{n}^{(t_{c})}I_{3}(n,k,h)-I_{3}(t_{c},k,h)\right]\right.$ (48) $\displaystyle+\left.v_{h}^{(t_{c})}\left[\sum_{n=1}^{K}v_{n}^{(t_{c})}I_{3}(n, h,k)-I_{3}(t_{c},h,k)\right]\right\}$ $\displaystyle-\sum_{k=1}^{K}\sum_{t=1}^{T}\hat{M}_{kt}\left[v_{k}^{(t_{c})} \sum_{n=1}^{K}v_{n}^{(t_{c})}I_{3}(n,k,t)-v_{k}^{(t_{c})}I_{3}(t_{c},k,t)\right]$ $\displaystyle+$ $\displaystyle\sum_{k=1}^{K}\hat{v}_{k}^{(t_{c})}\left[-\sum_{n=1}^{K}v_{n}^{(t _{c})}I_{2}(k,n)+I_{2}(k,t_{c})\right]\,.$ $$

Performing the minimization over $\eta$ first, we obtain

$$ $\eta^{*}(\alpha,t_{c})=-\frac{b}{2a}\,.$ (49) $$

The minimisation over $t_{c}$ yields

$$ $t_{c}^{*}(\alpha)=\begin{cases}1\,&\text{ if }\hat{\mathbb{Q}}(\alpha)^{\top} \left[f_{\mathbb{Q}}\left(\mathbb{Q}(\alpha),(1,\eta^{*}(\alpha,1))\right)-f_{ \mathbb{Q}}\left(\mathbb{Q}(\alpha),(2,\eta^{*}(\alpha,2))\right)\right]>0\\ 2\,&\text{ otherwise. }\end{cases}$ (50) $$

and hence

$$ $\eta^{*}(\alpha)=\eta^{*}(\alpha,t_{c}^{*}(\alpha))\,.$ (51) $$

Interestingly, we observe that the learning rate schedule has a different functional form depending on the current task $t_{c}$. This can be seen in Fig. 4 where the learning rate switches between two different schedules depending on the current task $t_{c}$.

## Appendix B Readout layer convergence properties

In this appendix, we examine the asymptotic behaviour of the readout layer weights during the late stages of training. In particular, we are interested in the convergence rate as a function of the task similarity $\gamma$. As in the main text, we consider the case $K=T=2$. From the overlap trajectories in Fig. 8 for $\gamma>0.3$, we observe that the cosine similarity quickly approaches unity, i.e., $|M_{kt}|/\sqrt{Q_{kk}}\approx\delta_{kt}$, which corresponds to perfect feature recovery. Therefore, the decrease in performance for $\gamma>0.3$ seen in Fig. 3 must be attributed to the dynamics of the second layer. Indeed, in Fig. 8, we observe a slowdown in the readout dynamics as $\gamma\to 1$.

Assuming perfect convergence of the feature layer to $\bm{w}_{1}=\bm{w}_{*}^{(1)}$ and $\bm{w}_{2}=\bm{w}_{*}^{(2)}$, we consider the dynamics of the readout layer while training on task $t=1$. We expect the corresponding readout layer to converge to the specialised configuration $\bm{v}^{(1)}=(v^{(1)}_{1},v^{(1)}_{2})=(1,0)^{\top}$ and we would like to compute the convergence rate as a function of $\gamma$. The dynamics of the readout layer reads

$$ $\displaystyle\frac{{\rm d}v^{(1)}_{1}}{{\rm d}\alpha}$ $\displaystyle=\eta\left[\frac{1}{3}(1-v^{(1)}_{1})-\frac{2}{\pi}\operatorname{ arcsin}\left(\frac{\gamma}{2}\right)v^{(1)}_{2}\right]\,,$ (52) $\displaystyle\frac{{\rm d}v^{(1)}_{2}}{{\rm d}\alpha}$ $\displaystyle=\eta\left[\frac{2}{\pi}\operatorname{arcsin}\left(\frac{\gamma}{ 2}\right)(1-v^{(1)}_{1})-\frac{1}{3}v^{(1)}_{2}\right]\,,$ $$

which can be rewritten as

$$ $\displaystyle\frac{{\rm d}}{{\rm d}\alpha}\begin{pmatrix}1-v^{(1)}_{1}\\ v^{(1)}_{2}\end{pmatrix}=\eta\bm{A}\begin{pmatrix}1-v^{(1)}_{1}\\ v^{(1)}_{2}\end{pmatrix}\,,$ (53) $$

where

$$ $\bm{A}=\begin{bmatrix}-1/3&a\\ a&-1/3\end{bmatrix}\;,$ (54) $$

and $a=2\operatorname{arcsin}\left(\gamma/2\right)/\pi$.
Note that $a<1/3$ for $0<\gamma<1$, hence $\bm{A}$ is negative definite, implying convergence to $\bm{v}^{(1)}=(1,0)^{\top}$. The rate of convergence is determined by the smallest eigenvalue (in absolute value): $a-1/3$. The associated convergence timescale is therefore

$$ $\alpha_{\rm conv}=\frac{3\pi}{\eta(\pi-6\operatorname{arcsin}\left(\gamma/2 \right))}\,,$ (55) $$

as anticipated in Eq. 10 of the main text.

## Appendix C Supplementary figures

Fig. 8 describes the dynamics of the optimal replay strategy for different values of task similarity in the same setting as Fig. 3 of the main text. In particular, the upper panel displays the evolution of the magnitude of the readout weights $|v^{(t)}_{k}|$, while the lower panel shows the trajectory of the cosine similarity $|M_{kt}|/\sqrt{Q_{kk}}$.

Fig. 4 compares the values of the loss at the end of training, averaged on both tasks, for different task-selection strategies. In particular, it highlights the performance gap between the four replay strategies at constant learning rate considered in the main text (no-replay, interleaved, optimal and pseudo-optimal) and the strategy that simultaneously optimise over task-selection and learning rate.

Figure: Figure 8: Overlap dynamics with optimal replay. We plot the absolute value of the task-dependent readout weights $|v^{(t)}_{k}|$ (upper panel) and the cosine similarity $|M_{kt}|/\sqrt{Q_{kk}}$ as a function of the training time $\alpha$. Different columns refer to different choices of task similarity $\gamma=0.1$, $0.3$, $0.6$, $0.9$.
Refer to caption: x8.png

Figure: Figure 9: Adopting an optimal learning rate schedule leads to major perfomance improvement. Average loss on both tasks at the end of the second training phase as a function of task similarity $\gamma$ under the same setting and parameters as Fig. 3 of the main text. The top four lines correspond to different strategies at constant learning rate $\eta=1$: no replay (purple crosses), optimal replay (red dots), interleaved (blue squares), pseudo-optimal replay (cyan dashed line). The bottom curve (brown plus signs) corresponds to jointly optimal replay and learning rate schedules (see Fig. 4).
Refer to caption: x9.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
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?
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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.
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
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)
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
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.
soundness: 4
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.
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.
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.

Meta-review: 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.

Decision: Accept (Poster)