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: **Summary**:  This paper analyses the statistical and computational properties of one particular type of meta-representation learning approaches, which alternatively optimises a shared feature extractor that produces linear representations, and solves inner task-specific linear models taking in the shared linear representations. Under the canonical settings when the feature extractor is linear, they show that the meta-representation learning algorithm can achieve near optimal error rate.  **Contributions**: - This paper analyses one family of widely adopted meta-representation learning algorithms they call AltMinGD, and prove a near optimal error rate, following the proof strategy for matrix sensing (which is still significantly different from the linear meta-representation learning problem). - This paper introduces a new variation of AltMinGD-S, motivated by their analysis. The new algorithm selects a subset of well-behaved tasks at each iteration, and requires a much smaller number of samples per task to estimate representation. - This paper also analyses a variation called AltMin (in appendix) where the outer loop feature extractor is optimised exactly, and proves slightly improved guarantees.