Paper: Meta-Learning Strategies through Value Maximization in Neural Networks (ICLR-2024)

Reviewer comment (Experiments):
A simple two-layer linear neural network is used, so it may not account for all the effects during training, and the results may not translate to more complex NNs.

Author response:
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.

Reviewer comment (Theory):
It's likely one will need to find a solution for every novel considered intervention.

Author response:
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.

Reviewer comment (Theory):
I don’t think all interventions can be modeled within this framework, even when restricted to the context of two-layer linear neural networks.

Author response:
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.

Reviewer comment (Theory):
Will the framework be able to find tractable solutions if we consider learning rules as a control signal (e.g. [1])?

Author response:
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.

Reviewer comment (Theory):
To what extent can we generalize control signals in a two-layer linear neural network while still maintaining solution tractability?

Author response:
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.

Reviewer comment (Experiments):
Do you think it is possible to develop task switching in your framework when one doesn't know what task will be next in advance?

Author response:
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.

Reviewer comment (Theory):
As pointed out by the authors themselves, a limitation is the assumption of linear models. Since the motivation behind the current work is to provide ways in improving the current neural networks, more analysis on non-linear systems is needed. Although very large neural networks are hard to analyze, simpler variations could be considered for non-linear settings to make this direction even more interesting.

Author response:
> W1 and W2: New experiment, use of linear networks, and extension to non-linear networks.

While the forward path of the networks used in this paper is linear, the gradient descent dynamics on the weights of multi-layer linear networks is non-linear if we think of them as a dynamical learning system. As we mentioned in Section 3, the dynamics are non-linear due to the nonlinear loss function, weight coupling and non-convex loss landscape. In other words, the optimized control signal is dealing with a non-linear system (the learning dynamics of the linear network). As also mentioned in App. A and B, a linear network’s non-linear dynamics resemble the dynamics of other more complex non-linear networks (see paragraph 3, App. A). To show this qualitative resemblance, and thus the relevance of our linear analysis, we included a new simulation of a non-linear neural network trained on the category engagement modulation experiment described in Sec. 5 (new results in K.6.2). The results for this non-linear neural network resemble the ones of the linear neural network, in terms of learning dynamics and optimal meta-learning control found. While this nonlinear experiment is feasible in this small-scale setting, we note that the computational cost increases, whereas our framework remains tractable.

We agree that the extension to non-linear networks is necessary, and we indicate related work to accomplish this in App. A and B. It is possible to just use batch training on a non-linear network to get the dynamics and optimize the control signal using our framework (as in the new experiment), but that requires storing all activations for the entire inner loop training period to do one gradient step on the control signal, as well as sacrificing mathematical tractability. On the other hand, writing closed-form ODEs available in linear networks (see referenced in App. A and B), allows for a compact description of the entire training of the network by using a few carefully chosen summary statistics, instead of saving every data point. Our goal is to show that this is a promising direction by exhibiting a large number of settings we can accommodate with this framework while keeping some of the analytical tools. We believe that this will motivate further mathematical analysis and techniques addressing the dynamics of optimal meta-learning. We included this discussion in App. K.6.2.

Reviewer comment (Experiments):
Optimization is a big challenge in neural networks. An analysis focused on the effect of different optimization techniques might be helpful when extending the work on non-linear settings, which is currently being analyzed in only simpler artificial settings.

Author response:
> W1 and W2: New experiment, use of linear networks, and extension to non-linear networks.

While the forward path of the networks used in this paper is linear, the gradient descent dynamics on the weights of multi-layer linear networks is non-linear if we think of them as a dynamical learning system. As we mentioned in Section 3, the dynamics are non-linear due to the nonlinear loss function, weight coupling and non-convex loss landscape. In other words, the optimized control signal is dealing with a non-linear system (the learning dynamics of the linear network). As also mentioned in App. A and B, a linear network’s non-linear dynamics resemble the dynamics of other more complex non-linear networks (see paragraph 3, App. A). To show this qualitative resemblance, and thus the relevance of our linear analysis, we included a new simulation of a non-linear neural network trained on the category engagement modulation experiment described in Sec. 5 (new results in K.6.2). The results for this non-linear neural network resemble the ones of the linear neural network, in terms of learning dynamics and optimal meta-learning control found. While this nonlinear experiment is feasible in this small-scale setting, we note that the computational cost increases, whereas our framework remains tractable.

We agree that the extension to non-linear networks is necessary, and we indicate related work to accomplish this in App. A and B. It is possible to just use batch training on a non-linear network to get the dynamics and optimize the control signal using our framework (as in the new experiment), but that requires storing all activations for the entire inner loop training period to do one gradient step on the control signal, as well as sacrificing mathematical tractability. On the other hand, writing closed-form ODEs available in linear networks (see referenced in App. A and B), allows for a compact description of the entire training of the network by using a few carefully chosen summary statistics, instead of saving every data point. Our goal is to show that this is a promising direction by exhibiting a large number of settings we can accommodate with this framework while keeping some of the analytical tools. We believe that this will motivate further mathematical analysis and techniques addressing the dynamics of optimal meta-learning. We included this discussion in App. K.6.2.

Reviewer comment (Presentation):
Figure 1 could be made a bit larger for better readability.

Author response:
> W4: Figure 1

Thank you, we will increase the size of the figure.

Reviewer comment (Evaluation):
Have the authors come across any unexpected results while analyzing the settings? For example, for curriculum learning where defining the curriculum itself might affect the learning in unexpected ways.

Author response:
> Q1: 

One surprise has been that we can find all of these intuitions from just maximizing value. This was not obligatory, as it could have been that certain intuitions arose from other aspects or considerations. We would also note that, while with the benefit of hindsight, the resulting control signals frequently are interpretable, before running these experiments we were often uncertain of the outcomes.

Reviewer comment (Experiments):
Is it possible to include some experiments from non-linear models (without the approximation) to address the concerns from the "weaknesses" section?

Author response:
> Q2: 

Thank you for the suggestion, we now include non-approximate results in nonlinear networks. See the answers above.