Meta-Learning Strategies through Value Maximization in Neural Networks

Thank you for your review. We are glad you found our writing style excellent, and our experiments well explained. We note that our work offers a path to get intuition on meta-learning strategies, and hopefully motivate further theoretical work on this tractable setting. Please, see below our response to your comments.

Relevance for modern machine learning models and practical advice

To show relevance in more complex models, we added a new experiment using a non-linear neural network trained on the semantic task, to then optimize the category engagement control signal to maximize value using our framework. The results of this new experiment show that the learning dynamics of linear networks resemble the ones from a non-linear network, including the optimal meta-learning signal found in this non-linear network (See qEL8 for more details).

Regarding practical advice to train modern architectures, we can suggest a few things from the conclusions of our work. For example, we recover the general idea of learning rate annealing which is generally used in modern architectures (See Figure 8 in Appendix F.2). This was not forced into the meta-learning system in any way, and just comes out from finding the optimal learning rate signal to maximize value through the learning trajectory. Another example is the idea of a curriculum from easy to hard examples. Besides the links with cognitive neuroscience mentioned in App. B, there is literature on machine learning developing tools around this concept (Bengio, et al. (2009). “Curriculum learning”). In Sinha et al. 2021 “Curriculum by smoothing”, the authors blur the feature maps of a convolutional neural network (large and complex architecture) to expose the network to lower frequency components on images at the early stages of the training, while exposing the network to higher frequency components in latest stages of learning. This is a way of curriculum where lower frequency information is arguably easier, and higher frequencies are harder. Exposing the network to easy components in early stages improves learning as we also concluded using our framework. We hope more of these types of interventions can be discovered, or perhaps, fully understood in simplified settings or through mathematical analysis. We will add these examples in the extended discussion. Thank you for this question.

Thanks for your review. We appreciate you found our paper well explained and with an extensive description of the framework and experiments. The goal of our simplified setting is to give a better intuition on how the framework is implemented, we are glad you found it useful. Please, find our answers to each of your comments below:

Non-linear networks

We further justify our use of linear networks by providing an extra experiment on a real non-approximated non-linear network, showing we obtain similar results as when using a linear network (see answer to reviewer qEL8).

Relation to other meta-learning algorithms

We find the connection between our framework and standard meta-learning algorithms very important. This is why we show that our framework can instantiate MAML, and bilevel programming (App F1 and F2). This is not a comparison but a way to show that they are mathematically equivalent under certain parameter configurations. Describing this connection is a way to motivate the use of our framework to understand meta-learning solutions to other problems as well.

We discussed the applicability of scalable meta-learning algorithms to our setting in Appendix F, second paragraph. Scalable meta-learning algorithms that can be applied to state-of-the-art architectures are either simplifying the meta objective (compared to the value function computed in our framework), restricting the optimized meta-variables, or approximating some of the relevant quantities (like meta-gradients). As we mentioned in this discussion, “each step in our outer loop considers the entire inner loop learning trajectory and computes the gradient of the meta-loss at all time steps to capture the effect of the meta-parameters across the entire learning trajectory, not just the last step as in the referenced work. In addition, our meta-parameters can be as complex as the (size of the network) times (inner loop training iterations), in other words, our meta-variables also depend on time, increasing the complexity of the optimization problem. We are solving the full complex meta-learning problem (which is the desired target in both references), by considering a simpler model, instead of approximating our computation.” This can be seen explicitly in App. C where we show the form of the meta-gradient to update the control signal. Further discussion of other meta-learning algorithms we are not covering with our framework can be found in Appendix F4. In addition, we included a new experiment in Appendix F3 where we meta-learn the learning rule, further showing the capabilities of our framework to relate to other meta-learning problems.

Collapsing the two layers into one:

This is an interesting question, and the answer reflects the complexity that linear network dynamics is bringing to our model.

First, you can collapse the forward path to a single matrix, as $\hat{Y} = W_{2}W_{1}X = WX$ because the network is linear. The best solution we can find using this network is linear regression, but the path it takes from the initial weights to the solution is given by the non-linear dynamics of $W_1$ and $W_2$, so it does not follow the same trajectory as a linear regression trained using gradient descent, due to the weight coupling between $W_{1}$ and $W_{2}$. The next step we could do is to compute gradient descent updating $W$ directly. We included the derivation of the learning dynamics of $W$ under a control signal which can be written in closed form (see Appendix G2). Each $W(t)$ depends on every previous value of the control signal g, giving a non-trivial relation that does not look like a common regularization (it might regularize in some way that is not explicit, such as finding sparser solutions as shown in the results of Section 6). The reason for this is that $W(t)$ reduces the loss instantly, while g reduces the cumulated loss. This can be also seen in the computation of the gradient of the control signal in App. C. Since we are optimizing the control signal in the gradient flow regime, there is no such thing as “single samples” from the dataset in the dynamics, but rather the summarized statistics of the data as $\Sigma_{x}$, $\Sigma_{xy}$ and $\Sigma_{y}$, hence instance weighting does not apply.

Thanks for your thoughtful review of our paper. We are pleased that you found our paper is well-motivated, offers useful intuition, contains many experiments in diverse settings, and usefully highlights links to cognitive neuroscience, which is a main purpose of our paper. Our goal is to show the flexibility of the framework to account for other learning settings that go beyond the experiments we presented here. Regarding your comments, while we agree that the use of linear networks is a limitation of the current paper, it is not an in-principle limitation of the framework itself. As in other parts of deep learning theory, we start with linearity as a useful prerequisite before tackling nonlinear extensions, and emphasize that deep linear networks still yield a nonconvex loss landscape and nonlinear dynamics that bears similarity to nonlinear networks. We further support the relevance of our deep linear network analysis by running our framework on a non-linear network, showing that similar qualitative results hold for non-linear networks. We also indicate the path forward to extensions of this framework to include non-linear network dynamics.

Thank you for clarifying the concerns raised, including concerns from other reviewers. I am updating my scores and will continue updating them after proof-reading the document again to see if all comments, including from other reviewers, have been addressed.

W1: Use of linear networks and comparison to non-linear networks.

As mentioned in the answer to reviewer qEL8, the dynamics of linear networks are non-linear with a non-convex landscape, resembling the dynamics of more complex non-linear neural networks. The part that is linear is the forward path, but the control signal is dealing with essentially a non-linear dynamical system due to the nonlinear loss function and weight coupling between the first and second layers as described in Sec 3. We run an additional experiment to show this, please see the answer to reviewer qEL8.

W2: It's likely one will need to find a solution for every novel considered intervention.

We agree with your statement. However, as shown in the appendix for each of the models, the math necessary to write the differential equations describing learning as a function of the control signal, in most cases, is standard gradient computation. Your intuition is right, changing the intervention of the control signal or the learning system, will change these equations. This is part of the challenge when extending to non-linear networks. Techniques to describe learning dynamics in non-linear networks exist (see paragraph 3, appendix A), but they have their own drawbacks in terms of mathematical and computational tractability, or required limits in input or hidden dimensionality to obtain closed-form differential equations describing the dynamics. This is a direction we are currently exploring and we can definitely see a path forward, but this is beyond the scope of the paper, since the general framework (Sec 2) presented here is the same.

W3, Q1 and Q2: Other interventions and learning rules as a control signal

As you mentioned in your review, the experiments presented in this work are vast, and due to the flexibility of our framework, we do not believe it is possible to cover all possible interventions in a single manuscript. The tractability of each control intervention will largely depend on the learning system, data distribution, and degrees of freedom of the control signal. This is all a design choice, for example, in App. H.1.1 we discussed the possibility of using a low-rank basis for the control signal, which might simplify computation and search space. We acknowledge that not all meta-learning phenomena can be instantiated with our framework. There is a literature in meta-learning that we are not covering with our framework, where there is no explicit variable that is meta-optimized. We discussed this in App. F4.

Thank you for the suggestion, using learning rules as control signals is possible, and we now include an experiment demonstrating one approach. In particular, we have considered the setting of “Cao, Summerfield, and Saxe et al. 2020, NeurIPS”, which describes a space of possible learning rules using two parameters, spanning gradient descent, Hebbian contrastive learning, quasi-predictive coding, Hebbian and anti-Hebbian rules. These rules instantiate several popular ideas in theoretical neuroscience. Concretely, for the deep linear network setting, the weight updates are

$\Delta W_{1} = \Delta W_{1}^{C}(\gamma) + \Delta W_{1}^{H}(\eta), \hspace{0.2cm} \Delta W_{2} = \Delta W_{2}^{C}(\gamma),$

with $\Delta W_{i}^{C}$ being a contrastive learning rule, and $\Delta W_{i}^{H}$ a Hebbian learning rule, which are controlled by the parameters $\gamma$ and $\eta$ respectively. Using our framework, we can find the optimal $\gamma(t)$ and $\eta(t)$ that maximizes value as described in Sec 2. By inferring these parameters, the system is deciding which learning rule to employ within this space over learning. We found that the optimal learning rule is consistent with the original Cao et al. 2020 paper. We provide a new experiment in App F3 which will be referenced in the main text. These methods could offer a route to examining how a learning agent might adapt its learning algorithm to better learn different tasks, for instance. We believe this additional experiment further shows the flexibility of our framework to accommodate different meta-learning settings.

W3: Switching without knowing the next task in advance

This is an important future direction. We note that in many cognitive science experiments, task switches happen at predictable times (after every 50 trials, say), and our analysis is directly applicable to this setting. To consider the case where task switches are unpredictable would require extensions: fundamentally, one could compute the expected loss over possible next tasks, either analytically in certain cases or through sampling. In essence, while we roll out a single learning trajectory corresponding to a predictable sequence of tasks, a sampling approach would roll out many trajectories with possible next tasks, and construct the loss function as the sum of these tasks. The control signal would then be optimized to perform well for the distribution of switches.

Thanks for your review. We are glad you found the experimental section vast and well explained. As we mentioned in other responses, our framework can accommodate a large set of interventions, mostly in the context of cognitive neuroscience, but we believe it might help the study of meta-learning strategies in machine-learning models as well. In this work, we have shown at least 10 (+ 2 more in the revised paper) different combinations of learning settings and interventions, describing them thoroughly. Our goal is to show that many more interventions can be studied using this framework, for example, the one you suggest on learning the learning rules [1]. We have added an experiment in this direction. We describe more details below.