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 studies overparametrized representation learning in meta-learning, i.e. representation dimensionality (R) is larger than number of samples ($n_2$) per task. It does so in the setting of linear representations, but goes beyond existing analyses that typically assume that the covariance of task regressors to be low rank and restrict representation dimensionality to be small. The paper characterizes the optimal representation for every dimensionality R (in an asymptotic regime), observes a double descent phenomenon w.r.t. R in the performance of the representation on test tasks, also concludes that the optimal dimensionality can be overparametrized in many cases. It then analyzes the sample complexity of learning such a representation with finite train tasks samples, and finds that choosing R adaptively with number of samples can strike a balance between bias and variance of the representation. An important phenomenon of alignment of feature covariance and task covariance is introduced and its effect on sample complexity and optimal dimensionality is studied. Experiments on simulated data verify some of findings and the presence of the alignment is observed on some real datasets (MNIST, CIFAR).