Meta-Learning Strategies through Value Maximization in Neural Networks
Biological and artificial learning agents face numerous choices about how to
learn, ranging from hyperparameter selection to aspects of task distributions
like curricula. Understanding how to make these meta-learning choices could
offer normative accounts of cognitive control functions in biological learners
and improve engineered systems. Yet optimal strategies remain challenging to
compute in modern deep networks due to the complexity of optimizing through the
entire learning process. Here we theoretically investigate optimal strategies
in a tractable setting. We present a learning effort framework capable of
efficiently optimizing control signals on a fully normative objective:
discounted cumulative performance throughout learning. We obtain computational
tractability by using average dynamical equations for gradient descent,
available for simple neural network architectures. Our framework accommodates a
range of meta-learning and automatic curriculum learning methods in a unified
normative setting. We apply this framework to investigate the effect of
approximations in common meta-learning algorithms; infer aspects of optimal
curricula; and compute optimal neuronal resource allocation in a continual
learning setting. Across settings, we find that control effort is most
beneficial when applied to easier aspects of a task early in learning; followed
by sustained effort on harder aspects. Overall, the learning effort framework
provides a tractable theoretical test bed to study normative benefits of
interventions in a variety of learning systems, as well as a formal account of
optimal cognitive control strategies over learning trajectories posited by
established theories in cognitive neuroscience.
Rodrigo Carrasco-Davis, Javier Masís, Andrew M. Saxe
# Meta-Learning Strategies through Value Maximization in Neural Networks

###### Abstract

Biological and artificial learning agents face numerous choices about how to learn, ranging from hyperparameter selection to aspects of task distributions like curricula. Understanding how to make these'meta-learning' choices could offer normative accounts of cognitive control functions in biological learners and improve engineered systems. Yet optimal strategies remain challenging to compute in modern deep networks due to the complexity of optimizing through the entire learning process. Here we theoretically investigate optimal strategies in a tractable setting. We present a _learning effort_ framework capable of efficiently optimizing control signals on a fully normative objective: discounted cumulative performance throughout learning. We obtain computational tractability by using average dynamical equations for gradient descent, available for simple neural network architectures. Our framework accommodates a range of meta-learning and automatic curriculum learning methods in a unified normative setting. We apply this framework to investigate the effect of approximations in common meta-learning algorithms; infer aspects of optimal curricula; and compute optimal neuronal resource allocation in a continual learning setting. Across settings, we find that control effort is most beneficial when applied to easier aspects of a task early in learning; followed by sustained effort on harder aspects. Overall, the learning effort framework provides a tractable theoretical test bed to study normative benefits of interventions in a variety of learning systems, as well as a formal account of optimal cognitive control strategies over learning trajectories posited by established theories in cognitive neuroscience.

## 1 Introduction

Deploying a learning system requires making many considered decisions about hyperparameters, architectures, and dataset properties. As learning systems have grown more complex, so have these decisions about how to learn. One approach to managing this complexity is to place these decisions under the control of the agent and meta-learn them. Building on this strategy, a range of meta-learning algorithms have been developed that are capable of fast adaptation to new tasks within a distribution (Finn et al., 2017; Nichol et al., 2018), continual learning (Parisi et al., 2019), and multitasking (Crawshaw, 2020). Meta-learning methods target diverse aspects of a learning system: they can adapt hyperparameters (Franceschi et al., 2018; Baik et al., 2020; Zucchet & Sacramento, 2022); learn weight initializations well-suited to a task distribution (Finn et al., 2017; Baik et al., 2020); manage different modules or architectural components (Andreas et al., 2017); enhance exploration (Gupta et al., 2018; Liu et al., 2021); and order tasks into a suitable curriculum (Stergiadis et al., 2021; Zhang et al., 2022). While this prior work has shown that meta-learning can bring important performance benefits, algorithms are often hand-designed for a specific intervention and a large gap remains in our theoretical understanding of how meta-learning operates (see App. A).
The aim of this paper is to develop a normative framework for investigating optimal meta-strategies in neural networks, implemented in biological and artificial agents. A core difficulty in computing optimal strategies is the complexity of optimizing through the learning process. To tackle this problem, we simplify the inner-loop learning dynamics using simpler tractable network models. We specifically study meta-learning dynamics in deep linear networks, which exhibit complex non-linear dynamics sharing properties of observed non-linear networks dynamics (Saxe et al., 2019; Braun et al., 2022). Examining this problem in a reduced setting, we derive optimal meta-learning strategies under various control designs and meta-learning scenarios. We concentrate on questions that are pertinent to the cognitive control literature, such as learning effort allocation, task switching, and attention to multiple tasks. The Expected Value of Control Theory (EVC, (Shenhav et al., 2013, 2017; Musslick et al., 2020; Masis et al., 2021)) has proposed answers to these questions. It posits that higher-level areas in the brain perform executive functions (cognitive control) over lower-level areas to maximize the cumulative return. The framework we present is a formal and computationally tractable example of the EVC theory that takes into account the impact of the control signal on future learning dynamics (see Appendix E).

**Main contributions**

\(\bullet\): We develop a computationally tractable _learning effort_ framework1 to study diverse and complex meta-learning interventions that normatively maximize value throughout learning.

Footnote 1: Anonymized Python package at [https://anonymous.4open.science/r/neuromod-6A3C/](https://anonymous.4open.science/r/neuromod-6A3C/) for reproducibility.

\(\bullet\): We fully solve learning dynamics as a function of control variables for simple models, and use this to derive efficient optimization procedures that maximize discounted performance throughout learning.

\(\bullet\): We express meta-learning algorithms such as Model Agnostic Meta-Learning (Finn et al., 2017) and Bilevel Programming (Franceschi et al., 2018) in our framework, studying the impact of approximations on their performance.

\(\bullet\): We compute optimized control strategies for a range of settings spanning continual learning, multi-tasking, and curriculum learning, and examine these normative strategies.

\(\bullet\): Due to this framework's normative goal of maximizing expected return, we draw qualitative connections to phenomena in cognitive neuroscience such as task engagement, mental effort, and cognitive control (Shenhav et al., 2013, 2017; Lieder et al., 2018; Masis et al., 2021).

## 2 Learning Effort Framework

We start by defining our framework in a general abstract way before turning to a simple example in Section 2.1. The generality of this description allows the framework to apply to a variety of different settings of interest spanning machine learning (Section 4) and cognitive control (Sections 5 and 6). Consider a learning model trained on a task \(\mathcal{T}\) for a period of time \(T\). Two equations define the learning model. We define the input-output mapping \(f\) and learning dynamics \(h\) as

\[\hat{Y}=f(X;w(t),g(t)),\quad\tau_{w}\frac{dw(t)}{dt}=h(w(t),g(t),\mathcal{T}) \tag{1}\]

respectively. In the first equation, \(f\) is a continuously differentiable function, \(X\) the input, and \(\hat{Y}\) the output. Here \(w(t)\) are the parameters of the learning model (e.g. weights in a neural network) during training with \(T\geq t\geq 0\). We introduce \(g(t)\) as an _effort signal_ (or _control signal_) that crucially will be chosen by the meta-learning optimization. This vector of control signals can model a number of interventions in the learning system, and will be chosen to maximize cumulative learning performance. The learning dynamics equation describes the evolution of the parameters during training and is given by a differential equation over the parameters of the learning model, \(h\) is a continuously differentiable function, and the evolution of the learned parameters \(w(t)\) (starting at \(w(0)\)) may depend on the control signal \(g(t)\) and task parameters \(\mathcal{T}\).
Given this setup, we can understand the control signal \(g(t)\) as a meta parameter that can be chosen in different ways, and which influences the network's input-output map and learning behavior. Regarding cognitive neuroscience literature, this could take the form of controlled attention or neural activity modulation. To determine how we choose \(g(t)\), we define a task performance metric during the learning period \(\mathcal{P}(t)\) (e.g. mean squared error during regression). Further, we assume that using the control signal \(g(t)\) is costly, according to a cost function \(C(g(t))\), as commonly used in control theory to describe, for instance, energy resource needed to exert control, or mental effort allocated to produced sustained engagement on a task. At any time during the learning of the task \(\mathcal{T}\) we consider an instant reward rate \(R(t)=\eta\mathcal{P}(t)\), where \(\eta\) is a constant that converts performance on the task \(\mathcal{P}(t)\) to reward/time units. We define the instant net reward rate as the difference between scaled performance and the cost of control \(v(t)=R(t)-C(g(t))\). The expected return or value function at the start of training can then be written as the cumulative discounted reward from learning and performing the task from time \(t=0\) to \(t=T\), with a discount factor \(1\geq\gamma>0\),

\[V=\int_{0}^{T}dt\gamma^{t}v(t)=\int_{0}^{T}dt\gamma^{t}\left[\eta\mathcal{P}(t )-C(g(t))\right]. \tag{2}\]

We emphasize this value function measures performance across the whole learning period. Finally, we posit that the goal of meta-learning is to choose \(g(t)\) to maximize the value function in equation 2. To find an approximately optimal \(g(t)\), we take gradient steps

\[g_{k+1}(t)=g_{k}(t)+\alpha_{g}\frac{dV}{dg(t)}. \tag{3}\]

for every \(0\geq t\geq T\), \(k\) being the iteration index. The optimal \(g(t)\) thus depends on a complex interplay of past and future values of the control signal, and how these interact with the whole trajectory of learning. Indeed, computing the gradient in equation 3 is computationally intractable in general. In the remainder of the paper, we carefully choose learning models and settings with rich dynamics but for which we have partial analytical tractability of the learning dynamics, such that efficient computation of the full control signal through time is possible. Further details on the algorithm implemented and estimation of involved quantities can be found in App. C and D.

Figure 1: Learning effort framework. A neural network is under the influence of a control signal \(g(t)\). This control signal is optimized iteratively by initializing \(g(t)\), then: (1) Solving learning dynamics in Eq. equation 1; (2) Computing the performance \(\mathcal{P}(t)\); (3) Integrating performance and control cost to compute the exact cumulative return \(V\) in Eq. 2; (4) Taking the gradient of \(V\) with respect to the control signal \(g(t)\) and update as in Eq. 3, then go back to (1). **(b)**: Multi-step MAML. **(c)**: Learning rate optimization as in Bilevel Programming. **(d)**: Task engagement, where the control signal determines the optimal amount of _engagement_ through time to multiple regression tasks. **(f)**: Category assimilation, where a model is trained to learn a classification task and can control the _engagement_ on each class \(c\) throughout training. **(e)**: Effort allocation, where the control signal (gain modulation of weights) is computed to maximize value throughout the learning of a single task. **(g)**: Task switching, where the gain modulation model is trained to switch tasks repeatedly and the control signal is computed throughout the switches.
By appropriate choice of how \(g(t)\) influences the network and learning dynamics, this general framework can accommodate a variety of possible interventions on a learning system. Some interventions correspond to other meta-learning algorithms such as Multi-Step MAML and Bilevel Programming (Fig. 1b and c). Further connection to the meta-learning literature in machine learning can be found in App. F. The results in subsequent sections investigate several scenarios illustrated in Fig. 1d to g. All of these experiments are variations on the influence of the control signal over the learning dynamics, keeping the rest of the framework as is.

### Single Neuron Example

Having described the general framework, we now turn to a simple case to illustrate it, yet with complex emergent solutions. We consider a _single neuron learning model_ trained on a _two-Gaussians regression task_ where the control signal acts as a _weight gain modulation_. This case offers insights regarding the dependence of the optimal control signal on task parameters and learning model hyperparameters.

**Two Gaussians regression task:** A dataset of examples \(i=1,\cdots,P\) is drawn as follows: A label \(y_{i}\) is first sampled as either \(+1\) or \(-1\) with probability \(1/2\). The input \(x_{i}\) is then sampled from a Gaussian \(x_{i}\sim\mathcal{N}(y_{i}\cdot\mu_{x},\sigma_{x}^{2})\). The task is to predict \(y_{i}\) from the value of \(x_{i}\). The intrinsic difficulty of the task is controlled by how much the Gaussians overlap, controlled by the relative value of \(\mu_{x}\) and \(\sigma_{x}\).

**Single neuron learning model:** The input-output mapping of our single neuron model is \(\hat{y}_{i}=x_{i}\cdot w(t)\left[1+g(t)\right]\), \(w(t)\) is our learned weight parameter, and \(g(t)\) is the control signal which acts as a multiplicative gain. The learning dynamics of \(w(t)\) are given by gradient descent on the loss function \(\mathcal{L}=\frac{1}{2}(y_{i}-\hat{y}_{i})^{2}+\frac{\lambda}{2}w(t)^{2}\). Taking the gradient flow limit (small learning rate (Saxe et al., 2019; Elkabetz & Cohen, 2021)), we find average learning dynamics for the weight described by

\[\tau_{w}\frac{dw}{dt}=-\left\langle\frac{\partial\mathcal{L}}{\partial w} \right\rangle=\mu_{x}\tilde{g}(t)-w(t)\left(\left\langle x^{2}\right\rangle \tilde{g}^{2}(t)+\lambda\right) \tag{4}\]

where \(\tilde{g}(t)=1+g(t)\), \(\left\langle\cdot\right\rangle\) denotes expectation over the data distribution, and \(\tau_{w}\) is the learning time scale of the weight. This gradient depends on \(g(t)\), making the learning dynamics of \(w(t)\) dependent on the control signal. For the single neuron model, we can find a closed form expression for \(w(t)\) as a function of the control signal \(g(t)\), giving us an expression for \(\left\langle\mathcal{L}(t)\right\rangle\) as well (see App. G.1). This tractability allows us to compute average dynamics and the necessary gradient efficiently.

**Control signal optimization**: As a performance and control measure for this model we use \(\mathcal{P}(t)=-\left\langle\mathcal{L}(t)\right\rangle\), \(C(g(t))=\beta g(t)^{2}\) respectively, meaning smaller loss leads to better performance and exerting control has a cost that is monotonic in the control signal magnitude, with cost per unit of control \(\beta\). Note that if \(g(t)=0\) for all \(T\geq t\geq 0\), then \(C(g(t))=0\), and \(\hat{y}_{i}=x_{i}\cdot w(t)\), which means that the weight is learned purely by gradient descent with no influence from the control signal, we call this is the _Baseline_ model. Having \(\mathcal{P}(t)\) and \(C(g(t))\), we can compute the value function in equation 2 and find the optimal \(g(t)\) by gradient ascent following equation 3 (algorithm described in App. D). In essence, this setting considers a simple learning scenario in which an agent can adjust the gain of the weights in a neural network that otherwise learns via gradient descent.

**Results:** In Fig. 2a, we show the difference in instant net reward \(v(t)\) for the baseline (\(g(t)=0\) for every \(t\)) and the control case (optimizing \(g(t)\)). The optimal meta-learning strategy that maximizes expected return in equation 2 invests more control at the start (Fig. 2b) of the learning period at the cost of some instant reward, with the result of faster learning (demonstrated in the lower loss for the control case in Fig. 2c). The control signal \(g(t)\) influences the instant net reward rate both at present \(t\) and future \(t^{\prime}>t\). The instant change in net reward rate \(v(t)\) will be caused by both the instant change on the \(\tilde{w}(t)=w(t)\cdot(1+g(t))\) (Fig. 2d) and \(C(g(t))\), making the effective weight \(\tilde{w}(t)\) closer to the solution at early stages. As expected, increasing the discount factor \(\gamma\) leads to higher levels of control, since future net reward will contribute more to the cumulative expected return, compensating the cost of increasing \(g(t)\) (Fig. 2e,f). Increasing the intrinsic noise of the task \(\sigma_{x}\) reduces the overall optimal control (Fig. 2g,h). Because it is not possible to
overcome this noise, the use of control will generate a cost that cannot be compensated by boosting learning. This inter-temporal choice of allocating effort based on the prospect of future reward has been widely studied in psychology and neuroscience (Masis et al., 2021; Keidel et al., 2021; Fromer et al., 2021; Masis et al., 2023) (App. A), and naturally arises from maximizing the discounted cumulative performance in Eq. 2. For more parameter variations see Fig. 9 in App. K.1.

## 3 Baseline Deep Linear Networks and Datasets

We now generalize the single neuron approach to more complex neural networks. In the case of a two-layer linear neural network, the corresponding input-output mapping in Eq. equation 1 is \(\hat{Y}=W_{2}(t)W_{1}(t)X\), where \(X\in\mathbb{R}^{I}\), \(\hat{Y}\in\mathbb{R}^{O}\), \(W_{1}(t)\in\mathbb{R}^{H\times I}\) and \(W_{2}(t)\in\mathbb{R}^{O\times H}\) are the first and second layer weights. Training a two-layer network to minimize MSE with weight regularization and taking gradient flow limit yields the learning dynamics equations (Saxe et al., 2019; Braun et al., 2022)

\[\tau_{w}\frac{dW_{1}}{dt}=W_{2}^{T}\left(\Sigma_{xy}^{T}-W_{2}W_{1}\Sigma_{x} \right)-\lambda W_{1},\ \ \tau_{w}\frac{dW_{2}}{dt}=\left(\Sigma_{xy}^{T}-W_{2}W_{1}\Sigma_{x} \right)W_{1}^{T}-\lambda W_{2} \tag{5}\]

where \(\Sigma_{xy}^{T}=\left\langle XY^{T}\right\rangle\), \(\Sigma_{x}^{T}=\left\langle XX^{T}\right\rangle\), \(\tau_{w}\) is a learning time-scale of the weights, and \(\lambda\) controls the weight regularization (see App. H). Learning is completely defined by the initial weights \(W_{1}(0)\), \(W_{2}(0)\), the task at hand and the hyperparameters, it follows non-linear dynamics due to weight coupling and a non-convex loss landscape while keeping computational tractability. Other frameworks such as the teacher-student setting (Goldt et al. 2019; Ye and Bors 2022) or mean field theory approximations (Mignacco et al. 2020; Bordelon and Pehlevan 2022) provide closed-form dynamics for non-linear networks, but relies on assumptions on task structure and taking architecture size limits (App. A). With the general framework and tractable models in hand, we now turn to probe the behavior of this two-layer network in a variety of settings. First, in Section 4, we draw out implications for standard meta-learning algorithms. Next, in Section 5 we turn to aspects of curriculum learning and the choice of which tasks to engage with. Finally, in Section 6 we study control interventions in the form of gain modulation throughout a network, of relevance to theories in neuroscience. In all of these sections, we compose meta-learning tasks from a base set of three datasets: (1) **Correlated Gaussian** regression, (2) **Semantic Tasks** with hierarchical concepts, and (3) **MNIST** (Details in App. J), from which we can determine the statistics needed to compute the learning dynamics (e.g. \(\Sigma_{x}\) and \(\Sigma_{xy}\)).

Figure 2: Results in single neuron model throughout the learning period \(0\geq t\geq T\). **(a)** Instant net reward \(v(t)\). **(b)** Loss \(\left\langle\mathcal{L}(t)\right\rangle\) for theoretical predictions (solid) and simulations using SGD (shaded). **(c)** Optimal control signal decreases through learning (Baseline \(g(t)=0\)). **(d)** Weight \(w(t)\) through learning for control and baseline case, \(\hat{w}(t)=w(t)\cdot(1+g(t))\). Dependence of optimal control signal on task parameters. **(e)** and **(g)**: optimal \(g(t)\) when varying discount factor \(\gamma\) and noise level \(\sigma_{x}\) respectively. **(f)** and **(h)**: Difference between instant net rewards \(v(t)\) between control and baseline when varying \(\gamma\) and \(\sigma_{x}\) respectively. Longer time horizons and less noisy tasks recruit more control.
Relation to Meta-Learning Algorithms in Machine Learning

The normative objective in Eq. equation 2 and the way it is maximized through gradient steps on the control signal \(g(t)\) can describe other meta-learning algorithms. Here we show the connections to two well-established algorithms, Model Agnostic Meta-Learning (MAML Finn et al., 2017) and Bilevel Programming (Franceschi et al., 2018).

MAML is an instance of our framework where the initial weights \(W_{1}(0)\) and \(W_{2}(0)\) in our deep linear network _are_ the control signal \(g(t)\). By defining the performance as the average loss per task indexed by \(\tau\), \(\mathcal{P}(t)=\sum_{\tau}\left\langle\mathcal{L}_{\tau}(t)\right\rangle\), this becomes the meta-objective in MAML when considering only one step ahead in the value function, this is \(V_{\text{MAML}}=\mathcal{P}(\delta t)\) with \(\delta t\) being the time after one gradient update on \(g(t)\) (See App F.1). Our framework can also optimize performance after multiple gradient steps, therefore obtaining Multi-Step MAML in a computationally tractable setting (Fig. 1b). We used the two-layer linear network (Section 3) and a set of 5 binary regression tasks with different pairs of digits from MNIST (App. F.1) to simulate Multi-Step MAML. Results in Fig. 3a and b show that the standard MAML loss \(V_{\text{MAML}}\) changes depending on how many steps ahead are considered during the initial weights optimization. \(V_{\text{MAML}}\) decreases when considering a few steps ahead, increasing the capacity to optimize the dynamics. After a certain number of steps considered in the optimization, \(V_{\text{MAML}}\) increases, sacrificing immediate performance to optimize the dynamics in steps further away, as shown in Fig. 3b. These multi-step results are only possible due to the tractability of our setting. We see that one-step MAML can substantially underperform Multi-Step MAML.

We also optimized hyperparameters of the network throughout training. Bilevel Programming optimization can compute this, with the main distinction being the reverse-hypergradient method used to update the meta-parameters (control signal) (Franceschi et al., 2017, 2018). We extend these methods by adding features with intuitive meaning under our normative frameworks, such as the discount factor \(\gamma\) and control cost \(C\) (App. F.2). We optimized the learning rate throughout time to maximize the cumulative reward in equation 2, and varied \(\gamma\) and \(\beta\) as in the single neuron example to illustrate the normative meaning in a hyperparameter optimization context (Fig. 1c). We observed qualitatively similar behavior as in the single neuron model, longer time horizons and less cost of increasing the learning rate recruit more control. Our work provides further utility to these meta-learning algorithms by interpreting them under a normative value-based framework. We note that there is a family of meta-learning phenomena that is not explicitly described by our framework. This being emergent meta-learning agents where no outer loop or meta-variable is optimized explicitly (Wang et al., 2017). We extend this discussion in App. F.3.

## 5 Engagement Modulation

Next we turn to the question of which tasks among many to engage with over time. We provide the model control over its _engagement_ on a set of available tasks, or class in a classification problem during learning. Selecting the optimal control signal in this setting involves improving multi-task capabilities and estimating optimal curriculum. Consider a set of \(N_{\tau}\) datasets, and a loss function \(\mathcal{L}(\hat{Y}_{\tau},Y_{\tau})\), where \(\hat{Y}_{\tau}\) is the estimation of a model and \(Y_{\tau}\) is the required target for dataset \(\tau\). The average loss for a set of datasets is \(\mathcal{L}=\sum_{\tau=1}^{N_{\tau}}\mathcal{L}(\hat{Y}_{\tau},Y_{\tau})+ \mathcal{R}(W)\) which is used to measure the performance \(\mathcal{P}(t)=-\left\langle\mathcal{L}(t)\right\rangle\) only, and we assume that weights are updated via gradient descent on the auxiliary loss \(\mathcal{L}_{\text{aux}}=\sum_{\tau=1}^{N_{\tau}}\psi_{\tau}(t)\mathcal{L}( \hat{Y}_{\tau},Y_{\tau})+\mathcal{R}(W)\), where \(\psi_{\tau}(t)\) are control signals we call _engagement coefficients_, and \(\mathcal{R}(W)\)

Figure 3: **(a)**: Single step MAML loss \(V=\mathcal{P}(\delta t)\) when considering more steps in the learning dynamics. **(b)**: Resulting learning dynamics from initial parameters found with Multi-Step MAML. **(c) and (d)**: Optimal learning rate when varying discount factor \(\gamma\) and cost coefficient \(\beta\).
is a weight decay regularize. We use this auxiliary loss just to obtain a learning dynamics equations that explicitly depends on the engagement coefficients \(\psi_{\tau}(t)\). Assuming the network receives inputs from all of the datasets at the same time (concatenated in \(X\)) and has specific outputs allocated to each dataset (concatenated in \(Y\)) as schematized in Fig. 1d, we can derive the learning dynamics equations for the weights as a function of \(\psi_{\tau}(t)\) giving

\[\tau_{w}\frac{dW_{1}}{dt} =\sum_{\tau}\psi_{\tau}(t)W_{2\tau}^{T}\left(\Sigma_{xy\tau}^{T}- W_{2\tau}W_{1}\Sigma_{x}\right)-\lambda W_{1},\] \[\tau_{w}\frac{dW_{2}}{dt} =\sum_{\tau}\psi_{\tau}(t)\left(\Sigma_{xy\tau}^{T}-W_{2\tau}W_{1 }\Sigma_{x}\right)W_{1}^{T}-\lambda W_{2}, \tag{6}\]

where \(W_{2\tau}\) denotes the weights of the neurons for the output to dataset \(\tau\) and \(\Sigma_{xy\tau}\) is \(\left\langle XY_{\tau}^{T}\right\rangle\), both padded with zeros to preserve dimension (see App. H.2). Each of the \(\psi_{\tau}(t)\) modulates the amount of learning of each dataset. The auxiliary loss to get a learning dynamics is to avoid the trivial solution of \(\psi_{\tau}=0\) to minimize the loss. We can find the optimal \(\psi_{\tau}(t)\) throughout learning by computing \(\mathcal{P}(t)\), using \(C(\psi(t))=\beta\|\mu_{\psi}-\bar{\psi}(t)\|^{2}\) (\(\bar{\psi}=(\psi_{1}(t),\psi_{2}(t),...)\)), then taking gradient steps on \(\psi_{\tau}(t)\) to maximize \(V\). Taking \(\mu_{\psi}=0\) means that to learn a dataset \(\tau\) (\(\psi_{\tau}(t)>0\)) the agent must pay a cost. We call this case _active engagement_. For \(\mu_{\psi}=1\), the agent must pay a cost to increase or suppress the learning signal from a specific dataset relative to a baseline. We call this case _attentive engagement_. In these cases, each of the elements in \(\bar{\psi}(t)\) are forced to stay in a certain range independently. Finally, we can force \(\bar{\psi}(t)\) to be of a fixed norm by making the cost \(C(\bar{\psi}(t))=\beta\left(\|\bar{\psi}(t)\|^{2}-\Psi\right)^{2}\), such that there is a fixed overall amount of engagement to distribute. We call this case _vector engagement_. For category engagement, which is focusing on particular subclasses in a classification problem, a similar set of equations can be derived (see App. H.3), where the engagement on class \(c\) through learning is denoted by \(\phi_{c}(t)\) (Fig. 1f). The meta-learning tasks used to train this model are the following:

**Task engagement**: Given a set of \(N_{\tau}\) datasets, and a total training period of \(T\), we trained the engagement modulation model described in Section 5. The idea of this task is to estimate the optimal learning curriculum (order of datasets presented in the neural network training) that maximizes expected return \(V\) during the time period \(T\). In this task, three binary MNIST classification datasets were used, specifically the digits \((0,1)\), \((7,1)\) and \((8,9)\) ordered by difficulty (easier to harder according to linear separability, see App. H.2).

Figure 4: Results for task engagement experiment. **(a)**, **(c)** and **(e)**: \(\mathcal{L}(t)\) for baseline and control case for Attentive, Active and Vector engagement. **(b)**, **(d)** and **(f)**: Engagement coefficients \(\psi_{\tau}(t)\) for each of the binary classification tasks Attentive, Active and Vector engagement. Mean and standard deviations from 5 independent trainings. **(h)** and **(j)**: Results for category engagement task, improvement in the loss function when using control for MNIST and Semantic dataset respectively. **(i)** and **(k)**: Optimal category engagement coefficients for MNIST and Semantic datasets. **(l)**: Class proportion experiment. **Uniform**: Loss when using uniform distribution for the abundance of classes in each batch. **Balanced**: Loss on a balanced batch, but using the inferred curriculum of classes in the batch to train. **Curriculum**: Loss on curriculum batch when using the curriculum. **(m)**: Loss per class using control (solid lines) and baseline (dashed lines).
**Category engagement**: For a classification task, there might be a better set of classes to learn during different stages of training. We trained the category engagement modulation model (described in Section 5 and App. H.3) to estimate the optimal _engagement_ or _attention_ to each of the categories in a classification task (Semantic and MNIST datasets) through learning. In addition, we trained the gain modulation model (next Section) in this same setting using a _neuron basis_ (see App. H.1).

### Results

**Task engagement**: We simultaneously gave the neural network inputs and targets for three datasets, as described in Section 5, each of them a different binary regression problem from MNIST. Each dataset used was chosen to vary on the level of difficulty to learn: the pair of numbers (0, 1) is easier to classify than (7, 1) and (8, 9) (based on the lowest loss achievable with linear regression in App. K.4). We computed the engagement coefficients \(\psi_{\tau}(t)\), one per dataset, that maximizes the expected return in Eq. 2. Learning curves and the evolution of engagement coefficients are depicted in Fig. 4; the baseline case corresponds to simultaneous training on all datasets at the same time (\(\psi_{\tau}(t)=1\) and \(C(\psi(t))=0\)). In the _attentive engagement_ agent, where \(\mu_{\psi}(t)=1\) (shown in Fig. 4a and 4b), the agent just needs to pay a cost to either amplify or suppress engagement on a dataset. In this setting, the agents amplify the engagement of all of the datasets, effectively increasing the learning rate per dataset, and achieving a lower \(\mathcal{L}_{C}(t)\) compared to \(\mathcal{L}_{B}(t)\). The order of learning each of the datasets goes from easier to harder, it is in the same order as in the _active engagement_ and _vector engagement_, and none of the datasets are engaged with \(\psi_{\tau}(t)<1\), presumably to avoid forgetting of early amplified datasets. In the case of _active engagement_, where \(\mu_{\psi}=0\) in \(C(\psi(t))=\beta\|\mu_{\psi}-\bar{\psi}(t)\|^{2}\) (shown in Fig. 4c and 4d), the agent must pay a cost to learn any of the tasks (\(\psi_{\tau}(t)>0\)). By distributing the learning between the tasks, the agent is capable of reaching \(\mathcal{L}_{C}(t)\) close to \(\mathcal{L}_{B}(t)\) as shown in the top panel of Fig. 4, without the need of fully engaging on all of the datasets at every time step. None of these datasets are fully disengaged at any point, possibly as a mechanisms to avoid catastrophic forgetting (Kirkpatrick et al., 2017) of datasets previously engaged during training. The engagement coefficients in the _vector_ case behave similarly. Since the control signal in this case is forced to keep a constant size of \(\Psi\), the agent is not able to fully engage in all of the datasets, and distributes this _attention_ resource on each dataset from easier to harder, as in the _active_ case. The meta-learning strategy found in our setting of keep re-visiting previous tasks to keep performance is well studied in psychology (Ericsson and Harwell, 2019; Eglington and Pavlik Jr, 2020), and it is also related to memory _replay_ theories as a value-based mechanism that avoids catastrophic forgetting (Mattar and Daw, 2018; Agrawal et al., 2022).

**Category engagement**: In some classification tasks, it might be better to learn some categories first and others later during training. We trained the engagement modulation model to control _engagement_ or _attention_ on categories of a classification dataset. In Fig. 4 we show the results of this model trained on the Semantic dataset, and MNIST dataset classifying all digits. The engagement coefficients \(\phi_{c}(t)\) describe the focus on class \(c\) in the classification problem, which basically scales the error signal for that specific class through training (see App. H.3). Fig. 4h and 4j show the improvement in the loss when optimizing for categoric assimilation coefficients for both datasets. Fig. 4i and 4k depict the engagement coefficients per class \(\phi_{c}(t)\). In the Semantic task, the engagement coefficients are clustered depending on the level of the hierarchy for the respective output. Higher coefficients are spent on categories in higher levels of the hierarchy, as well as earlier during learning. Because we kept \(\beta\) high for this experiment (\(\beta=5.0\)), the cost of deviating from a control vector of size \(C\) is high (where \(C\) is the number of classes); therefore the amplification of engagement in some categories goes along with suppression for other categories to keep the control with constant size. For the MNIST dataset, each \(\phi_{c}(t)\) corresponds to a specific digit, and the order of assimilation that maximizes value shows a consistent order of digits among different runs, being ordered as \((0,1,7,6,4,2,3,9,8,5)\), which is roughly the same as the average linear separation per digit (see App. K.4). As in the task engagement results, we found that it is optimal to assimilate easier elements first, allocating higher \(\phi_{c}(t)\) and more concentrated in the early stages of learning. More difficult categories are assimilated later, allocating a smaller maximum \(\phi_{c}(t)\) compared to easier classes, but with sustained engage over time. The benefits of learning from easier to harder aspects of tasks have been shown in cognitive science (Krueger and Dayan, 2009; Wilson et al., 2019) and machine learning (Parisi et al., 2019; Saglietti et al.,
2022; Zhang et al., 2022), and we are able to reproduce this finding in the task engagement and category engagement experiments within our normative framework. The engagement level per class amplifies the error signal of learning a particular class through time, which can be roughly controlled by modifying the proportion of classes in the batch through training. To show this, we trained the baseline network on MNIST (no control, only backpropagation), and used \(\phi_{c}(t)\) to modify the proportion of classes in the batch throughout the training (App. H.3). This gives a better curriculum than sampling each class uniformly to populate the batch, as shown in Fig. 4l and 4m.

## 6 Gain Modulation

Motivated by studies of neuromodulation (Lindsay and Miller, 2018; Ferguson and Cardin, 2020), we finally address a neuroscience inspired model (Shenhav et al., 2013, 2017) where the _learning effort control signals_\(G_{1}(t)\in\mathbb{R}^{H\times I}\) and \(G_{2}(t)\in\mathbb{R}^{O\times H}\) modulate the gain of each layers weights as \(\tilde{W}_{i}(t)=(1+G_{i}(t))\circ W_{i}(t)=\tilde{G}_{i}(t)\circ W_{i}(t)\) where \(\circ\) denotes element-wise multiplication. This control signal will modify the input-output mapping of the network to \(\hat{Y}=\tilde{W}_{2}(t)\tilde{W}_{1}(t)X\). Given the control signals, we assume the weights are learned using gradient descent, yielding the learning dynamics equations

\[\tau_{w}\frac{dW_{1}}{dt} =\left(\tilde{W}_{2}^{T}\Sigma_{xy}^{T}\right)\circ\tilde{G}_{1 }-\left(\tilde{W}_{2}^{T}\tilde{W}_{2}\tilde{W}_{1}\Sigma_{x}\right)\circ \tilde{G}_{1}-\lambda W_{1},\] \[\tau_{w}\frac{dW_{2}}{dt} =\left(\Sigma_{xy}^{T}\tilde{W}_{1}^{T}\right)\circ\tilde{G}_{2 }-\left(\tilde{W}_{2}\tilde{W}_{1}\Sigma_{x}\tilde{W}_{1}^{T}\right)\circ \tilde{G}_{2}-\lambda W_{2}. \tag{7}\]

The control signal \(G_{i}(t)\) effect is _similar_ to a time-varying learning rate, except (1) it is weight specific (i.e. with coupling between the elements of the control matrix), (2) it does not change the weight decay rate which is originally controlled by \(\lambda\) and \(\tau_{w}\), and (3) \(G_{i}(t)\) also changes the input-output mapping. Solving the learning dynamics gives \(\mathcal{P}(t)=-\left\langle L(t)\right\rangle\), using \(C(G(t))=\exp\left(\beta\left(\|G_{1}(t)\|_{F}^{2}+\|G_{2}(t)\|_{F}^{2}\right) \right)-1\), to then estimate \(dV/dG_{i}(t)\) as in App. C, and find the control trajectory that maximizes cumulative reward in Eq. 2 (More details in App. H.1, we provide an exact solution of the learning dynamics in a single-layer network given a control signal \(G(t)\) in App. G.2). In addition, we simulated this model in a non-linear network using approximations (see App. I). The meta-learning tasks used to train this model are the following:

**Effort Allocation**: We train the gain modulation model separately on each of the three datasets for a time period of \(T\), and estimate the control signal that maximizes the expected return \(V\) in Eq. 2.

**Task Switch**: We defined two different Gaussian datasets (App. L). We sequentially train the network on each dataset for a time period \(T_{s}\). The expected reward \(V\) is computed for the whole training period \(T>T_{s}\) of the gain modulation model, and maximized through gradient updates on \(G_{i}(t)\).

### Results

**Effort Allocation**: This setting is similar to the single neuron setting of Section 2.1, but with a two layers network instead of just one neuron, where every weight in the network has its own gain signal as described in Eq. 7 and schematized in Fig. 1e. The results of the baseline training and controlled training using gain modulation are presented in Fig. 5. In the gain modulation model, we can see the same qualitative behavior as in the single-neuron model when varying parameters of the learning model and control optimization. The control signal that maximizes expected return reduces the instant net reward rate by the use of control in the early stages of learning, to get better performance later as shown in the lower loss for the controlled case (Fig. 5a and b). Through both optimizing the learning and minimizing \(C(G(t))\) at the same time, the gain modulation is not only more efficient by getting a more sparse solution (\(L_{1}\) norm in Fig. 5b), using fewer weights than when no control is used, it also learns faster (Fig. 5c, more details in App. K.2). There are times during learning when it is more effective to apply control. As can be seen in the \(L_{2}\) norm of the control matrices \(G_{1}(t)\) and \(G_{2}(t)\), and the absolute value of the time derivative of the loss \(d_{t}\mathcal{L}(t)=|d\mathcal{L}(t)/dt|\) for the baseline and control case (Fig. 5d), the control signal is larger early in training and near the stages of learning when
the increase in performance (\(d_{t}\mathcal{L}(t)\)) is larger (Fig. 5a). The control signal shifts the peaks in \(d_{t}\mathcal{L}(t)\) earlier in learning, leading to better performance and higher reward earlier, compensating for the momentarily increased cost of control. Similar results are obtained when training on the other two datasets (see Fig. 12 and 13 in App. K.2). Neuromodulators are known to be involved in high-level executive tasks such as engagement in learning (Shenhav et al., 2013, 2017; Lieder et al., 2018; Grossman et al., 2022), and some of them are believed to act as gain modulation (Lindsay et al., 2020; Ferguson and Cardin, 2020) (App. A). We provide a testable and tractable setting where neuromodulators could manage performing and learning tasks to maximize cumulative reward.

**Task Switch**: The task is schematized in Fig. 1g. In Fig. 5e, each peak in the loss is a task switch (every 1800 time steps), and as expected, the baseline loss \(\mathcal{L}_{B}(t)\) is higher than the loss with control \(\mathcal{L}_{C}(t)\) almost at every point throughout learning. After each switch, the control signal manages to iteratively drive the learning dynamics to places in parameter space \(W\) where each switch is less costly (Fig. 5f). Since the linear network is over-parametrized, the drive to adjust for the next task can be done without meaningfully changing the solution for the current task. The control signal starts acting before the switch (Fig. 5g) to amortize the loss peak at the time of the switch, and to speed up the approach of the weight to the solution, skipping the plateau in the loss. In addition, the sparsity of the weights is higher compared to the baseline case, the cost of using control to switch is transferred to the size of the weights, making it easier to move the effective weight \(\tilde{W}(t)\) by a large amount when changing \(G(t)\) (See App. K.3). This setting poses meta-learning and gain modulation as a neural implementation of task/context switching in real scenarios (Puigbo et al., 2020; Ben-Iwhiwhu et al., 2022).

## 7 Discussion

We present a flexible computationally tractable _learning effort_ framework to study optimal meta-learning with neural network dynamics in a variety of settings, where control signals can influence both learning and performance. Our framework optimizes control signals on a fully normative objective: discounted cumulative performance throughout learning. The aim of our framework is to aid the evaluation of possible interventions in engineered systems, and provide formal underpinnings for cognitive control theories in neuroscience (see App. B). While a limitation of this work is its use of linear network models, we study non-linear network dynamics approximations in App. I. We hope our work will contribute to a greater understanding of how agents should act to maximize their learning abilities, based on their own prospects of learning.

Figure 5: Results of the gain modulation model trained on an MNIST classification task. **(a)**: Instant net reward \(v(t)\), baseline vs controlled. **(b)**: L1 and L2 norms of the weights. **(c)**: Loss \(\mathcal{L}(t)\) throughout learning. **(d)**: normalized \(d_{t}\mathcal{L}(t)\), and normalized L2 norm of the control signal \(G_{1}(t)\) and \(G_{2}(t)\). **(e)**: Results on the task switch meta-task. Comparison of \(\mathcal{L}(t)\) for the baseline and control case. **(f)**: Values of \(\mathcal{L}(t)\) at switch times, along with the normalized cost of control \(C(t)\) at switch times (green line). **(g)**: Zoom of \(\mathcal{L}(t)\) in the top panel, along with the normalized cost of control.