Statistically and Computationally Efficient Linear Meta-representation Learning
Year: 2021

Abstract:
In typical few-shot learning, each task is not equipped with enough data to be learned in isolation. To cope with such data scarcity, meta-representation learning methods train across many related tasks to find a shared (lower-dimensional) representation of the data where all tasks can be solved accurately. It is hypothesized that any new arriving tasks can be rapidly trained on this low-dimensional representation using only a few samples. Despite the practical successes of this approach, its statistical and computational properties are less understood. Recent theoretical studies either provide a highly suboptimal statistical error, or require many samples for every task, which is infeasible in the few-shot learning setting. Moreover, the prescribed algorithms in these studies have little resemblance to those used in practice or they are computationally intractable. To understand and explain the success of popular meta-representation learning approaches such as ANIL [43], MetaOptNet [36], R2D2 [9], and OML [33], we study a alternating gradient-descent minimization (AltMinGD) method (and its variant alternating minimization (AltMin) in the Appendix) which underlies the aforementioned methods. For a simple but canonical setting of shared linear representations, we show that AltMinGD achieves nearly-optimal estimation error, requiring only Ω(polylog d) samples per task. This agrees with the observed efficacy of this algorithm in the practical few-shot learning scenarios.

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

heading: 1.1 Related work
text: 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/ε)).

heading: 2 Problem Formulation: Meta-learning of Shared Representation
text: 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.

heading: 3 Alternating minimization
text: We focus on AltMinGD from [47], which learns a shared parameterized representation fU (·) as in (1). Several variations of this algorithm are widely used, for example [43, 36, 9]. However, we note that these previous works neither referred to this algorithm as AltMinGD, nor explicitly related it to the Alternating Minimization (AltMin) framework [30]. To highlight this connection we follow the notations from the latter [30]. AltMinGD alternately updates the matrix of regression parameters V using exact minimization with fixed U , and updates the representation parameter U using the standard gradient descent step. Concretely,
v(i) ∈ arg min v∈Rr ∑ j∈[m] `(〈v, fU (x(i)j )〉, y (i) j ) ,∀i ∈ [t] ,
U ← U − η ∑ i∈[t] ∑ j∈[m] ∇U `(〈v(i), fU (x(i)j )〉, y (i) j ) .
As `(·, ·) is typically a convex function, we can estimate v(i) efficiently for a fixed U . Note that many methods used in practice (ANIL [43] and MetaOptNet [36]) back-propagate through their respective inner (potentially inexact) optimization step. However, since we do exact inner minimization with respect to V , by a generalization of the Danskin’s theorem [10], back-propagating through the inner minimization is equivalent to computing the gradient ∇UL(U, V ) with respect to U and then setting V = V ∗(U), where V ∗(U) ∈ arg minV L(U, V ) is the minimizer for the current U . That is, ∇U minV L(U, V ) = ∇UL(U, V ∗(U)). For the linear representation learning problem specified in Section 2, the above updates reduce to the following:
v(i) ∈ arg min v ∑ j (y (i) j − 〈x (i) j , Uv〉) 2 , for all i ∈ [t] ,
U ← U − η∇UL(U, V ) = U + η ∑ i,j (y (i) j − 〈x (i) j , Uv (i)〉)x(i)j (v (i))> .
Given U , we can efficiently estimate each of the low r-dimensional regression parameters v(i)’s separately and in parallel using standard least squares regression. Our analysis requires that when we update V for current U , U should be independent from the training points. Similarly, during the update for U , V should be independent of the data points. We ensure the independence using two strategies: (a) similar to standard online meta-learning settings [20], we randomly select (previously unseen) tasks to update U and V , (b) within each task, we divide the datapoints into two sets to update V and U separately. But in our experiments, we re-used all the samples at each iteration. Algorithm 1 presents a pseudo-code of AltMinGD applied to Problem (4). Note that in Algorithm 1, we apply QR-decomposition on U after every U update to ensure that magnitude of U and V does not stray far away from that of true U∗ and V ∗, respectively. Otherwise, the sample complexity requirements of the algorithm increase in the condition number factors.
Run-time and memory usage: Exact update for v(i) has a time complexity of O(mr2 + r3), which can be brought down to O(m · r) by using gradient descent for solving the least squares. Our analysis shows that under the sample complexity assumptions of Theorem 1, each of the least squares problem has a constant condition number. So, the total number of iterations for this update scale as log 1 to achieve error. If we set = 1/poly(t, σ), then using standard error analysis, we should be able to
Algorithm 1 AltMinGD : Meta-learning linear regression parameters via alternating minimization gradient descent
Required: Data: {(x(i)j ∈ Rd, y (i) j ∈ R)}mj=1 for all 1 ≤ i ≤ t, K: number of steps, η: stepsize.
1 Initialize U ← Uinit 2 Randomly shuffle the tasks {1,. . . ,t}
for 1 ≤ k ≤ K do 3 Tk ← [1 + t(k−1)K , tk K ] for i ∈ Tk do 4 v(i) ← arg min
v̂∈Rr ∑ j∈[m/2] ( y (i) j − 〈 Uv̂, x (i) j 〉)2 end
5 U ← U + η ∑ i=Tk m∑ j=1+m2 ( y (i) j − 〈 Uv(i), x (i) j 〉) x (i) j (v (i))> 6 U ← QR(U)
end 7 return U
Algorithm 2 AltMinGD-S : Meta-Learning regression parameters via AltMinGD over task subsets
Required: Data: {(x(i)j ∈ Rd, y (i) j ∈ R)}mj=1 for all 1 ≤ i ≤ t, K: number of steps, η: stepsize.
Use the same steps as AltMinGD (Algorithm 1), but replace Line 3 with: 3 Tk ← { i ∈ [1 + t(k−1)K , tk K ] ∣∣ σmax(U>S(i)U) ≤ 2; σmin(U>S(i)U) ≥ 12 ;
where S(i) = 2m ∑ j∈[m/2] x (i) j (x (i) j ) > }
obtain the optimal error rate in Theorem 9. The gradient descent update for U requires O(mt · dr) time assuming large enough mt. Furthermore space complexity of AltMinGD is O(dr + t · r2). We provide an estimate of the statistical efficiency of the AltMinGD in Theorem 1. We also provide an analysis of the traditional Alternating Minimization algorithm (AltMin) which uses exact minimization for updating U in the Appendix A. We obtain a slightly improved statistical guarantee for AltMin in terms of the condition number but its run-time is slower than that of AltMinGD.

heading: 3.1 Subset Selection
text: Algorithm 1 operates over all the tasks in a batch, each of which are generated using a random process. Now, if the number of tasks t is large, then there is a non-trivial probability that some of the tasks are outliers, i.e., they have a large amount of error. This might lead to an arbitrary poor solution due to the outlier tasks. This is reflected in our analysis of AltMinGD (see Theorem 1), where the number of samples per task grows logarithmically with t which is non-intuitive as typically larger number of tasks should not hurt the sample complexity.
In more general representation learning problems, when the number of tasks t is large, there is more chance that some of them are outlier tasks. Ideally we want to design an estimator for shared representation U that is robust to a few outlier tasks. For the linear representation learning problem, we observe that to ensure small error for a task, we require the Hessian to be well-conditioned. So, we compute the eigenvalues of the Hessian U>S(i)U for each task with the current U , and select only the tasks whose eigenvalues lie close to the expected repeated eigenvalue 1. Algorithm 2 applies this criteria to select tasks in each iteration, and then use the standard AltMinGD updates on those selected tasks. This leads to an improved dependence on t as we show in Theorem 3.
Run-time: On top of the run-time complexity of AltMinGD, the subset selection scheme adds an additiveO(mt ·dr+t ·r3) term. This arises due toO(r3) eigen-decompositions of Ũ>S(i)Ũ ∈ Rr×r for each task.

heading: 4 Statistical guarantees for Alternating Minimization algorithms
text: We reiterate that Õ and Ω̃ hide logarithmic terms in d and r and other problem parameters. We analyze the AltMinGD and AltMinGD-S algorithms using the rescaled Frobenius norm error: ‖(I− U∗(U∗)>)U‖F / √ r ∈ [0, 1] between the rank-r subspaces corresponding to the true U∗ and the output of the algorithm U . We first provide our main results analyzing these algorithms, and present detailed comparisons to previous results in Section 4.1
AltMinGD: We first present our main result for the AltMinGD method (Algorithm 1), applied to the linear representation learning problem described in Section 2. Theorem 1 (Simplified version of Theorem 9 in Appendix C). Let there be t linear regression tasks, each with m samples satisfying Assumptions 1, 2. Let κ := λ∗1/λ ∗ r and let,
m ≥ Ω̃(r2 + r log t+ κ · (σ/ √ λ∗r)
2r2 log t), t ≥ Ω̃(κ · µ2r3), and mt ≥ Ω̃(κ · µdr2 + κ3 · µdr2(σ/ √ λ∗r) 2).
Then AltMinGD (Algorithm 1), initialized at Uinit s.t. ‖(I − U∗(U∗)>)Uinit‖F ≤ min(21/121, Õ(1/κ)) and run for K = Ω(dκ log(mt/(κ · µdr · (σ/ √ λ∗r))e) iterations with the stepsize η = r/t2λ∗1 , outputs U so that the following holds (w.p. ≥ 1−K/(dr) 10):
‖(I− U∗(U∗)>)U‖F√ r
≤ Õ (√ κ ( σ√ λ∗r )√ µd r m t ) . (5)
Remark 1 (Initialization): Our result holds if the initial point Uinit is reasonably accurate. One choice of initialization is to use the Method-of-Moments (MoM) [48]. Due to sub-optimality of MoM approach ([48, Theorem 3], also provided in Theorem 12 in Appendix), we get an additional sample complexity requirement of mt ≥ Ω̃(κ2dr2 (µκ+ r(σ/ √ λ∗r)
4). Note that this does not degrade the asymptotic error rate, Õ( √ dr/mt) when ε = Õ( √ dr/mt)→ 0. In our experiments, we observed that random initialization works just as well. Such a requirement of a good initialization is common in theoretical analyses of alternating update methods [30, 38], where it has been widely observed that random initialization works well in practice.
Remark 2 (Generalization in few-shot learning): Learning a shared representation helps in generalizing to new arriving tasks in few-shot learning. Suppose we run Algorithm 1, under the conditions of Theorem 1 to get an estimated subspace U . Let a new task, whose task specific regression parameter v∗+ lie in U∗, be introduced with m+ samples. Now, we can apply the step 4 of Algorithm 1, with U and the new samples, to meta-learn an estimate v+ of v∗+. Then the mean-squared-error (MSE) of the estimated parameter is Õ((σ/ √ λ∗r)(µdr
2/mt+ r/m+)). Therefore, as long as mt is large enough, we only need m+ = Ω(r) additional samples to get an arbitrarily small MSE, as opposed to m+ = Ω(d) of the trivial baseline of solving the new task by itself. We also improve upon other baselines from [48] in terms of dependence on σ and t; see Section 4.1 and Table 1 for more details.
Remark 3 (Near-optimality of the error rate): We note that our error rate matches – up to poly(κ, µ) factors – the information theoretic lower bound given in Corollary 2. Corollary 2. [48, Theorem 5] Let r ≤ d/2 and mt ≥ r(d− r), then for all V ∗, w.p. ≥ 1/2
inf Û sup U∈Grr,d ‖(I− U∗(U∗)>)Û‖F√ r ≥ Ω ( 1 κ σ√ λ∗r √ d r m t ) , (6)
where Gr,d is the Grassmannian manifold of r-dimensional subspaces in Rd, the infimum for Û is taken over the set of all measurable functions that takes mt samples in total from the model in Section 2 satisfying Assumption 1 and 2.
However, the sufficient conditions onmt in Theorem 1 has a factor r gap from the necessary condition above, which we discuss with a concrete example in the next remark.
Remark 4 (Gaussian example): Let us interpret our result using a concrete example. Consider independent Gaussian parameters v∗(i) ∼ N (0, (1/r)Ir×r) such that the signal-to-noise ratio (i.e., xTU∗v∗(i)/σ2) is independent of r. Then with high probability ‖v∗(i)‖ = Θ̃(1) and λ∗1 = λ∗r =
Θ̃(1). It follows that as per Assumption 2 the condition number κ = Θ̃(1) and µ = Θ̃(1). To estimate U∗ up to an ε error, AltMinGD needs a total of mt = Õ(dr2 + σ2dr/ε2) samples. The second term is dominant for small ε and is optimal, which follows from the near-optimality in Remark 2. However, it is an open question if the first term is necessary, as the best known lower bound in the noiseless case will require mt = Ω(dr). In this well-behaved Gaussian case, AltMinGD requires m ≥ Ω̃(r2 + (1 + σ2)r log t) per task samples. Remark 5 (Dependence on the minimum eigenvalue): Notice that in the limit of λ∗r → 0, V ∗ is rank deficient, thus making it impossible to recover the entire subspace of U∗. This is reflected in our Theorem 3 where the error-rate approaches the maximum possible value of one as λ∗r approaches zero (the LHS of Eq. (5) is at most one). However, for prediction error, smaller rank of V ∗ implies smaller dimensional representation to be learned, thus the prediction error bound should improve with lower λ∗r (and also smaller rank of V
∗). Proving a tight guarantee in the prediction error is more challenging and most of the existing results in matrix sensing literature [29] only provide guarantees in parameter estimation error. On the contrary, the lower-bound in (6) becomes zero as λ∗r decreases, implying that the lower-bound is significantly weaker in λ∗r . This is expected since the lower-bound is derived through a lower-bound for the corresponding subspace regression loss. Intuitively when λ∗r = 0 the tasks become less diverse (more homogeneous), and therefore the regression becomes easier. Such condition number mismatch in upper and lower-bounds are common in low-rank literature [30].
Task subset selection (AltMinGD-S): One downside of Algorithm 1 is that m needs to increase with t (i.e., m = Ω(log t)). We introduce AltMinGD-S in Algorithm 2 to study a fundamental question of whether this log t dependence is necessary. We show that when the noise is sufficiently small, AltMinGD-S achieves a per task sample complexity that does not increase with t. Theorem 3 (Simplified version of Theorem 11 in Appendix D). Consider the setting of Theorem 1. Let κ := λ∗1/λ ∗ r .
m ≥ Ω̃(r2 + κ · (σ/ √ λ∗r)
2r2 log t), t ≥ Ω̃(κ · µ2r3), and mt ≥ Ω̃(κ · µdr2 + κ3 · µdr2(σ/ √ λ∗r) 2).
Then AltMinGD (Algorithm 2), initialized atUinit s.t. ‖(I−U∗(U∗)>)Uinit‖F ≤ min(21/121, Õ( 1κ )) and run forK = Ω(dκ log(mt/(κ · µdr · (σ/ √ λ∗r))e) iterations using the stepsize η =
r/t 2λ∗1 , outputs U so that the following holds (w.p. ≥ 1−K/(dr)10):
‖(I− U∗(U∗)>)U‖F√ r
≤ Õ (√ κ ( σ√ λ∗r )√ µd r m t ) . (7)
Remark 7 (Bias of Subset Selection): One may observe that this scheme may introduce a bias in the training data at each iteration. However, we control this bias by adding a new requirement that the number of tasks should be at least t ≥ Ω̃(κ · µ2r3) (Theorem 6). This ensures that the only a small O(1/µr) fraction of the tasks are discarded at each step (Lemma D.1 in the Appendix D.1), and this leads to a low bias. This requirement may be insignificant in our regime of interest where the number of tasks may be exponentially large, so that AltMinGD-S can provide a gain over AltMinGD.
Remark 6 (When noise is small enough): Note that when the noise variance σ is small enough or when there are large number of tasks, i.e. σ2 O(1/log t), AltMinGD-S only needs m ≥ Ω̃(r2) samples per task, assuming suitable initialization (see Remark 1). Furthermore, since AltMin-S selects a fraction of tasks to perform updates and the selection process requires onlyO(mt ·dr+t ·r3), the time-complexity of the method remains same as that of AltMinGD, up to constant factors. Next we see that AltMinGD-S removes the dependence of m on t completely, in the noiseless setting. Corollary 4. Let there be t linear regression tasks, each with m samples satisfying Assumptions 1, 2,
m ≥ Ω̃(r2), t ≥ Ω̃(µ2r3K), and mt ≥ Ω̃(µdr2K). Additionally assume that the observations are noiseless, i.e. σ = 0. Then AltMinGD-S (Algorithm 2), initialized at Uinit s.t. ‖(I− U∗(U∗)>)Uinit‖F ≤ min(21/121, Õ(1/κ)) and run for K iterations using the stepsize η = 12λ∗1 , outputs U so that the following holds (w.p. ≥ 1−K/(dr) 10):
‖(I− U∗(U∗)>)U‖F√ r ≤ (1− 1 6κ )KÕ (κ) . (8)
The above corollary shows that in the noiseless setting, the per-task sample complexity for AltMinGDS does not grow with t, and is nearly optimal. Also note that that even for noiseless setting, techniques like Method-of-Moments (MoM) still incur error of √ dr/mt, ignoring κ terms. In contrast,
AltMinGD-S when initialized using MoM method (see Remark 1)), incurs just Õ(exp(−t/κ)) error. Proofs of Theorems 1 & 3 are in Appendix C.1 & Appendix D.1.

heading: 4.1 Sample complexity comparison
text: To the best of our knowledge, Theorems 1 and 3 presents the first analysis of an efficient method for achieving optimal error rate in σ, d and r. [48] is most relevant that analyzes the landscape of the Empirical Risk Minimization (ERM) with Burer-Monteiro factorization. It shows that ERM can achieve a rescaled Frobenius norm error of ε with t tasks (assuming t ≥ d), when m ≥ Ω̃(r4 log(t) + r2 log(t)σ2/ε2). We stress that this is highly sub-optimal as for small estimation ε, more tasks do not help improve the per-task sample complexity. This also does not reconcile with practice where more tasks tend to help accuracy and helps overcome small number of samples per-task. In contrast, AltMinGD requires m ≥ Ω̃(r2(1 + σ2) log(t) + (r2σ2/ε2)(d/t)) where small error ε can be achieved by collecting more tasks and increasing t. [19] studies the global minimizer of the non-convex ERM optimization in Eq. (4), without providing an efficient algorithm to solve it. The authors show that non-convex ERM achieves a small generalization error if m = Ω̃(d), which is impractical in the few-shot learning setting.
Another approach is Method-of-Moments (MoM), which estimates U by finding the principal directions of a particular 4th moment of the data [48, 35]. MoM can indeed trade-off smaller error ε by increasing the number of tasks t. But the algorithm is inexact, i.e., even for σ = 0, we need m→∞ to achieve exact recovery of U∗; see Appendix G. This is in a stark contrast with AltMinGD and AltMinGD-S where for noiseless case, we can find U∗ exactly, as long as m = O(r log t+ r2) and t = O(dr); see Figure 1a for an illustration. We consolidate these comparisons in Table 1.
Finally, a concurrent and independent work by [16] also analyzes AltMinGD but only for the special case when there is no noise, i.e., σ = 0. We show tighter results that are more generally applicable: (i) our analysis applies to general noise σ that is not necessarily zero, (ii) even in the noiseless case, our analysis of AltMinGD is tighter and shows a smaller sample complexity, and (iii) we present novel AltMinGD-S that further improves the sample complexity. Precisely, in the noiseless case, [16] proves that m = Ω̃(κ2 · r3 log t) is sufficient for finding U∗ with a large enough t. Our tighter analysis shows that m = Ω̃(r log t+ r2) (Theorem 1) is sufficient with no dependence in κ. Note that the condition number κ > 1 and can be arbitrarily large depending on the problem instance. Further, we present a novel algorithm, AltMinGD-S , that only requires m ≥ Ω̃(r2) (Corollary 4).

heading: 5 Experimental results
text: In this section we empirically compare the performance of AltMinGD (Algorithm 1) and its exact minimization variant AltMin (Algorithm 3 in Appendix), two different versions of Method-of-
Moments (MoM [48], MoM2 [35]), and simultaneous gradient descent on (U, V ) using the BurerMonteiro factorized loss (4) (BM-GD [48]). We omit AltMinGD-S here because the logarithmic gain (1/ log(t)) of AltMinGD-S will only be observed when we have an exponentially large number of tasks (t). This is challenging to simulate using our modest computing hardware. However, for big-data scenarios similar subset select schemes may be useful. In all the figures, the magenta dashed line with square marker represents AltMinGD, the blue straight line with circular marker denotes the AltMin , the red dotted line with downwards pointing triangular marker denotes the MoM, the yellow dotted line with upwards pointing triangular marker represents the MoM2, and the green dashed and dotted line with diamond marker represents the BM-GD. In all the figures we plot the subspace estimation error of the output U of the algorithms. The error is calculated using the rescaled Frobenius norm ‖(I− U∗(U∗)>)U‖F / √ r, which takes a value in the interval [0, 1]. All results are averaged over multiple runs, and more experiments details and plots are provided in Appendix H.
Figure 1a plots subspace distance against the standard deviation σ of the regression noise, ε(i)j ∼ N (0, σ2); see (2). Clearly, as predicted by Theorems 1 and 5 (in Appendix), the methods we consider—AltMinGD and AltMin —achieve smaller error than MoM methods for small noise regime. Here the error of AltMinGD and AltMin are linearly proportional to σ. However as predicted the error of MoM and MoM is a constant multiple of √ dr3/mt = √ r for all values of σ, and it does not improve when σ decreases (see Table 1). While BM-GD does not have any known algorithmic guarantees, it still performs better than MoM methods. However, BM-GM becomes unstable and challenging to tune at low noise regime, even at a lower or comparable step-size than AltMinGD. To highlight this, the individual trials for each algorithm in this plot are plotted in Figure 2. Figure 1b plots the subspace error against the number of tasks t. In Figure 1c, we plot the the error against
the number samples per tasks m. In both of these figures, we observe that, AltMinGD, AltMin and BM-GD achieve much smaller subspace error than the MoM and MoM2. Furthermore, as predicted, the squared error of all these methods decrease linearly in m and t. We again note that BM-GD is unstable and hard to tune, especially for large t. The individual trials for each algorithm on these two plots are plotted in the Appendix H. ERM-based BM-GD performs poorer than Alternate Minimization-based AltMin and AltMinGD. This might be due to the presence of many bad local minima in the optimization landscape of the ERM problem jointly over (U, V ) [48, Theorem 2].
In Figure 3, we plot the subspace estimation error against the number of iterations of AltMinGD, AltMin, and BM-GD for varying levels of task diversity µ (Assumption 2). We observe that AltMin takes significantly fewer iterations to converge than AltMinGD and BM-GD, and AltMinGD converges earlier than BM-GD. However, each iteration of AltMin is very slow as it needs O(d3) operations, where as AltMinGD and BM-GD need only O(d) operations per iteration. Therefore, AltMinGD could be the fastest in practical high-dimensional setting. BM-GD seems to be slower than AltMinGD because BM-GD seems to need a smaller stepsize than AltMinGD to stabilize its convergence. While all the methods perform worse when the task diversity decreases ((a) → (b) → (c)), we see that AltMin is more robust than others. This may be attributable to AltMin’s tighter dependence on the condition number κ (Theorem 5, in Appendix) when compared to AltMinGD (Theorem 1).

heading: 6 Conclusion
text: When learning a shared representation for multiple tasks, a common approach is to alternate between finding the best linear model for each task on the current representation, and taking one gradient descent step to update the shared representation. This algorithm, AltMinGD, has been widely used in meta-representation learning with little theoretical understanding. We provide insights into the empirical success of AltMinGD by studying it in the canonical problem of linear meta-learning. We showed that, AltMinGD provides a nearly optimal error rate, along with nearly optimal per-task and overall sample complexities in their dependence in the dimensionality d of the data. To the best of our knowledge, this is the first such optimal error rate that scales appropriately with the noise in observations, while still ensuring per-task sample complexity to be nearly independent of the dimensionality d. Latter is a key requirement in meta-learning as individual tasks are data-starved. The limitations of our results are: (i) the analysis does not extend to non-linear representations, (ii) the dependence on the rank r of the shared subspace, the incoherence µ, and the condition number κ may not be tight; and (iii) our analysis is “local” and requires a good initialization. We also proposed and analyzed a task subset selection-based method (AltMinGD-S) that further improves the per-task sample complexity and ensures that it is independent of the number of tasks in small noise or large number of tasks regime. However, the subset selection scheme heavily relies on the linearity of the shared representation. Therefore, this scheme cannot be directly applied to more practical neural network training. It also remains an open question if it is possible to achieve a per-task sample complexity that does not depend on the number of tasks t, even in the large noise setting.
Our work leads to several interesting future directions and questions. For the non-linear version of the problem, ensuring optimal error rate with optimal per-task sample complexity is an interesting open question. Finally, analyzing alternating minimization methods with stochastic gradients and streaming tasks is another promising direction. Our proof techniques could be combined with that of recent results in efficient one-pass SGD [31] to design a nearly optimal stochastic algorithm.

heading: Acknowledgements
text: Oh acknowledges funding from Google faculty research award, NSF grants IIS-1929955, CCF1705007, CNS-2002664, CCF 2019844 as a part of Institute for Foundation of Machine Learning, and CNS-2112471 as a part of Institute for Future Edge Networks and Distributed Intelligence. We also thank anonymous reviewers for their reviews and suggestions for improving our manuscript.

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

This paper provides sample complexity and error bounds for “alternating gradient descent minimization” (AltMinGD), an algorithm used for meta-learning. It shows that AltMinGD applied to linear regression is superior to competitors such as Empirical Risk Minimization (ERM) and Method of Moments (MoM). Furthermore, they show that the bounds are tight and basically cannot be improved upon. 

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.