Towards Sample-efficient Overparameterized Meta-learning

1 Introduction . In a multitude of machine learning ( ML ) tasks with limited data , it is crucial to build accurate models in a sample-efficient way . Constructing a simple yet informative representation of features is a critical component of learning a model that generalizes well to an unseen test set . The field of meta-learning dates back to [ 8 , 4 ] and addresses this challenge by transferring insights across distinct but related tasks . Usually , the meta-learner first ( 1 ) learns a feature-representation from previously seen tasks and then ( 2 ) uses this representation to succeed at an unseen task . The first phase is called representation learning and the second is called few-shot learning . Such information transfer between tasks is the backbone of modern transfer and multitask learning and finds ubiquitous applications in image classification [ 14 ] , machine translation [ 6 ] and reinforcement learning [ 17 ] . Recent literature in ML theory has posited that overparameterization can be beneficial to generalization in traditional single-task setups for both regression [ 28 , 38 , 3 , 32 , 29 ] and classification [ 31 , 30 ] problems . Empirical literature in deep learning suggests that overparameterization is of interest for both phases of meta-learning as well . Deep networks are stellar representation learners despite containing many more parameters than the sample size . Additionally , overparameterization 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . is observed to be beneficial in the few-shot phase for transfer-learning in Figure 1 ( a ) . A ResNet-50 network pretrained on Imagenet was utilized to obtain a representation of R features for classification on CIFAR-10 . All layers except the final ( softmax ) layer are frozen and are treated as a fixed feature-map . We then train the final layer of the network for the downstream task which yields a linear classifier on pretrained features . The figure plots the effect of increasing R on the test error on CIFAR-10 , for different choices of training size n2 . For each choice of n2 , increasing R beyond n2 is seen to reduce the test-error . These findings are corroborated by [ 17 ] ( MAML ) and [ 37 ] , who successfully use a transfer learning method that adapts a pre-trained model , with 112980 parameters , to downstream tasks with only 1-5 new training samples . In Figure 1 ( b ) , we consider a sequence of linear regression tasks and plot the few-shot error of our proposed projection and eigen-weighting based meta-learning algorithm for a fixed few-shot training size , but varying dimensionality of features . The resulting curve looks similar to Figure 1 ( a ) and suggests that the observations regarding overparameterization for meta-learning in neural networks can , to a good extent , be captured by linear models , thus motivating their detailed study . This aligns with trends in recent literature : while deep nets are nonlinear , recent advances show that linearized problems such as kernel regression ( e.g. , via neural tangent kernel [ 20 , 16 , 23 , 34 , 12 ] ) provide a good proxy to understand some of the theoretical properties of practical overparameterized deep nets . However , existing analysis of subspace-based meta-learning algorithms for both the representation learning and few-shot phases of linear models have typically focused on the classical underparameterized regime . These works ( see Paragraphs 2-3 of Sec . 1.2 ) consider the case where representation learning involves projection onto a lower-dimensional subspace . On the other hand , recent works on double descent shows that an overparameterized interpolator beats PCA-based method . to build upon these results to develop a theoretical understanding of overparameterized meta-learning1 . 1.1 Our contributions . This paper studies meta-learning when each task is a linear regression problem , similar in spirit to [ 36 , 22 ] . In the representation learning phase , the learner is provided with training data from T distinct tasks , with n1 training samples per task : using this data , it selects a matrix Λ ∈ Rd×R with arbitrary R to obtain a linear representation of features via the map x→ Λ > x . In the few-shot learning phase , the learner faces a new task with n2 training samples and aims to use the representation Λ > x to aid prediction performance . We highlight that obtaining the representation consists of two steps : first the learner projects x onto R basis directions , and then performs eigen-weighting of each of these directions , as shown in Figure 2 ( b ) . The overarching goal of this paper is to propose a scheme to use the knowledge gained from earlier tasks to choose Λ that minimizes few-shot risk . This goal enables us to engage with important questions regarding overparameterization : 1The code for this paper is in https : //github.com/sunyue93/Rep-Learning . Q1 : What should the size R and the representation Λ be to minimize risk at the few-shot phase ? Q2 : Can we learn the Rd dimensional representation Λ with N Rd samples ? The answers to the questions above will shed light on whether overparameterization is beneficial in few-shot learning and representation learning respectively . Towards this goal , we make several contributions to the finite-sample understanding of linear meta-learning , under assumptions discussed in Section 2 . Our results are obtained for a general data/task model with arbitrary task covariance Σβ and feature covariance ΣF which allows for a rich set of observations . Optimal representation for few-shot learning . As a stepping stone towards the goal of characterizing few-shot risk for different Λ , in Section 3 we first consider learning with known covariances ΣT and ΣF respectively ( Algorithm 1 ) . Compared to projection-only representations in previous works ( see Paragraphs 2-3 of Sec . 1.2 ) , our scheme applies eigen-weighting matrix Λ∗ to incentivize the optimizer to place higher weight on promising eigen-directions . This eigen-weighting procedure has been shown in the single-task case to be extremely crucial to avail the benefit of overparameterization [ 5 , 29 , 32 ] : it captures an inductive bias that promotes certain features and demotes others . We show that the importance of eigen-weighting extends to the multi-task case as well . Canonical task covariance . Our analysis in Section 3 also reveals that , the optimal subspace and representation matrix are closed-form functions of the canonical task covariance Σ̃T = Σ 1/2 F ΣTΣ 1/2 F , which captures the feature saliency by summarizing the feature and task distributions . Representation learning . In practice , task and feature covariances ( and hence the canonical covariance ) are rarely known apriori . However , we can estimate the principal subspace of the canonical task covariance Σ̃T ( which has a degree of freedom ( DoF ) of Ω ( Rd ) ) from data . In Section 4 we first present empirical evidence that feature covariance ΣF is “ positively correlated ” with Σ̃T . Then we propose an efficient algorithm based on Method-of-Moments ( MoM ) , and show that the sample complexity of representation learning is well below O ( Rd ) due to the inductive bias . Our sample complexity bound depends on interpretable quantities such as effective ranks ΣF , Σ̃T and improves over prior art ( e.g. , [ 22 , 36 ] ) , even though the prior works were specialized to low-rank Σ̃T and identity ΣF ( see Table 2 ) . End to end meta-learning guarantee . In Section 5 , we consider the generalization of Section 3 , where we have only estimates of the covariances instead of perfect knowledge . This leads to an overall meta-learning guarantee in terms of Λ∗ , N and n2 and uncovers a bias-variance tradeoff : As N decreases , it becomes more preferable to use a smaller R ( more bias , less variance ) due to inaccurate estimate of the weak eigen-directions of Σ̃T . In other words , we find that overparameterization is only beneficial for few-shot learning if the quality of representation learning is sufficiently good . This explains why , in practice , increasing the representation dimension may not help reduce few-shot risk beyond a certain point ( see Fig . 5 ) . 1.2 Related work . Overparameterized ML and double-descent The phenomenon of double-descent was first discovered by [ 5 ] . This paper and subsequent works on this topic [ 3 , 32 , 31 , 29 , 10 ] emphasize the importance of the right prior ( sometimes referred to as inductive bias or regularization ) to avail the benefits of overparameterization . However , an important question that arises is : where does this prior come from ? Our work shows that the prior can come from the insights learned from related previously-seen tasks . Section 3 extends the ideas in [ 33 , 38 ] to depict how the optimal representation described can be learned from imperfect covariance estimates as well . Theory for representation learning Recent papers [ 22 , 21 , 36 , 15 ] propose the theoretical bounds of representation learning when the tasks lie in an exactly r dimensional subspace . [ 22 , 21 , 36 ] discuss method of moment estimators and [ 36 , 15 ] discuss matrix factorized formulations . [ 36 ] shows that the number of samples that enable meaningful representation learning is O ( dr2 ) . [ 22 , 21 , 36 ] assume the features follow a standard normal distribution . We define a canonical covariance which handles arbitrary feature and task covariances . We also show that our estimator succeeds with O ( dr ) samples when n1 ∼ r , and extend the bound to general covariances with effective rank defined . Subspace-based meta learning With tasks being low rank , [ 22 , 21 , 36 , 18 , 15 ] do few-shot learning in a low dimensional space . [ 39 , 40 ] study meta-learning for linear bandits . [ 26 ] gives information theoretic lower and upper bounds . [ 7 ] proposes subspace-based methods for nonlinear problems such as classification . We investigate a representation with arbitrary dimension , specifically interested in overparameterized case and show it yields a smaller error with general task/feature covariances . Related work [ 15 ] provides results on overparameterized representation learning , but [ 15 ] requires number of samples per pre-training task to obey n1 & d , whereas our results apply as soon as n1 & 1 . Mixed Linear Regression ( MLR ) In MLR [ 41 , 24 , 11 ] , multiple linear regression are executed , similar to representation learning . The difference is that , the tasks are drawn from a finite set , and number of tasks can be larger than d and not necessarily low rank . [ 25 , 9 , 27 ] propose sample complexity bounds of representation learning for mixed linear regression . They can be combined with other structures such as binary task vectors [ 2 ] and sparse task vectors [ 1 ] . 2 Problem Setup . The problem we consider consists of two phases : 1 . Representation learning : Prior tasks are used to learn a suitable representation to process features . 2 . Few-shot learning : A new task is learned with a few samples by using the suitable representation . This section defines the key notations and describes the data generation procedure for the two phases . In summary , we study linear regression tasks , the features and tasks are generated randomly , i.i.d . from their associated distributions DT and DF , and the two phases share the same feature and task distributions.The setup is summarized in Figure 2 ( a ) .

Summary: This paper considers the representation learning for multi-task linear regression problem in the setting with arbitrary feature covariance, arbitrary task covariance, and arbitrary representation dimension. In this setting, it aims to address two questions: (1) given a set of $T$ linear regression tasks with $N$ total samples, what is the optimal representation $\Lambda \in \mathbb{R}^{d\times R}$ for a downstream few-shot linear regression task? and (2) what is the sample complexity for finding this representation? To address these questions, the authors find a closed-form solution for the optimal representation when the fine-tuning phase involves an eigen-weighting procedure, which they argue introduces an inductive bias that enables few-shot success. They show that the key to estimating this optimal representation is estimating the "canonical task covariance", which they show necessitates a number of samples growing with the effective ranks (traces) of the task and feature covariances, rather than the ambient dimension $d$. The authors also provide an overall risk bound that shows it is preferable to have smaller $R$ for smaller $N$.