Optimal Protocols for Continual Learning via Statistical Physics and Control Theory
Authors: Francesco Mori
Stefano Sarao Mannelli
Francesca Mignacco
Venue: ICLR 2025 Poster (2025)
Keywords: machine learning theory
statistical physics
online learning
continual learning
optimal control theory
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: This theoretical work combines statistical physics and control theory to design optimal task-selection protocols mitigating catastrophic forgetting while preserving performance in continual learning, validated on synthetic and real-world data.

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.

Published as a conference paper at ICLR 2025
OPTIMAL PROTOCOLS FOR CONTINUAL LEARNING VIA STATISTICAL PHYSICS AND CONTROL THEORY
Francesco Mori', Stefano Sarao Mannelli Francesca Mignacco 4.5
ABSTRACT
Artificial neural networks often struggle with forgettnni when learn- ing 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 the- oretical foundation assessing their optimality. In this paper, 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- con-trivial yet interpretable strategies for mitigating catastrophic forgetting, shedding light on how optimal learning protocols modu- late established effects, such as the influence of task similarity on forgetting. Fi- nally, we validate our theoretical findings with experiments on real-world data.
INTRODUCTION
Mastering a diverse range of problems is crucial for both artificial and biological systems. In the context of training neural network on series of tasks-a.k.a. multi-task learning (Caruana 1993
1994b a 1997) tth ability to learn new tasks can be improved by leveraging knowledge from pre- vious ones (Suddarth Kergosien 1990p. However, this process can lead catastrophic forgetti, where learning new tasks degrades performance on older ones. This phenomenon has been observed in both theoretical neuroscience (McCloskey Cohen 1989 Ratcliff 1990) and machine learning (Srivastava et al. 20133 Goodfellow et al.] 20144). and occurs when the network parameters encod- ing older tasks are overwritten while training on a new task. Several mitigation strategies have been proposed (French 1999. Kemker et al. 20188). including semi-distributed representations (French)
1991 1992). regularisation methods (Kirkpatrick et al. 2017 Zenke et al.) 2017; Li & Hoiem, 2017 dynamical architectures (Zhou et al.) 2012 Rusu et al. 2016), and others (see e.g. Parisi reviews). A common strategy, known as replay, is for thorough De Lange et al. thooouugh
ct al.
present the network with examples from the old tasks while training on the new one minimise forgetting (Shin et al. 2017 Draelos et al. 2017 Rolnick et al. 2019).
On the theoretical side, [axter (2000) pioneered the research 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 al. 2020 Zhang & Yang [021] Zhang & Gao| 2022). However, these results focused on worst-case anal lsis, offering bounds that may not reflect the typical performance of algorithms. In contrast, Dhifal- lah & Lu (222) began investigating the typical-case scenario, providing precise characterisation of transfer learning in simple neural network models. [Graace et al. (2022): [Ingrosoo al.) (2024) extended this analysis to more complex architectures and generative models, allowing for better
description of the relation between tasks. Finally, Lee et al. 2021 2022 proposed a theoretical
Rudolf Peierls Centre for Theoretical Physics, University of Oxford, Oxford OX1 3PU, United Kingdom
" "Data Science and AL, Computer Science and Engineering, Chalmers University of Technology and University of Gothenburg
"School of Computer Science and Applied Mathematics, University of the Witwatersrand GGraduate Center, City University of New York, New York, NY 10016, USA "Joepph Henry I Laboratories of Physics, Princeton University, Princeton, NJ 08544, USA

---

Published as a conference paper at ICLR 2025
framework for the study of the dynamies of continual learning with focus catastrophic forget- ting. Their work provided a theoretical explanation for the surprising empirical results of Ramasesh al. (2020), which revealed non non-monotonic relation between forgetting and task similarity, where
maximal forgetring occurs at intermediate task similarity. Analogously, [Shan et al. (2024) 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 prototypical mitigation strategy. We use optimal control theory to determine the optimal training protocol that maximises performance across different tasks.
Our contribution. In this work, we combine dimensionnliit-rrduucioi techniques from statistical physics (Saad & Solla, 1995a b. Biehl & Schwarze, 1995) and Pontryagin's 's maximum principle from control theory (Feldbaum) 1955 Pontryagin 1957 Koppl 1962). This approach allows to derive optimal task-selection protocols for the training dynamics of a neural network in a continual learning setting. Pontryagin's 's principle works efficiently in low-dimensional deterministic systems. Hence, applying it to neural networks requires the statistical physics approach (Engell 2001). which reduces the evolution of high-dimensional stochastic systems erned by ordinary differential equations (ODEs) (Saad & Solla 1995a b Biehl & Schwarze 1995).
few key order parameters gov-
Specifically, we consider the teacher-student framework of Lee al. (2021) prototype continual learning setting amenable to analytic characterisation. Our main contributions are:
We leverage the ODEs for the learning curves of online SGD derive dlosed-form formu- lae for the optimal training protocols. In particular, we provide equations for the optimal task-selection protocol and the optimal learning rate schedule, as function of the task similarity 7 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 range of problem parameters and find highly structured protocols. Interestingly, we are nonetheless able to interpret these strategies 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 revision phase, where old tasks are replayed.
We clarify the impact of task similarity catastrophic forgetting. At variance with what previously observed (Ramasesh et al.) [220) Lee et al. [202]] 2022), 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 con- tinual learning task using the Fashion-MNIST dataset. Here, the pseudo-optimal strategy effectively interpolates between simple heuristics depending on the problem's parameters.
Further related works. Recent theoretical works on online dynamics in one-hidden-layer neural networks have addressed various learning problems, including over-parameterisatio (Goldt et al.) 2019), algorithmic analysis (Retinetti al. 2021; [Srinivasan et al., 2024). and learning strate-
gies Lee et al., 2021 Saglietti et al.] [2022] Sarao Mannelli et al. 2024). However, these studies
have not explored the problem from an optimal control perspective.
Early works addressed the optimality of hyperparameters in high- high-dimensional online Icarning for committee machines via control theory. These studies focused on optimising the learning rate (Saad & Rattray 1997, |Ratray& & Saad 1998, [Schlösser et al.) 1999), the regularisation (Saad & Rattray) 1998), and the learning rule (Rattray & Saad, 1997). However, to the best of our knowledge, the problem of optimal task selection has not been explored yet. Carrasco- Davis et al. (2023) and Li et al. (2024) applied optimal control to the dynamics of connectionist models of behaviour, but their analysis was limited to low-dimensional and minite-dmmensinall settings. Urbani (221 extended the Bellman equation to nigh-dimensional mean-field dynamical systems, though without considering learning processes.
2

---

Published as a conference paper at ICLR 2025
Teacher
Student gen
My VN
Task (I)
x- wil)
go
Loss
Teacher
Catastrophic forgetting of Task
X-
Task (2)
y() g+
a)
b)
Train on Task
Train on Task 2
(raining step
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 learn both tasks. However, sequential training results in catastrophic forgetting, where the performance on a previously learned task significantly deteriorates when a new task introduced (panel b). Parameters: K T
Several other works have combined ideas from machine learning and optimal control. Notably, Han et al. (2019) interpreted deep learning as an optimal control problem on dynamical system, where the control variables correspond to the network parameters, Chen & Hazan (2024) formulated contaooptiizztiinn 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 model-based approach investigate the supervised learning of multiple tasks. Follow- ing Lee et al. (2021 2022). we consider a teacher-student framework (Gardner & Derrida 1989) "student" neural network is trained on synthetic inputs E R.N. drawn i.i.d. from standard (N0 1). The labels for each task f = 1...., T are generated by single- layer "teacher" networks: (()) g. (x w. /VN), where W. (w. (1) ......._?") ( T denote the corresponding teacher vectors, and g- the activation function. The student is two-layer neural network with K hidden units, first-layer weights W
Gaussian distribution,
E R activation
function g, and second-layer weights v E R It outputs the prediction:
K
(z; W. v)
Σ
Uk 9
Wk
VN
Following standard multi-headed approach to continual learning Zenke et al.
2017 Farquhar
Gal 2018), we allow for task-dependent readout weights: V = (v())
")
7''x
Specifically, the readout for task is updated only when that task is presented. 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. [ Training is performed via Stochastic Gradient Descent (SGD) the squared loss of y/ and g j(æ; W, o()). We consider the online regime, where at each training step the algorithmic update computed using a new sample (ær, (9). The generalisation error of the student on task is given by
)
(W. V. W.) = 2 -1))))))))) =
w.
Ú(æ; W x(1)) w())))
F
g"
(2)
N
The angular brackets (-) denote the expectation over the input distribution for given set of teacher and student weights. Crucially, the error depends on the input data only through the preactivations
(8) w.
d
x. Wkk
1= ..... K
and
x
t 1,...,T.
(3)
VN
VN
3

---

Published as a conference paper at ICLR 2025
Eq. 3 defines jointly Gaussian variables with zero mean and second moments given by
MR == Een (Acd
WE we
N
= Wk Whh
Qua = Ex [
(4)
NN
(8)
Sw = Ex N!
w
w
N
called overlaps in the statistical physics literature. Therefore, the dynamics of the generalisation er- ioris entirely captured by the evolution of the student readouts V and the overlaps. As shown in Leee et al. [2021: 2022), we can track the evolution of the generalisation error in the high-dimensional. We leverage this 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 derive equations governing the dynamics of the overlaps and readouts under given task-selection protocol. These equations fully determine the evolution of the generalization error. For the remainder of the paper, we consider K = T to guarantee that the student network has enough capacity learn all tasks perfectly. Teacher vectors and the task similarity is tuned by a parameter so that S = ót,-' + 7(1 - ōt,r). For simplicity, useful to encode all the relevant degrees of freedom-namee weights- in the same vector. We use the shorthand notation Q (vec(Q), vec(M), vec(V)) e R 22*+2KT As further discussed in Appendix A in the limit of large input dimension N with K, T ~ Ox(,), the training dynamics is described by a set of ODEs
normalised,
the overlaps and the readout
dQ(a) = So (Q(a), u(a)) da
with a E (0,
(5)
The parameter a denotes the effective training time-the ratio between training steps and input dimension N. The vector u encodes the dynamical variables that we want to control optimally. In particular, we study the optimal schedules for task-selection (a( and learning rate 1)(a). Here, ta(a) E (1..... ( indicates on which task the student is trained at time a. The specific form of the functions fa is derived in Appendix A The initial condition 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 stress that Eq. 5 is low-dimensional deterministic equation that fully captures the high-dimensional stochastic dynamics of SGD) as N - xo. This dimensionality reduction is crucial apply the optimal control techniques presented 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 linear combination of the generalisation errors on the different tasks at the final training time o
T
T
(((aF)) = - (Q((F))
with C 2 O and Σ ct 1==1,
(6)
where the coefficients c, identify the relative importance of different tasks and e, denotes the infinite- dimensional limit of the average generalisation error on task I, as defined in Eq. 2 Crucially, we
have analytic expression for Ez, derived in Appendix [ In the remainder of the paper, assume
equally important tasks a 1/T'. As customary in optimal control theory (Pontryagin 1957). we adopt a variational approach to solve the problem. We define the cost functional
FIQ. Q uj h (Q(ar)) + [ da Q(a)) [-
dQ(a) da
+ fo (Q(a).(())))
where the conjugate order parameters namics in the training interval (0,) Finding the optimal protocol amounts to minimising the
= (vec Q), vM) veo(()) enforce the training dy-
4

---

Published as a conference paper at ICLR 2025
protocol
No replay
Interleaved
Optimal
Task
Task
Task 2
10"1
10
1311
10-1
10
12
Loss
Task
10-1
13"4
Average
time step input dimension
Figure 2: How to learn new task without forgetting the old one? The student trained on task until convergence during the first phase (a E [0, 1000)), then task 2 is introduced. During the second phase (a (1000, 1025)). ) task may be replayed to prevent forgetting. For better visibility, we only display the regions E [o, 20] U [990, 1025]. We compare three strategies: a) no replay, b) interleaved replay, i.e., alternating between the two tasks, e) the optimal strategy derived in Sec. 2 Crosses mark numerical simulations of single trajectory at N = 20000, lines mark the solution of Eq. 5
Colour bars represent the protocol le- Parameters: 7 = 0.3,
2, and 1.
cost functional F with respect to Q, Q. and u. We defer the details of this variational procedure to Appendix[A] [ and present only the main steps here. For general introduction to the control methods adopted here, see, e.g., Kirk (2004). The minimisation with respect to Q provides a set of equations for the backward dynamics the conjugate parameters
dQ(a) ((a) Volo(Q(a) (a))
with (0,
da
The final condition for the dynamics is given by
T
((aF) = vah(Qr)= = Σ Vcsx(((()))
(9)
t=1
The optimal control curve u" (a) is obtained as the solution the minimisation:
u (a)
argmin wEll
{Q(()) fa(Q(a).
(10)
where L/ is the set of allowed controls, that can be either continuous or discrete. For instance, for task the notation la((a) to indicate
selection take u((a) (ax) and 21,2 T), where the current task, i.c. the task on which the student is trained at time a. When optimising over both task selection and learning rate schedule we take u (te. 17) and U/ 21,2... Th} x R+ Crucially, the optimal control equations 5) g and TO must be iterated until convergence, starting from an initial guess on u, which can be for instance taken random. Let us stress that the space U of possible controls is high-dimensional and hence it is not feasible to explore it via greedy search strategies.
RESULTS AND APPLICATIONS
The theoretical framework of Sec. 2 extremely flexible and can be applied to variety of settings. We will present a technical paper focused on the method in specialised venue. In this section, we focus on specific settings 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 the optimal strategies in synthetic settings to applications real datasets.
3.1 EXPERIMENTS ON SYNTHETIC DATA
We formulate the problem of continual learning as follows. During first training phase, the student learss perfectly task i 1. Then, the goal is learn new task 2 without forgetting the
5

---

Published as a conference paper at ICLR 2025
Loss protocel
0.1
Task
Average loss
Taak Task Awerage
Loss 11....33 g
Task protocol
replay
Optimel rep ay
Interleaved
Pseudo Optimal
Task similarity
a)
2013
02
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 function of the task similarity y under the replay setting from Fig. 2 Different lines correspond to different strategies: no replay (purple crosses), optimal replay (red dots), interleaved (blue squares), aseudo-optimal replay (cyan dashed line). b-e) Optimal replay strategies for different values of y = 0.1, 0.3, 0.6, 0.9. Colour bars represent the protocol ((()).
previous one. A given time window of duration oF is assigned for the second training phase. In particular, we investigate the role of replay-i.ce -i... presenting samples from task 1 during the second training phase- andsee-adiee the structure the optimal replay strategy.
We use the equations derived in Sec. to study optimal replay during the second phase of training. To this end, we take the task-selection variable as our control u(a) = te (ax) E 21, while we set te = during the first training phase. The result of the optimisation in Eq. T0 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 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 ay. 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) function of the training time a. Numerical simulations marked by cross symbols-are are in excellent agreement with our theory. Deviations are smaller than 1/VN, compatible with finite-size effects.
The student is trained exclusively on task until a = 1000, when the task is perfectly learned with loss ~ 10-6 Then, the student trained on a combination of new and old tasks for a training time of duration or = 25. A colour bar above each plot illustrates the associated task-selection strategy ta).. 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 heuristic "interleaved" strategy, where training alternates one sample from the new task to one sample from the old one. As observed in Lee al. (2022), the interleaved strategy already provides performance gain, demonstrating the relevance of replay to mitigate catastrophic forgetting. Panel e) of Fig. Zishows the loss dynamics for the optimal replay strategy. Notably, the optimal strategy exhibits complex structure and displays 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 range of task similarities.
The impact of task similarity. To understand the structure of the optimal strategy, we examine its performance in relation task similarity 7. Fig. 3a) depicts the average loss at the end of training as function of y. For the no-replay strategy, we reproduce the findings from previous works (Lee et al.) 2021 y022): the highest error occurs at intermediate task similarity. [ee et al. (2022) explained this trade-off between node re-use and node activation. Specifically, for small y, there is minimal interference between tasks. One hidden neuron predominantly aligns with the new task, while the other neuron retains the knowledge of the old task, leading task specialisation. At large 7, features from task 1 are reused for task 2, avoiding forgetting. However, at intermediate y, interference is maximal: both neurons quickly align with task 2, and task 1 is forgotten. Remarkably,
non-monotonicity

---

Published as a conference paper at ICLR 2025
Fig. Bl) shows that replay reverses this trend, with the minimal enror occurring at intermediate yy. 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 and displays a highly structured protocol. We can interpret this structure a posteriori: an initial focus phase without replay 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. Ga). In this case, 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 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 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 y under the optimal protocol. First, as shown in Fig. g of Appendix [ the optimal protocol achieves a good level of node specialisation across all values of yy Thus, replay prevents the task interference that typically causes performance deterioration at intermediate y. The non- monotonic behaviour of the optimal curve in Fig. 3aa arises from different origin, involving two opposing effects related to the first-layer weights and the readout. The initial decrease of the loss with y is quite intuitive, as only minimal knowledge can be transferred from task 1 to task 2 when Y is small. Consequently, the focus phase on task must be longer for smaller 7, leaving less time to revise task 1, thereby reducing performance. On the other hand, the performance decrease ob- served in Fig. 3 for y > 0.3 is more subtle and is related to the readout layer. Once the two hidden neurons have specialised specialised-each aligning with
of the teacher vectors expect the readout weights corresponding to the incorrect teacher to be suppressed. Specifically, if w- = ! and w2 w. (2) the learning dynamics should drive the readout weights (1)) = (u) : u2 (
and (2 (m) 01 ( towards (1)) (1, TT and 22)) (0, T) to achieve full recovery of
the teacher networks. As shown in Fig. g of Appendix C the time required to suppress the off-
diagonal weights (1) and of²) increases as y + 1. This is intuitive, as higher task similarity ? reduces the distinction between tasks, slowing the suppression of the off-diagonal weights. In Ap-
pendix we derive analytically the convergence timescale ('ronv of the readout layer, showing that
Task protocol
3π
aconv 7( 6 arcsin (1/2))
(11)
where 1) is the learning rate. This timescale is a monotonically increasing function of y and di- verges as - with accoav ~ //3m//22((1 7))- This result explains the performance decrease of the optimal strategy as ? - 1. In summary, the performance decrease for 7 - is due to the first-layer weights, while for y - 1 it is related to the readout weights.
20.0
Task 1
Learning rate
Task 2
125
10.0
7.5
5.0
2.5
0.0
Optimal learning rate schedules. While the above results are obtained with a constant learn- ing rate, annealing schedules are widely used in machine learning. Thus, it is relevant to study the optimal interplay between replay and learn- ing rate. Optimal learning rate dynamics have been studied with a similar approach in (aad & Rattray (1997). We jointly optimise the task- selection protocol together with the learning rate to investigate its impact on continual learning. Fig. [ shows the optimal learning rate schedule for task similarity 7 0.3 in the second training
1000.0
1012.5
1025
time step / input dimension
Figure 4: Joint optimisation of task-selection and learning-rate. Optimal learning rate as function of training time a for the same parame- ters as Fig. There is a single optimal learning rate curve, but for visibility purposes we show it as solid orange line when training on task 1 and a dashed green line on task 2. The task-selection protocol (a) is illustrated in the colour bar.
7

---

Published as a conference paper at ICLR 2025
Tank protocol
7 0.1
0.3
Tosk Task
Average loss
Task protocol
0=6.6
= 0.9
2x30
Optimal
Sequential
Loss
02
Task similarity
a)
time stop input dimarsion
Figure 5: The impact of task similarity multi-task learning. a) Average loss function of task similarity at the end of training (a.r = 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 = 0.1, 0.6, 0.9. Parameters: K = and 7= 10.
phase of duration op = 25. Similarly as for constant learning rate (panel e) of Fig. 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 enter- ing the revision phase, the optimal learning rate schedule exhibits a highly nontrivial structure (see Fig. 4. Indeed, although the optimal learning rate curve is unique, we find that effectively it can be seen two different curves, associated the respective tasks. In practice, the optimal learning rate curve "jumps" between these two curves according to the task selected at given training time. performance. Fig. of Appendix showos that the joint optimization of learning rate and task selec- tion outperforms all other protocols, including exponential and power-law learning rate schedules combined with interleaved replay.
Multi-task learning from scratch. We also consider a multi-task setting (see Fig. 5 where both tasks must be learned from scratch within fixed number of steps, corresponding time or. We first consider sequential learning, i.e., training only on task for a < ap/2. /2, then only task 2, or vice versa. As shown in Fig. (5a) sequential learning leads catastrophic forget- ting, with the worst performance observed at intermediate task similarity. In contrast, a randomly interleaved strategy, where examples from both tasks are presented in equal proportion but in ran- dom order, shows significant improvement. This approach can exploit task similarity, leading to monotonic decrease average loss as y increases. The optimal strategy, displayed in Fig. (5D-e) for various values of y, 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 strategy. This observation aligns with our pseudo-optimal strategy, which suggests employing interleaved replay once performance on both tasks becomes comparable.
total training
3.2 EXPERIMENTS ON REAL DATA
We consider the experimental framework established in Ramasesh al.) (2020); Lee al. (2022) for the study of task similarity in relation to catastrophic forgetting. We use the FaShion-MNIST dataset (iiao al. 2017) generate upstream and downstream tasks. The upstream dataset-
D, (a) y/ (1) consists in a pair of classes from the standard dataset. The downstream dataset
is generated by linear interpolation of the upstream dataset with second auxiliary dataset-
D [8, }. contaaiinnnnnngg a new pair of classes.
D2 =x .ii
(2)
{r20 (1) (1 7/ ) ry, (1) + +1 7)0:).
(12)
8

---

Published as a conference paper at ICLR 2025
where the parameter " 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 multi-heat architecture (Zhou al. [012 Rusu et al. 2016) 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, strategy with interleaved replay, and a "pseudo-optimal"" strategy. Recall that the latter is inspired by the optimal protocol derived 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 phase of interleaved replay. Crucially, this protocol can be easily implemented in practice, only requires an estimate of the generalisation error on the two tasks, which can be obtained in real-world settings.
No Reploy
Interlaaeed
Psoudo-Optimal
Average loss
10
Replay
Interfeaved
Poeuto-Optimal
Lo3
10-1
1.3k
20k
10000
20000
10000
20000
10000
20000
10000
20000
Step
Step
Step
Step
Figure 6: Training dynamics. Training curves on the modified fashion MNIST task at similarity 0.5. The network is trained for 10.000 steps on the first task before switching to the second task and being trained for additional 10.000 steps. 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.
Fig. 6hoows the training loss under the different training protocols for y = 0.5. While the no-replay strategy appears to be successful for small downstream datasets (i.e., a few training steps in the online framework) in the longer run it leads to strong forgetting and high average loss. This behavior is intuitive: for small datasets, the initial loss the new task is high, leading to substantial decrease in loss early on, which temporarily outweighs the decline in performance on the previous task. 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 better performance over the entire trajectory.
500 samples
1000 samples
2000 samples
4000 samples
8000 samples
Raplay
10
Inter rerieeeed
Pocudo Opemmal
Averopa loss
10-9
0.25 0.50 0.75
0.25 0.50 0.75
0.25 0.50 0.75
0.25 0.50 0.75
0.25 0.50 0.75
Figure 7: Average loss comparison. The figure focuses on the average loss and shows the final loss achieved by the three strategies we increase the size of task 2 (from left to right: 500, 1.000, 2.000, 4.000, and 8,000 samples) while task has always 10.000 samples. Individual panels show the performance of the three strategies as we span the value of y form 0.05 to 0.95.
This result is not limited to the specific value of 7. Fig. 7 shows the average final loss as a function of 7. From left to right, different panels correspond to an increasing number of available samples for the downstream task, comparable to o in the theoretical model. For small downstream tasks, the no-replay strategy is optimal, as shown in the second and third panels. As the size of the down- stream task increases, the interleaved strategy approaches optimal performance while the no-replay becomes suboptimal. The pseudo-optimal strategy combines the advantages of both approaches, in-
9

---

Published as a conference paper at ICLR 2025
terpolating between no-replay and interleaved strategies to automatically adapt the computational budget. This results in the best overall performance across regimes. We confirm this observation with additional experiments tween the synthetic and real settings-such as data structure- -te pseudo-optimal strategy remains effective on real-world data, demonstrating its robustness and broad applicability.
CIFARIO in Appendix C.2 Notably, despite the differences be-
4 DISCUSSION
Conclusion. In this work, we introduce systematic approach for identifying and interpreting opti- mal task-selection strategies in synthetic learning settings. We consider teacher-student scenario as 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 control-theod -heoryy framework that allows 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 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 takes first step toward understanding the theory behind optimal training protocols for neural networks. In the following, we discuss current limitations and outline promising directions for future research. First, Pontryagin's 's maximum principle provides necessary condition for optimality but does not guarantee a global optimum. Nevertheless, the strategies derived from this approach in the settings under consideration perform significantly better than previously proposed heuristics. Additionally, Pontryagin's 's principle does not easily extend to stochastic problems. This limitation overcome in the high-dimensional limit where concentration results provide deterministic dynamical equations. For simplicity, we focus i.i.d. Gaussian inputs, but our analysis can be extended to more structured data models (Goldt et al. 2020 Loureiro et al.) 2021 Adomaityte et al. (023) to study how input distribution affects task se 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 2020; Saglietti al. 2022 Cornacchia & Mossel] 2023 Abbe et al., 2023). Furthermone, it would be interesting to go beyond the study of online dynamics to understand the impact of memorisation in batch learning settings (Sagawa al. 2020
Existing results in the spurious correlations (Ye et al. 202T) and fairness (Ganesh et al. 2023
literature suggest a strong dependence of the classifie's bias on the presentation order in batch learning. Our method can be applied to mean-field models llke (Mannelli et al.) [024] Jain et al.)
2024 to theoretically investigate this phenomenon. An interesting extension of our work involves
applying revenlly statistical physics methods to the study of deeper networks and more
complex learning architectures (Bordelon & Pehlevan [002b]]a al Rende et al. [0024] Tiberi et al.
2024). Another interesting direction concerns finding optimal protocols for shaping, where task
order significantly impacts both animal learning and neural networks (Skinner 1938 Tong et al,
2023
Lee al. 2024).
REPRODUCIBILITY STATEMENT
The code to reproduce the numerical experiments is available in the supplementary material.
ACKNOWLEDGMENT
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).
REFERENCES
Emmanuel Abbe, Elisabetta Cornacchia, and Aryo Lotfi. Provable advantage of curriculum learning on parity targets with mixed inputs. Advances in Neural Information Processing Systems, 36:
10

---

Published as a conference paper at ICLR 2025
24291-24321, 2023.
Urte Adomaityte, Gabriele Sicuro, and Pierpaolo Vivo. Classification of superstatistical features in high dimensions. In 2023 Conference on Neural Information Procecessing Systems, 2023.
Jonathan Baxter. A model of inductive bias learning. Journal of artificial intelligence research, 12: 149-198, 2000.
John Bechhoefer. Control theory for physicists. Cambridge University Press, 2021.
Michael Bichl and Holm Schwarze. Learning by on-line gradient descent. Journal of Physics A: Mathematical and general, 28(3):643, 1995.
Blake Bordelon and Cengiz Pehlevan. The influence of learning rule on representation dynamics in wide neural networks. In The Eleventh International Conference Learning Representations, 2022a.
Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution in wide neural networks. Advances in Neural Information Processing Systems, 35:32240-32256, 2022b.
Rodrigo Carrasco-Davis, Javier Masís, and Andrew M Saxe. Meta-learning strategies through value maximization in neural networks. arXiv preprini arXiv:2310. 19919, 2023.
Caruana. Multitask learning: A knowledge- source of inductive bias. In Proceedings of the Tenth International Conference on Machine Learning, pp. 41-48. Citeseer, 1993.
Rich Caruana. Learning many related tasks at the same time with backpropagation. Advances in neural information processing systems, 7, 1994a.
Rich Caruana. Multitask learning. Machine learning, 28:41-75, 1997.
Richard A Caruana. Multitask connectionist learning. In Proceedings of the 1993 connectionist models summer school, pp. 372-379. Psychology Press, 1994b.
Xinyi Chen and Elad Hazan. Online control for meta-optimization. Advances in Neural Information Processing Systems, 36, 2024.
Lenaic Chizat and Francis Bach. On the global convergence of gradient descent for over- parameterized models using optimal transport. Advances in neural information processing sys- lems, 31, 2018
Elisabetta Cornacchia and Elchanan Mossel. A mathematical model for curriculum learning for parities. International Conference on Machine Learning, pp. 6402-6423, PMLR, 2023.
Matthias De Lange, Rahaf Aljundi, Mare Masana, Sarah Parisot, Xu Jia, Aleš Leonardis, Gregory Slabaugh, and Tinne Tuytelaars. A continual learning survey: Defying forgetting classification tasks. IEEE transactions on pattern analysis and machine intelligence, 44(7):3366-33) 2021.
Oussama Dhifallah and Yue M Lu. Phase transitions in transfer learning for high-dimensional per- ceptrons. Entropy, 23(4):400, 2021.
Timothy J Draclos, Nadine E Miner, Christopher C Lamb, Jonathan A Cox, Craig M Vineyard, Kristofor Carlson, William M Severa, Conrad D James, and James B Aimone. Neurogenesis deep learning: Extending deep networks to accommodate new classes. In 20/7 international joint conference on neural networks (IJCNN), pp. 526-5 533-533. IEEE, 2017.
Andreas Engel. Statistical mechanics of learning. Cambridge University Press, 2001.
Sebastian Farquhar and Yarin Gal. Towards robust evaluations of continual learning. arXiv preprint arXiv: 1805.09733 2018.
AA Feldbaum. On the synthesis of optimal systems with the aid phase space. Avtomatika i Telemehanika, 16(2):129-149) 1955.
11

---

Published as a conference paper at ICLR 2025
Robert M French. Using semi-distributed representations to overcome catastrophic forgetting in connectionist networks. In Proceedings of the 13th anmual cognitive science society conference, volume 1. pp. 173-178, 1991.
Robert French. Semi-distributed representations and catastrophic forgetting connectionist net- works. Connection Science, 4(3-4):3553777, 1992.
Robert M French. Catastrophic forgetting in connectionist networks. Trends in cognitive sciences, 3(4): 135, 1999.
Prakhar Ganesh, Hongyan Chang, Martin Strobel, and Reza Shokri. On the impact of machine learning randomness on group fairness. In Proceedings of the 2023 ACM Conference on Fairness, Accountability, and Transparency, pp. 1789-1800, 2023.
Elizabeth Gardner and Bernard Derrida. Three unfinished works on the optimal storage capacity of networks. Journal Physics A: Mathematical and General, 22(12): 1989.
Federica Gerace, Luca Saglietti, Stefano Sarao Mannelli, Andrew Saxe, and Lenka Zdeborová. Probing transfer learning with model of synthetic correlated datasets. Machine Learning: Sci- ence and Technology, 3(1)001, 2022.
Sebastian Goldt, Madhu Advani, Andrew M Saxe, Florent Krzakala, and Lenka Zdeborová. Dy- namics of stochastic gradient descent for two-layer neural networks in the teacher-student setup. Advances in neural information processing systems, 32, 2019.
Sebastian Goldt, Marc Mézard, Florent Krzakala, and Lenka Zdeborová. Modeling the influence of data structure on learning in neural networks: The hidden manifold model. Physical Review X, 10(4):041044, 2020.
lan Goodfellow, Mehdi Mirza, Xia Da, Aaron C. Courville, and Yoshua Bengio. An empiri- cal investigation of catastrophic forgeting in gradient-basee based neural networks. In Yoshua Bengio and Yann LeCun (eds.), 2nd International Conference on Learning Representations, ICLR 2014, Banff. AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014.
Jiequn Han, Qianxiao Li, al. mean-fiele optimal control formulation of deep learning. Research in the Mathematical Sciences, 6(1): 2019.
Alessandro Ingrosso, Rosalba Pacelli, Pietro Rotondo, and Federica Gerace. Statistical mechan- ics of transfer learning in fully-connected networks in the proportional limit. arXiv preprint MXXi:22077.07688 2024.
Anchit Jain, Rozhin Nobahari, Aristide Baratin, and Stefano Sarao Mannelli. Bias motion: Theo- retical insights into the dynamics of bias in sgd training. arXiv preprint arXiv:2405.18299 2024.
Ronald Kemker, Mare McClure, Angelina Abitino, Tyler Hayes, and Christopher Kanan. Measuring catastrophic forgetting in neural networks. In Proceedings of the AAAI conference artificial intelligence, volume 32, 2018
Donald E Kirk. Optimal control theory: an introduction. Courier Corporation, 2004.
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinskka et al. Overcom- ing catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114(13):3521-35) 2017.
Richard E Kopp. Pontryagin maximum principle. In Mathematics in Science and Engineering, volume 5, pp. 255-279. Elsevier, 1962.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Toronto, ON, Canada, 2009.
Jin Hwa Lee, Stefano Sarao Mannelli, and Andrew Saxe. Why do animals need shaping? theory of task composition and curriculum learning. arXiv preprint arXiv:2402.1 140.62402./8311, 2024.
12

---

Published as a conference paper at ICLR 2025
Sebastian Lee, Sebastian Goldt, and Andrew Saxe. Continual learning in the teacher-student setup: Impact of task similarity. In International Conference on Machine Learning, pp. 6109-6119, PMLR. 2021.
Scbastian Lee, Stefano Sarao Mannelli, Claudia Clopath, Sebastian Goldt, and Andrew Saxe. Maslow's hammer in catastrophic forgetting: Node re-use vs. node activation. In International Conference on Machine Learning, pp. 12455-12477. PMLR, 2022.
Yujun Li, Rodrigo Carrasco-Davis, Younes Strittmatter, Stefano Sarao Mannelli, and Sebastian Mus- slick. A meta-learning framework for rationalizing cognitive fatigue in neural systems. In Pro- ceedings of the Annual Meeting of the Cognitive Science Society, volume 46, 2024.
Zhizhong Li and Derek Hoiem. Learning without forgetting. IEEE transactions on pattern analysis and machine intelligence, 40(12):2935-294) 2017.
Bruno Loureiro, Cedric Gerbelot, Hugo Cui, Sebastian Goldt, Florent Krzakala, Marc Mezard, and Lenka Zdeborová, Learning curves of generic features maps for realistic datasets with a teacher- student model. Advances in Neural Information Processing Systems, 34:18137-18151, 2021.
Stefano Sarao Mannelli, Federica Gerace, Negar Rostamzadeh, and Luca Saglictti. Bias-inducing
geometries: exactly solvable data model with fairness implications. In ICMI. 2024 Workshop on
Geooeery-grounded Representation Learning and Generative Modeling, 2024. URL https:
///opezeview.co net. wwt.:/ett.oouao wostt//orummmioo wttp///romm../oooll..ooo/oo zzpMpy
Michael McCloskey and Neal J Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of learning and motivation, volume 24, pp. 109-165. Elsevier, 1989.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 115(33):E7665--777), 2018.
Genman Parisi, Ronald Kemker, Jose Part, Christopher Kanan, and Stefan Wermter. Continual lifelong learning with neural networks: A review. Neural networks, 113:54-71, 2019.
LS Pontryagin. Some mathematical problems arising in connection with the theory of optimal au- tomatic control systems. Proc. Conf. on Basic Problems in Automatic Control and Regulation, 1957.
Vinay V Ramasesh, Ethan Dyer, and Maithra Raghu, Anatomy of catastrophic forgetting: Hidden representations and task semantics. arXiv preprint arXiv:2007.07400) 2020.
Roger Ratcliff. Connectionist models of recognition memory: constraints imposed by learning and forgetting functions. Psychological review, 97(2):285, 1990.
Magnus Rattray and David Saad, Globally optimal on-line learning rules for multi-layer neural networks. Journal of Physics A: Mathematical and General, 30(22):1.771, 1997.
Magnus Rattray and David Saad. Analysis of on-line training with optimal learning rates. Physical Review E, 58(5):6379, 1998.
Maria Refinetti, Stéphane A'scolii Ruben Ohana, and Sebastian Goldt. Align, then memorise: the dynamics of learning with feedback alignment. In International Conference on Machine Learning, pp. 8925-8935. PMLR, 2021.
Riccardo Rende, Federica Gerace, Alessandro Laio, and Sebastian Goldt. Mapping of attention mechanisms generalized potts model. Physical Review Research, 6(2):023057, 2024.
David Rolnick, Arun Ahuja, Jonathan Schwarz, Timothy Lillicrap, and Gregory Wayne. Experience replay for continual learning. Advances in neural information processing systems, 32, 2019.
Grant Rotskoff and Eric Kanden-Eijnder Trainability and accuracy of artificial neural networks: An interacting particle system approach. Communications on Pure and Applied Mathematics, 75 (9):1889-1933 2022.
13

---

Published as a conference paper at ICLR 2025
Andrei A. Rusu, Neil C. Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick, Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell. Progressive neural networks. CoRR, abs/1606.04B5 04671, 2016.
David Saad and Magnus Rattray. Globally optimal parameters for on-line learning in multilayer neural networks. Physical review letters, 79(13):5578, 1997.
David Saad and Magnus Rattray. Learning with regularizers in multilayer neural networks. Physical Review E, 57(2):2170, 1998.
David Saad and Sara Solla. Dynamics of on-line gradient descent learning for multilayer neural networks. Advances in neural information processing systems, 8, 1995a.
David Saad and Sara Solla. On-line learning in soft committee machines. Physical Review E, 52 (4):4225, 1995b.
Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameterization exacerbates spurious correlations. International Conference on Machine Learning, pp.
PMLR, 2020.
Luca Saglietti, Stefano Mannelli, and Andrew Saxe. An analytical theory of curriculum learning in teacher-student networks. Advances in Neural Information Processing Systems, 35:21113-21127, 2022.
Stefano Sarao Mannelli, Yaraslau Ivashynka, Andrew M Saxe, Luca Saglietti, et al. Tilting the odds at the lottery: the interplay of overparameterisation and curricula neural networks. In Proceedings of the 41st International Conference Machine Learning, 2024.
E Schlösser, D Saad, and M Biehl. Optimization of on-line principal component analysis. Journal of Physics A: Mathematical and General, 32(22))4061, 1999.
Haozhe Shan, Qianyi Li, and Haim Sompolinsky. Order parameters and phase transitions of contin- ual learning in deep neural networks. arXiv preprint 10Xiv:2407.10315. 2024.
Hanul Shin, Jung Kwon Lee, Jachong Kim, and Jiwon Kim. Continual learning with deep generative replay. Advances in neural information processing systems, 30, 2017.
Burrhus Frederic Skinner. The behavior of organisms: An experimental analysis. BF Skinner Foundation, 1938.
Ravi Francesco Srinivasan, Francesca Mignacco, Martino Sorbaro, Maria Refinetti, Avi Cooper, Gabriel Kreiman, and Giorgia Dellaferrera. Forward learning with top-down feedback: Empirical and analytical characterization. In The Twelfth International Conference Learning Represen- tations, 2024. URL htLlhtttss:::pperreetee wttp:e/wwwweeeeeteooueeeeeeeeee.
net/ forum? L
Rupesh K Srivastava, Jonathan Masci, Sohrob Kazerounian, Faustino Gomez, and Jürgen Schmid- huber. Compete compute. Advances in neural information processing systems, 26, 2013.
Steven C Suddarth and YL Kergosien. Rule-injection hints means of improving network perfor- mance and learning time. In European association for signal processing workshop, pp. 120-129. Springer, 1990.
Lorenzo Tiberi, Francesca Mignacco, Kazuki Irie, and Haim Sompolinsky. Dissecting the in- terplay attention paths a statistical mechanics theory of transformers. arXiv preprint arXiv:2405. 159266 2024.
William L Tong, Anisha Iyer, Venkatesh N Murthy, and Gautam Reddy. Adaptive algorithms for shaping behavior. bioRxiv, 2023.
Pierfrancesco Urbani. Disordered high- high-dimensional optimal control. Journal of Physics A: Mathe- matical and Theoretical, 54(32):32401, 2021.
Rodrigo Veiga, Ludovic Stephan, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Phase diagram stochastic gradient descent in digennimensional two-layer neural networks. Advances 2022.
in Neural Information Processing Systems,
14

---

Published as a conference paper at ICLR 2025
Yaqing Wang, Quanming Yao, James T Kwok, and Lionel Mi.. Generalizing from few examples: A survey on few-shot learning. ACM computing surveys (csur), 53(3):1-34, 2020.
Daphna Weinshall and Dan Amir. Theory of curriculum learning, with convex loss functions. Jour- nal of Machine Learning Research, 21(222):1-19, 2020.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist. a novel image dataset for benchmark- ing machine learning algorithms. arXiv preprint arXiv:1708.0771 2017.
Han-Jia Ye, De-Chuan Zhan, and Wei-Lun Chao. Procrustean training for imbalanced deep learning. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 92-102, 2021.
Friedemann Zenke, Ben Poole, and Surya Ganguli. Continual learning through synaptic intelligence. In International conference on muchine learning, pp. 3987-3995, PMLR, 2017.
Lei Zhang and Xinbo Gao. Transfer adaptation learning: A decade survey. IEEE Transactions Neural Networks and Learning Systems, 2022.
Yu Zhang and Qiang Yang. A survey on multi-task learning. IEEE transactions on knowledge and data engineering, 34(12):5586-566)) 2021.
Guanyu Zhou, Kihyuk Sohn, and Honglak Lee. Online incremental feature learning with denoising autoencoders. In Artificial intelligence and statistics, pp. 1453-1461. PMLR, 2012.
A DETAILS ON THE THEORETICAL DERIVATIONS
In this appendix, we provide detailed derivations of the equations in Sec. 2.c.2e 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., (021)) We remind that inputs are -dimensional vectors æ E RN with independent identically distributed (i.i.d)) standard Gaussian entries I; A 1), while the labels are generated by single-layer teacher
networks: (()) = go (x w. /VN), = .,..., T, with a different teacher for each task. The student a one-hidden layer network that outputs the prediction:
K
(e)
We
(13)
In the main text, focus on the case K T. where the student has, in principle, the capacity to perfectly solve the problem and represent all teachers. Specifically, there exists a configuration of the student's parameters that achieves perfect recovery. This configuration corresponds to aligning each of the student's hidden neurons with a specific teacher/task. Explicitly, this configuration is
(k) and Uk (t) 6k,1, where ō denotes the Kronecker delta. However, our theory remains valid for arbitrary and T.
given by wA w
We focus the online (on one one-pass) setting, so that at each training step the student network is presented with a fresh example x" / = 1,.... P. and P/N ~ (1)) The weights of the student are updated through gradient descent on 1 (g(t)) - (0)) following the task-selection protocol t.:
w p-1
1/" A (oo) uk (LdH of (N
x
VN
k 2k (())
Mp-1
11th
A
(0P g(X
N
(14)
K
A (t)) = g((l - ())
(( - go (x (2(t)))
where 17th 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,
20t uD (1)
No
1000
(15)
VN
VN
15

---

Published as a conference paper at ICLR 2025
Notice that, due to the online-learning setup, each training step the input 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:
MR == Ex (xad1']
ut
QkA Ee [A-An]
WR Whh
(16)
N
(e') w
Sw = En (
N
A.1 GENERALISATION ERROR AS A FUNCTION OF THE ORDER PARAMETERS
We write the generalisation error (Eq.
2 of the main text) as average the local fields:
Et (W, V. W.)
(0) "h
E.A. [0(A6)9(A)) +
2 Ea.d, [
k,h
(17)
U
E.A.
[o(Ao)
where the expectation computed over the multivariate Gaussian distribution
(-)
P.A.
exp
((2T (2Tm)++T[C
(18)
Q
M
C=
M
S
From now on, we adopt the unified notation
(2( Ex.x. (a=(Aa)g)(Nn)]
(19)
where p,p refer both to the indices of the student weights k, h the tasks 6, U. We then rewrite the generalisation error as
Eσ (, V. W.)
[2 h) [2(111
Σ (00 12k,)).
(20)
In all the results presented in Sec. we consider g = go (z) = erf (z/v2). In this case, there is an analytic expression for the integral l2 (Saad & Solla 1995a):
q3p
(3,
arcsin
(21)
V1
+ 089 / + 4pp
and we use the symbolg 9 to denote generically an overlap from Eq. 16 according to the choice of indices 3, p, e.g., 9sh Quh 9kt MRts and qtt, Stta- In this special case, the generalisation error can be written explicitly as a function of the overlaps
e, (W, W.)
Σ k
(0 arcsin
Qkh
+ aresin
Se
π
v1 + Q + Qn
π
1 S.c
(22)
MR.
Uk
aresin
VT QuiVI+ St.
A.2 ORDINARY DIFFERENTIAL EQUATIONS FOR THE FORWARD TRAINING DYNAMICS
Given that the generalisation error depends only on the overlaps, order to characterise the learning curves need to compute the equations of motion for the overlaps from the SGD dynamics of the weights given in Eq. 14 The order parameter Su associated to the teachers is constant in time. We
obtain an ODE for MRt by multiplying both sides of the first of Eq. 14 by !9 and dividing by N:
w
++1
w
w.
17"
(t. (te) (((N))
(23)
N
N
16

---

Published as a conference paper at ICLR 2025
where stress the difference between te, the task selected for training at step μ4, and t, the task for which we compute the overlap. We define a "training time" a = n/N and take the infinite- dimensional limit N + 30. The parameter becomes continuous and Mat concentrates to the solution of the following ODE:
dMs da - (te) Ex..
(A = fM,At
(24)
where the expectation is computed over the distribution in Eq. similarly from Eq.
18
The ODE for Qkh is obtained
14
u+t w
w whth
17"
(1.) (to) g' (A
T A ( (el (A
o
N
N
N
N
(25)
2 + (17t))² (A (1tt)e)
1 (ta)" g'( (2
(t))
x
2k
N
In the infinite-dimensional limit, we find
dQua
77 (te) (1c) E...
A A - not no E...
A
(26)
da
[(A)))) g' (Ax)) (An) = fQuu
+ u Ex..
(27)
Finally, taking the infinite dimensional limit of the second Eq. 14, we find the ODE for the readout:
_ (t)
da 17 BAA. A m(Ax)) = Sv.
(28)
is useful to write this system of ODEs more compact form. With the shorthand notation Q = (Qe,(Q), vec(M), vec( fo = (vec(fq), vec(fM), vec(fv)) we can write dQ(a) fo (Q(a), u(a))
a E (0, ar.
(29)
da
The initial condition for Q(0) is chosen to reproduce the random initialisation of the SGD algo- rithm. 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 respectively. A thorough analysis of the validity of this ODE description is provided in Veiga et al. (2022). where the authors study the crossover be- tween narrow and infinitely wide networks, clartyyig the connection with the so-called mean-field or hydrodynamic regime (Mei et 2018 Chizat & Bach 2018 Rotskoff & Vanden-Fijndenn 2022).
It is useful to write explicit expressions for the integrals involved in fo (Lee al. 2021
First,
expanding the terms in A (t), we can write
K
SQ.h = -
k,), - Ta k,)
n
K
(tc)
(to) Iz(n, - h,)k)
-70h
(30)
2
Uk
(tc) (te)
un
Σ
Up
(ote 1, k, k, (teer
Ln,m=1
K
2
Ia(n, ter h))
==1
fM = -nu
(tz)
y (t) (t) 12 k, 1oh La(to,
(31)
n=1
K
fv. 17
( (te) Iz(, n) (2kttt) Si..
(32)
17

---

Published as a conference paper at ICLR 2025
Similarly in Eq. 19
adopt the unified notation for the integrals
Ts(B, = Ex.x. [Asgs(A)) (. = Ex.x (gn
(33)
where B, S, can refer both to the indices of the student weights k. n, the tasks In the special case g(z) 9x(2) erf(2/ /2), the integrals have explicit expressions function of the overlaps
2qpc (1 + QaR) 208p03c
Ta(B,p,C)
T / + 988)
(34)
T.(3. T² / aresin Ao
4
VAAA²
the symbol y denotes generically an overlap from Eq. T6 and
Ag = Aa Q3p - 98r Apr (1 + (cc) - 93c UpC (1 + (rr) + 9er 986 9pr + 9cr OoC 937
A = A4 (1 + (aa) - 43r (1 + (cc) - 4 (1 + @rr) ) + 20cr 93c 4gr
A2 Aa (1 + app) Apr (1 + 9cc) qpc (1 + 9rr) + 29c-9pc99p
(35)
A3 (1 11--0 (38) (1 App) 4sp
A4 = (1 - <<<) (1 + (-r) Re
A.3 INFORMAL DERIVATION OF PONTRYAGIN MAXIMUM PRINCIPLE
Let us consider the augmented cost functional
u u h (Q(ar)) + /"d da Q(a)T
for
[- dQ(a) + (Q(a).
(36)
da
where the conjugate variables Q(a) act Lagrange multipliers, enforcing the dynamics at time a. Setting variations with respect to (()) results in the forward dynamics
FFQQQ.u) 0 = 8Q(a)
dQ(a) So (Q(a) da
(37)
Integrating by parts, we find
F h ( da @(a) fo (Q(a),
f da dQ(a) Q(a)
(38)
da
Q())((0)
Setting to zero variations with respect to Q(a) for a < ap, find the backward dynamics
dQ(a)) ((a) Vala TVaaa u(a)),
(39)
da
while for a ar we get the final condition
((ar) = voh(Q(ar))
(40)
Note that we do consider variations with respect to Q0)) as this quantity is fixed by the initial condition (0) Qo. Finally, minimizing the cost functional with respect to the control we get the optimality condition Eq. of the main text.
A.4 OPTIMAL CONTROL FRAMEWORK
To determine the optimal control, iterate Eqs. 5 and 10 of the main text until convergence
[Bechhoefer] 2021). Let us consider first the case where the control is the current task (a). such that (,) if the network is trained on task (1 (1....; T) training time For simplicity,
18

---

Published as a conference paper at ICLR 2025
focus on the case T 2, but the following discussion is easily generalised to any T. In particular, can be written as
since here u(a) = (a) the evolution equation 5
dQ(a) fo (Q(a)) (0)) - Qo
(41)
de
Similarly, the backward dynamics reads
dQ(a)
Q(a) (Q(a))((a)),
(42)
da
with final condition
((FF) 2 1 (Qsi(((aF)) (ee2(((F))
(43)
The optimality equation 10 yields
(a) argmin (Q(a) so (Q(o), te) (1,2
(44)
Therefore, we find the explicit formula for the optimal task protocol
{ 2
if (a)) Ifa (Q(a).ta(a) = 2) fo (Q(a)) = =)] 10
(5(a)=
(45)
otherwise.
Then, we start from a guess for the control variable (e(()). We integrate Eq. 41 forward, obtaining the trajectory (aa) for a ar,) Then, we integrate the backward equation 42
starting from the final condition 43 obtaining the trajectory Q(aa) for a a0. Then, the control variable can be updated using Eq. 45 and used in the next iteration of the algorithm. These equations 41 42 and 45
are iterated until convergence.
We next consider the joint optimisation of the learning rate schedule naa) and the task protocol te(a). The optimality condition 10 can be written as
(C(a), = argmin +,((1.))
(Q(a) (Q(a), (t(o)) (lo:n))
(46)
Crucially, the function Q fe(Q, (tc, 7)) turns out to be quadratic in m.. Explicitly,
fQ (te, n)) an² + by.
(47)
where
a Qunuke) (t) (tt)
Σ
o (to) (te) La(n, m, k,,) + La to h)
(48)
k,h=1
nmm=1
2
Talm,ter
and
K
K
Quh
Uk
(tc)
Σ
(t) knn,k,h) Io(to, k.))
(49)
k
Is(tes h, h 1 }
+ 2
(tc)
Σ
(te)
K
T
K
Ma
U (nn,k (()) Iste,
1=1
K
K
Ult 12 (k, m)) + 12k
n=l
19

---

Published as a conference paper at ICLR 2025
Performing the minimization over 1) first, we obtain
1f (a,t)))) = b
(50)
2a
The minimisation over Le yields
{ 2 otherwise.
(e (e((C))) 1))) So (200) ())))
(a))
(51)
and hence
(a,))""(a
(52)
Interestingly, we observe that the learning rate schedule has a different functional form depending on the current task te. This can be seen in Fig. 4 where the learning rate switches between two different schedules depending on the current task to.
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 function of task similarity 7. As the main text, we consider the case K 7 = 2. From the overlap trajectories in Fig. 8 for y > 0.3, observe that the cosine similarity quickly approaches unity, ic., |Mxl// Q x Ōge, which corresponds to perfect feature recovery. Therefore, the decrease in performance for y > 0.3 seen in Fig. must be attributed the dynamics of the second layer. Indeed, in Fig. 8 we observe a slowdown in the readout dynamics as + 1.
w. we consider the 1. We expect the corresponding readout layer to converge to the specialised configuration v!)) = (uf") (b()) = (1. and we would like
Assuming perfect convergence of the feature layer to wi (1) and
dynamics of the readout layer while training on task
to compute the convergence rate as a function of 7. The dynamics of the readout layer reads
de
da
3
1
(1)
arcsin
2/
(1)
(53)
dug
2
aresin
?
(1 2
(1)
1
U2
da
2
3
which can be rewritten as
d
1
(1)
1- v
(1)
nA
(54)
do
2
U22
where
- 1/3
A =
(55)
1/3
and a = 2 arcsin (7/2) /r. Note that a < f/r for 0 < 7 < 1, hence A is negative definite, implying convergence to (1) = (1,0)T The rate of convergence determined by the smallest eigenvalue (in absolute value): a 1/3. The associated convergence timescale is therefore
Oconv
ア, - 6 arcsin (1/2))
(56)
anticipated in Eq. I of the main text.
C SUPPLEMENTARY FIGURES
C. ADDITIONAL RESULTS IN THE SYNTHETIC FRAMEWORK
Fig. describes the dynamics of the optimal replay strategy for different values of task similarity in the same setting Fig. of the main text. In particular, the upper panel displays the evolution of
20

---

Published as a conference paper at ICLR 2025
Y = 0.1
y 0.3
y 0.6
y 0.9
M CAn
/ dimension
Figure 8: Overlap dynamics with optimal replay. We plot the absolute value of the task dese-depn readout weights (0) (upper panel) and the cosine similarity |Mael/v as a function of the training time a. Different columns refer to different choices of task similarity y 0.1, 0.3, 0.6, 0.9.
the magnitude of the readout weights 10/') while the lower panel shows the trajectory the cosine similarity / VQ
Fig. 9 compares the values of the loss at the end of training. averaged on both tasks, for different lask-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. Additionally, consider exponential and power law learning rate schedules, in combination with interleaved replay protocol. For each value of task similarity % we optimize over the schedule parameters via grid search. We still find performance gap with respect the optimal strategy, which highlights the relevance of the joint optimization of training protocols.
Fig. T0 illustrates continual learning setting with T 3 tasks. The student is two-layer neural networks with K = 3 hidden units, and a different readout is trained for each task. In the initial training phase, the student is trained task up time a 1000, when the loss reaches 10-6 In the second phase, the student must learn tasks 2 and 3 without forgetting task 1. Panel a) shows the losses of the optimal strategy as function of time for the three tasks and their average, during the second training phase. The optimal strategy is represented by a colour bar in the upper panel. consists in an initial phase where only the new tasks 2 and 3 are presented to the network, and second phase where task is replayed. Despite its complicated structure, it shares some similarities with the pscudo-optimal strategy described in the main text. Specifically, task is replayed only when its loss is comparable to the losses on the other two tasks. Panel b) shows the average loss as a function of time, comparing different mesk-selection protocols. In particular, we consider
interleaved protocols, where two all three tasks are alternated during training;
sequential protocols, where tasks presented in distinct blocks without being replayed;
shuffled protocol, that preserves the relative fraction of samples from each tasks obtained from the optimal strategy but presents them in a randomised order.
The final performance of the optimal strategy surpasses all the aforementioned approaches. Notably, as the number of tasks grows, the number of possible heuristic strategies expands significantly, mak- ing difficult to identify effective solutions through intuition alone. This highlights the importance of theoretical framework for systematically determining the optimal strategy.
C.2 ADDITIONAL RESULTS ON REAL DATASET
We run additional experiments on real data using the simulation setup detailed in Sec. 3.2 We consider the CIFAR10 dataset Krizhevsky et al. (2009) and create classification tasks taking as task the classes with odd labels and the others as task 2. Fig. shows the results of the training curve for the two tasks according the three strategies and the final figure compares the average
21

---

Published as a conference paper at ICLR 2025
ruplay nterlesved
Optimal replay Psoudo Optimal
Average loss
1072.
nnerlesved,
Optimal replay and
10-3
Task similarity
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 function of task similarity ? under the same setting and parameters as Fig. 3 of the main text. The top four lines correspond to different strategies constant learning rate 1j 1: no replay (purple crosses), optimal (red dots), interleaved (blue squares), pseudo-optimal (cyan dashed line). The bottom three curves correspond different annealing schemes for the learning rate. Up-facing orange triangles correspond inter- leaved replay with power law annealing 17(a) no a/a,) ), while down-facing green triangles indicate exponential annealing v7(a) = Too sxp(-a/ao). For both annealing schedules, the scalar parameters y/. 3, and og optimized by grid search for each value of % Finally, the brown plus signs correspond to jointly optimal replay and learning rate schedules (see Fig. 4
Loss function protacol
Interleaved (1-2-3)
0.150
Task 2
Task 1
10a1
Interteaved (2
Shuffed
0.125
Average loss
Scquential
Sequential (3
0.300
Average
Sequential (3-2.1)
Optimal
0.075
0.050
0.025
0.000
a)
b)
1005
10..0
3015
1020
1025
Figure 10: Optimal replay schedule for = = tasks. The student network has k hidden units. In the initial phase, [0, 1000] the network is trained on task I until convergence (loss is 10 10-6). In the second phase = [1000, 1025], we determine the optimal replay strategy. In both phases the learning
is constant / 1. The tasks are chosen such that the overlaps S1,2 0.5. S2.3 0.5, and S1,3 = 0. Panel a) shows the optimal task protocol and the evolution of the loss over the three tasks. The result is complicated replay strategy, though shares some similarities with the pseudo-optimal strategy described in the main text. Specifically, task is replayed only when its loss is comparable the losses on the other two tasks. Panel b) compares the optimal strategy with different heuristics, all at 1. "Interleaved (1 2 3)" is an interleaved strategy containing tasks equal proportion. "Interleaved (2 3)" is an interleaved strategy containing tasks (2 - 3) in equal proportion. "Shuffled" contains the same task proportions as in the optimal strategy, but in random order (showing that the structure of the optimal replay strategy matters). "Sequential (2 3)" corresponds to the sequential strategy with t = 2 in the first half and 3 in the second. Similarly for "Sequential (3 - 2)". "Sequential (3 of the replay sequence, then f in the second third and f in the last third.
1)" has t = 3 the first third
loss throughout the training steps. Notice that the observation reported in the controlled Fashion MNIST experiment still valid in this scenario.
22

---

Published as a conference paper at ICLR 2025
Replay
Optimall
Average
10a1
task
Interleaved
Retuia-Ottimal
Loss
8x-0
7.0
6.9
7x 10
6.8
15k
20k
10000
20000
10000
20000
10000
20000
10000
20000
Figure 11: Continual learning on CIFARIO. Training curves of No-Replay, Interleaved, and Pseudo-Optimal learning strategies (from left right) on CIFARIO. The two tasks are obtained partitioning the dataset according the parity of the labels and training using online stochastic gra- dient descent. The final panel (rightmost) shows the average loss for the three strategies. The inset zooms the final stage of learning.
As already observed the main text, the performance of the Seeudo-Optimal learning strategies appears to interpolate between No-Replay and Interleaved. In Fig. 12 we highlight the differences between Pseudo-Opp Sestimal and the alternative strategies. We see a up to 3% improvement in perfor- mance during learning. Contrarily the main text, observe also region where Pseudo-Optime appears sub-performing with respect to the Interleaved strategies, however this limited to a less than 0.5% decrease in loss.
Interleaved Pseudo-
No Replay Pseudo-Optimal
Loss difference
10000
12000
14000
16000
18000
20000
Step
Figure 12: Learning difference CIFAR10. This figure complement Fig. differences between the strategies.
11 highlighting the
23


Reviews:
[{"review_id": "SauxVReo8B", "reviewer": "Reviewer_vgxp", "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.", "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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.\n- 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.\n- 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).\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ojSmwfsqZo", "reviewer": "Reviewer_PcHB", "summary": "This paper provides a theoretical framework for optimal task selection and step-size optimization in continual learning.\nThe 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.\nUsing 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.\nThe ideas are thoroughly evaluated in a toy setting with simulated data, with some additional experiments on Fashion MNIST and MLPs.", "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?\n\n- 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.\n\n- 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?\n\n- \"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.\n\n- 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?\n\n- 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?\n\n- 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?\n\n- 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)\n\n- 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?\n\n\n** Minor Comments\n- 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.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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.\n\n- 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.\n\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "g1zHFrdPTL", "reviewer": "Reviewer_pAPx", "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.\n\n**Contributions:**\n- Derive optimal task-selection protocols and learning rate schedules as functions of task similarity and problem parameters.\n- Proposes a \"pseudo-optimal\" replay strategy with distinct focus and revision phases to minimize forgetting.", "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?\n- 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.\n- (see weaknesses for more suggestions)", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- Figures are clear, well explained, and easy to understand.\n- The paper uses an original and refreshing approach with strong theoretical backing.\n- Demonstrates robust alignment between theoretical predictions and experimental results.\n- Complex mathematical formulations are well-supported.", "weaknesses": "- Please try to use shorter sentences. The sentences tend to be overly long and complex.\n- The mathematics and equations are dense and could limit accessibility for readers unfamiliar with optimal control or statistical physics.\n- 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.\n- 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.\n- 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.\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MMmXHVAiZD", "reviewer": "Reviewer_4ows", "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.", "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.\n- 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper is well-written and organized, with a comprehensive discussion of related work.\n- 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.\n- The experiments are detailed, especially for a theoretical work.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)