Towards Sample-efficient Overparameterized Meta-learning
Year: 2022

Abstract:
An overarching goal in machine learning is to build a generalizable model with a small number of samples. To this end, overparameterization has been the subject of immense interest to explain the generalization ability of deep nets even when the size of the dataset is smaller than that of the model. While prior literature focuses on the classical supervised setting, this paper aims to demystify overparameterization for meta-learning. Here we have a sequence of linear-regression tasks and we ask: (1) Given earlier tasks, what is the optimal linear representation of features for a new downstream task? and (2) How many samples do we need to build this representation? This work shows that surprisingly, overparameterization arises as a natural answer to these fundamental meta-learning questions. Specifically, for (1), we first show that learning the optimal representation coincides with the problem of designing a task-aware regularization to promote inductive bias. This inductive bias explains how the downstream task actually benefits from overparameterization, in contrast to prior works on few-shot learning. For (2), we develop a theory to explain how feature covariance can implicitly help reduce the sample complexity well below the degrees of freedom and lead to small estimation error. We then integrate these findings to obtain an overall performance guarantee for our metalearning algorithm. Numerical experiments on real and synthetic data verify our insights on overparameterized meta-learning.

heading: 1 Introduction
text: 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.

heading: 1.1 Our contributions
text: 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).

heading: 1.2 Related work
text: 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].

heading: 2 Problem Setup
text: 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).

heading: 2.1 Data generation
text: Definition 1 (Task and feature distributions) Throughout, DT and DF denote the distributions of tasks βi and features xij respectively. These distributions are subGaussian, zero-mean with corresponding covariance matrices ΣT and ΣF .
Definition 2 (Data distribution for a single task) Given a specific realization of task vector β ∼ DT , the corresponding label/input distribution (y,x) ∼ Dβ is obtained via y = x>β + ε where x ∼ DF and ε is zero-mean subgaussian noise with variance σ2.
Data for Representation Learning (Phase 1). We have T tasks, each with n1 training examples. The task vectors (βi)Ti=1 ⊂ Rd are drawn i.i.d. from the distribution DT . The data for ith task is given by (yij ,xij)n1j=1 i.i.d.∼ Dβi . In total, there are N = T × n1 examples.
Data for Few-Shot Learning (Phase 2). Sample task β? ∼ DT . Few-shot dataset has n2 examples (yi,xi) n2 j=1
i.i.d.∼ Dβ? . We use representation learning data to learn a representation of feature-task distribution, called eigen-weighting matrix Λ in Def. 3 below. The matrix Λ is passed to few-shot learning stage, helping learn β? with few data.

heading: 2.2 Training in Phase 2
text: We will define a weighted representation, called eigen-weighting matrix, and show how it is applied for few-shot learning. The matrix is learned during representation learning using the data from the T tasks. DenoteX ∈ Rn2×d whose ith row is xi, and y = [y1, ..., ym]>. We are interested in studying the weighted 2-norm interpolator defined below for overparameterization regime R ≥ n2.
Definition 3 (Eigen-weighting matrix and Weighted `2-norm interpolator) Let the representation dimension be R, where R is any integer between 1 and d. We define an eigen-weighting matrix Λ ∈ Rd×R and the associated weighted `2-norm interpolator
β̂Λ = arg min β ‖Λ†β‖2 s.t. y = Xβ and β ∈ range_space(Λ).
The solution is equivalent to defining α̂Λ = Λ†β̂Λ and solving an unweighted minimum 2-norm regression with featuresXΛ. This corresponds to our few-shot learning problem
α̂Λ = arg min α ‖α‖2 s.t. y = XΛα
from which we obtain β̂Λ = Λα̂Λ. When there is no confusion, we can replace β̂Λ with β̂. One can easily see that β̂ = Λ(XΛ)†y. We note that Definition 3 is a special case of the weighted ridge regression discussed in [38], as stated in Observation 1. An alternative equivalence between min-norm interpolation and ridge regression can be found in [32].
Observation 1 LetX ∈ Rn2×d and y ∈ Rn2 , define β̂1 = lim
t→0 argminβ‖Xβ − y‖22 + tβ>(ΛΛ>)†β, β ∈ column space of Λ. (2.1)
We have that β̂1 = β̂.

heading: 3 Canonical Covariance and Optimal Representation
text: In this section, we ask the simpler question: if the covariances ΣT and ΣF are known, what is the best choice of Λ to minimize the risk of the interpolator from Definition 3? In general, the covariances are not known; however, the insights from this section help us study the more general case in Section 5. Define the risk as the expected error of inferring the label on the few-shot dataset,
risk(Λ,ΣT ,ΣF ) = Ex,y,β(y − x>β̂Λ)2 = Eβ(β̂Λ − β)>ΣF (β̂Λ − β) + σ2. (3.1) The natural choice of optimization for choosing Λ would be to choose the weighting that minimizes the eventual risk of the learned interpolator.
Λ∗ = arg min Λ′∈Rd×R risk(Λ′,ΣT ,ΣF ) (3.2)
Since the label y is bilinear in x and β, we introduce whitened features x̃ = Σ−1/2F x and associated task vector β̃ = Σ1/2F β. This change of variables ensures x
Tβ = x̃T β̃; now, the task covariance in the transformed coordinates takes the form
Σ̃T = Σ 1/2 F ΣTΣ 1/2 F ,
which we call the canonical task covariance; it captures the joint behavior of feature and task covariances ΣF ,ΣT . Below, we observe that the risk in Equation (3.1) is invariant to the change of co-ordinates that we have described above i.e it does not change when Σ1/2F ΣTΣ 1/2 F is fixed and we vary ΣF and ΣT .
Algorithm 1 Constructing the optimal representation
Require: Projection dimension R, noise level σ, canonical covariance Σ̃T , task covariance ΣF .
1: function COMPUTEOPTIMALREP(R,ΣF , Σ̃T , σ, n2) 2: U1,ΣRF , Σ̃ R T , σR = COMPUTEREDUCTION(R,ΣF , Σ̃T , σ) 3: Optimization: Get θ∗ from (OPT-REP). 4: Map to eigenvalues: Set diagonal Λ∗R ∈ RR×R with entries Λ∗R,i = (1/θ∗i − 1)−2. 5: Lifting and feature whitening: Λ∗ ← U1(ΣRF )−1/2Λ∗R. 6: return Λ∗
7: function COMPUTEREDUCTION(R,ΣF ,Σ̃T , σ) 8: Get eigen-decomposition Σ̃T = UΣU>. 9: Principal eigenspace U1 ∈ Rd×R = the first R columns of U .
10: Top eigenvalues: Set Σ̃RT = U > 1 Σ̃TU1,Σ R F = U > 1 ΣFU1 11: Equivalent noise level: σ2R ← σ2 + tr(Σ̃T )− tr(Σ̃RT ). 12: return U1,ΣRF , Σ̃RT , σR
Observation 2 (Equivalence to problem with whitened features) Let data be generated as in Phase 1. Denote Σ̃T = Σ 1/2 F ΣTΣ 1/2 F . Then risk(Σ −1/2 F Λ,ΣT ,ΣF ) = risk(Λ, Σ̃T , I).
This observation can be easily verified by substituting the change-of-coordinates into Equation (3.1) and evaluating the risk.
The risk in (3.1) quantifies the quality of representation Λ; however it is not a manageable function of Λ that can be straightforwardly optimized. In this subsection, we show that it is asymptotically equivalent to a different optimization problem, which can be easily solved by analyzing KKT optimality conditions. Theorem 1 characterizes this equivalence; the COMPUTEREDUCTION subroutine of Algorithm 1 calculates key quantities that are used in specifying the reduction, and the COMPUTEOPTIMALREP subroutine of Algorithm 1 uses the solution of the simpler problem to obtain a solution for the original. Assumption 1 (Bounded feature covariance) There exist positive constants Σmin, Σmax such that ΣF is lower/upper bounded as follows: 0 ≺ ΣminI ΣF ΣmaxI .
Assumption 2 (Joint diagonalizability) ΣF and ΣT are diagonal matrices.2
Assumption 3 (Double asymptotic regime) We let the dimensions and the sample size grow as d,R, n2 →∞ at fixed ratios κ̄ := d/n2 and κ := R/n2.
Assumption 4 The joint empirical distribution of the eigenvalues of ΛR and Σ̃RT is given by the average of Dirac δ’s: 1R ∑R i=1 δΛR,i, √ RΣ̃RT,i . It converges to a fixed distribution as d→∞.
With these assumptions, we can derive an analytical expression to quantify the risk of a representation Λ. We will then optimize this analytic expression to obtain a formula for the optimal representation.
Theorem 1 (Asymptotic risk equivalence) Suppose Assumptions 1, 2, 3, 4 hold. Let ξ > 0 be the unique number obeying n2 = ∑R i=1 ( 1 + (ξΛ2i ) −1)−1. Define θ ∈ RR with entries θi = ξΛ2i1+ξΛ2i and calculate Σ̃RT , σR using the COMPUTEREDUCTION procedure of Algorithm 1. Then, define the analytic risk formula
f(θ, Σ̃RT , n2) = 1
n2 − ‖θ‖22
( n2
R∑ i=1 (1− θi)2Σ̃RT,i + (‖θ‖22 + 1)σ2R
) . (3.3)
We have that lim
n2→∞ f(θ, Σ̃RT , n2) = lim n2→∞ risk(Σ−1/2F Λ,ΣT ,ΣF ) (3.4)
The proof of Theorem 1 applies the convex Gaussian Min-max Theorem (CGMT) in [35] and can be found in the Appendix B.2.We show that as dimension grows, the distribution of the estimator β̂ converges to a Gaussian distribution and we can calculate the expectation of risk.
2This is equivalent to the more general scenario where ΣF and ΣT are jointly diagonalizable.
Theorem 1 provides us with a closed-form risk for any linear representation. Now, one can solve for the optimal representation by computing (OPT-REP) below. In order to do this, we propose an algorithm for the optimization problem in Appendix B.5 via a study of the KKT conditions for the problem 3.
θ∗ = arg min θ f(θ,ΣT ,ΣF ), s.t. 0 ≤ θ < 1, R∑ i=1 θi = n2 (OPT-REP)
The optimal representation is4 Λ∗R,i = ((1/θ ∗ i − 1)ξ)−2. The subroutine COMPUTEOPTIMALREP in Algorithm 1 summarizes this procedure.
Remark 1 Thm. 1 states that risk(Σ−1/2F Λ,ΣT ,ΣF ) can be arbitrarily well-approximated by f(θ, Σ̃RT , n2) if n2 is sufficiently large. In Fig. 1(b), we set ΣF = I100, ΣT = diag(I20, 0.1I80), n2 = 40. The curves in Fig1(b) are the finite dimensional approximation of f (LHS of (3.4)); the dots are empirical approximations of the risk (RHS of (3.4)). We tested two cases when Λ is the optimal eigen-weighting or projection matrix with no weighting. Our theorem is corroborated by the observation that the dots and curves are visibly very close. The approximation is already accurate for the finite dimensional problem with just n2 = 40.
The benefit of overparameterization. Theorem 1 leads to an optimal eigen-weighting strategy via asymptotic analysis. In Figure 3, we plot the effect on the risk of increasing R for different shapes of task covariance; the parameter ι controls how spiked ΣT is, with a smaller value for ι indicating increased spikedness. For the underparameterized problem, the weighting does not have any impact on the risk. In the overparameterized regime, the eigen-weighted learner achieves lower few-shot error than its unweighted (Λ = I) counterpart, showing that eigen-weighting becomes critical.
The eigen-weighting procedure can introduce inductive bias during few-shot learning, and helps explain how optimal representation minimizing the few-shot risk can be overparameterized with R n2. We note that, an R dimensional representation can be recovered by a d dimensional representation matrix of rank R, thus the underparameterized case can never beat d dimensional case in theory. The error with optimal eigen-weighting in overparameterized regime is smaller than the respective underparameterized counterpart. The error is lower with smaller ι. It implies that, while Σ̃T gets closer to low-rank, the excess error caused by choosing small dimension R (equal to the gap σ2R − σ2 in Algo 1) is not as significant.
Low dimensional representations zero out features and cause bias. By contrast, when Σ̃T ∈ Rd×d is not low rank, every feature contributes to learning with the importance of the features reflected by the weights. This viewpoint is in similar spirit to that of [19] where the authors devise a misspecified linear regression to demonstrate the benefits of overparameterization. Our algorithm allows arbitrary representation dimension R and eigen-weighting.

heading: 4 Representation Learning
text: In this section, we will show how to estimate the useful distribution in representation learning phase that enables us to calculate eigen-weighting matrix Λ∗. Note that Λ∗ depends on the canonical covariance Σ̃T = Σ 1/2 F ΣTΣ 1/2 F . Learning the R-dimensional principal subspace of Σ̃T enables us 5 to calculate Λ∗. Denote this subspace by S̃T .
3In Sec. 5 the constraint is θ ≤ θ ≤ 1− d−n2 n2
θ for robustness concerns. 4In the algorithm, ξ = 1 and ΛR,i = (1/θ∗i − 1)−2, because cΛ∗ for any constant c gives the same β̂. 5We also need to estimate ΣF for whitening. Estimating ΣF is rather easy and incurs smaller error compared
to Σ̃T . The analysis is provided in the first part of Appendix B.
Subspace estimation vs. inductive bias. The subspace-based representation S̃T has degrees of freedom= Rd. When Σ̃T is exactly rank R and features are whitened, [36] provides a samplecomplexity lower bound of Ω(Rd) examples and gives an algorithm achieving O(R2d) samples. However, in practice, deep nets learn good representations despite overparameterization. In this section, recalling our Q2, we argue that the inductive bias of the feature distribution can implicitly accelerate learning the canonical covariance. This differentiates our results from most prior works such as [22, 21, 36] in two aspects:
1. Rather than focusing on a low dimensional subspace and assuming N & Rd, we can estimate Σ̃T or S̃T in the overparameterized regime N . Rd. 2. Rather than assuming whitened features ΣF = I and achieving a sample complexity of R2d, our learning guarantee holds for arbitrary covariance matrices ΣF ,ΣT . The sample complexity depends on effective rank and can be arbitrarily smaller than DoF. We showcase our bounds via a spiked covariance setting in Example 1 below.
For learning Σ̃T or its subspace S̃T , we investigate the method-of-moments (MoM) estimator. Definition 4 (MoM Estimator) For 1 ≤ i ≤ T , define b̂i,1 = 2n−11 ∑n1/2 j=1 yijxij , b̂i,2 =
2n−11 ∑n1 j=n1/2+1 yijxij . Set
M̂ = n−11 T∑ i=1 (bi,1b > i,2 + bi,2b > i,1),
The expectation of M̂ is equal toM = ΣFΣTΣF .
Inductive bias in representation learning: Recall that canonical covariance Σ̃T = Σ 1/2 F ΣTΣ 1/2 F is the attribute of interest. However, feature covariance Σ1/2F term implicitly modulates the estimation procedure because the population MoM is not Σ̃T but M = Σ 1/2 F Σ̃TΣ 1/2 F . For instance, when estimating the principle canonical subspace S̃T , the degree of alignment between ΣF and Σ̃T can make or break the estimation procedure: If ΣF and Σ̃T have well-aligned principal subspaces, S̃T will be easier to estimate since ΣF will amplify the S̃T direction withinM .
We verify the inductive bias on practical image dataset, reported in Appendix A. We assessed correlation coefficient between covariances Σ̃T ,ΣF via the canonical-feature alignment score defined as the correlation coefficient
ρ(ΣF , Σ̃T ) :=
〈 ΣF , Σ̃T 〉 ‖ΣF ‖F ‖Σ̃T ‖F = trace(M) ‖ΣF ‖F ‖Σ̃T ‖F .
Observe that, the MoM estimator M naturally shows up in the alignment definition because the inner product of Σ̃T ,ΣF is equal to trace(M). This further supports our inductive bias intuition. As reference, we compared it to canonical-identity alignment defined as trace(Σ̃T )√
d‖Σ̃T ‖F (replacing ΣF with
I). The canonical-feature alignment score is higher than the canonical-identity alignment score. This significant score difference exemplifies how ΣF and Σ̃T can synergistically align with each other (inductive bias). This alignment helps our MoM estimator defined below, illustrated by Example 1 (spiked covariance).
In the following subsections, let N = n1T refer to the total tasks in representation-learning phase. Let rF = tr(ΣF ), rT = tr(ΣT ), and r̃T = tr(Σ̃T ). Define the approximate low-rankness measure of feature covariance by6
sF = min s ′ F , s.t. s ′ F ∈ {1, ..., d}, s′F /d ≥ λs′F+1(ΣF )
We have two results for this estimator.
1. Generally, we can estimateM with O(rF r̃2T ) samples. 2. Let n1 ≥ sT , we can estimateM with O(sF r̃T ) samples. Paper [36] has sample complexity O(dr2) (r is exact rank). Our sample complexity is O(rF r̃2T ). rF , r̃T can be seen as effective ranks and our bounds are always smaller than [36]. We will discuss later in Example 1. Our second result says when n1 ≥ sT , our sample complexity achieves theO(dr) which is proven a lower bound in [36].
6The (sF + 1)-th eigenvalue is smaller than sF /d. Note the top eigenvalue is 1.
Theorem 2 Let data be generated as in Phase 1. Assume ‖ΣF ‖, ‖ΣT ‖ = 1 for normalization7.

heading: 1. Let n1 be a even number. Then with probability at least 1−N−100,
text: ‖M̂ −M‖ . (r̃T + σ2) √ rF N + √ rT T .
2. Assume T ≥ sF . If n1 & r̃T + σ2, then with probability at least 1− CT−100
‖M̂ −M‖ . ( (r̃T + σ 2)/n1 )1/2 .
Denote the top-R principal subspaces ofM ,M̂ byMtop,M̂top and assume the eigen-gap condition λR(M) − λR+1(M) > 2‖M̂ −M‖. Then a direct application of Davis-Kahan Theorem [13] bounds the subspace angle as follows
angle(Mtop,M̂top) . ‖M̂ −M‖/(λR(M)− λR+1(M)).
Estimating eigenspace of canonical covariance. Note that if ΣF and ΣT are aligned, (e.g. Example 1 below with sF = sT = R), then Mtop = S̃T is exactly the principal subspace of Σ̃T . Theorem 2 indeed gives estimation error for the principal subspace of Σ̃T . Note that, such alignment is and more general requirement compared to related works which require whitened features [36, 22].
Example 1 (Spiked Σ̃T , Aligned principal subspaces) Suppose the spectra of ΣF and Σ̃T are bimodal as follows ΣF = diag(IsF , ιF Id−sF ), ΣT = diag(IsT , ιT Id−sT ). Set statistical error ErrT,N := √ r2T rF /N + √
rT /T . When ιT , ιF < 1, sF ≥ sT , the recovery error of Σ̃T and its principal subspace S̃T are bounded as
angle(M̂top, S̃T ) . ErrT,N + ι2F ιT and ‖M̂ − Σ̃T ‖ . ErrT,N + ιF ιT .
The estimation errors for Σ̃T , S̃T are controlled in terms of the effective ranks and the spectrum tails ιF , ιT . Typically sF sT & n1 so √ r2T rF /N term dominates the statistical error in practice. In Fig. 4 we plot the error of estimating M (whose principal subspace coincides with Σ̃T ). ΣF = diag(I30, ιI70), ΣT = diag(I30,070). T = N = 100. We can see that the error increase with ι .

heading: 5 Robustness of Optimal Representation and Overall Meta-Learning Bound
text: In Section 3, we described the algorithm for computing the optimal representation with known distributions of features and tasks. In Section 4, we proposed the MoM estimator in representation learning phase to estimate the unknown covariance matrices. In this section, we study the algorithm’s behaviors when we calculate Λ using the estimated canonical covariance, rather than the fullinformation setting of Section 3.
7This is simply equivalent to scaling yij , which does not affect the normalized error ‖M̂ −M‖/‖M‖. In the appendix we define S = max{‖ΣF ‖, ‖ΣT ‖} and prove the theorem for general S.
Armed with the provably reliable estimators of Section 4, we can replace Σ̃T and ΣF in Algorithm 1 with our estimators. In this section, we inquire: how does the estimation error in covariance-estimation in representation learning stage affect the downstream few-shot learning risk? That says, we are interested in8 risk(Λ,ΣT ,ΣF )− risk(Λ∗,ΣT ,ΣF ).
Let us replace the constraint in (OPT-REP) by θ ≤ θ ≤ 1− d−n2n2 θ. This changes the “optimization" step in Algorithm 1. Theorem 3 does not require an explicit computation of the optimal representation by enforcing θ. Instead, we use the robustness of such a representation (due to its well-conditioned nature) to deduce its stability. That said, for practical computation of optimal representation, we simply use Algorithm 1. We can then evaluate θ after-the-fact as the minimum singular value of this representation to apply Theorem 3 without assuming an explicit θ.
Let Λθ(R) = COMPUTEOPTIMALREP(R,ΣF ,M̂ , σ, n2) denote the estimated optimal representation and Λ∗θ(R) = COMPUTEOPTIMALREP(R,ΣF , Σ̃T , σ, n2) denote the true optimal representation, which cannot be accessed in practice. Below we present the bound of the whole meta-learning algorithm. It shows that a bounded error in representation learning leads to a bounded increase on the downstream few-shot learning risk, thus quantifying the robustness of few-shot learning to errors in covariance estimates.
Theorem 3 Let Λθ(R), Λ∗θ(R) be as defined above, and rF = tr(ΣF ), rT = tr(ΣT ), r̃T = tr(Σ̃T ). The risk of meta-learning algorithm satisfies9
risk(Λθ(R),ΣT ,ΣF )− risk(Λ∗θ(R),ΣT ,ΣF ) . n22
d(R− n2)(2n2 −Rθ)θ
[ (r̃T + σ 2) √ rF N + √ rT T ] .
Notice that as the number of previous tasks T and total representation-learning samples N observed increases, the risk of the estimated Λθ(R) approaches that of the optimal Λ∗θ(R) as we expect. The result only applies to the overparameterized regime of interest R > n2. The expression of risk in the underparameterized case is different, and covered by the second case of Equation(4.4) in [38]. We plot it in Fig 1(b) on the left side of the peak as a comparison.
Risk with respect to PCA level R. In Fig. 5, we plot the error of the whole meta-learning algorithm. We simulate representation learning and get M̂ , use it to compute Λ and plot the theoretical downstream risk (experiments match, see Fig. 1 (b)). Mainly, we compare the behavior of Theorem 3 with differentR. WhenR grows, we search Λ in a larger space. The optimal Λ in a feasible subset is
always no better than searching in a larger space, thus the risk decreases with R increasing. At the same time, representation learning error increases with R since we need to fit a matrix in a larger space. In essence, this result provides a theoretical justification on a sweet-spot for the optimal representation. d = R is optimal when N = ∞, i.e., representation learning error is 0. As N decreases, there is a tradeoff between learning error and truncating small eigenvalues. Thus choosing R adaptively with N can strike the right bias-variance tradeoff between the excess risk (variance) and the risk due to suboptimal representation.

heading: 6 Conclusion
text: In this paper, we study the sample efficiency of meta-learning with linear representations. We show that the optimal representation is typically overparameterized and outperforms subspace-based representations for general data distributions. We refine the sample complexity analysis for learning arbitrary distributions and show the importance of inductive bias of feature and task. Finally we provide an end-to-end bound for the meta-learning algorithm showing the tradeoff of choosing larger representation dimension v.s. robustness against representation learning error.
8Note that Sec.6 of [38] gives the exact value of risk(Λ∗,ΣT ,ΣF ) so we have an end to end error guarantee. 9The bracketed expression applies first conclusion of Theorem 3. One can plug in the second as well.

heading: Acknowledgements
text: This work is supported in part by the NSF TRIPODS II grant DMS 2023166, NSF TRIPODS CCF 1740551, NSF CCF-2046816, NSF CCF 2007036, Army Research Office grant W911NF-21-1-0312, and the Moorthy Professorship at UW ECE.

Review summaries:
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).

This paper studies the role of overparameterization in meta representation learning.  Meta-learning studies how to learn from previous tasks to adapt quickly to new ones, while overparameterization studies the settings where the number of training points is smaller than the number of features. Recent works studied the generalization properties of meta-learning a linear representation but did not focus on the overparameterized setting, where the number of training samples per task is lower than the dimension of the meta-learned representation.  This work presents asymptotic results on the expected risk in the overparameterized regime. Furthermore they propose an algorithm to learn a linear representation together with sample complexity bounds. They also provide some empirical results validating their analysis. 

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$. 

This paper presents an analysis of the role of representation width when learning a shared feature extractor across multiple tasks. The authors argue that wider representations improve transfer, even when only a small number of labelled examples are accessible to fine-tune the last linear layer of a neural network. Their theoretical analysis focuses on linear regression with a linear feature extractor, with most experiments also on this setting.