Statistically and Computationally Efficient Linear Meta-representation Learning

1 Introduction . Common real world tasks follow a long tailed distribution where most of the tasks only have a small number of labeled examples [ 51 ] . Collecting more clean labels is often costly ( e.g. , medical imaging ) . As each task does not have enough examples to be learned in isolation under this few-shot learning scenario , meta-learning attempts to jointly learn across a large number of tasks to exploit some structural similarities among those tasks . One popular approach is to learn a shared representation , where new arriving tasks can be solved accurately [ 45 ] . The premise is that ( i ) there is a shared low-dimensional representation fU ( x ) ∈ Rr represented by a task-independent meta-parameterU and ( ii ) a simple linear model of 〈vi , fU ( x ) 〉 can make accurate prediction on the i-th task with a task-specific parameter vi . Once the representation fU has been learnt , we can rapidly adapt to new arriving tasks as the representation dimension r is much smaller than the dimension d of the input data . This approach is becoming increasingly popular with a growing list of recent applications [ 33 , 36 , 9 , 42 , 43 , 27 , 47 , 14 , 13 , 18 ] and has been empirically shown to achieve the state-of-the-art performances on benchmark few-shot learning datasets [ 47 , 14 , 43 ] . ∗†thekump2 @ illinois.edu , ‡prajain @ google.com , ¶pnetrapalli @ google.com ( Part of the work was done while at Google Research India ) , and ±sewoong @ cs.washington.edu . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . These successes rely on a simple but effective training algorithm which alternately updates U and { vi } which we call AltMinGD ( Alternating Minimization and Gradient Descent ) . Suppose we are given t tasks , and the i-th task is associated with a dataset { ( x ( i ) j ∈ Rd , y ( i ) j ) } mj=1 of size m. In this paper , we closely follow the formulation of [ 47 ] , which solves for a function fU : Rd → Rr ( typically a deep neural network ) and a task-specific linear model vi ∈ Rr on a choice of a loss ` ( · , · ) : min U { ∑ i∈ [ t ] min vi∈Rr ∑ j∈ [ m ] ` ( 〈vi , fU ( x ( i ) j ) 〉 , y ( i ) j ) } , ( 1 ) by alternately applying a ( stochastic ) gradient descent step of U in the outer loop ( for given vi ’ s ) and numerically finding the optimal solution vi in the inner loop ( for a given U ) . Several closely related algorithms have been proposed , including separating training-set used for the inner loop and the validation-set used for the outer-loop [ 43 , 36 , 9 , 5 ] , early stopping the inner-loop [ 33 ] , applying to datasets with imbalanced data sizes [ 42 , 14 ] , and proposing new architectures and regularizers [ 27 ] . There is an increasing list [ 47 , 14 , 43 ] of numerical evidences showing that these meta representation learning improves upon competing approaches including MAML [ 20 ] and its variants [ 23 , 32 , 39 ] . Further , [ 43 ] provides experimental evidences that shared representation is the dominant component in the efficacy of MAML [ 20 ] , even though MAML does not explicitly seek a shared representation . In this paper , we analyze the computational and statistical properties of AltMinGD and its variant AltMin under the simple but canonical setting of learning a shared linear representation for linear regression tasks [ 48 ] . The fundamental question of interest is : as the number of tasks grow , does AltMinGD learn the underlying r-dimensional shared representation ( subspace ) more accurately , and consequently make more accurate predictions on new tasks ? This question is critical in explaining the empirical success in few-shot learning where the number of tasks in the training set is large while each of those tasks is data starved . Further , in settings like crowdsourcing or bioinformatics , collecting more data on new tasks is easier than collecting more data on existing tasks . Contributions . We analyze the widely adopted AltMinGD and prove a nearly optimal error rate . We show that AltMinGD requires only m = Ω ( log t+ log log ( 1/ε ) ) samples per task to achieve an error of ε in estimating the representation U when we have a large enough number t of tasks in the training data and assuming a constant dimensionality r = O ( 1 ) of the representation . Under this condition , AltMinGD achieves an error decaying as Õ ( σ √ d/mt ) , which nearly matches the fundamental lower bound . Together , these analyses imply that AltMinGD is able to compensate for having only a few samples per task ( small m ) by having many few-shot tasks ( large t ) , significantly improving the state-of-the-art ( see Table 1 ) . Note that the log log ( 1/ε ) dependence of m is hidden in the Ω̃ notation and is not explicitly visible from our main theorems or the table . A fine grained analysis showing this dependence is provided in Theorem 9 in Appendix C. We follow the proof strategy of alternating minimization algorithms for matrix sensing [ 30 , 38 ] , but there are important differences making the analysis challenging . First , the meta-learning dataset does not satisfy Restricted Isometry Property ( RIP ) central in the existing matrix sensing analysis , and hence none of the technical lemmas can be directly applied . We leverage on the task diversity property in Assumption 2 , to prove all necessary concentration bounds . Next , there is an inherent asymmetry in the problem ; we require accurate estimation of U for generalization to new arriving tasks ( which is the primary goal of meta-learning ) , but we do not necessarily require accurate estimation of vi ’ s . We exploit this to ensure accurate estimation of U with a small m. Our analysis of AltMinGD leads to a fundamental theoretical question : is the condition m = Ω ( log t ) necessary ? We introduce a variation AltMinGD-S , which at each iteration selects a subset of tasks that are well-behaved ( covering the r-dimensional subspace of current estimated U ) and uses ( the empirical risk of ) only those tasks in the update . While log t dependence is unavoidable if we require all t tasks to be well-behaved , ensuring a large fraction to be well-behaved requires smaller m. When the noise is sufficiently small with variance O ( 1/ log t ) , we show that AltMinGD-S requires only m = Ω ( log log ( 1/ε ) ) ( with no dependence on t ) to estimate the shared representation accurately . Inspired by a long line of successes in matrix completion and matrix sensing [ 30 ] , we also analyze a variation of AltMinGD that alternately applies minimization for U and { vi } updates , which we call AltMin , and prove a slightly improved guarantees at an extra computational cost of a factor of dr2 . Notations : [ n ] = { 1 , 2 , . . . , n } . ‖A‖ and ‖A‖F denote the spectral and Frobenius norms of a matrix A . 〈A , B〉 denotes the inner product . A† is the Moore-Penrose pseudoinverse . x ∼ N ( 0 , Id×d ) means that x is a d dimensional standard isotropic Gaussian random vector . Õ , Ω̃ and Θ̃ hide logarithmic terms in dimension d , rank r , tolerance ε and other problems parameters . 1.1 Related work . There is a large body of work in meta-learning since the seminal work in learning to learn [ 46 ] , inductive bias learning [ 7 ] , and multitask learning [ 12 ] . One popular approach starting from [ 28 , 6 ] is to learn a shared low-dimensional representation for a set of related tasks . This is becoming increasingly popular with empirical successes in the few-shot learning scenarios [ 33 , 36 , 9 , 42 , 43 , 27 , 47 , 14 ] . Linear representation learning . In this paper , we show that the popular AltMinGD algorithm for solving meta representation learning , achieves near-optimal error rate and sample complexity when applied to recovering linear representations , i.e . fU ( x ) = UTx . This problem has been studied in [ 3 , 44 , 40 ] and Nuclear-norm minimization approaches are proposed in [ 4 , 24 , 2 , 41 ] but they do not provide subspace/generalization error guarantees and suffer from large training time . Closest to our work are [ 48 , 35 , 34 , 19 ] which propose new algorithms with statistical guarantees . We also point out a concurrent and independent work [ 16 ] , which also analyzes AltMinGD but for a special case of the noiseless setting . Authors empirically showed that AltMinGD performs better than other baseline federated learning algorithms for neural meta-representation learning on some datasets . We compare these results with our guarantees in Section 4.1 . Competing against meta-representation learning approaches listed above are the bi-level optimization based methods . A pioneer in this direction is MAML [ 20 ] , which is analyzed under linear regression tasks in [ 15 , 21 , 8 , 22 ] . [ 15 ] and [ 21 ] identify that MAML outperforms a simple Empirical Risk Minimization ( ERM ) when tasks are heterogeneous in their respective level of difficulty . [ 8 ] shows that , perhaps surprisingly , negative learning rate is optimal for MAML applied to linear regression tasks , where zero learning rate corresponds to the standard ERM . Matrix sensing . Starting from matrix sensing and completion problems [ 11 , 37 , 30 ] , recovering a low-rank matrix from linear measurements has been a popular topic of research . Linear metalearning is a special case of matrix sensing , but with a non-standard sensing operator of the form A ( UV T ) = [ A1 ( UV T ) , . . . , Amt ( UV T ) ] where Aij ( UV T ) = 〈xije > i , UV > 〉 . This operator does not satisfy restricted isometry property in general , and existing matrix sensing results do not apply . Similar sensing operators have been studied in [ 29 , 52 ] which gives m = Ω ( d ) . We provide a significantly tightened analysis to require only m = Ω ( log t+ log log ( 1/ε ) ) . 2 Problem Formulation : Meta-learning of Shared Representation . We focus on the meta-learning problem with a shared linear representation for linear regression tasks . Let t denote the number of tasks . The i-th task is associated with m samples { ( x ( i ) j ∈ Rd , y ( i ) j ∈ R ) } mj=1 . We assume there is a common low-dimensional representation ( U∗ ) Tx of each data point x , parameterized by U∗ ∈ Rd×r where r d. The corresponding observation y is sampled by regressing over the low-dimensional representation ( U∗ ) Tx . Now , in general , learning U∗ is NP-hard [ 25 ] . Instead , similar to [ 48 ] , we study the problem in the following tractable random design setting . Assumptions 1 . Let U∗ ∈ Rd×r be an orthonormal matrix . For a task i ∈ [ t ] , with task specific parameter vector v∗ ( i ) ∈ Rr and j-th example x ( i ) j ∼ N ( 0 , Id×d ) , its observation is : y ( i ) j = 〈x ( i ) j , U ∗v∗ ( i ) 〉+ ε ( i ) j , ( 2 ) where ε ( i ) j ∼ N ( 0 , σ2 ) is the measurement noise which is independent of x ( i ) j . We denote by ṽ∗ ( i ) = U∗v∗ ( i ) the model parameter vector for each regression task in d-dimensions . We denote the matrix of these parameters as : Ṽ ∗ = U∗ ( V ∗ ) T where ( V ∗ ) T = [ v∗ ( 1 ) , . . . , v∗ ( t ) ] . The difficulty of estimating U∗ still depends on the diversity or incoherence of the tasks . Assumptions 2 . Let λ∗1 and λ∗r denote the largest and smallest eigenvalues of the task diversity matrix ( r/t ) ( V ∗ ) TV ∗ ∈ Rr×r respectively . Let κ = λ∗1/λ∗r . We say that V ∗ is µ-incoherent if max i∈ [ t ] ‖v∗ ( i ) ‖2 ≤ µλ∗r . ( 3 ) To estimate the subspace U , we minimize the empirical risk of the t tasks in the training data , over the meta-parameter U ∈ Rd×r and the task-specific model parameters V = [ v ( 1 ) , . . . , v ( t ) ] T ∈ Rt×r : L ( U , V ) = t∑ i=1 m∑ j=1 1 2 ( y ( i ) j − 〈 Uv ( i ) , x ( i ) j 〉 ) 2 . ( 4 ) The problem is non-convex due to the bi-linearity of Uand V . We are interested in the few-shot learning setting where the goal is to learn the representation accurately despite a small number of samples per task in the training data . Now , even if the representation U∗ is known a priori , we would require O ( r ) samples per task to learn the parameter v. Furthermore , information theoretically the total number of samples m · t should scale at most linearly with the data dimension d .

Summary: The paper analyses the problem of **linear** meta-representation learning, and more particularly the *alternating minimization and gradient descent* algorithm (AltMinGD) that has become empirically popular for meta-learning methods. The authors prove statistical guarantees for this algorithm.    They also propose an alternate version, AltMinGD-S, which achieves better per-task sample complexity, but is less tractable in practice because it is based on a task subset selection which relies on the linearity of the representations. Then, they show with experiments than the AltMin algorithms achieve smaller errors than the Method of Moments and similar error than the Burer-Monteiro ERM, while being easier to tune. They also study the number of iterations required, and show that the AltMin algorithms (AltMin, AltMinGD, AltMin-S) are the fastest ones.