Paper: First-order ANIL provably learns representations despite overparametrisation (ICLR-2024)

Reviewer comment (Theory):
In the high level, the theoretical results in this paper show that overparameterization/more task-specific data provably _hurts_ generalization in their simplified meta-learning setting. I am not sure if this is consistent with the growing empirical evidence that overparameterization/more data usually _helps_ generalization in practical deep learning. Does this mismatch suggest that the linear model considered in this paper is somewhat over-simplified?

Author response:
The negative impact of using more task-specific data occurs when there are fewer samples during testing. This situation often leads to overfitting when the features lack a regularized, low-rank structure. The intuition is that the network, having been trained on tasks with a larger number of samples, may become overconfident in its adaptability. Note that this setting is somewhat atypical. Typically, if pretraining involves tasks with a large number of samples, one would expect the test tasks to also have a large number of samples. Moreover, it is important to note that this issue arises only with a single-step gradient descent at test time. As detailed in Appendix F, FO-ANIL has optimal performance in this setting when the number of test-time samples matches the number of training samples.

In more complex fine-tuning scenarios, the asymptotic low-rank behaviour that is beneficial for generalisation applies to any finite sample size. However, whether a larger number of samples is advantageous remains unclear. This question is theoretically challenging, even in the context of this simple setting.

On the other hand, a slowdown caused by overparameterisation has been proven in supervised learning [1]. Our findings seem to support the ubiquity of this type of slowdown. However, our perspective on overparameterisation is more optimistic: feature learning through meta-learning pretraining is effective, regardless of the size of the hidden dimension. This eliminates the need to estimate the correct dimension, which can be hard in real-world scenarios, and allows the use of overparameterised networks regardless of the hidden dimension. A more general result on adaptability is relevant as it supports the use of the most expressive model feasible for solving a problem. And as the reviewer already suggested, the experimental evidence from practice seems to support the idea that overparametrisation usually helps. Our work represents an initial step towards a theoretical understanding of such adaptive behaviour in the context of model-agnostic meta-learning.

[1] Xu, Weihang, and Simon Du. "Over-parameterization exponentially slows down gradient descent for learning a single neuron." The Thirty Sixth Annual Conference on Learning Theory. PMLR, 2023.

Reviewer comment (Experiments):
Will it change the learned representations if we add l2 weight decay regularization to the meta-training objective to avoid overfitting (with finite/infinite samples per task)?

Author response:
Following the reviewer's query, we have examined the impact of regularization. We must confirm this rigorously, but our preliminary understanding is that L2 weight decay regularization (applied in the outer loop) accelerates training and enhances the regularization effect. The adaptation performance remains similarly reliant on $k$. Depending on the magnitude of the L2 coefficient and the particular regime of interest—whether $m_{\mathrm{test}} > k$ or $m_{\mathrm{test}} < k$—the regularization could either hinder or aid generalization.

We reiterate the main focus in this work is to capture the learning dynamics of FO-ANIL in a well-understood, simple theoretical setting to improve our understanding of model-agnostic meta-learning. This has shown the good “implicit bias” of FO-ANIL despite overparameterisation.

Reviewer comment (Theory):
I do not think enough is done to thoroughly guide a reader through the actual theory of this work. While I appreciate that presenting theory in a page limit is difficult—and I do not propose for full derivations or anything of the sort be moved out of the appendix—there is in general little discussion or proof sketches presented. As a result, the theorems and propositions are stated and it is not clear how the statement is derived from the premise. This is better for Proposition 1 than Theorem 1, where Proposition 1 at least has some discussion in the paragraph following immediately, where it is stated that the error is decomposed into three terms and then bounded by concentration inequalities. What is two lines of writing then lead to deeper insight and more of this is needed.

Author response:
Our presentation of Theorem 1 was indeed brief and lacked necessary details, largely due to space limitations. We have now added a paragraph that details the proof strategy and explains how specific assumptions are utilised in the proof.

Reviewer comment (Presentation):
While the in-depth discussion of this work is appreciated, I am not certain that it should be prioritised over more technical detail. Particularly, the long comparison to [1] seems excessive for the main paper and could be deferred to the appendix along with much of this section. Similarly, the statement "Theorem 1 answers the conjecture of Saunshi et al. (2020)." is unhelpful because the conjecture is not stated, thus bringing no benefit to the reader.

Author response:
We have done our best to condense the discussion. The first three paragraphs are retained as each highlights a unique aspect of our work. The “Infinite tasks models” paragraph, to which the reviewer referred, remains since it emphasises an essential modelling decision which contrasts with prior works. However, we have revised the language to reflect the reviewer's suggestions. Lastly, most of the technical comparisons with [1] have been moved to the Appendix.

Reviewer comment (Writing):
Phrases such as "More importantly, the infinite samples idealisation does not reflect the initial motivation of meta-learning, which is to learn tasks with a few samples. Interesting phenomena are thus not observed in this simplified setting." are conversational and speculative for the main body of work, come across as combative, and seem to belittle prior work. The discussion should be compressed, retaining mainly the limitations component and the introduction of Burer-Monteiro (since it is used in the experiments), to make significant space for Appendix A in the main paper.

Author response:
We have done our best to condense the discussion. The first three paragraphs are retained as each highlights a unique aspect of our work. The “Infinite tasks models” paragraph, to which the reviewer referred, remains since it emphasises an essential modelling decision which contrasts with prior works. However, we have revised the language to reflect the reviewer's suggestions. Lastly, most of the technical comparisons with [1] have been moved to the Appendix.

Reviewer comment (Experiments):
I am concerned that the experiment summarised in Table 1 does not adequately assess or demonstrate the truth of Proposition 1. Since a scaling law is given in Proposition 1, it should be possible to experimentally show how tight of a bound this is, rather than just demonstrating that meta-learning can outperform ridge regression on the input features but not ridge regression on the ground-truth representation.

Author response:
We have included additional experiments in the appendix I.5. to verify the scaling described in Proposition 1. Specifically, we have conducted controlled experiments to separately assess the impact of the first term (that does not depend on $m_{\mathrm{test}}$) and the last two terms (that scale with $\sqrt{k/m_{\mathrm{test}}}$) in the bound and show that the scaling in Proposition 1 is tight.

Reviewer comment (Presentation):
I suggest including a figure summarising the setup to help aid clarity. While the setting is clear with the notation, such a figure would be helpful and could significantly improve clarity. Space could be made for this by compressing the Introduction. If inclusion in the main paper is not possible, adding the figure in the appendix would still add value.

Author response:
As suggested by the reviewer, we have added a figure that summarises the meta-learning task considered in the paper, and highlights the task distribution, its decomposition into two-layer teachers and the misspecification/overparameterisation in the student network.

Reviewer comment (Theory):
Can Equation 7 not be simplified further since $\Sigma_*$ is proportional to the identity matrix and $B_*$ is orthogonal?

Author response:
The covariance $\Sigma_\star$ could be replaced by identity, but $B_\star B_\star^\top \in \mathbb{R}^{d \times d}$ is a rank $k$ projection that cannot be further simplified. We have kept $\Sigma_\star$ in the presentation to enhance clarity, ensuring that Equation 7 is preserved with arbitrary covariances.

Reviewer comment (Experiments):
The paper clearly motivates the need to study a simplified linear, first-order, single gradient-step version of ANIL. Nevertheless, a small-scale empirical verification to the extent by which violating some of these assumptions affect the results could have further strengthened the paper.

Author response:
We have provided additional experiments in Appendix I.6. with 2 and 3-layer ReLU networks and checked the scaling of generalisation as $d$ varied. Our experiments show that the adaptation performance of FO-ANIL pretraining with ReLU networks does not scale with $d$. This is evidence supporting our main insight in ReLU networks: model-agnostic meta-learning pretraining adapts to the implicit problem complexity despite overparametrisation.

Reviewer comment (Theory):
In your theoretical analysis you assume matching depth in the model architecture. Since this is a linear model, is it possible to say something about the case where we have a mismatch in terms of the number of layers? (given that in the linear case this does not change the expressiveness of the model)

Author response:
Thanks for this great question. Firstly, in the linear case, a teacher with multiple layers is effectively equivalent to a two-layer formulation. Secondly, with students that have multiple layers, controlling the weight dynamics could be significantly more challenging. As discussed in Appendix F, one key insight persists: the global minima maintains the form $(B, w)$ form, where $B$ is now the product of initial layers up to the last layer. This structural property reveals that despite its nonconvexity, the problem has good global minima with rank deficiencies. It may suggest that our insights (regarding the low-rank regularised asymptotic limit and the dependency on $k$ during adaptation) from our paper could extend to multi-layer linear students. We leave such explorations for future work.

Reviewer comment (Novelty):
Novelty: the authors emphasize that their analysis is in the finite-sample regime, which they successfully argue is significantly more relevant than infinite-sample one. However they are not the first ones to study this setting: as they mention, Collins et al., 2022 already consider it, but Ji et al., 2022 (Theoretical Convergence of Multi-Step Model-Agnostic Meta-Learning) and Fallah et al., 2021 (Generalization of Model-Agnostic Meta-Learning Algorithms: Recurring and Unseen Tasks) also do, and provide convergence guarantees and generalization bounds for MAML, a more generic algorithm than ANIL. Neither of these works are mentioned in the paper, so I’d encourage the authors to discuss the novel insights they bring to the meta-learning community which are not already available in these works.

Author response:
Our work departs from the studies by Ji et al. (2022) and Fallah et al. (2021), which analyse the convergence of MAML for general loss functions without addressing representation learning. Ji et al. (2022) show convergence to a local minimum, and Fallah et al. (2021) study meta-learning under a strong convexity assumption (which does not apply to networks with 2 or more layers, even with linear activation). In contrast, our work focuses on feature learning, drawing inspiration from the insights of Raghu et al. (2020) (Rapid Learning or Feature Reuse? Towards Understanding the Effectiveness of MAML). In the same vein as Collins et al. (2022), we explore feature learning within a specific teacher-student framework. We highlight that, despite these methods not being explicitly designed for learning shared representations, such learning occurs if the tasks share underlying structures. We have included a new paragraph in Appendix A to further explain the connection between our work and these two prior works. We will expand on this brief literature review in order to better contrast our work with other orthogonal theoretical works on meta-learning.

The novelties of our work compared to the setting of Collins et al. (2022) are detailed thoroughly in the discussed section. Most notably, we consider a misspecified setting where the hidden layer is overparametrised. This is an important aspect that is important to capture in theory, and our work shows successful representational learning irrespective of the degree of overparametrisation. Although Collins et al. (2022) address the finite task and sample regime, it should be noted that the analysis is based on the concentration of the finite sample update around the infinite sample case. This regime, as argued in our work, does not capture the additional unlearning behaviour that is explained by our analysis. Another key improvement in our work is the initialisation, which does not require any weak alignment. In $k \ll d$ settings, the usual full rank initialisation assumed by Collins et al. (2022) fails, even when $k' = k$, indicating that their analysis might only describe the latter stages of the training process (i.e., training phase after the initial harder alignment phase is completed).

Reviewer comment (Theory):
While ANIL is typically motivated as an approximation to MAML, it much more closely resembles metric-based few-shot learning such as prototypical networks, metaoptnet, or R2D2 where the main difference is the algorithm to solve the linear head. How difficult would it be to extend your results for ANIL to these methods?

Author response:
Thank you for raising this excellent question. We have also identified a related question for future investigation: what would happen if the inner loop involved running multiple steps of gradient descent? Providing a precise answer to these questions is challenging without studying a teacher-student setting under a form of idealisation, such as infinite tasks, individually for each meta-learning algorithm. For algorithms like R2D2 and MetaOptNet, the multi-task framework could be used (one could take linearly separable classification tasks for SVM and logistic regression).

More broadly, we expect to see regularisation in FO-ANIL from the few-shot inner step to transfer to such methods with more intricate inner finetuning steps. This regularisation would then be beneficial to obtain a simple structure that is better at adaptation with fewer samples at hand. A priori, it is harder to comment on learning dynamics. One of our interesting findings is that FO-ANIL does not require any form of alignment at initialisation. An interesting related question is if this applies to algorithms with more sophisticated finetuning.

Reviewer comment (Theory):
In section 4, you discuss overfitting at the end of the “infinite task models” paragraphs. It isn’t obvious to me how unlearning the orthogonal space reduces overfitting and why there’s no need to unlearn it with infinite samples — could you expand on what is meant? I’m trying to understand how this observation may translate to the other methods mentioned above.

Author response:
In the presence of a finite number of samples, the excess risk of the corresponding linear regression is known to scale in the square root of the dimension. Even worse can happen when the number of samples is smaller than the dimension: in that case, the linear regression can overfit the training set and yields a very bad performance on new, unobserved, test data.

Unlearning the orthogonal space here means that we decrease the initial representation dimension from $k’$ to $k$, which first improves the scaling of the excess risk, but also avoids this potential bad overfitting regime when the number of samples ($m_{\mathrm{test}}$) is between $k$ and $k’$.

On the other hand, when having an infinite number of samples in linear regression, the excess risk will be exactly $0$ (i.e., we will recover exactly the model parameters). As a consequence, there is no benefit from reducing the representation dimension in that case, which is why it is not done by the algorithm.

On the pretraining dynamics, this translates in the outer step. The outer step indeed plays a regularisation role, as it uses validation data (with $m_{\mathrm{out}}$ samples per task): if the inner step tends to overfit, this overfitting, that will happen on the orthogonal subspace, will be mitigated during the outer loop. This mitigation will be done by unlearning, i.e., reducing the weight given to this orthogonal subspace.
All of this does not happen with $m_{\mathrm{in}}=\infty$, as there is no overfitting to mitigate.