An overarching goal in machine learning is to build a generalizable model with few 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 the 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. We leverage this inductive bias to explain 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 meta-learning algorithm. Numerical experiments on real and synthetic data verify our insights on overparameterized meta-learning.

§ INTRODUCTION

In a multitude of machine learning (ML) tasks with limited data, it is crucial to build accurate models in a sample-efficient way. Constructing a simple yet informative representation of features is a critical component of learning a model that generalizes well to an unseen test set. The field of meta-learning dates back to <cit.> 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 <cit.>, machine translation <cit.> and reinforcement learning <cit.>.

Recent literature in ML theory has posited that overparameterization can be beneficial to generalization in traditional single-task setups for both regression <cit.> and classification <cit.> 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 is observed to be beneficial in the few-shot phase for transfer-learning in Figure <ref>(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 $\nfs$. For each choice of $\nfs$, increasing $R$ beyond $\nfs$ is seen to reduce the test-error. These findings are corroborated by <cit.> (MAML) and <cit.>, 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 <ref>(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 <ref>(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 <cit.>) 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. <ref>) 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. This motivates us to build upon these results to develop a theoretical understanding of overparameterized meta-learning[The code for this paper is in <https://github.com/sunyue93/Rep-Learning>.].

§.§ Our contributions

This paper studies meta-learning when each task is a linear regression problem, similar in spirit to <cit.>. In the representation learning phase, the learner is provided with training data from $\ntask$ distinct tasks, with $\nspt$ training samples per task: using this data, it selects a matrix $\hbLa\in\R^{d\times R}$ with arbitrary $R$ to obtain a linear representation of features via the map  $\x \to \hbLa^\top \x$. In the few-shot learning phase, the learner faces a new task with $n_2$ training samples and aims to use the representation $\hbLa^\top \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 <ref>.
The overarching goal of this paper is to propose a scheme to use the knowledge gained from earlier tasks to choose $\hbLa$ that minimizes few-shot risk. This goal enables us to engage with important questions regarding overparameterization:

Q1: What should the size $R$ and the representation $\hbLa$ be to minimize risk at the few-shot phase?

Q2: Can we learn the $Rd$ dimensional representation $\hbLa$ with $\ntot\ll \Rr d$ 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 <ref>. Our results are obtained for a general data/task model with arbitrary task covariance $\bSt$ and feature covariance $\bSi$ 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 $\hbLa$, in Section <ref> we first consider learning with known covariances $\B$ and $\bSi$ respectively (Algorithm <ref>). Compared to projection-only representations in previous works (see Paragraphs 2-3 of Sec. <ref>), our scheme applies eigen-weighting matrix $\bLa$ 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 <cit.>: 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 <ref> also reveals that, the optimal subspace and representation matrix are closed-form functions of the canonical task covariance $\Bbar = \sbSi\B\sbSi$, 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 $\Bbar$ (which has a degree of freedom (DoF) of $\Omega(Rd)$) from data. In Section <ref> we first present empirical evidence that feature covariance $\bSi$ is “positively correlated” with $\Bbar$.    Then we propose an efficient algorithm based on Method-of-Moments (MoM), and show that the sample complexity of representation learning is well below $\order{Rd}$ due to the inductive bias. Our sample complexity bound depends on interpretable quantities such as effective ranks $\bSi,\Bbar$ and improves over prior art (e.g., <cit.>), even though the prior works were specialized to low-rank $\Bbar$ and identity $\bSi$ (see Table <ref>).

End to end meta-learning guarantee. In Section <ref>, we consider the generalization of Section <ref>, where we have only estimates of the covariances instead of perfect knowledge. This leads to an overall meta-learning guarantee in terms of $\bLa$, $\ntot$ and $\nfs$ and uncovers a bias-variance tradeoff: As $\ntot$ decreases, it becomes more preferable to use a smaller $\Rr$ (more bias, less variance) due to inaccurate estimate of the weak eigen-directions of $\Bbar$. 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. <ref>).

(a) Steps of the meta-learning algorithm. (b)
Our representation-learning algorithm has two steps:  projection and eigen-weighting. We focus on the use of overparameterization$+$weighting matrix (Def. <ref>), and compare this with overparameterization with simple projection (no eigen-weighting), and underparameterization (for which eigen-weighting has no impact and is equivalent to projection).  <cit.> study underparameterized projections only. To distinguish from eigen-weighting, we will refer to simple projections as subspace-based representations.

§.§ Related work

Overparameterized ML and double-descent
The phenomenon of double-descent was first discovered by <cit.>. This paper and subsequent works on this topic <cit.> 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 <ref> extends the ideas in <cit.> to depict how the optimal representation described can be learned from imperfect covariance estimates as well.

Theory for representation learning
Recent papers <cit.> propose the theoretical bounds of representation learning when the tasks lie in an exactly $r$ dimensional subspace.  <cit.> discuss method of moment estimators and <cit.> discuss matrix factorized formulations.
<cit.> shows that
the number of samples that enable meaningful representation learning is $\order{dr^2}$. <cit.> 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 $\order{dr}$ samples when $\nspt\sim r$, and extend the bound to general covariances with effective rank defined.

Subspace-based meta learning With tasks being low rank, <cit.> do few-shot learning in a low dimensional space.
<cit.> study
meta-learning for linear bandits. <cit.>
gives information theoretic lower and upper bounds.
<cit.> 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 <cit.> provides results on overparameterized representation learning, but <cit.> requires number of samples per pre-training task to obey $\nspt\gtrsim d$, whereas our results apply as soon as $\nspt\gtrsim 1$.

Mixed Linear Regression (MLR) In MLR  <cit.>, 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. <cit.> propose sample complexity bounds of representation learning for mixed linear regression. They can be combined with other structures such as binary task vectors  <cit.> and sparse task vectors <cit.>.

§ PROBLEM SETUP

The problem we consider consists of two phases:

* Representation learning: Prior tasks are used to learn a suitable representation to process features.

* 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 $\Dct$ and $\Dcf$, and the two phases share the same feature and task distributions.The setup is summarized in Figure <ref>(a).

§.§ Data generation

Throughout, $\Dct$ and $\Dcf$ denote the distributions of tasks $\bt_i$ and features $\x_{ij}$ respectively. These distributions are subGaussian, zero-mean with corresponding covariance matrices $\Sigt$ and $\Sigf$.

Given a specific realization of task vector $\bt \sim \Dct$, the corresponding label/input distribution $(y,\x)\sim \Dc_{\bt}$ is obtained via $y=\x^\top\bt+\eps$ where $\x\sim\Dcf$ and $\eps$ is zero-mean subgaussian noise with variance $\sige^2$.

Data for Representation Learning (Phase 1). We have $\ntask$ tasks, each with $\nspt$ training examples. The task vectors $(\bt_i)_{i=1}^\ntask\subset\R^d$ are drawn i.i.d. from the distribution $\Dct$. The data for $i$th task is given by $(y_{ij},\x_{ij})_{j=1}^{n_1}\distas \Dc_{\bt_i}$. In total, there are $\ntot=\ntask\times \nspt$ examples.

Data for Few-Shot Learning  (Phase 2).  Sample task $\bts\sim\Dct$. Few-shot dataset has $\nfs$ examples $(y_{i},\x_{i})_{j=1}^{\nfs}\distas \Dc_{\bts}$.

We use representation learning data to learn a representation of feature-task distribution, called eigen-weighting matrix $\hbLa$ in Def. <ref> below. The matrix $\hbLa$ is passed to few-shot learning stage, helping learn $\bts$ with few data.

§.§ Training in Phase 2

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 $\ntask$ tasks.

Denote $\X\in\R^{\nfs\times d}$ whose $i^{\rm th}$ row is $\x_i$, and $\y = [y_1,...,y_m]^\top$. We are interested in studying the weighted 2-norm interpolator defined below for overparameterization regime $R\ge\nfs$.

Let the representation dimension be $R$, where $R$ is any integer between $1$ and $d$. We define an eigen-weighting matrix $\hbLa\in\R^{d\times R}$ and the associated weighted $\ell_2$-norm interpolator

\begin{align*}
\hat \bt_{\hbLa} = \arg\min_{\bt} \tn{\hbLa^{\dag} \bt} \quad \text{s.t.}\quad\y=\X \bt~~~\text{and}~~~ \bt\in\mathrm{range\_space}(\hbLa).
\end{align*}

The solution is equivalent to defining $\hat {\vct{\alpha}}_{\hbLa}= \hbLa^{\dag} \hat \bt_{\hbLa}$ and solving an unweighted minimum 2-norm regression with features $\X\hbLa$. This corresponds to our few-shot learning problem

\begin{align*}
\hat {\vct{\alpha}}_{\hbLa} = \arg\min_{\vct{\alpha}} \tn{\vct{\alpha}} \quad \text{s.t.}\quad\y=\X\hbLa\vct{\alpha}
\end{align*}

from which we obtain $\hat \bt_{\hbLa} = \hbLa\hat {\vct{\alpha}}_{\hbLa}$. When there is no confusion, we can replace $\hat \bt_{\hbLa}$ with $\hat \bt$. One can easily see that $\hat \bt = \hbLa(\X\hbLa)^\dag\y$. We note that Definition <ref> is a special case of the weighted ridge regression discussed in <cit.>, as stated in Observation <ref>. An alternative equivalence between min-norm interpolation and ridge regression can be found in <cit.>.

Let $\X \in\R^{\nfs\times d}$ and $\y\in\R^{\nfs}$, define

\begin{align}
\hat \bt_1 &= \lim_{t\rightarrow 0} \argmin_{\bt} \|\X \bt - \y\|_2^2 + t\bt^\top (\hbLa\hbLa^\top)^\dag\bt,\ \bt\in\mathrm{column~space~of~}\hbLa. \label{eq: obs 1 ridge}
\end{align}

We have that $\hat \bt_1 = \hat \bt$.

§ CANONICAL COVARIANCE AND OPTIMAL REPRESENTATION

In this section, we ask the simpler question: if the covariances $\B$ and $\bSi$ are known, what is the best choice of $\hbLa$ to minimize the risk of the interpolator from Definition <ref>? In general, the covariances are not known; however, the insights from this section help us study the more general case in Section <ref>. This is equivalent to selecting the penalty of the ridge term in (<ref>), in accordance of the feature and task distributions to weight the solution $\hat\beta$.
Define the risk as the expected error of inferring the label on the few-shot dataset,

\begin{align}
&\risk(\hbLa,\B, \bSi) = \mtx{E}_{\x,y,\bt}(y - \x^\top\hat\bt_{\hbLa})^2
= \Eb_{\bt} (\hat \bt_{\hbLa}-\bt)^\top \bSi(\hat \bt_{\hbLa}-\bt) + \sigma^2.\label{eq:bSi_test_risk}
\end{align}

The natural choice of optimization for choosing $\hbLa$ would be to choose the weighting that minimizes the eventual risk of the learned interpolator.

\begin{equation}
\label{eq:opt_rep}
\bLa=\arg\min_{{\hbLa}'\in\R^{d\times R}} \risk( {\hbLa}',\B, \bSi)
\end{equation}

Since the label $y$ is bilinear in $x$ and $\beta$, we introduce whitened features $\tilde{\x}=\sbSiinv\x$ and associated task vector $\tilde{\bt}=\sbSi \bt$. This change of variables ensures $\x^T\bt=\tilde{\x}^T\tilde{\bt}$; now, the task covariance in the transformed coordinates takes the form

\[
\Bbar=\sbSi\B\sbSi,
\]

which we call the canonical task covariance; it captures the joint behavior of feature and task covariances $\bSi, \B$.
Finding the above sentence a bit confusing i.e what does it mean to capture joint behavior?Below, we observe that the risk in Equation (<ref>) is invariant to the change of co-ordinates that we have described above i.e it
does not change when $\sbSi\B\sbSi$ is fixed and we vary $\bSi$ and $\B$.
[Equivalence to problem with whitened features]
Let data be generated as in Phase 1. Denote $\Bbar=\sbSi\B\sbSi$. Then
$\risk({\sbSiinv\hbLa},\B, \bSi) = \risk(\hbLa, \Bbar, \Ib)$.

This observation can be easily verified by substituting the change-of-coordinates into  Equation (<ref>) and evaluating the risk.

The risk in (<ref>) quantifies the quality of representation $\hbLa$; however it is not a manageable function of $\hbLa$ 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 <ref> characterizes this equivalence; the computeReduction subroutine of Algorithm <ref> calculates key quantities that are used in specifying the reduction, and the computeOptimalRep subroutine of Algorithm <ref> uses the solution of the simpler problem to obtain a solution for the original.

[Bounded feature covariance]
There exist positive constants $\Sigma_{\min}$, $\Sigma_{\max}$ such that $\bSi$ is lower/upper bounded as follows: $\mathbf{0}\prec\Sigma_{\min}\Ib\preceq \bSi \preceq\Sigma_{\max}\Ib$.

[Joint diagonalizability]
$\bSi$ and $\B$ are diagonal matrices.[This is equivalent to the more general scenario where $\bSi$ and $\B$ are jointly diagonalizable.]

[Double asymptotic regime]

We let the dimensions and the sample size grow as $d,R,\nfs\rightarrow\infty$ at fixed ratios $\bar{\kappa}:= {d}/{\nfs}$ and $\kappa:= {R}/{\nfs}$.

The joint empirical distribution of the eigenvalues of $\hbLa_R$ and $\BbarR$ is given by the average of Dirac $\delta$'s: $\frac{1}{
R}\sum_{i=1}^R\delta_{\hbLa_{R,i},\sqrt{R}\BbarRi}$. It converges to a fixed distribution as $d\rightarrow\infty$.

With these assumptions, we can derive an analytical expression to quantify the risk of a representation $\hbLa$. We will then optimize this analytic expression to obtain a formula for the optimal representation.

Suppose Assumptions <ref>, <ref>, <ref>, <ref> hold.
Let $\xi>0$ be the unique number obeying $
\nfs = \sum_{i=1}^R\big({1+(\xi{\hbLa}^2_i)^{-1}}\big)^{-1}$.
Define $\btheta \in \R^R$ with entries $\btheta_i=\frac{\xi\hbLa^2_i}{1+\xi\hbLa^2_i}$ and calculate $\BbarR, \sigma_R$ using the ComputeReduction procedure of Algorithm <ref>. Then, define the analytic risk formula

\begin{align}
f(\btheta, \BbarR, \nfs)
= \dfrac{1 }{\nfs - \|\btheta\|^2_{\red{2}}}
\left(\nfs\sum_{i=1}^R (1 - \btheta_i)^2{\BbarRi} + (\|\btheta\|^2_{\red{2}}+1)\sigma_{R}^2\right).
\label{optimal_lambda}
\end{align}

We have that

\begin{align}\label{eq: asym thm}
\lim_{\nfs \to \infty} f(\btheta, \BbarR, \nfs) = \lim_{\nfs \to \infty} \risk(\sbSiinv\hbLa, \B, \bSi)
\end{align}

The proof of Theorem <ref> applies the convex Gaussian Min-max Theorem (CGMT) in <cit.> and can be found in the Appendix B.2.
We show that as dimension grows, the distribution of the estimator $\hat \beta$ converges to a Gaussian distribution and we can calculate the expectation of risk.

Theorem <ref> provides us with a closed-form risk for any linear representation. Now, one can solve for the optimal representation by computing (<ref>) 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 [In Sec. <ref> the constraint is $\uth\le\btheta\le1-\frac{d-\nfs}{\nfs}\uth$ for robustness concerns.].

\begin{align}
\btheta^* = \arg\min_{\btheta}\ f(\btheta, \B, \bSi),\ \mbox{s.t.}\ 0\le \btheta< 1, \sum_{i=1}^R\btheta_i = \nfs  \label{eq:theta opt}\tag{OPT-REP}
\end{align}

Theoretical risk of optimal representation. $\bSi=\Ib_{100}$, $\B=\diag{\Ib_{20},\iota\Ib_{80}}$, $\nfs=40$.

The optimal representation is[In the algorithm, $\xi=1$ and $\hbLa_{R,i} = (1/\btheta_i^*-1)^{-2}$, because $c\bLa$ for any constant $c$ gives the same $\hat \bt$. ] $\bLa_{R,i} = ( (1/\btheta_i^*-1) \xi)^{-2}$. The subroutine computeOptimalRep in Algorithm <ref> summarizes this procedure.

Thm. <ref> states that $\risk(\sbSiinv\hbLa, \B, \bSi)$ can be arbitrarily well-approximated by $f(\btheta, \BbarR, \nfs)$ if $n_2$ is sufficiently large. In Fig. <ref>(b), we set $\bSi=\Ib_{100}$, $\B=\diag{\Ib_{20},0.1\Ib_{80}}$, $\nfs=40$. The curves in Fig<ref>(b) are the finite dimensional approximation of $f$ (LHS of (<ref>)); the dots are empirical approximations of the risk (RHS of (<ref>)). We tested two cases when $\hbLa$ 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 $n_2 = 40$.

The benefit of overparameterization. Theorem 1 leads to an optimal eigen-weighting strategy via asymptotic analysis.
In Figure <ref>, we plot the effect on the risk of increasing $R$ for different shapes of task covariance; the parameter $\iota$ controls how spiked $\B$ is, with a smaller value for $\iota$ indicating increased spiked-ness. 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 ($\hbLa = \Ib$) 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\gg \nfs$. 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 $\iota$. It implies that, while $\Bbar$ gets closer to low-rank, the excess error caused by choosing small dimension $R$ (equal to the gap $\sigma_R^2-\sigma^2$ in Algo <ref>) is not as significant.

Figure <ref>(b) depicts the importance of eigen-weighting: in the overparameterized regime, the blue eigen-weighted learner achieves markedly lower test-error than its unweighted ($\hbLa = \Ib$) counterpart. This 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\gg \nfs$. In contrast, for the underparameterized problem, the weighting does not have any impact.
in overparameterized regime, the smaller error corresponding to optimal eigen-weighting compared to no weighting, i.e., the eigen-weighting matrix being identity, shows that
eigen-weighting becomes critical (for underparameterized problem, weighting has no impact).

Low dimensional representations zero out features and cause bias. By contrast, when $\Bbar\in\R^{d\times 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 <cit.> where the authors devise a misspecified linear regression to demonstrate the benefits of overparameterization. Our algorithm allows arbitrary representation dimension $R$ and eigen-weighting. think where this should appearIn the next section, we will acquire prior information by learning $\Bbar$. Thm. <ref> and (<ref>) infuses this information to design the best eigen-weighting.

§ REPRESENTATION LEARNING

In this section, we will show how to estimate the useful distribution in representation learning phase that enables us to calculate eigen-weighting matrix $\bLa$. Note that $\bLa$ depends on the canonical covariance $\Bbar=\sbSi\B\sbSi$.  Learning the $R$-dimensional principal subspace of $\Bbar$ enables us[We also need to estimate $\bSi$ for whitening. Estimating $\bSi$ is rather easy and incurs smaller error compared to $\Bbar$. The analysis is provided in the first part of Appendix B.] to calculate $\bLa$. Denote this subspace by $\red{\Sbar}$.

Subspace estimation vs. inductive bias. The subspace-based representation $\Sbar$ has degrees of freedom$=Rd$. When $\Bbar$ is exactly rank $R$ and features are whitened, <cit.> provides a sample-complexity lower bound of $\Omega(Rd)$ examples and gives an algorithm achieving $\order{R^2d}$ 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 <cit.> in two aspects:

* Rather than focusing on a low dimensional subspace and assuming $N\gtrsim Rd$, we can estimate $\Bbar$ or $\Sbar$ in the overparameterized regime $N\lesssim Rd$.

* Rather than assuming whitened features $\bSi = \Ib$ and achieving a sample complexity of $R^2d$, our learning guarantee holds for arbitrary covariance matrices $\bSi,\B$. 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 <ref> below.

For learning $\Bbar$ or its subspace $\red{\Sbar}$, we investigate the method-of-moments (MoM) estimator.

For $1\leq i\leq T$, define $\hat \bb_{i,1} = 2\nspt^{-1}\sum_{j=1}^{\nspt/2}\yij\xij$, $\hat \bb_{i,2} = 2\nspt^{-1}\sum_{j=\nspt/2+1}^{\nspt}\yij\xij$. Set

\begin{align*}
\hat \bM &= \nspt^{-1}\sum_{i=1}^{\ntask} (\bb_{i,1}\bb_{i,2}^\top + \bb_{i,2}\bb_{i,1}^\top),
\end{align*}

The expectation of $\hat \bM$ is equal to $\bM = \bSi\B\bSi$.

Inductive bias in representation learning: Recall that canonical covariance $\Bbar=\sbSi\B\sbSi$ is the attribute of interest. However, feature covariance $\sbSi$ term implicitly modulates the estimation procedure because the population MoM is not $\Bbar$ but $\bM = \sbSi\Bbar\sbSi$. For instance, when estimating the principle canonical subspace $\Sbar$, the degree of alignment between $\bSi$ and $\Bbar$ can make or break the estimation procedure: If $\bSi$ and $\Bbar$ have well-aligned principal subspaces, $\Sbar$ will be easier to estimate since $\bSi$ will amplify the $\Sbar$ direction within $\M$.

We verify the inductive bias on practical image dataset, reported in Appendix A. We assessed correlation coefficient between covariances $\Bbar,\bSi$ via the canonical-feature alignment score defined as the correlation coefficient

\[
\rho(\bSi,\Bbar):=\frac{\big<\bSi,\Bbar\big>}{\|\bSi\|_F\|\Bbar\|_F}=\frac{\text{trace}(\bM)}{\|\bSi\|_F\|\Bbar\|_F}.
\]

Observe that, the MoM estimator $\bM$ naturally shows up in the alignment definition because the inner product of $\Bbar,\bSi$ is equal to $\text{trace}(\bM)$. This further supports our inductive bias intuition.
As reference, we compared it to canonical-identity alignment defined as  $\frac{\text{trace}(\Bbar)}{\sqrt{d}\|\Bbar\|_F}$ (replacing $\bSi$ with $\Ib$). The canonical-feature alignment score is higher than the canonical-identity alignment score. This significant score difference exemplifies how $\bSi$ and $\Bbar$ can synergistically align with each other (inductive bias).
This alignment helps our MoM estimator defined below, illustrated by Example <ref> (spiked covariance).

In the following subsections, let $\ntot = \nspt \ntask$ refer to the total tasks in representation-learning phase.
Let $\TX = \tr(\bSi)$, $\Tbt = \tr(\B)$, and $\TbtX = \tr(\Bbar)$.
Define the approximate low-rankness measure of feature covariance by[The $(\rf+1)$-th eigenvalue is smaller than $\rf/d$. Note the top eigenvalue is $1$.]

\begin{align*}
\rf = \min \ \rf',\ \mbox{s.t.~} \rf'\in\{1,...,d\},\  \rf'/d \ge \lambda_{\rf'+1}(\bSi)
\end{align*}

We have two results for this estimator.

* Generally, we can estimate $\bM $ with $\Order(\TX\TbtX^2)$ samples.

Let $\nspt\ge \rt$, we can estimate $\bM$ with $\Order(\rf\TbtX)$ samples.

Paper <cit.> has sample complexity $\Order(dr^2)$ ($r$ is exact rank). Our sample complexity is $\Order(\TX\TbtX^2)$. $\TX , \TbtX$ can be seen as effective ranks and our bounds are always smaller than $d,r$ of the setting in <cit.>. We will discuss later in Example <ref>. Our second result says when $\nspt\ge \rt$, our sample complexity achieves the $\order{dr}$ which is proven a lower bound in <cit.>.

feature cov     3|c|$\bSi = \Ib$, $\B=\diag{\Ib_{\rt},{\bf 0}}$     3|c|
$\bSi = \diag{\Ib_{\rf},\iotaF\Ib_{d-\rf}}$,

estimator     sample $\ntot$     sample $\nspt$     error     sample $\ntot$     sample $\nspt$     error

MoM     $d\rt^2$    $1$     $(d\rt^2/\ntot)^{1/2}$
$\barrf \barrt^2$     $1$     $(\barrf \barrt^2/\ntot)^{1/2}$

MoM     $d\rt$      $\rt$     $(\rt/\nspt)^{1/2}$     $\barrf\barrt$     $\barrt$     $(\barrt/\nspt)^{1/2}$

Right side: Sample complexity and error of MoM estimators. $\rf$ ($\rt$) is the dimension of the principal eigenspace of the feature (task) covariance. $\barrf = \rf + \iotaF(d-\rf)$, $\barrt = \rt + \iotaT(d-\rt)$ are the effective ranks. Left side: This is the well-studied setting of identity feature covariance and low-rank task covariance. Our bound in the second row is the first result to achieve optimal sample complexity of ${\cal{O}}(ds_{T})$ (cf. <cit.>).

Let data be generated as in Phase 1. Assume $\|\bSi\|, \|\B\|\ = 1$ for normalization[This is simply equivalent to scaling $\yij$, which does not affect the normalized error $\|\hat\bM - \bM\| / \|\bM\|$. In the appendix we define $\Snorm = \max\{\|\bSi\|, \|\B\|\}$ and prove the theorem for general $\Snorm$.].

* Let $\nspt$ be a even number. Then with probability at least $1 - \ntot^{-100}$,

\begin{align*}
\| \hat \bM - \bM \| \lesssim (\TbtX + \sigma^2)\sqrt{\frac{\TX}{\ntot}} + \sqrt{\frac{\Tbt}{\ntask}}.
\end{align*}

* Assume $T\ge \rf$. If $\nspt \gtrsim \TbtX+\sigma^2$,
then with probability at least $1 - CT^{-100}$

\begin{align*}
\| \hat \bM - \bM \| \lesssim \left((\TbtX+\sigma^2)/\nspt\right)^{1/2}.
\end{align*}

Denote the top-$R$ principal subspaces of $\bM, \hat \bM$ by $\bMt, \hatbMt$ and assume the eigen-gap condition $\lambda_R(\bM) - \lambda_{R+1}(\bM) > 2\| \hat \bM - \bM \|$. Then a direct application of Davis-Kahan Theorem <cit.> bounds the subspace angle as follows
\[\text{angle}( \bMt, \hatbMt) \lesssim \| \hat \bM - \bM \|/(\lambda_R(\bM) - \lambda_{R+1}(\bM)).\]

Error of MoM estimator How is there still an estimation error with $\iota=0$?

Estimating eigenspace of canonical covariance. Note that if $\bSi$ and $\B$ are aligned, (e.g. Example <ref> below with $\rf=\rt=R$), then $\bMt=\red{\Sbar}$ is exactly the principal subspace of $\Bbar$. Theorem <ref> indeed gives estimation error for the principal subspace of $\Bbar$. Note that, such alignment is a more general requirement compared to related works which require whitened features <cit.>.

[Spiked $\Bbar$, Aligned principal subspaces]Suppose the spectra of $\bSi$ and $\Bbar$ are bimodal as follows $\bSi = \diag{\Ib_{\rf},\iotaF\Ib_{d-\rf}}$, $\B=\diag{\Ib_{\rt},\iotaT\Ib_{d-\rt}}$. Set statistical error $\text{Err}_{T,N}:=\sqrt{{\barrt^2\barrf}/{\ntot}} + \sqrt{{\barrt}/{\ntask}}$.
When $\iotaT, \iotaF<1$, $\rf\geq \rt$, the recovery error of $\Bbar$ and its principal subspace $\Sbar$ are bounded as

\[
\text{angle}(\hatbMt, \Sbar)
%\lesssim \sqrt{\frac{(\rt + \iotaT(\rf-\rt) + \iotaF\iotaT(d-\rf))^2\barrf}{\ntot}} + \sqrt{\frac{\barrt}{\ntask}}
\lesssim \text{Err}_{T,N}+ \iotaF^2\iotaT\quad  \mbox{and}\quad \|\hat \bM -  \Bbar\| \lesssim  \text{Err}_{T,N} + \iotaF\iotaT.
\]

The estimation errors for $\Bbar,\Sbar$ are controlled in terms of the effective ranks and the spectrum tails $\iota_F,\iota_T$. Typically $\rf\rt\gtrsim \nspt$ so $\sqrt{{\barrt^2\barrf}/{\ntot}}$ term dominates the statistical error in practice. In Fig. <ref> we plot the error of estimating $\bM$ (whose principal subspace coincides with $\Bbar$). $\bSi=\diag{\Ib_{30},\iota \Ib_{70}}$, $\B=\diag{\Ib_{30},{\bf 0}_{70}}$. $\ntask = \ntot = 100$. We can see that the error increase with $\iota$ Please check why $\iota=0$ has error if you are choosing $R=30$.

§ ROBUSTNESS OF OPTIMAL REPRESENTATION AND OVERALL META-LEARNING BOUND

In Section <ref>, we described the algorithm for computing the optimal representation with known distributions of features and tasks. In Section <ref>, 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 $\hbLa$ using the estimated canonical covariance, rather than the full-information setting of Section <ref>.

Armed with the provably reliable estimators of Section <ref>, we can replace $\Bbar$ and $\bSi$ in Algorithm <ref> 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 in[Note that Sec.6 of <cit.> gives the exact value of $\risk(\bLa,\B,\bSi)$ so we have an end to end error guarantee.]  $\risk(\hbLa,\B,\bSi) - \risk(\bLa,\B,\bSi)$.

Let us replace the constraint in (<ref>) by $\uth\le\btheta\le1-\frac{d-\nfs}{\nfs}\uth$. This changes the “optimization" step in Algorithm <ref>.
Theorem <ref> does not require an explicit computation of the optimal representation by enforcing $\uth$. 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 <ref>. We can then evaluate $\uth$ after-the-fact as the minimum singular value of this representation to apply Theorem <ref> without assuming an explicit $\uth$.

Let $\hbLa_{\uth}(R) = \textproc{ComputeOptimalRep}(R, \bSi, \hat\bM, \sigma, \nfs)$ denote the estimated optimal representation and $\bLa_{\uth}(R) = \textproc{ComputeOptimalRep}(R, \bSi, \Bbar, \sigma, \nfs)$ 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.

End to end learning guarantees. $d=100, \nfs=40, \ntask=200$, $\B=(\Iden_{20},0.05\cdot\Iden_{80})$, $\bSi=\Iden_{100}$.

Let $\hbLa_{\uth}(R)$, $\bLa_{\uth}(R)$ be as defined above, and $\TX = \tr(\bSi)$, $\Tbt = \tr(\B), \TbtX = \tr(\Bbar)$. The risk of meta-learning algorithm satisfies[The bracketed expression applies first conclusion of Theorem <ref>. One can plug in the second as well.]

\begin{align*}
\risk(\hbLa_{\uth}(R), \B, \bSi) - \risk(\bLa_{\uth}(R),\B, \bSi) \lesssim
\frac{\nfs^2}{d(R-\nfs)(2\nfs - R\uth)\uth} \left[(\TbtX + \sigma^2)\sqrt{\frac{\TX}{\ntot}} + \sqrt{\frac{\Tbt}{\ntask}}\right].
\end{align*}

Notice that as the number of previous tasks $T$ and total representation-learning samples $N$ observed increases, the risk of the estimated $\hbLa_{\uth}(R)$ approaches that of the optimal $\bLa_{\uth}(R)$ as we expect.  The result only applies to the overparameterized regime of interest $R > \nfs$. The expression of risk in the underparameterized case
is different, and covered by the second case of Equation(4.4) in <cit.>. We plot it in Fig <ref>(b) on the left side of the peak as a comparison.

Risk with respect to PCA level $R$. In Fig. <ref>, we plot the error of the whole meta-learning algorithm. We simulate representation learning and get $\hat \bM$, use it to compute $\hbLa$ and plot the theoretical downstream risk (experiments match, see Fig. <ref> (b)). Mainly, we compare the behavior of Theorem  <ref> with different $R$. When $R$ grows, we search $\hbLa$ in a larger space. The optimal $\hbLa$ 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 $\ntot=\infty$, i.e., representation learning error is $0$. As $\ntot$ decreases, there is a tradeoff between learning error and truncating small eigenvalues. Thus choosing $R$ adaptively with $\ntot$ can strike the right bias-variance tradeoff between the excess risk (variance) and the risk due to suboptimal representation.

§ NUMERICAL EXPERIMENTS

In this section, we will verify by experiments the three main contributions proposed before: (1) We apply the method of moment estimator for retrieving the covariance of the task or the top eigenspace.
We show that when feature space aligns with the task space, we only need $\order{\rf\rt^2}$
instead of $\order{d\rt^2}$
samples. (2) In Theorem <ref> part 2, we introduce task average estimator and argue that when each task contains $\Omega(\rt)$ samples, we can retrieve the subspace of tasks with $\order{\rf\rt}$ samples.
(3) In Def. <ref>, we present the optimal representation matrix that minimizes the excess risk and in Theorem <ref> we show the overall risk of meta learning algorithm.

*Error of representation learning with respect to task feature alignment. We depict this in Fig. <ref>(a). After applying method of moments estimators we do $\rt$-SVD truncation.
We can see from Fig. <ref>(a) that, MoM's learn the subspaces accurately when $1 - \iota$ is large, (i.e., with feature-task alignment), whereas behave worse when $\bSi = \Ib$ Specifically, The MoM-F estimator works better only when $\iota$ is small, i.e., high task-feature alignment.

*Task average helps learning with fewer samples. We show the sample efficiency of the task average estimator in Fig. <ref>(b). We fix the total number of samples $\ntot$, and vary the number of samples per task $\nspt$. In Theorem <ref> part 2, we argue that it estimates the task subspace with $\order{d\rt}$ samples when $\nspt \ge \rt$. We can see from the figure that, when $\nspt$ is small, this estimator degenerate to $\hat\Qb$ which asks for $\order{d\rt^2}$ samples, so the subspace estimation is bad. When $\ntask$ is small, the sampled task features cannot approximate the true task distribution, which also cause large error. We need $\nspt$ and $\ntask$ both reasonably large.

*Optimal representation. In Fig. <ref>(c), we know $\B$ and do $R$-SVD truncation and apply optimal shaping for varying $R$. When the problem becomes overparameterized ($\nfs<R$), the estimation of downstream task depends on the shaping matrix. The risk corresponding to the optimal shaping matrix is smaller than using raw features for solving the downstream task.

*End to end behavior of meta-learning algorithm. We run the whole meta-learning process. The normalized risk is plotted in Fig. <ref>(d).
We use the optimal shaping matrix $\hbLa$, calculated as a function of $\hB$, for solving the downstream task.
Both dimension reduction ($\nfs>R$) or optimal overparameterized problem ($\nfs<R$) end up with small risk, whereas it becomes hard to learn when $\nfs\approx R$. The risk of overparameterized case can be smaller than using low dimensional representation (typically applied in <cit.>).
As discussed in Remark <ref>, because we cannot learn $\B$ perfectly, the smallest risk with finite data happens when the chosen representation dimension $R<d$, unlike the case with known $\B$.

§ CONCLUSION

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.

* For all authors...

* Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

* Did you describe the limitations of your work?

* Did you discuss any potential negative societal impacts of your work?

* Have you read the ethics review guidelines and ensured that your paper conforms to them?

* If you are including theoretical results...

* Did you state the full set of assumptions of all theoretical results?

* Did you include complete proofs of all theoretical results?
In supplementary file

* If you ran experiments...

* Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

* Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

* Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

* Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?

* If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

* If your work uses existing assets, did you cite the creators?

* Did you mention the license of the assets?

* Did you include any new assets either in the supplemental material or as a URL?

* Did you discuss whether and how consent was obtained from people whose data you're using/curating?

* Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?

* If you used crowdsourcing or conducted research with human subjects...

* Did you include the full text of instructions given to participants and screenshots, if applicable?

* Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

* Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

§ NUMERICAL VERIFICATION OF INDUCTIVE BIAS FOR REPRESENTATION LEARNING

(a) Alignment of feature-task on image classification models. For MNIST, we train 45 linear pairwise classifiers between each two classes. We apply the pretrained ResNet classification model on the other three datasets, compute the (last layer) feature/task covariances and get the alignments. The alignment is a measure of correlation which is denoted by $\rho$ here. (b) We use the cifar100 dataset, take the pretrained ResNet18 network and vary the number of tasks (i.e., varying the number of output classes of the neural net, also equivalent to number of rows of the last layer matrix $B$ defined below). The alignments increase with number of tasks.

We add a figure with experiments on a few image datasets. We take the pretrained ResNet18 neural network, and feed the images into it. For every image, we take the last (closest to output) layer output as the feature $\x$, which is of dimension $d=512$. The weights of the last layer are the tasks, which is a $\ntask\times d$ matrix (We call it $B$). $\ntask=1000$, each row of $B$ is a task vector. Then $Bx\in\R^{\ntask}$ generates the label, whose each entry corresponds to each class. We calculate the feature and task covariance, as well as the alignments defined in Sec. <ref>. We can clearly see the inductive bias of every dataset.

§ ANALYSIS OF OPTIMAL REPRESENTATION

§.§ Proof of Observation <ref> and equivalent noise

Let $\hbLa\in\R^{d\times R}$, $\X \in\R^{\nfs\times d}$ and $\y\in\R^\nfs$, and define

\begin{align}
\hat \bt &= \hbLa(\X \hbLa)^\dag \y,\\
\hat \bt_1 &= \lim_{t\rightarrow 0} \argmin_{\bt} \|\X\bt - \y\|^2 + t\bt^\top (\hbLa\hbLa^\top)^\dag\bt
\end{align}

Then $\hat \bt_1 = \hat \bt$.

Denote the SVD $(\X\hbLa)^\top  = \Ub\bSigma \Vb^\top$, where
$\Ub\in\R^{R\times R},\bSigma \in\R^{R\times \nfs}, \Vb\in\R^{\nfs\times \nfs}$.

\begin{align*}
\hat \bt_1 &= \lim_{t\rightarrow 0} \argmin_{\bt} \|\X \bt - \y\|^2 + t\bt^\top (\hbLa\hbLa^\top)^\dag\bt\\
&= \lim_{t\rightarrow 0} (\X^\top\X + t(\hbLa\hbLa^\top)^\dag)^{-1} \X \y\\
&= \lim_{s\rightarrow \infty} s \hbLa (s \hbLa^\top \X^\top\X \hbLa + I)^{-1} \hbLa^\top \X^\top \y\\
&= \lim_{s\rightarrow \infty} s \hbLa (s \Ub\bSigma \Vb^\top \Vb\bSigma^\top \Ub^\top + I)^{-1} \Ub\bSigma \Vb^\top \y\\
&= \lim_{s\rightarrow \infty} s \hbLa (s \Ub\diag{\bSigma^\top \bSigma + I_{\nfs}, I_{R-\nfs}}\Ub^\top)^{-1} \Ub\bSigma \Vb^\top \y\\
&= \lim_{s\rightarrow \infty} \hbLa \Ub (\diag{\bSigma^\top \bSigma, I_{R-\nfs}/s})^{-1} \bSigma \Vb^\top \y.\\
&= \hbLa(\X \hbLa)^\dag \y
\end{align*}

The risk of $\hat \bt$ is given by

\begin{align*}
\text{risk}(\hat \bt) = \mtx{E}(y - \x^\top\hat\bt) &= \Eb (\hat \bt-\bt)^\top \bSi(\hat \bt-\bt) + \sigma^2.
\end{align*}

In Sec. <ref>, we study the asymptotic optimal representation. Below, we characterize the properties of the problem for fixed $\bt$ and arbitrary input covariance $\bSi$. We first go over this and then discuss how to obtain the optimal representation $\bLa$ minimizing test risk.

Projection onto $R$ dimensional subspace.
For the remaining proof after this part, we will mainly analyze the relation between $\hbLa_R$ and $\btheta$ in Thm. <ref>, which lie in an $R$ dimensional subspace. Here we will build the connection from the $d$ dimensional problem to $R$ dimensional, mainly computing the equivalent noise below. The equivalent noise consists of original noise and the extra noise caused by PCA truncation.

Let $\x_{R}$ be the projection of $\x$ onto the $R$-dimensional subspace spanned by columns of $\mtx{U}_1$, and $\x_{R^\perp}$ is the projection of $\x$ onto the orthogonal complement. Namely, $\x_{R}=\mtx{U}_1^\top\x \in \R^{R}$ and $\x_{R^\perp}=\mtx{U}_2^\top\x\in \R^{(d-R)}$. Similarly we can define $\bt_{R}$ and $\bt_{R^\perp}$. Thus,

\begin{align}
y = \x^\top\bt + \eps = \x_{R}^\top\bt_{R} + \x_{R^\perp}^\top\bt_{R^\perp} + \eps
\end{align}

We can treat $\eps_{R}=\x_{R^\perp}^\top\bt_{R^\perp} + \eps$ as the new noise, and try to solve for $\bt_{R}$. Then define $\mathbf{\Sigma}_{T,R^\perp}$ as the matrix containing the same eigenvectors as $\B$ and the top $R$ eigenvalues are zeroed out, our noise variance becomes $\sigma_R^2=\sigma ^2 + \Eb(\|\x_{R^\perp}\|^2\|\bt_{R^\perp}\|^2)= \sigma^2 + \tr(\Bbar)-\tr(\Bbar^R)$ in our algorithm. If we are still in overparameterized regime, namely $R>\nfs$, then we define optimal representation on top of it.

In summary, the $R$-SVD truncation reduces the search space of $\hbLa$ into $R$ dimensional space, where the covariance of the noise in $\y$ increases from $\sigma^2 \Ib$ to $\sigma^2_R\Ib$.

§.§ Distributional characterization of least norm solution

In this part, for simplicity of discussion, we focus on the $R$ dimensional space while omitting the projection step, and the equivalence of a diagonal eigen-weighting matrix $\hbLa_R\in\R^{R\times R}$ and $\btheta\in\R^R$ in Thm. <ref>. Here, we assume a truncated feature matrix $\Xasym\in\R^{n \times R}$ where the feature is projected into an $R$ dimensional space.

Define $\Xasym\in \R^{n \times R}, \yasym\in\R^n$. We study the following least norm solution of the least squares problem

\begin{align}
\hat \bt = \arg\min_{\bt'}\ \|\bt'\|,\quad \mbox{s.t.,}~\Xasym\bt' = \yasym
\end{align}

Assume the rows of $\tilde \X$ are independently drawn from $\Nn(0, \bSiasym)$.
We focus on a double asymptotic regime where $R,n\rightarrow\infty$ at fixed overparameterization ratio $\kappa:=R/n>0$.

assumptionasstwo The covariance matrix $\bSiasym$ is diagonal and there exist constants $\Sigma_{\min},\Sigma_{\max}\in(0,\infty)$ such that:
0\prec\Sigma_{\min}\Ib \preceq \bSiasym \preceq \Sigma_{\max}\Ib.

The joint empirical distribution of $\{(\lambda_i(\bSiasym),\bt_i)\}_{i\in[R]}$ converges in Wasserstein-k distance to a probability distribution $\mu$ on $\R_{>0}\times\R$ for some $\ntask\geq 4$. That is
\frac{1}{R}\sum_{i\in[R]}\delta_{(\lambda_i(\bSiasym),\bt_i)} \stackrel{W_k}{\Longrightarrow} \mu.

<cit.> Let random variables $(\Sigma,B)\sim \mu$ (where $\mu$ is defined in Assumption <ref>) and fix $\kappa>1$. Define parameter $\xi$ as the unique positive solution to the following equation

\begin{align}\label{eq:ksi}
\E_{\mu}\Big[ \big({1+(\xi\cdot\Sigma)^{-1}}\big)^{-1} \Big] = {\kappa^{-1}}\,.
\end{align}

Define the positive parameter $\gamma$ as follows:

\begin{align}
\label{eq:gamma}
\hspace{-0.1in}\gamma :=
\Big({\sigma^2 + \E_{\mu}\Big[\frac{B^2\Sigma}{(1+\xi\Sigma)^2}\Big]}\Big)\Big/\Big({1-\kappa\E_{\mu}\Big[\frac{1}{\left(1+(\xi\Sigma)^{-1}\right)^2}\Big]}\Big).
\end{align}

With these and $H\sim\Nn(0,1)$, define the random variable

\begin{align}
X_{\kappa,\sigma^2}(\Sigma,B,H) := \Big(1-\frac{1}{1+ \xi\Sigma}\Big) B + \sqrt{\kappa}\frac{\sqrt{\gamma}\,\Sigma^{-1/2}}{1+(\xi\Sigma)^{-1}} H, \label{eq:X}
\end{align}

and let $\Pi_{\kappa,\sigma^2}$ be its distribution.

[Asymptotic distribution characterization – Overparameterized linear Gaussian problem]theoremmainthm
<cit.> Fix $\kappa>1$ and suppose Assumptions <ref> and <ref> hold. Let

\[
\frac{1}{R}\sum_{i=1}^{R}\delta_{\sqrt{R}{\hat \bt}_i,\sqrt{R}\bt_i,\bSiasymi}
\]
be the joint empirical distribution of $({\sqrt{R}\hat \bt},\sqrt{R}\bt,\bSiasym)$ and it converges to a fixed distribution as dimension grows. Let $f:\R^3\rightarrow\R$ be a function in $\rm{PL}(2)$. We have that

\begin{align}\label{eq:thm}
\hspace{-0.1in}\frac{1}{R} \sum_{i=1}^{R} f(\sqrt{R}\hat\bt_i,\sqrt{R}\bt_i,\bSiasymi) \rP \E\left[f(X_{\kappa,\sigma^2},B,\Sigma) \right].
\end{align}

In particular, the risk is given by

\begin{align}
\text{risk}({\hat \bt}_n)&\rP \E[\Sigma (B-X_{\kappa,\sigma^2})]+\sigmaze^2\\
&=\E[\frac{\Sigma}{(1+\xi\Sigma)^2} B^2+\frac{\kappa \gamma}{(1+(\xi\Sigma)^{-1})^2}] + \sigmaze^2.
\end{align}

§.§ Finding Optimal Representation

Now, for simplicity (and actually without losing generality) assume $\bSiasym=\Iden$. This means that empirical measure of $\bSi$ trivially converges to $\Sigma=1$. With the representation $\bLa$ with asymptotic distribution $\Lambda$, the ML problem has the following mapping

\[
\bt\rightarrow \hbLa_R^{-1}\bt\quad\text{and}\quad \bSiasym\rightarrow \hbLa_R\bSiasym\hbLa_R.
\]

This means the empirical measure converges to the following mapped distributions

\[
B\rightarrow \bar{B}=\Lambda^{-1}B\quad \text{and}\quad \Sigma=1\rightarrow \bar{\Sigma}=\Lambda^2\Sigma=\Lambda^2.
\]

Our question: Craft the optimal distribution $\Lambda$ to minimize the representation learning risk. Specifically, for a given $(B,\Lambda)$ pair, we know from the theorem above that

\begin{align}
\text{risk}^{\hbLa_R}({\hat \bt}_n) &\rP\E[\frac{\bar{\Sigma}}{(1+\xi\bar{\Sigma})^2} \bar{B}^2+\frac{\kappa \gamma}{(1+(\xi\bar{\Sigma})^{-1})^2}]+ \sigmaze^2\\
&=\E[\frac{B^2}{(1+\xi\Lambda^2)^2}+\frac{\kappa \gamma}{(1+(\xi\Lambda^2)^{-1})^2}]+ \sigmaze^2.
\end{align}

Thus, the optimal weighting strategy (asymptotically) is given by the distribution

\[
\Lambda^*=\arg\min_{\Lambda}\E[\frac{B^2}{(1+\xi\Lambda^2)^2}+\frac{\kappa \gamma}{(1+(\xi\Lambda^2)^{-1})^2}],
\]

where $\gamma,\xi$ are strictly positive scalars that are also functions of $\Lambda$.

§.§ Non-asymptotic Analysis (for simpler insights)

We apply the discussion iin Sec. <ref> non-asymptotically in few-shot learning. Remember we define $\X\in \R^{\nfs \times R}, \y\in\R^{\nfs}$, each row of $\X$ is independently drawn from $\Nn(0, \bSi)$. We study the following least norm solution of the least squares problem

\begin{align}
\hat \bt = \arg\min_{\bt'}\ \|\bt'\|,\quad \mbox{s.t.,}~\X\bt' = \y.
\end{align}

Set $\kappa=R/{\nfs}>1$. Given $\sigmaze>0$, covariance $\bSi$ and latent vector $\bt$ and define the unique non-negative terms $\xi,\gamma,\z\in\R^R$ and $\bg\in\R^R$ as follows:

\begin{align}
&\xi>0\quad\text{is the solution of}\quad \kappa^{-1}={R^{-1}}\sum_{i=1}^R\big({1+(\xi\bSii)^{-1}}\big)^{-1},\nn\\
\end{align}

Let $\h\sim\Nn(0,\Id/R)$. The non-asymptotic distributional prediction is given by the following random vector

\[
{\hat \bt}(\bSi)= \frac{1}{1+(\xi\bSi)^{-1}}\odot\bt+\frac{\sqrt{\kappa\gamma}\sbSiinv}{1+(\xi\bSi)^{-1}}\odot \h.
\]

Note that, the above formulas can be slightly simplified to have a cleaner look by introducing an additional variable $\z=\frac{1}{1+(\xi\bSi)^{-1}}$.

Also note that, the terms in the non-asymptotic distribution characterization and asymptotic distribution characterization have one to one correspondence. Non-asymptotic distribution characterization is essentially a discretized version of asymptotic DC where instead of expectations (which is integral over pdf) we have summations.

Now, we can use this distribution to predict the test risk by using Def. <ref> in the risk expression.

Going back to representation question, without losing generality, assume $\bSi=\Iden$ and let us find optimal $\hbLa_R$. Then

\[
{\hat \bt}=\hbLa_R\left[\frac{1}{1+(\xi\hbLa_R^2)^{-1}}\odot\hbLa_R^{-1}\bt+\frac{\sqrt{\kappa\gamma}\hbLa_R^{-1}}{1+(\xi\hbLa_R^2)^{-1}}\odot \h\right].
\]

The risk is given by (using $\h\sim\Nn(0,\Iden_p)$)

\begin{align}
\text{risk}^{\hbLa_R}({\hat \bt}_n) - \sigmaze^2
&=\E[({\hat \bt}-\bt)^\top \bSi ({\hat \bt}-\bt)]\\
&=\sum_{i=1}^R \frac{\Bi}{(1+\xi(\hbLaRi)^2)^2}+\sum_{i=1}^R \frac{\kappa\gamma}{(1+(\xi(\hbLaRi)^2)^{-1})^2}.
\end{align}

Here, note that $\xi$ is function of $\bLa$ and $\gamma$ is function of $\bt,\bLa$. If we don't know $\B$, we use the estimation from representation learning $\hB$ instead.

To find the optimal representation, we will solve the following optimization problem that minimizes the risk.

\begin{equation}\label{eq: thm1 obj f}
\begin{array}{rrclcl}
\displaystyle \min_{\bLa} & \multicolumn{3}{l}{\sum_{i=1}^R \dfrac{\bt_i^2}{(1+\xi(\hbLaRi)^2)^2}+\sum_{i=1}^R\dfrac{\kappa\gamma}{(1+(\xi(\hbLaRi)^2)^{-1})^{2}}}\\
\\
\textrm{s.t.} & \kappa^{-1}=\dfrac{1}{R}\sum_{i=1}^R (1+(\xi(\hbLaRi)^2)^{-1})^{-1}\\
\\
& \gamma=\dfrac{\sigmaze^2+\sum_{i=1}^R \frac{\bt_i^2}{(1+\xi(\hbLaRi)^2)^2}}{1-\frac{\kappa}{R}\sum_{i=1}^R (1+(\xi(\hbLaRi)^2)^{-1})^{-2}}.
\end{array}
\end{equation}

So we plug in the expression of $\gamma$ and get

\begin{align}
\kappa\gamma&=\dfrac{\sigmaze^2+\frac{1}{R}\sum_{i=1}^R \frac{\bt_i^2}{(1+\xi(\hbLaRi)^2)^2}}{\kappa^{-1}-\frac{1}{R}\sum_{i=1}^R (1+(\xi(\hbLaRi)^2)^{-1})^{-2}}
= \dfrac{R\sigmaze^2+\sum_{i=1}^R \frac{\bt_i^2}{(1+\xi(\hbLaRi)^2)^2}}{\sum \frac{\xi(\hbLaRi)^2}{(1 + \xi(\hbLaRi)^2)^2} }.
\end{align}

Let $\btheta_i=\frac{\xi(\hbLaRi)^2}{1+\xi(\hbLaRi)^2}$, then the objective function becomes

\[
\sum_{i=1}^R\Bi(1-\btheta_i)^2+(\sum_{i=1}^R \btheta_i^2)\dfrac{R\sigmaze^2+\sum \Bi(1-\btheta_i)^2}{\sum_{i=1}^R \btheta_i(1-\btheta_i)} = \dfrac{\nfs(\sum_{i=1}^R\Bi(1-\btheta_i)^2)+R\sigmaze^2(\sum_{i=1}^R \btheta_i^2)}{{\nfs}-\sum_{i=1}^R\btheta_i^2}
\]
such that $0\leq\btheta_i<1$ and $\sum_{i=1}^R \btheta_i = \frac{R}{\kappa}={\nfs}$. This quantity is same as the objective (<ref>). We divide this quantity by $d$ to get the risk function, which is same as the definition of $f$ in (<ref>).

§.§ Solving the optimization problem.

Here, we propose the algorithm for minimizing $f(\btheta)$. We explore the KKT condition for its optimality.

The objective function is

\begin{align}
f(\btheta) = \sum_{i=1}^R\Bi(1-\btheta_i)^2+(\sum_{i=1}^R \btheta_i^2)\dfrac{R\sigmaze^2+\sum \Bi(1-\btheta_i)^2}{\sum_{i=1}^R \btheta_i(1-\btheta_i)}. \label{eq:risk_asy}
\end{align}

Let $C,S,V\in\R$. Define

\begin{align*}
\phi(\Bi;C,V,S) := \frac{Cp(R-{\nfs}-S)^2}{2\nfs(V + R\sigmaze^2 + (R-{\nfs}-S)\Bi^2)}
\end{align*}

and we find the root of the following equations:

\begin{align*}
\sum_{i=1}^R \phi(\Bi;C,V,S) &= R-{\nfs},\\
\sum_{i=1}^R \phi^2(\Bi;C,V,S) &= S - (2\nfs-R),\\
\sum_{i=1}^R \Bi\phi^2(\Bi;C,V,S) &= V.
\end{align*}

Let $\btheta_i  = 1 - \phi(\Bi;C^*,V^*,S^*)$ where $C^*,V^*,S^*$ are the roots, then

\begin{align*}
\btheta = \arg\min_{\btheta'}\ f(\btheta'),\quad \mbox{s.t.,}\ 0\le \btheta'<1,\ \sum_{i=1}^R \btheta'_i = \nfs.
\end{align*}

Define $s = \sum_{i=1}^R \btheta_i^2$, $\phi_i = 1-\btheta_i$.
Define $Q = \dfrac{1}{R}\sum_{i=1}^R\Bi\phi_i^2$.

\begin{align*}
f(\phi) &= \sum_{i=1}^R\Bi\phi_i^2 + \frac{s}{{\nfs}-s} (R\sigmaze^2 + \sum_{i=1}^R\Bi\phi_i^2)\\
&= R(Q + \frac{s}{{\nfs}-s}(\sigmaze^2 + Q))\\
&= \frac{{R\nfs}}{R-{\nfs}-\sum_{i=1}^R\phi_i^2} (Q + \sigmaze^2).
\end{align*}

The last line uses

\begin{align*}
s &= \sum_{i=1}^R (1 - \phi^2) = R - 2\sum_{i=1}^R\phi_i + \sum_{i=1}^R\phi_i^2= R - 2(R-{\nfs}) + \sum_{i=1}^R\phi_i^2 = 2\nfs - R + \sum_{i=1}^R\phi_i^2.
\end{align*}

Now define $\sum_{i=1}^R\phi_i^2 = S$, and we compute the gradient of $f$, we have

\begin{align*}
\frac{df}{R\phi_i} &= \left( 2\nfs(\sum_{j=1}^R \B_j \phi_j^2 + (R-{\nfs}-s)\Bi )+  2R\nfs\sigmaze^2 \right)\phi_i.
\end{align*}

Suppose $0<\phi_i<1$, then we need $ \frac{df}{R\phi_i}$ equal to each other for all $i$. Suppose $ \frac{df}{R\phi_i} = C$, and denote $\sum \B_j \phi_j^2 = V$, we can solve for $\phi_i$ from $ \frac{df}{R\phi_i} = C$ as

\begin{align}
\phi_i = \frac{Cd(R-{\nfs}-S)^2}{2\nfs(V + R\sigmaze^2 + (R-{\nfs}-S)\Bi^2)} := \phi(\Bi;C,V,S).\label{eq:phi}
\end{align}

We define the function $\phi(\Bi;C,V,S)$ as above, and use the fact that

\begin{align*}
\sum_{i=1}^R \phi(\Bi;C,V,S) &= R-{\nfs},\\
\sum_{i=1}^R \phi^2(\Bi;C,V,S) &= S - (2\nfs-R),\\
\sum_{i=1}^R \Bi\phi^2(\Bi;C,V,S) &= V.
\end{align*}

We can solve[For the root of $3$-dim problem, the worst case we can grid the space and search with time complexity $\Order(\eps^{-3})$.] $C,V,S$ and retrieve $\phi_i$ by (<ref>). $\btheta_i = 1-\phi_i$.

§ ANALYSIS OF MOM ESTIMATORS

§.§ Covariance estimator

We will first present the estimation error of the feature covariance $\bSi$, which is not covered in the main paper due to limitation of space. Note that if $\bSi$ is fully aligned with $\B$, e.g., $\bSi = \B$, then estimating $\bSi$ is enough for getting optimal representation, and we will show it has lower sample complexity and error compared to estimating canonical covariance $\Bbar$. That is a naive case, if it does not work, this intermediate result will help in our latter proof.

We will use the following Bernstein type concentration lemma, generalized from <cit.>:

Let $\Zb\in\R^{n_1\times n_2}$. Choose $T_0, \sigma^2$ such that

* $\mtx{P}(\|\Zb\|\ge C_0T_0 + t) \le \exp(-c\sqrt{t/T_0})$.

* $\|\Eb(\Zb\Zb^\top)\|, \|\Eb(\Zb^\top \Zb)\|\le \sigma^2$.

Then with probability at least $1-(n T_0)^{-c}$, $c>10$,

\begin{align*}
\|\frac{1}{n}\sum_{i=1}^n \Zb_i - \Eb(\Zb_i)\| \lesssim \log(n T_0)\left(\frac{T_0\log(n T_0)}{n} + \frac{\sigma}{\sqrt{n}}\right).
\end{align*}

Define $K = \log^2(C_KnT_0)$ for $C_K>0$, $\Zb' = \Zb\ones(\|\Zb\| \le KT_0)$, then

\begin{align*}
\|\mtx{E}(\Zb - \Zb')\| &\le \int_{KT_0}^\infty \exp(-c\sqrt{t/T_0}) dt \lesssim (1+\sqrt{K})\exp(-c\sqrt{K})T_0 \\
&\lesssim (1+\log(C_KnT_0))(n T_0)^{-C}.
\end{align*}

We can choose $C_K$ large enough so that $C>10$. We will use <cit.>. Set $R = \log^2(C_KnT_0)T_0 + C_0T_0$, $\Delta = (1+\log(C_KnT_0))(n T_0)^{-C}$, $t = C_t \log(n T_0)(\frac{T_0\log(n T_0)}{n} + \frac{\sigma}{\sqrt{n}})$ for some $C_t>0$, plugging in the last inequality of <cit.>, the LHS is smaller than $(n T_0)^{-c}$ for some $c$. We can also check $\mtx{P}(\|\Zb\|\ge R) \le (n T_0)^{-c}$ for some $c$, thus we prove the lemma.

*Feature Covariance. We can directly estimate the covariance of features by

\begin{align}\label{eq:bsi_est}
\hbSi = \frac{1}{\ntot} \sum_{j=1}^{\nspt}\sum_{i=1}^{\ntask} \xij\xij^\top,
\end{align}

The mean of this estimator is $\bSi$ and we can estimate the top $r$ eigenvector of $\bSi$ with $\tilde \Order(r)$ samples.

As we have defined in Phase 1, features $\xij$ are generated from $\Nn(0,\bSi)$. We aim to estimate the covariance $\bSi$. Although there are different kinds of algorithms, such as maximum likelihood estimator <cit.>, to be consistent with the algorithms in the latter sections, we study the sample covariance matrix defined by (<ref>).

Suppose $\x_i$, $i=1,...,\ntot$
are generated independently from $\Nn(0, \bSi)$. We estimate (<ref>), then when $\ntot\gtrsim\TX$, with probability $1 - \Order((\ntot \tr(\bSi))^{-C})$,

\begin{align*}
\|\hbSi - \bSi\| \lesssim \sqrt{\frac{\|\bSi\|\tr(\bSi)}{\ntot}}.
\end{align*}

Denote the span of top $\rf$ eigenvectors of $\bSi$ as $\W$ and the span of top $\rf$ eigenvectors of $\hbSi$ as $\hat \W$. Let $\delta_{\lambda} = \lambda_{\rf}(\bSi) - \lambda_{\rf+1}(\bSi)$. Then if $\ntot \gtrsim \frac{\|\bSi\|\tr(\bSi)}{\delta_{\lambda}^2}$, we have

\begin{align*}
\sin(\angle\W,\hat\W) \lesssim \sqrt{\frac{\|\bSi\|\tr(\bSi)}{\ntot\delta_{\lambda}^2}}
\end{align*}

When $\bSi = \diag{\mtx{I}_{\rf}, 0}$, we have $\sin(\angle\W,\hat\W) \lesssim \sqrt{\frac{\rf}{\ntot}}$.

Lemma <ref> gives the quality of the estimation of the covariance of features $\x$. When the condition number of the matrix $\bSi$ is close to $1$, we need $\ntot\gtrsim d$ to get an estimation with error $\Order(1)$. However, when the matrix $\bSi$ is close to rank $\barrf$, the amount of samples to achieve the same error is smaller, and we can use $\ntot \gtrsim \barrf$ samples to get $\Order(1)$ estimation error.

We will use Bernstein type concentration results to bound its error, and a similar technique will be used for $\hat \bM$ in the next sections.

First we observe that, the features $\xij$ among different tasks are generated i.i.d. from $\Nn(0,\bSi)$. So we can rewrite (<ref>) as

\begin{align}\label{eq:bsi_est_2}
\hbSi = \frac{1}{\ntot }\sum_{i=1}^{\ntot} \x_i\x_i^\top
\end{align}

where $\x_i\sim\Nn(0,\bSi)$. The error of $\hbSi$ depends on $\ntot$ regardless of $\ntask$ and $\nspt$ respectively.

First, we know by concentration inequality

\begin{align}\label{eq:norm_x_bd}
\mtx{P}(\|\x\x^\top\| - \tr(\bSi) \ge t) = \mtx{P}(\|\x\|^2 - \tr(\bSi) \ge t) \le \exp(-c\min\{\frac{t^2}{\tr(\bSisq)}, \frac{t}{\|\bSi\|}\}).
\end{align}

We will use the fact $\sqrt{\tr(\bSisq)} \le \tr(\bSi)$.
Define $K = C_0\log(\ntot \tr(\bSi)) \tr(\bSi)$, $\Zb = \x\x^\top$, $\Zb' = \Zb\cdot\ones\{\|\Zb\|\le K\}$ where $\ones$ means indicator function ($\ones(\mathrm{True}) = 1, \ones(\mathrm{False}) = 0$), for some positive number $C_0$. Then

\begin{align*}
\|\Eb(\Zb - \Zb')\| &\le \int_{t=K}^\infty (\exp(-c \frac{t^2}{\tr^2(\bSi)}) + \exp(-c\frac{t}{\|\bSi\|})) dt\\
& \le \int_{t=K}^\infty (\exp(-c \frac{t}{\tr(\bSi)}) + \exp(-c\frac{t}{\|\bSi\|})) dt\\
&\le 2\frac{\tr(\bSi)}{c} \exp(-c \frac{K}{\tr(\bSi)})\\
&\le \frac{\sqrt{K\tr^2(\bSi)}}{c}\exp(-\frac{cK}{\tr(\bSi)})\\
&\lesssim(\ntot \tr(\bSi))^{-C}
\end{align*}

where $C\ge C_0 - 3/2$.
Then we compute $(\x\x^\top)^2 = \|\x\|^2\x\x^\top$. Let $\bSi$ be diagonal (the proof is invariant from the basis. In other words, if $\bSi$ is not diagonal, then we can make the eigenvectors of $\bSi$ as basis and the proof applies). Then

\begin{align}
\Eb(\|\x\|^2\x\x^\top)_{ij} =
\begin{cases}
\bSi_{ii}(\tr(\bSi) + 2\bSi_{ii}), & i=j,\\
0, & i\ne j. \label{eq: x pow 4}
\end{cases}
\end{align}

So $\|\Eb(\|\x\|^2\x\x^\top)\| \le \|\bSi\|(\tr(\bSi) + 2\|\bSi\|) \approx \|\bSi\|\tr(\bSi)$. $\approx$ means $\gtrsim$ and $\lesssim$.

Using  Lemma <ref>, with (<ref>) and the inequality above, we get that with probability $1 - \Order((\ntot \tr(\bSi))^{-C})$,

\begin{align}\label{eq:bSi err}
\|\hbSi - \bSi\| \lesssim \log(\ntot \tr(\bSi))\left(\frac{\log(\ntot \tr(\bSi))\tr(\bSi)}{\ntot } + \sqrt{\frac{\|\bSi\|\tr(\bSi)}{\ntot }} \right).
\end{align}

If the number above is smaller than $\lambda_{r} - \lambda_{r+1}$, we have that

\begin{align}\label{eq:xxt}
\ntot  \gtrsim \frac{\|\bSi\|\tr(\bSi)}{(\lambda_{r} - \lambda_{r+1})^2}
\end{align}

which is $\Order(r)$ if condition number is $1$.

The bound of the angle of top $R$ eigenvector subspace is a direct application of the following lemma.

<cit.> Let $\A$ be a square matrix. Let $\hat\W$, $\W$ denote the span of top $r$ singular vectors of $\hat A$ and $\A$. Suppose $\|\hat\A - \A\| \le \Delta$, and $\sigma_r(\A) - \sigma_{r+1}(\A) \ge \Delta$, then

\begin{align*}
\sin(\angle \W, \hat \W) \le \frac{\Delta}{\sigma_r(\A) - \sigma_{r+1}(\A) - \Delta}.
\end{align*}

So that the error of principle subspace recovery of feature covariance is upper bounded by $\frac{\|\hbSi - \bSi\|}{\sigma_{r}(\bSi) - \sigma_{r+1}(\bSi) - \|\hbSi - \bSi\|}$, where $\|\hbSi - \bSi\|$ is calculated in (<ref>).

§.§ Method of moment

This section contains three parts. We first bound the norm of task vectors. Then we analyze the second result of Thm. <ref>, where $\nspt$ is lower bounded by effective rank. Last we prove the first result of Thm. <ref> which is a generalization of <cit.>.

§.§.§ Property of task vectors

We first study the property of the tasks $\bt_1,...,\bt_{\ntask}$. We know that, for any $\bt\sim\Nn(0,\B)$,

\begin{align*}
\mtx{P}(\|\bt\|^2 - \tr(\B)\ge t) \le \exp(-c\min\{\frac{t^2}{\tr(\Bsq)}, \frac{t}{\|\B\|}\}).
\end{align*}

So that with probability at least $1-\delta$, we have

\begin{align}
\|\bt_i\|^2 &\lesssim \tr(\B) + \sqrt{ (\log(1/\delta) + \log(\ntask))\tr(\Bsq)} + (\log(1/\delta) + \log(\ntask))\|\B\|\notag\\
&\lesssim \tr(\B) + \log(\ntask/\delta)\sqrt{\tr(\Bsq)}\lesssim \tr(\B)\log(\ntask/\delta),\ \forall i=1,...,\ntask.
\label{eq: bt i norm}
\end{align}

With similar technique we know that with probability at least $1-\delta$,

\begin{align}
\|\bSi\bt_i\|^2 &\lesssim \tr(\bSi\B\bSi) + \log(\ntask/\delta)\sqrt{\tr((\bSi\B\bSi)^2)},\ \forall i=1,...,\ntask. \label{eq: bSibt}\\
\|\sbSi\bt_i\|^2 &\lesssim \tr(\sbSi\B\sbSi) + \log(\ntask/\delta)\sqrt{\tr((\sbSi\B\sbSi)^2)},\ \forall i=1,...,\ntask. \label{eq: sqrbSibt}
\end{align}

We will use $\delta = \ntask^{-c}$ for some constant $c$ so that $\log(\ntask/\delta) = (c+1)\log(\ntask) \approx \log(\ntask)$. Later, we will use the norm bounds of above quantities which happen with probability at least $1 - T^{-c}$.

Now we will prove Theorem <ref> by proving the almost same theorem below.

We generate tasks $\bt_i\in\R^d\sim\Nn(0, \B)$, $i=1,...,\ntask$. We generate features and labels with each task by $(\xij,\yij)$ where $\yij = \xij^\top \bt_i + \eps_j$, where $\xij\sim\Nn(0, \bSi)$, $\eps_j\sim\Nn(0, \sigma)$. There are $\ntask$ tasks and each task has $\nspt$ data, let $\ntot  = \nspt \ntask$. Define $\barB = \frac{1}{\ntask}\sum_{j=1}^{\ntask} \bt_j\bt_j^\top$, $\bar \bM = \bSi\barB\bSi$, $\bar \Qb = \bSi\barB\bSi + \tr(\barB\bSi)\bSi$. The moment estimator above satisfies

\begin{align*}
\|\hat \bM - \bar \bM\| &\lesssim
\sqrt{\frac{\tr(\bSi)\tr^2(\B\bSi)\|\bSi\| +  (\tr(\B\bSi)+\sigma^2\tr(\bSi))\cdot\sigma^2 \|\bSi\|}{\ntot }}\\
&\quad + \frac{\tr^2(\bSi)\tr(\B\bSi) + (\sigma+\sqrt{\tr(\B)})\sigma\tr(\bSi) }{\ntot }.
\end{align*}

We define $\barB = \frac{1}{\ntask}\sum_{i=1}^{\ntask} \bt_i\bt_i^\top$. The proof contains these steps:

* First, we show that the norm of $\|\bt_i\|, \|\bt_i\|_{\bSi}, \|\bt_i\|_{\bSisq}$, $\forall i=1,...,\ntask$ are all around their expected value.

* For any $i\in\{1,...,\ntask\}$, we study the covariance of

\begin{align*}
\frac{1}{2}y_{ij}y_{i(j+\nspt/2)}\cdot(\x_{ij}\x_{i(j+\nspt/2)}^\top + \x_{i(j+\nspt/2)}\x_{ij}^\top) - \bSi\bt_i\bt_i^\top\bSi
\end{align*}

Then use Bernstein inequality to analyze the average of them as compared to $\barB$.

* Finally we bound the distance of $\barB$ and $\B$ via the same technique in Section <ref>.

Then we focus on each term $y_{ij}y_{i(j+\nspt/2)} \x_{ij}\x_{i(j+\nspt/2)}^\top$ or $y_{ij}y_{i(j+\nspt/2)} \x_{i(j+\nspt/2)}\x_{ij}^\top$ and mainly bound their variance. For simplicity of notation, denote a pair of data as $(\x, y)$, $(\x', y')$ and we use $yy'\x(\x')^\top$ to simplify either $y_{ij}y_{i(j+\nspt/2)} \x_{ij}\x_{i(j+\nspt/2)}^\top$ or $y_{ij}y_{i(j+\nspt/2)} \x_{i(j+\nspt/2)}\x_{ij}^\top$.
Then we study the quantity $yy'\x(\x')^\top$. We study a fixed $i$ thus we simplify $\bt_i$ as $\bt$ in our proof. In the following paragraphs, we denote $y = \x^\top \bt + \eps$ and $y' = (\x')^\top \bt + \eps'$, and take the expectation only over $\x,\x',\eps,\eps'$. That means we treat $\bt$ as fixed vectors. In Sec. <ref> we will further use the bounds about $\bt$ above.

First, we compute the expectation of $\hat \bM$.

\begin{align*}
\mtx{E}(y\x) &= \Eb(\x\x^\top\bt) = \bSi\bt.
\end{align*}

So the expectation of $\bM$ is

\begin{align}
\Eb_{\x}(\hat \bM) = \Eb_{\x} (yy'\x(\x')^\top) = [\mtx{E}(y\x)][\mtx{E}(y\x)]^\top = \bSi\bt\bt^\top\bSi. \label{eq:exp_yyxx}
\end{align}

Next we use the Bernstein type inequality to bound $\hat \bM - \Eb_{\x}(\hat \bM)$.

Denote $\hat y = \bt^\top \x$, we study the estimator $\hat y\x$ and $\hat y\hat y'\x(\x')^\top$. This is equivalent to the noiseless case $\sigma = 0$.

First we bound $\|\hat y\x\|$. With probability $1-\delta$, $\|\x\|$ is bounded by (<ref>), and $|\hat y| = |\bt^\top\x| \lesssim \|\sbSi\bt\|(1 + \sqrt{\log(1/\delta)})$. Then with probability $1-\delta$

\begin{align*}
\|\hat y\hat y'\x(\x')^\top\| \le  (1 + \log^2(1/\delta))\tr^2(\bSi)\|\sbSi\bt\|^2.
\end{align*}

In other words,

\begin{align}
\mtx{P}(\|\hat y\hat y'\x(\x')^\top\| > \tr^2(\bSi)\|\sbSi\bt\|^2 + t) \le \exp(-c\frac{\sqrt{t}}{\tr(\bSi)\|\sbSi\bt\|^2} )\label{eq:yyxx_prob}
\end{align}

Next we need to compute the expectation of $(\hat y\hat y'\x(\x')^\top)(\hat y\hat y'\x'\x^\top)$.

$\sbSi\bt = \bb$, let $\x = \sbSi\z$ so that $\z$ is standard normal.

\begin{align}
&\quad \mtx{E}_{\x}(\hat y\hat y'\x(\x')^\top) = [\mtx{E}_{\x}(\hat y\x)][\mtx{E}_{\x}(\hat y\x)]^\top\\
&\quad \mtx{E}_{\x}(\hat y\x) = \Eb_{\x}(\bt^\top \x \x) = \Eb_{\x}(\bb^\top \z \sbSi \z) = [\Eb_{\x}(\bb^\top \z\z^\top \sbSi)]^\top = \sbSi\bb.\\
&\Rightarrow \mtx{E}_{\x}(\hat y\hat y'\x(\x')^\top) = \sbSi\bb\bb^\top\sbSi = \sbSi^2\bt\bt^\top(\sbSi^2) = \bSi\bt\bt^\top\bSi.
\end{align}

Denote any vector $\bar\vb\in\R^d$, and let $\vb = \sbSi\bar\vb$.

\begin{align}
&\quad \mtx{E}_{\x}(\bar\vb^\top(\hat y\hat y'\x(\x')^\top)(\hat y\hat y'\x(\x')^\top)^\top\bar\vb)\\
&= \Eb_{\x}[ (\bt^\top \x_1)^2 (\bt^\top \x_2)^2 (\x_1^\top \bar\vb\bar\vb^\top \x_1)(\x_2^\top \x_2)].\\
&= \Eb_{\x}[ (\bb^\top \z)^2 (\z^\top \vb\vb^\top \z)] \Eb_{\x}[ (\bb^\top \z)^2 \z^\top\bSi \z].
\end{align}

We compute $\Eb_{\x}[ (\bb^\top \z)^2 (\z^\top \vb\vb^\top \z)]$ and $\Eb_{\x}[ (\bb^\top \z)^2 \z^\top\bSi \z]$ separately.

\begin{align}
\Eb_{\x}[ (\bb^\top \z)^2 (\z^\top \vb\vb^\top \z)] = \|\bb\|^2 \|\vb\|^2 + 2\sum_{i=1}^{d} \bb_i^2 \vb_i^2.
\end{align}

To bound $\Eb_{\x}[ (\bb^\top \z)^2 \z^\top\bSi \z]$, we denote $\bSi = \sum_{i=1}^d \la_i {\ub}_i{\ub}_i^\top$, then

\begin{align}
\Eb_{\x}[ (\bb^\top \z)^2 \z^\top\bSi \z] &= \sum_{i=1}^d \la_i (\|\bb\|^2 + 2\sum_{j=1}^{d} \bb_j^2 \ub_{ij}^2)\\
& \le 3\tr(\bSi) \|\sbSi\bt\|^2. \label{eq:yyxx_mid}
\end{align}

So that $\mtx{E}_{\x}(\hat y\hat y'\x(\x')^\top) = \bSi\bt\bt^\top\bSi$, and

\begin{align}
\mtx{E}_{\x}(&\bar\vb^\top(\hat y\hat y'\x(\x')^\top)(\hat y\hat y'\x(\x')^\top)^\top\bar\vb) \\
&= (\|\sbSi\bt\|^2 \|\sbSi\bar\vb\|^2 + 2\sum_{i=1}^{d} (\sbSi\bt)_i^2 (\sbSi\bar\vb)_i^2) \cdot (\sum_{i=1}^d \la_i (\|\sbSi\bt\|^2 \|{\ub}_i\|^2 + 2\sum_{j=1}^{d} (\sbSi\bt)_j^2 \ub_{ij}^2)).
\end{align}

The right hand side is upper bounded as

\begin{align}
\mtx{E}_{\x}(\|(\hat y\hat y'\x(\x')^\top)(\hat y\hat y'\x'\x^\top)\|) \lesssim \tr(\bSi)\|\sbSi\bt\|^4\|\bSi\| \label{eq:exp_yyxx_var}
\end{align}

§.§.§ The variance caused by noise.
Suppose $y = \x^\top\bt + \eps$ where $\eps$ is a noise with variance $\sigma^2$. Then we have

\begin{align}
\Eb_{\x}(yy'\x(\x')^\top) &= \Eb_{\x}((\hat y + \eps)(\hat y'+ \eps')\x(\x')^\top)\\
&= \Eb_{\x}(\hat y(\hat y')\x(\x')^\top) + 2\Eb_{\x}(\eps(\x^\top\bt)\x'\x^\top) + \Eb_{\x}(\eps\eps'\x(\x')^\top). \label{eq:y noisy init}
\end{align}

And with simple computation, we have the following theorem (compare to Lemma <ref>)

Let $c>0$ be a constant. The following inequalities hold:

\begin{align}
\mtx{P}(\|\eps(\x^\top\bt)\x'\x^\top\| \ge \sigma\|\sbSi\bt\|\tr(\bSi) + t) &\le \exp(-c\sqrt{\frac{t}{\sigma\|\bt\|\tr(\bSi)}}), \label{eq:exp_yyxx_err1}\\
\| \Eb_{\x}(\eps(\x^\top\bt)\x'\x^\top)(\eps(\x^\top\bt)\x(\x')^\top)\| &\lesssim \sigma^2\tr(\bSi)\|\bSi\| \|\sbSi \bt\|^2,\label{eq:exp_yyxx_err2}\\
\mtx{P}(\|\eps\eps'\x'\x^\top\| \ge \sigma^2\tr(\bSi) + t) &\le \exp(-c\sqrt{\frac{t}{\sigma^2\tr(\bSi)}}),\label{eq:exp_yyxx_err3}\\
\| \Eb_{\x}((\eps\eps'\x'\x^\top)(\eps\eps'\x(\x')^\top))\| & \lesssim \sigma^4 \tr(\bSi)\|\bSi\| .\label{eq:exp_yyxx_err4}
\end{align}

Proof of Lemma <ref>. By concentration of Gaussian random vecotr, we know that $|\eps|\le \sigma (1+\sqrt{\log(1/\delta)})$, $|\x^\top\bt|\le \Order(\|\sqrt{\bSi}\bt\|(1+\sqrt{\log(1/\delta)}))$ and $\|\x\|^2\le \Order(\tr(\bSi) + \tr(\bSi)\log(1/\delta))$ all with probability $1-\delta$, so that we get (<ref>) and (<ref>).

For (<ref>),

\begin{align}
\|\Eb_{\x}(\eps(\x^\top\bt)\x'\x^\top)(\eps(\x^\top\bt)\x(\x')^\top)\|=\|\Eb_{\x} \sigma^2 (\x^\top\bt)^2\x^\top\x\bSi\|
\end{align}

Here we can apply (<ref>). At the end we get
\begin{align}
\|\Eb_{\x}(\eps(\x^\top\bt)\x'\x^\top)(\eps(\x^\top\bt)\x(\x')^\top)\| \lesssim \sigma^2\tr(\bSi)\|\bSi\| \|\sbSi \bt\|^2
\end{align}

For (<ref>),

\begin{align}
\|\Eb_{\x}(\eps\eps'\x'\x^\top)(\eps\eps'\x(\x')^\top)\| = \| \Eb_{\x} \sigma^4(\x^\top\x)\x'(\x')^\top\|
\end{align}

which easily ends up with (<ref>).

§.§.§ Overall error of MoM

We will assemble the useful bounds of noiseless case and noise together and apply Lemma <ref>. We look at the data in the $j$th batch. From (<ref>)

\begin{align*}
\Eb_{\x} (\frac{1}{2}y_{ij}y_{i(j+\nspt/2)} (\x_{ij}\x_{i(j+\nspt/2)}^\top + \x_{i(j+\nspt/2)}\x_{ij}^\top) - \bSi\bt_i\bt_i^\top\bSi) = 0
\end{align*}

From (<ref>) and (<ref>),

\begin{align}
&\mtx{P} (\|\frac{1}{2}y_{ij}y_{i(j+\nspt/2)} (\x_{ij}\x_{i(j+\nspt/2)}^\top + \x_{i(j+\nspt/2)}\x_{ij}^\top) - \bSi\bt_i\bt_i^\top\bSi\|\\
&\quad \ge C\log(\ntask)(\sqrt{\tr((\bSi\B\bSi)^2)} + \tr^2(\bSi)\sqrt{\tr((\sbSi\B\sbSi)^2)}) + \sigma^2\tr(\bSi) + t) \\
&\le \exp\left(-c \sqrt{\frac{t}{\max\left\{ \log(\ntask) \tr(\bSi)\sqrt{\tr((\sbSi\B\sbSi)^2)} ,\sigma^2\tr(\bSi)\right\}}}\right)
\end{align}

Note that

\begin{align*}
&\quad \|\frac{1}{2}y_{ij}y_{i(j+\nspt/2)} (\x_{ij}\x_{i(j+\nspt/2)}^\top + \x_{i(j+\nspt/2)}\x_{ij}^\top) - \bSi\bt_i\bt_i^\top\bSi\| \\
&\le \|\frac{1}{2}y_{ij}y_{i(j+\nspt/2)} (\x_{ij}\x_{i(j+\nspt/2)}^\top + \x_{i(j+\nspt/2)}\x_{ij}^\top) \| + \|\bSi\bt_i\bt_i^\top\bSi\| \\
&\lesssim \|\frac{1}{2}y_{ij}y_{i(j+\nspt/2)} (\x_{ij}\x_{i(j+\nspt/2)}^\top + \x_{i(j+\nspt/2)}\x_{ij}^\top) \| + \log^2(\ntask)\tr((\bSi\B\bSi)^2)
\end{align*}

Summing up the upper bound of (<ref>) and (<ref>) in (<ref>)

\begin{align}
&\quad \|\Eb_{\x} (\frac{1}{2}y_{ij}y_{i(j+\nspt/2)} (\x_{ij}\x_{i(j+\nspt/2)}^\top + \x_{i(j+\nspt/2)}\x_{ij}^\top) )^2\| \\
&\lesssim \log^2(\ntask)\tr(\bSi)\tr^2((\sbSi\B\sbSi)^2)\|\bSi\| + \sigma^4 \tr(\bSi) \|\bSi\|.
\end{align}

Thus we can denote

\begin{align}
T_0 &=  (\log(\ntask) \sqrt{\tr((\B\bSi)^2)}  + \sigma^2)\tr(\bSi),\\
\sigma &= (\log^2(\ntask)\tr(\bSi)\tr^2((\B\bSi)^2)\|\bSi\| + \sigma^4 \tr(\bSi) \|\bSi\|)^{1/2},
\end{align}

and apply Lemma <ref> by plugging in $T_0$ and $\sigma$,

\begin{align}
\log^{-1}(\ntot T_0) \|\frac{1}{2n} \sum_{i=1}^{\nspt}&(y_iy_i'(\x_i(\x_i')^\top + \x_i'\x_i^\top)) - \frac{1}{\ntask}\bSi (\sum_{i=1}^{\ntask} \bt_i\bt_i^\top) \bSi \| \le \Order(\frac{T_0\log(\ntot T_0)}{\nspt} + \frac{\sigma}{\sqrt{\nspt}}). \label{eq:mom_bd}
\end{align}

\begin{align*}
\|\frac{1}{2n}\sum_{i=1}^{\nspt}&(y_iy_i'(\x_i(\x_i')^\top + \x_i'\x_i^\top)) - \Eb(yy'\x'\x^\top)\| \le\\
&\tilde O\left(\sqrt{\frac{{\color{blue}\tr(\bSi)\|\sbSi\bt\|^4\|\bSi\|} + {\color{red} (\|\sbSi\bt\|^2+\sigma^2\tr(\bSi))\cdot\sigma^2 \|\bSi\|}}{\nspt}}\right)\\
&+ \tilde O\left(\frac{{\color{blue}\tr^2(\bSi)\|\sbSi\bt\|^2} + {\color{red}(\sigma+\|\bt\|)\sigma\tr(\bSi) }}{\nspt}\right).
\end{align*}

Blue part is the concentration of $(\x^\top\bt)\x$ and red part is the concentration of noise. Taking expectation over $\bt$ yields the bound in the Lemma.

If $\bSi = \diag{I_r,0}$, $\sigma=0$, then the error is
$\tilde \Order(\sqrt{\frac{r\|\bt^r\|^4}{\nspt}}) + \tilde \Order(\frac{r\|\bt^r\|^2}{\nspt})$.
The empirical covariance estimator of $\bt$ has the same concentration property as Lemma <ref>, which we will omit the proof of.
(of Lemma <ref>)
Suppose $\bt_i$, $i=1,...,\ntask$ are generated independently from $\Nn(0, \B)$. Then with probability $1 - \Order((\ntot\tr(\B))^{-C})$,

\begin{align*}
&\|\barB - \B\| \lesssim \log(\ntask\tr(\B))(\frac{\log(\ntot\tr(\B))\tr(\B)}{\ntask} + \sqrt{\frac{\|\B\|\tr(\B)}{\ntask}})
\end{align*}

Ensembling Corollary <ref> with Theorem <ref> we can prove Theorem <ref>.

§.§ Method of moment - another estimator

In this section, we will define another moment estimator, which as well potentially recovers the subspace we want. Define

\begin{align*}
\hat \Qb = \frac{1}{\ntask} \sum_{i=1}^{\ntask} \frac{1}{\nspt} \sum_{j=1}^{\nspt} y_{ij}^2\xij\xij^\top.
\end{align*}

The mean of estimator $\hat \Qb$ is

\begin{align*}
\Eb(\Qb) := \Qb = \bSi\B\bSi + \tr(\B\bSi)\bSi %\tr(\sbSi\B\sbSi)\bSi.
\end{align*}

\begin{align*}
\bar \Qb = \bSi\barB\bSi + \tr(\barB\bSi)\bSi %\tr(\sbSi\B\sbSi)\bSi.
\end{align*}

with the same assumption as Theorem <ref>, we have

\begin{align*}
\|\hat \Qb - \bar \Qb\| &\lesssim
\sqrt{\frac{\tr(\bSi)\tr^2(\B\bSi)\|\bSi\| +  (\tr(\B\bSi)+\sigma^2\tr(\bSi))\cdot\sigma^2 \|\bSi\|}{\ntot }}\\
&\ + \frac{\tr^2(\bSi)\tr(\B\bSi) + (\sigma+\sqrt{\tr(B)})\sigma\tr(\bSi) }{\ntot }.
\end{align*}

We first compute the expectation of $\hat \Qb$. We simplify the notation the same way as before: fix $\bt$ and let $y = \x^\top\bt + \eps$, and study $y\x$ and $y^2\x\x^\top$. Define $\z_{i}\sim\Nn(0,1)$, and $\vct{\alpha} = \sbSi\bt$.

\begin{align}
\Qb = \Eb(\hat \Qb) &= \Eb y^2\x\x^\top = \Eb( (\bt^\top \x + \eps)^2\x\x^\top ) = \Eb \left[ \sqrt{\sigma_i\sigma_j}\z_i\z_j ((\sum_{k=1}^d \z_k\alpha_k)^2)\right] + \sigma^2\bSi\\
\Eb( \sigma_i\z_i^2 ((\sum_{k=1}^d \z_k\alpha_k)^2) ) &= \Eb(\|\vct{\alpha}\|^2 + 2\alpha_i^2 )\sigma_i\\
\Eb(\sqrt{\sigma_i\sigma_j}\z_i\z_j ((\sum_{k=1}^d \z_k\alpha_k)^2) &= 2\Eb\sqrt{\sigma_i\sigma_j}\alpha_i\alpha_j,\ i\ne j.
\end{align}

\begin{align}
\Qb = \bSi\B\bSi + \tr(\sbSi\B\sbSi)\bSi + \sigma^2\bSi \label{eq:Q}
\end{align}

Now we follow the method in the last section, we will bound

\begin{align}
\mtx{P}(\|y^2\x\x^\top\|\ge t)~\mathrm{and}~ \Eb(\|y^4(\|\x\|^2\x\x^\top)\|).
\end{align}

* $\|\hat y^2\x\x^\top\|$. Let $\z\sim\Nn(0, \Ib_d)$. We first write $\|\hat y^2\x\x^\top\| = \|\bt^\top\sbSi\z\|^2 \|\sbSi\z\|^2$. We know that with probability $\Order(\delta)$

\begin{align}
\|\sbSi\z\|^2 &\ge \tr(\bSi) + \tr(\bSi)\log(1/\delta),\\
\hat y^2 = (\bt^\top\sbSi\z)^2 &\ge \|\sbSi\bt\|^2(1 + \log(1/\delta))
%\|\sbSi\bt\|^2 &\ge \tr(\sbSi\B\sbSi) + \sqrt{\tr((\sbSi\B\sbSi)^2)}\log(1/\delta)
\end{align}

So with probability $1 - \Order(\delta)$ we have

\begin{align}
\|\hat y^2\x\x^\top\| \lesssim  \tr(\bSi)\|\sbSi\bt\|^2 + \tr(\bSi)\|\sbSi\bt\|^2\log^2(1/\delta)
\end{align}

* $\Eb(\|\hat y^4(\|\x\|^2\x\x^\top)\|)$. With the same derivation above we have that with probability $1 - \Order(\delta)$

\begin{align}
\|\hat y^4(\x\x^\top)^2\| = \|\hat y^4\|\x\|^2(\x\x^\top)\| \le c\left( \tr(\bSi)^2\|\sbSi\bt\|^4 + \tr^2(\bSi)\|\sbSi\bt\|^4\log^4(1/\delta)\right)
\end{align}

So for all $K\ge 1$, the expectation is smaller than

\begin{align}
&\quad K\tr(\bSi)\|\sbSi\bt\|^4 + \int_{ K\tr(\bSi)\|\sbSi\bt\|}^\infty \exp(-c(\frac{t}{\tr^2(\bSi)\|\sbSi\bt\|^4})^{1/4})dt\\
&\lesssim K\tr(\bSi)\|\sbSi\bt\|^4 + K^{3/4}\exp(-cK^{1/4})\tr^2(\bSi)\|\sbSi\bt\|^4.
\end{align}

* Noise. Suppose $y = \x^\top\bt + \eps$. Then we have

\begin{align*}
y^2\x\x^\top &= (\x^\top\bt + \eps)^2\x\x^\top \preceq (\x^\top\bt)^2\x\x^\top + \eps^2\x\x^\top
\end{align*}

The first term was computed in the noiseless case, so in this part we focus on $\eps^2\x\x^\top$. And

\begin{align*}
y^4(\|\x\|^2\x\x^\top) &= (\x^\top\bt + \eps)^4(\|\x\|^2\x\x^\top) \preceq  (\x^\top\bt)^4\|\x\|^2\x\x^\top + \eps^4\|\x\|^2\x\x^\top
\end{align*}

The first term is bounded by the last part, we focus on $\eps^4\|\x\|^2\x\x^\top$.

* $\eps^2\x\x^\top$. With probability at least $1-\delta$, we know that $\eps^2 \le \Order(\sigma^2 (1 + \log(1/\delta))$ and $\|\x\x^\top\| \le \tr(\bSi) + \Order(\tr(\bSi)\log(1/\delta))$. Thus with probability at least $1-\delta$ we have

\begin{align*}
\|\eps^2\x\x^\top\| \lesssim \sigma^2\tr(\bSi) + \sigma^2\tr(\bSi)\log^2(1/\delta)
\end{align*}

* $\eps^4(\x\x^\top)^2$. We know that

\begin{align*}
\| \Eb (\eps^4\|\x\|^2\x\x^\top) \| &= 3\sigma^4 \| \Eb(\|\x\|^2\x\x^\top)\| \lesssim \sigma^4\|\bSi\|\tr(\bSi)
\end{align*}

So in the noisy case, we know that

\begin{align}
&\|y^2\x\x^\top\| \lesssim
(\|\sbSi\bt\|^2 + \sigma^2)\tr(\bSi) (\|\sbSi\bt\|^2 + \sigma^2)\tr(\bSi)\log^2(1/\delta),\ \mbox{w.p. } 1 -  \delta\\
&\Eb(\|y^4(\x\x^\top)^2\|) \lesssim (\|\sbSi\bt\|^4 + \sigma^4\|\bSi\|)\tr(\bSi)
\end{align}

We can exactly define the same quantities as in (<ref>), and get a similar result to (<ref>), let $\barB = \frac{1}{\ntask}\sum_{i=1}^{\ntask} \bt_i\bt_i^\top$, then

\begin{align}
\log^{-1}(\ntot T_0) \|\frac{1}{\nspt} \sum_{i=1}^{\nspt}&y_i^2\x_i\x_i^\top  - (\bSi\barB \bSi + \tr(\bSi\barB)\bSi ) \| \lesssim \frac{T_0\log(\ntot T_0)}{\nspt} + \frac{\sigma}{\sqrt{\nspt}}. \label{eq:mom_bd_2}
\end{align}

§.§.§ Estimating with fewer samples when each task contains enough samples

In this part we will prove Theorem <ref>, which is the second case of Theorem <ref>. First we will give a description of standard normal features, then prove the general version.
(Standard normal feature, noiseless)
Let data be generated as in Phase 1, let $\Snorm = \max\{\|\bSi\|, \|\B\|\}$ in this theorem and the following section[in the paper we assume $\Snorm=1$ for simplicity.], $\TbtX = \tr(\B\bSi)$,
$\TX = \tr(\bSi)$, $\Tbt = \tr(\B)$. Suppose $\sigma = 0$, $\bSi = \Ib$, and suppose the rank of $\B$ is $\rt$.
Define $\hat \bt_i = \nspt^{-1}\sum_{j=1}^{\nspt}\yij\xij$,  $\mfkb = [\bt_1,...,\bt_{\ntask}]$, and $\hat \mfkb = [\hat \bt_1,...,\hat \bt_{\ntask}]$. Let $\nspt>c_1\Tbt\lambda^{-1}_{\rt}(\B)$,
with probability $1 - \Order(\ntask^{-C})$, where $C$ is constant,

\begin{align*}
\sigma_{\max}(\hat \mfkb - \mfkb) \lesssim \sqrt{\frac{\ntask \Tbt}{\nspt}}.
\end{align*}

Denote the span of top $\rt$ singular column vectors of $\hat \mfkb$ and $\B$ as $\hat \W, \W$, then

\begin{align*}
\sin(\angle\hat \W, \W) \lesssim \sqrt{\frac{\Tbt}{\nspt\lambda_{\rt}(\B)}}.
\end{align*}

For example, if $\B = \diag{\Ib_{\rt},0}$, then $\sin(\angle\hat \W, \W) \lesssim \sqrt{\rt/\nspt}$.

We first estimate $\bt_i$ with

\begin{align*}
\hat \bt_i = \frac{1}{\nspt}\sum_{j=1}^{\nspt} \yij\xij.
\end{align*}

Then we fix $\bt_i$ and compute the covariance of $\yij\xij$ (its mean is $\bt_i$).

\begin{align*}
\mathrm{Cov}(\yij\xij - \bt_i) &= \Eb(\xij\xij^\top\bt_i\bt_i^\top\xij\xij^\top) - \bt_i\bt_i^\top \precsim \|\bt_i\|^2\Ib.
\end{align*}

The first term is similar to (<ref>), where the bound can is in <cit.>.  The vector $\hat \bt_i$ is the average of $\yij\xij$ over all $j$. With concentration we know that

\begin{align}
\mathrm{Cov}(\hat \bt_i - \bt_i) \precsim \frac{\|\bt_i\|^2}{\nspt}\Ib. \label{eq:cov_mfkb}
\end{align}

Let $\mfkb = [\bt_1,...,\bt_{\ntask}]$, and $\hat \mfkb = [\hat \bt_1,...,\hat \bt_{\ntask}]$. Then we know the covariance of each column of $\hat \mfkb - \mfkb$ is bounded by (<ref>). Thus with a constant $c$ and probability $1 - \exp(-c\ntask^2)$,

\begin{align}
\sigma_{\max}^2(\hat \mfkb - \mfkb) \lesssim \frac{\ntask\|\bt_i\|^2}{\nspt}.\label{eq:max_mfkb}
\end{align}

We have proved in (<ref>) that $\|\bt_i\|^2 \le \log(\ntask)\tr(\B)$ with probability $1-\ntask^{-c}$. The columns of $\mfkb$ is generated from $\Nn(0,\B)$, so that

\begin{align*}
\sigma_{\max}(\hat \mfkb - \mfkb) \lesssim \sqrt{\frac{\ntask\log(\ntask)\tr(\B)}{\nspt}}.
\end{align*}

Now we study $\mfkb$. We know that $\Eb(\mfkb\mfkb^\top) = \Eb(\sum_{i=1}^{\ntask} \bt_i\bt_i^\top) = \ntask\B$. $\mfkb$ is a matrix with independent columns. Thus let $\nspt>c_1\tr(\B)\lambda^{-1}_{\rt}(\B)$,  $\ntask>\max\{c_2d, \frac{\|\B\|\tr(\B)}{\lambda_{\rt}^2(\B)}\}$, then with Lemma <ref>, for Gaussian matrix with independent columns <cit.>, with probability at least $1 - \Order(\ntask^{-c_3} + (\ntask\tr(\B))^{-c_4} + \exp(-c_5\ntask^2)) = 1 - \Order(\ntask^{-C})$, where $c_i$ are constants,

\begin{align*}
\sigma_{\rt}(\mfkb) \ge \sqrt{\ntask\lambda_{\rt}(\B) - \Order(\sqrt{\ntask\|\B\|\tr(\B)})}.
\end{align*}

Denote the span of top $\rt$ singular vectors of $\hat \mfkb$ and $\B$ as $\hat \W, \W$,
with Lemma <ref>,

\begin{align*}
\sin(\angle\hat \W, \W) \le \sqrt{\frac{\log(\ntask)\tr(\B)}{\nspt\lambda_{\rt}(\B)}}.
\end{align*}

Next, we will propose a theorem with general feature covariance and noisy data, which is a generalization of Theorem <ref>.

Let data be generated as in Phase 1.
Suppose $\hat \bb_i = \nspt^{-1}\sum_{j=1}^{\nspt}\yij\xij$, $\mfkb = \bSi[\bt_1,...,\bt_{\ntask}]$, and $\hat \mfkb = [\hat \bb_1,...,\hat \bb_{\ntask}]$. Let $\delta_\lambda = \lambda_{\rt}(\bSi\B\bSi) - \lambda_{{\rt}+1}(\bSi\B\bSi))$, suppose $\bSi$ is approximately rank $\rf$,

\begin{align*}
\nspt &\gtrsim (\tr(\B\bSi)+\sigma^2)\|\bSi\|,\\
\ntask &\gtrsim \max\{\rf, \frac{d\lambda_{\rf+1}(\bSi)}{\|\bSi\|} \},
\end{align*}

then with probability $1 - \Order(\ntask^{-C})$, where $C$ is constant,

\begin{align*}
\sigma_{\max}(\hat \mfkb - \mfkb) \lesssim \sqrt{\frac{\ntask(\tr(\B\bSi)+\sigma^2)\|\bSi\|}{\nspt}}.
\end{align*}

Denote the span of top ${\rt}$ singular vectors of $\hat \mfkb$ and $\bSi\B\bSi$ as $\hat \W, \W$, if further we assume $\ntask \gtrsim \frac{\|\bSi\B\bSi\|\tr(\bSi\B\bSi)}{\delta^2_\lambda}$, then

\begin{align*}
\sin(\angle\hat \W, \W) \lesssim \sqrt{\frac{(\tr(\B\bSi)+\sigma^2)\|\bSi\|}{\nspt\delta^2_{\lambda}}}.
\end{align*}

Suppose $\bSi = \diag{\Ib_{\rf},\iota\Ib_{d-\rf}}$, and $\B = \diag{\Ib_{\rt}, 0}$, $\sigma = 0$. Suppose $\iota d < \rf$. Then with $\ntask \gtrsim \rf$, $\nspt \gtrsim \rt$ so that $\ntot \gtrsim \rf\rt$,

\begin{align*}
\sin(\angle\hat \W, \W) \lesssim \sqrt{\rt/n}.
\end{align*}

We let $\xij\sim\Nn(0, \bSi)$.
For the $i$th task, let

\begin{align*}
\hat \bb_i = \frac{1}{\nspt}\sum_{j=1}^{\nspt} \yij\xij.
\end{align*}

We fix $\bt_i$ and compute

\begin{align}
\mtx{E}(\yij\xij) &\precsim \Eb(\xij\xij^\top\bt_i) = \bSi\bt_i, \label{eq:hbb mean}
\end{align}

\begin{align}
\mathrm{Cov}(\yij\xij - \bSi\bt_i) &\precsim (\bt_i^\top\bSi\bt_i)\bSi + \sigma^2\bSi. \label{eq: cov xijyij}
\end{align}

To get the bound above, we can adopt the technique in <cit.> such that, write $\xij = \sbSi \z$, and reduce to $\Eb((\z^\top\sbSi\bt_i)^2\sbSi\z\z^\top\sbSi) $. The proof of <cit.> gives the explicit bound of $\|\Eb((\z^\top\vct{\alpha})^2\z\z^\top)\|$ for any $\vct{\alpha}$ that equals above. The vector $\hat \bb_i$ is the average of $\yij\xij$ over all $j=1,...,\nspt$. With concentration we know that

\begin{align}
\mathrm{Cov}(\hat \bb_i - \bSi\bt_i) \precsim \frac{\bt_i^\top\bSi\bt_i + \sigma^2}{\nspt}\bSi. \label{eq:cov_mfkb_gen}
\end{align}

Suppose $\mfkb = \bSi[\bt_1,...,\bt_{\ntask}]$, and $\hat \mfkb = [\bb_1,...,\bb_{\ntask}]$. $\hat \mfkb - \mfkb$ is a matrix with independent columns. Suppose $\X$ is approximately rank $\rf$, Let $\Vb_{\rf}\in\R^{d\times d}$ be the projection onto the top-$R$ sigular vector space of $\bSi$ and $\Vb_{\rf^\perp}\in\R^{d\times d}$ be the projection onto the $\rf+1$ to $d$th sigular vector space of $\bSi$.
With $\ntask$ columns and $\ntask\ge \rf$, we know that

\begin{align*}
\sigma_{\max}(\Vb_{\rf}(\hat \mfkb - \mfkb)) &\lesssim  \frac{\ntask(\max_i\bt_i^\top\bSi\bt_i + \sigma^2)\|\bSi\|}{\nspt}\\
\sigma_{\max}(\Vb_{\rf^\perp}(\hat \mfkb - \mfkb)) &\lesssim \frac{\max\{\ntask,d\}(\max_i\bt_i^\top\bSi\bt_i + \sigma^2)\lambda_{{\rt}+1}(\bSi)}{\nspt}
\end{align*}

With similar argument as before, with probability $1 - \exp(-c\ntask^2)$ for constant $c$,

\begin{align}
\sigma_{\max}^2(\hat \mfkb - \mfkb) \lesssim \frac{\colsmfkb(\max_i\bt_i^\top\bSi\bt_i + \sigma^2)\|\bSi\|}{\nspt}.\label{eq:max_mfkb_gen}
\end{align}

We know in (<ref>) that $\|\sbSi\bt_i\|^2 \le \Order(\log(\ntask)\tr(\B\bSi))$ with probability $1 - \ntask^{-c}$ for constant $c$. So that

\begin{align}
\sigma_{\max}(\hat \mfkb - \mfkb) \lesssim \sqrt{\frac{\colsmfkb(\log(\ntask)\tr(\B\bSi)+\sigma^2)\|\bSi\|}{\nspt}}. \label{eq:mfkb err}
\end{align}

Now we study $\mfkb$. $\Eb(\mfkb\mfkb^\top) = \Eb(\bSi(\sum_{i=1}^{\ntask} \bt_i\bt_i^\top)\bSi) = \ntask\bSi\B\bSi$.

Thus let

\begin{align*}
\nspt &> C_1(\log(\ntask)\tr(\B\bSi)+\sigma^2)\|\bSi\|.
\end{align*}

Now apply the concentration of Gaussian matrix with independent columns <cit.>. With probability $1 - \Order(\ntask^{-C_1} + (\ntask\tr(\bSi\B\bSi))^{-C_2} + \exp(-C_3\ntask^2))$, where $C_i$ are constants (the probability can be simplified as $1 - \Order(\ntask^{-C})$),

\begin{align*}
\sigma_{\rt}(\mfkb) \ge \sqrt{\ntask(\lambda_{{\rt}}(\bSi\B\bSi) - \lambda_{{\rt}+1}(\bSi\B\bSi)) - \Order(\sqrt{\ntask\|\bSi\B\bSi\|\tr(\bSi\B\bSi)})}.
\end{align*}

Denote the span of top $\rt$ singular vectors of $\hat \mfkb$ and $\bSi\B\bSi$ as $\hat \W, \W$, let

\begin{align}
\ntask \gtrsim \max\{\rf, \frac{d\lambda_{\rf+1}(\bSi)}{\|\bSi\|},\frac{\|\bSi\B\bSi\|\tr(\bSi\B\bSi)}{(\lambda_{{\rt}}(\bSi\B\bSi) - \lambda_{{\rt}+1}(\bSi\B\bSi))^2} \}
\end{align}

we plug in (<ref>) and Lemma <ref>,

\begin{align*}
\sin(\angle\hat \W, \W) &\lesssim \sqrt{(\frac{d\lambda_{\rf+1}(\bSi)}{\ntask\|\bSi\|} + 1) \cdot \frac{(\tr(\B\bSi)+\sigma^2)\|\bSi\|}{\nspt(\lambda_{\rt}(\bSi\B\bSi) - \lambda_{{\rt}+1}(\bSi\B\bSi))}}\\
&\approx \sqrt{\frac{(\tr(\B\bSi)+\sigma^2)\|\bSi\|}{\nspt(\lambda_{\rt}(\bSi\B\bSi) - \lambda_{{\rt}+1}(\bSi\B\bSi))}}.
\end{align*}

§.§.§ Method of moments with arbitrary $\nspt$

In this subsection we will analyze $\hat\mfkb$ with any $\nspt$, and propose the error of MoM estimator.

First, suppose there are at least two samples per task, we can separate the samples into two halves, and compute the following estimator.

Let data be generated as in Phase 1, and let $\nspt$ be a even number. Define $\hat \bb_{i,1} = 2\nspt^{-1}\sum_{j=1}^{\nspt/2}\yij\xij$, $\hat \bb_{i,2} = 2\nspt^{-1}\sum_{j=\nspt/2+1}^{\nspt}\yij\xij$. Define

\begin{align*}
\hat \bM &= \nspt^{-1}\sum_{i=1}^{\ntask} (\bb_{i,1}\bb_{i,2}^\top + \bb_{i,2}\bb_{i,1}^\top),\\
\bM &= \bSi\B\bSi.
\end{align*}

Then there is a constant $c>10$, with probability $1 - \ntot^{-c}$,

\begin{align*}
\| \hat \bM - \bM \| \lesssim (\TbtX + \sigma^2)\sqrt{\frac{\TX}{\ntot}} + \sqrt{\frac{\Tbt}{\ntask}}.
\end{align*}

For simplicity of notation, we will define a random vector $\x$ with zero mean and covariance $\bSi$, a random vector $\bt$ with zero mean and covariance $\B$, a random variable $\eps$ with zero mean and covariance $\sigma$, and they are subGaussian[We remove the subscripts when there is no confusion.]. Let $y = \x^\top\bt + \eps$.
We first estimate the mean of $\hat \bM$.

Note that if we fix $\bt$, $\hat \bb_{i,1}, \hat \bb_{i,2}$ are i.i.d., so

\begin{align*}
\Eb_{\x,\eps}(\hat \bb_{i,1}) &= \Eb_{\x,\eps}(y\x)
= \Eb_{\x,\eps}((\x^\top\bt + \eps)\x)= \bSi\bt,\\
\Eb_{\x,\eps}(\hat \bM) &= \frac{1}{2} (\Eb_{\x,\eps}(\hat \bb_{i,1}) \Eb_{\x,\eps}(\hat \bb_{i,2})^\top + \Eb_{\x,\eps}(\hat \bb_{i,2}) \Eb_{\x,\eps}(\hat \bb_{i,1})^\top)
\\&= \Eb_{\x,\eps}(\hat \bb_{i,1})\Eb_{\x,\eps}(\hat \bb_{i,1})^\top = \frac{1}{\ntask}\bSi(\sum_{i=1}^{\ntask} \bt_i\bt_i^\top) \bSi.
\end{align*}

We take expectation over $\bt_i$ and get $\bM$. We define the right hand side as $\bar \bM$ for the proof below.

Next, we will bound $\|\hat\bM-\bM\|$.

<cit.> proposes that, with probability $1-\delta$,

\begin{align*}
\|\xij\|^2 &\lesssim \log(1/\delta)\tr(\bSi),\\
(\xij^\top\bt_i)^2 &\lesssim \log(1/\delta) \tr(\bSi\B),\\
\eps_{ij}^2 &\lesssim \log(1/\delta)\sigma^2.
\end{align*}

If we enumerate $i=1,...,\ntask$ and $j=1,...,\nspt$, there are in total $\ntask\nspt=\ntot$ terms. So we set $\delta = \ntot^{-c+1}$ for a constant $c>1$, then with probability $1-\ntot^{-c}$, for all $i,j$ we have

\begin{align*}
\|\yij\xij\| = \|(\xij\bt_i+\eps_{ij})\xij\| \lesssim \log^{3/2}(\ntot) \sqrt{(\tr(\bSi\B) + \sigma^2)\tr(\bSi)}.
\end{align*}

Define $\ddelta_{i,l} = \hat\bb_{i,l} - \bSi\bt_i$ for $l=1,2$ (we will use $l=1$ below, the result for $l=2$ is the same). Note that $\ddelta_i$ is zero mean. With <cit.>  we have with probability $1 - \ntot^{-c}$,

\begin{align}
\|\ddelta_{i,1}\| \lesssim \nspt^{-1/2}\log^{5/2}(\ntot) \sqrt{(\tr(\bSi\B) + \sigma^2)\tr(\bSi)} \label{eq:ddelta bd}
\end{align}

\begin{align*}
\Zb_i &= \hat\bb_{i,1}\hat\bb_{i,2}^\top - \Eb_{\x,\eps} (\hat\bb_{i,1}\hat\bb_{i,2}^\top)\\
&= (\bSi\bt_i + \ddelta_{i,1})(\bSi\bt_i + \ddelta_{i,2})^\top - \Eb_{\x,\eps} (\hat\bb_{i,1}\hat\bb_{i,2}^\top)\\
&= \ddelta_{i,1}(\bSi\bt_i)^\top + \bSi\bt_i\ddelta_{i,2}^\top + \ddelta_{i,1}\ddelta_{i,2}^\top - \Eb_{\x,\eps}(\ddelta_{i,1}\ddelta_{i,2}^\top).
\end{align*}

\begin{align}
\|\Eb \Zb_i\Zb_i^\top\| &\le \|\Eb (\bSi\bt_i\ddelta_{i,2}^\top + \ddelta_{i,1}(\bSi\bt_i)^\top)(\bSi\bt_i\ddelta_{i,2}^\top + \ddelta_{i,1}(\bSi\bt_i)^\top)^\top\| \notag\\
&\quad + \|\Eb \ddelta_{i,1}\ddelta_{i,2}^\top\ddelta_{i,2}\ddelta_{i,1}^\top\|. \label{eq: Zbi pause 2}
\end{align}

Then we can use (<ref>) and (<ref>) to bound the first term by

\begin{align*}
\nspt^{-1}\log^{6}(\ntot)(\tr(\bSi\B) + \sigma)\tr(\bSi)\tr(\bSisq\B)\|\bSi\|^2.
\end{align*}

\begin{align*}
\Eb_{\x,\eps} \ddelta_{i,1}\ddelta_{i,2}^\top\ddelta_{i,2}\ddelta_{i,1}^\top &= (\Eb_{\x} \ddelta_{i,2}^\top\ddelta_{i,2})\|\Eb_{\x} \ddelta_{i,1}\ddelta_{i,1}^\top\|\\
&\lesssim \nspt^{-2}(\Eb_{\x,\eps} (\x^\top\bt+ \eps)^2 \x^\top\x) \|\Eb_{\x,\eps} (\x^\top\bt+ \eps)^2 \x\x^\top\|\\
&\lesssim \nspt^{-2}(\tr^2(\bSi\B) + \sigma^4)\tr(\bSi)\|\bSi\|.
\end{align*}

The second line is due to the fact that $\ddelta_{i,l}$ is the difference of $(\x^\top\bt+ \eps)\x$ and its mean, and covariance is upper bounded by variance (not subtracting the mean). The $\nspt^{-2}$ factor comes from the average over $\nspt$ terms. The reasoning of the last line is same
as (<ref>). Now we can go back to (<ref>) and get

\begin{align*}
\|\Eb \Zb_i\Zb_i^\top\| \lesssim \nspt^{-1}\log^{6}(\ntot)(\tr(\bSisq\B) + \tr(\bSi\B) + \sigma^2)^2\tr(\bSi)\|\bSi\|^2.
\end{align*}

Next we need to bound the norm of $\Zb_i$. We use (<ref>) and (<ref>), with probability $1 - \ntot^{-c}$,

\begin{align*}
\|\Zb_i\| &\le \nspt^{-1/2}\log^{3}(\ntot)(\tr(\bSisq\B) + \tr(\bSi\B) + \sigma^2)\sqrt{\tr(\bSi)}\|\bSi\| \\
&\quad + \nspt^{-1}\log^{5}(\ntot) (\tr(\bSi\B) + \sigma^2)\tr(\bSi).
\end{align*}

Define the upper bound for $\|\Eb \Zb_i\Zb_i^\top\|, \|\Zb_i\|$ as $Z_1, Z_2$ (the right hand side of two above inequalities).
Now we apply Bernstein type inequality (Lemma <ref>), with probability $1 - \ntot^{-c}$,

\begin{align*}
&\quad \| \hat \bM - \bar\bM\| \\
\| \ntask^{-1} \sum_{i=1}^{\ntask} \Zb_i - \Eb_{\x} \Zb_i\|\\
&\lesssim \log(\ntask Z_2) \left( \ntask^{-1/2}\log(\ntot) Z_1^{1/2} + \ntask^{-1}Z_2\log(\ntask Z_2) \right)\\
&\lesssim \log(\ntask Z_2) \Big(\sqrt{\frac{\log^{6}(\ntot)(\tr(\bSisq\B) + \tr(\bSi\B) + \sigma^2)^2\tr(\bSi)\|\bSi\|^2}{\nspt\ntask}} \\
&\quad + \frac{\log^{3}(\ntot)(\tr(\bSisq\B) + \tr(\bSi\B) + \sigma^2)\sqrt{\tr(\bSi)}\|\bSi\|}{\nspt^{1/2}\ntask}\\
&\quad + \frac{\log^{5}(\ntot) (\tr(\bSi\B) + \sigma^2)\tr(\bSi)}{\ntask}\Big)\\
&= \log(\ntask Z_2)\cdot \Big(\log^{3}(\ntot)\|\bSi\|(\tr(\bSisq\B) + \tr(\bSi\B) + \sigma^2)\sqrt{\frac{\tr(\bSi)}{\ntot}}\\
&\quad + \frac{\log^{5}(\ntot)(\tr(\bSisq\B) + \tr(\bSi\B) + \sigma^2)\sqrt{\tr(\bSi)}\|\bSi\|}{\ntot^{1/2}\ntask^{1/2}}\Big).
\end{align*}

The term

\begin{align*}
\|\bSi\|(\tr(\bSisq\B) + \tr(\bSi\B) + \sigma^2)\sqrt{\frac{\tr(\bSi)}{\ntot}}
\end{align*}

is the dominant term as shown in the theorem.

The following method of moment estimator is used in <cit.>, where $\nspt\ge 1$. In other words, if there is one sample per task, one can use the following estimator.

Let data be generated as in Phase 1. Define $\hat \bb_i = \nspt^{-1}\sum_{j=1}^{\nspt}\yij\xij$, $\mfkb = \bSi[\bt_1,...,\bt_{\ntask}]$, and $\hat \mfkb = [\hat \bb_1,...,\hat \bb_{\ntask}]$.

\begin{align*}
\hat \BB &= \hat \mfkb \hat \mfkb^\top = \ntask^{-1}\sum_{i=1}^{\ntask} \hat \bb_i\hat \bb_i^\top, \\
\BB &= \Eb(\hat \mfkb \hat \mfkb^\top) = \bSi\B\bSi + \nspt^{-1}(\bSi\B\bSi + \tr(\B\bSi)\bSi + \sigma^2\bSi ) ,\\
\barB &= \sum_{i=1}^{\ntask} \bt_i\bt_i^\top,\\
\bBB &= \bSi\barB\bSi + \nspt^{-1}(\bSi\barB\bSi + \tr(\barB\bSi)\bSi + \sigma^2\bSi )
\end{align*}

With probability $1 - \ntot^c$,

\begin{align*}
\|\hat \BB - \bBB\| \lesssim
\|\bSi\|(\tr(\bSisq\B) + \tr(\bSi\B) + \sigma^2)\sqrt{\frac{\tr(\bSi)}{\ntot}}.
\end{align*}

First, we compute the expectation of $\hat \BB$.

\begin{align}
\Eb_{\x,y,\eps}\hat \BB &= \Eb_{\x,y,\eps} \ntask^{-1}(\sum_{i=1}^{\ntask} \hat\bb_i\hat\bb_i^\top),\notag\\
\Eb_{\x,y,\eps}\hat\bb_i\hat\bb_i^\top &= \Eb_{\x,y,\eps} \left(\nspt^{-1} \sum_{j=1}^{\nspt} (\bt_i^\top\xij + \eps_{ij})\xij\right)\left(\nspt^{-1} \sum_{j=1}^{\nspt} (\bt_i^\top\xij + \eps_{ij})\xij\right)^\top\notag \\
&= \nspt^{-1}\sigma^2\bSi + \Eb_{\x}  (\nspt^{-1}\sum_{j=1}^{\nspt}\xij\xij^\top\bt_i)(\nspt^{-1}\sum_{j=1}^{\nspt}\xij\xij^\top\bt_i)^\top. \label{eq: BBi mean}
\end{align}

Now we will study the second term. (<ref>) states that $\Eb_{\x,y,\eps}(\hat\bb_i) = \bSi\bt_i$. And $\hat\bb_i$ is an average of $\nspt$ terms, we use the expression of the covariance of sample means to get

\begin{align}
\mathbf{Cov}(\hat\bb_i) &= \nspt^{-1}\mathbf{Cov}(\x\x^\top\bt_i), \label{eq: BBi cov} \\
\Eb_{\x,y,\eps}\hat\bb_i\hat\bb_i^\top &=
\Eb_{\x}  (\nspt^{-1}\sum_{j=1}^{\nspt}\xij\xij^\top\bt_i)(\nspt^{-1}\sum_{j=1}^{\nspt}\xij\xij^\top\bt_i)^\top \notag \\
&= \bSi\bt_i\bt_i^\top\bSi + \nspt^{-1} \mathbf{Cov} (\x\x^\top\bt_i) \label{eq: BBi mean 2}
\end{align}

Now we study $\mathbf{Cov}(\x\x^\top\bt_i)$.

\begin{align*}
\mathbf{Cov}(\x\x^\top\bt_i) &= \Eb_{\x}(\x\x^\top\bt_i - \bSi\bt_i)(\x\x^\top\bt_i - \bSi\bt_i)^\top\\
&= \Eb_{\x} (\x\x^\top\bt_i)(\x\x^\top\bt_i)^\top - \bSi\bt_i\bt_i^\top\bSi
\end{align*}

Let $\x = \sqrt{\bSi}\z$ so that $\z\sim\Nn(0,\Ib)$. Let two indices $k,l\in[\d]$. When $k\ne l$,

\begin{align*}
\Eb_{\x} [(\x\x^\top\bt_i)(\x\x^\top\bt_i)^\top]_{kl}
&= \Eb_{\z}(\sum_{j=1}^d\bt_{i,j}\sigma_j\z_j)^2 \sigma_k\z_k\sigma_l\z_l\\
&= 2\sigma_k^2\sigma_l^2\bt_{i,k}\bt_{i,l}
\end{align*}

\begin{align*}
\Eb_{\x} [(\x\x^\top\bt_i)(\x\x^\top\bt_i)^\top]_{kk}
&= \Eb_{\z}(\sum_{j=1}^d\bt_{i,j}\sigma_j\z_j)^2 \sigma_k^2\z_k^2\\
&= \tr(\bt_i^\top\bSi\bt_i)\sigma_k^2 + 2\sigma_k^4\bt_{i,k}^2.
\end{align*}

So that

\begin{align*}
\Eb_{\x} (\x\x^\top\bt_i)(\x\x^\top\bt_i)^\top &= 2\bSi\bt_i\bt_i^\top\bSi + \tr(\bt_i^\top\bSi\bt_i),\\
\mathbf{Cov}(\x\x^\top\bt_i) &= \Eb_{\x} (\x\x^\top\bt_i)(\x\x^\top\bt_i)^\top - \bSi\bt_i\bt_i^\top\bSi \\
&= \bSi\bt_i\bt_i^\top\bSi + \tr(\bt_i^\top\bSi\bt_i)\bSi.
\end{align*}

We plug it back into (<ref>) and (<ref>) and get

\begin{align*}
\Eb_{\x,y,\eps}\hat\bb_i\hat\bb_i^\top &= \bSi\bt_i\bt_i^\top\bSi + \nspt^{-1}(\bSi\bt_i\bt_i^\top\bSi + \tr(\bt_i^\top\bSi\bt_i)\bSi + \sigma^2\bSi ).
\end{align*}

Define $\barB = \frac{1}{\ntask}\sum_{j=1}^{\ntask} \bt_j\bt_j^\top$. So that

\begin{align*}
\Eb_{\x,y,\eps}\hat \BB &= \Eb_{\x,y,\eps} \ntask^{-1}(\sum_{i=1}^{\ntask} \hat\bb_i\hat\bb_i^\top) \\
&= \bSi\barB\bSi + \nspt^{-1}(\bSi\barB\bSi + \tr(\barB\bSi)\bSi + \sigma^2\bSi ) := \bBB.\\
\Eb_{\bt}\hat \BB &= \BB.
\end{align*}

We fix all $\bt_i$ and study $\Eb_{\x,y,\eps}\hat \BB$. Now we need to show how fast $\hat\BB$ converges to $\bBB$.

\begin{align*}
\Zb_i &= \hat\bb_i\hat\bb_i^\top - \Eb_{\x} (\hat\bb_i\hat\bb_i^\top)\\
&= (\bSi\bt_i + \ddelta_i)(\bSi\bt_i + \ddelta_i)^\top - \Eb_{\x}(\bSi\bt_i + \ddelta_i)(\bSi\bt_i + \ddelta_i)^\top\\
&= \bSi\bt_i\ddelta_i^\top + \ddelta_i(\bSi\bt_i)^\top + \ddelta_i\ddelta_i^\top - \Eb_{\x}(\bSi\bt_i\ddelta_i^\top + \ddelta_i(\bSi\bt_i)^\top + \ddelta_i\ddelta_i^\top).
\end{align*}

\begin{align*}
\|\Eb \Zb_i^2\| &\le \|\Eb (\bSi\bt_i\ddelta_i^\top + \ddelta_i(\bSi\bt_i)^\top)^2\| + \|\Eb \ddelta_i\ddelta_i^\top\ddelta_i\ddelta_i^\top\|.
\end{align*}

Then we can use (<ref>) and (<ref>) to bound the first term

\begin{align}
\|\Eb \Zb_i^2\| \lesssim \nspt^{-1}\log^{6}(\ntot)(\tr(\bSi\B) + \sigma)\tr(\bSi)\tr(\bSisq\B)\|\bSi\|^2 + \|\Eb \ddelta_i\ddelta_i^\top\ddelta_i\ddelta_i^\top\| \label{eq: Zbi pause}
\end{align}

So we need to bound $\|\Eb \ddelta_i\ddelta_i^\top\ddelta_i\ddelta_i^\top\|$. Note that $\ddelta_i$ is the average of $\xij(\xij^\top\bt_i + \eps_{ij})$ with respect to index $j=1,...,\nspt$. So we just let $\x\sim \Nn(0,\bSi)$ and study $\x(\x^\top\bt_i+ \eps_{ij})$. Denote it by $\ub_i$.

\begin{align*}
\|\Eb_{\x} \ub_i\ub_i^\top\ub_i\ub_i^\top\| &= \| \Eb_{\x} (\x^\top\bt_i+ \eps_{ij})^4 \x\x^\top\x\x^\top\|\\
&\lesssim \|\Eb_{\x} ((\x^\top\bt_i)^4 + \sigma^4)\x\x^\top\x\x^\top\|\\
&\lesssim (\tr^2(\bSi\B) + \sigma^4)\tr(\bSi)\|\bSi\|.
\end{align*}

So that

\begin{align*}
\|\Eb \ddelta_i\ddelta_i^\top\ddelta_i\ddelta_i^\top\| \lesssim \nspt^{-2}(\tr^2(\bSi\B) + \sigma^4)\tr(\bSi)\|\bSi\|.
\end{align*}

Now we can go back to (<ref>) and get

\begin{align*}
\|\Eb \Zb_i^2\| \lesssim \nspt^{-1}\log^{6}(\ntot)(\tr(\bSisq\B) + \tr(\bSi\B) + \sigma^2)^2\tr(\bSi)\|\bSi\|^2.
\end{align*}

Next we need to bound the norm of $\Zb_i$. We use (<ref>) and (<ref>), with probability $1 - \ntot^{-c}$,

\begin{align*}
\|\Zb_i\| \le \nspt^{-1/2}\log^{3}(\ntot)(\tr(\bSisq\B) + \tr(\bSi\B) + \sigma^2)\sqrt{\tr(\bSi)}\|\bSi\| + \nspt^{-1}\log^{5}(\ntot) (\tr(\bSi\B) + \sigma^2)\tr(\bSi).
\end{align*}

Define the upper bound for $\|\Eb \Zb_i^2\|, \|\Zb_i\|$ as $Z_1, Z_2$ (the right hand side of two above inequalities).
With Bernstein type inequality (Lemma <ref>),with probability $1 - \ntot^{-c}$,

\begin{align*}
&\quad \| \hat \BB - \bBB\| \\
\| \ntask^{-1} \sum_{i=1}^{\ntask} \Zb_i - \Eb_{\x} \Zb_i\|\\
&\lesssim \log(\ntask Z_2) \left( \ntask^{-1/2}\log(\ntot) Z_1^{1/2} + \ntask^{-1}Z_2\log(\ntask Z_2) \right)\\
&\lesssim \log(\ntask Z_2) \Big(\sqrt{\frac{\log^{6}(\ntot)(\tr(\bSisq\B) + \tr(\bSi\B) + \sigma^2)^2\tr(\bSi)\|\bSi\|^2}{\nspt\ntask}} \\
&\quad + \frac{\log^{3}(\ntot)(\tr(\bSisq\B) + \tr(\bSi\B) + \sigma^2)\sqrt{\tr(\bSi)}\|\bSi\|}{\nspt^{1/2}\ntask}\\
&\quad + \frac{\log^{5}(\ntot) (\tr(\bSi\B) + \sigma^2)\tr(\bSi)}{\ntask}\Big)\\
&= \log(\ntask Z_2)\cdot \Big(\log^{3}(\ntot)\|\bSi\|(\tr(\bSisq\B) + \tr(\bSi\B) + \sigma^2)\sqrt{\frac{\tr(\bSi)}{\ntot}}\\
&\quad + \frac{\log^{5}(\ntot)(\tr(\bSisq\B) + \tr(\bSi\B) + \sigma^2)\sqrt{\tr(\bSi)}\|\bSi\|}{\ntot^{1/2}\ntask^{1/2}}\Big).
\end{align*}

§ PROOF OF ROBUSTNESS OF OPTIMAL REPRESENTATION

Suppose the data is generated as Phase 2, $\hbLa$ and $\uth$ are defined in Def. <ref> and the estimated task is obtained as (<ref>). Let the upper bound of $\|\hat \bM - \bM\|$ be $\calE$. The risk of meta-learning algorithm satisfies

\begin{align*}
\risk(\hbLa_{\uth}(R), \B, \bSi) - \risk(\bLa_{\uth}(R),\B, \bSi) \lesssim
\frac{\nfs^2\cdot\calE}{d(R-\nfs)(2\nfs - R\uth)\uth}.
\end{align*}

In the proof below, we use $\hbLa$ and $\bLa$ to replace $\hbLa_{\uth}(R),\bLa_{\uth}(R)$ for simplicity.
We first decompose the risk as

\begin{align*}
&\quad \risk(\hbLa, \B, \bSi) - \risk(\bLa, \B, \bSi)\\
&= \underbrace{\risk(\hbLa, \hB,\bSi) - \risk(\bLa, \hB,\bSi)}_{\le 0} \\
&\quad + [\risk(\hbLa, \B, \bSi)-\risk(\hbLa, \hB,\bSi)] + [\risk(\bLa, \hB,\bSi)-\risk(\bLa, \B, \bSi)].
\end{align*}

We know $\risk(\hbLa, \hB,\bSi) - \risk(\bLa, \hB,\bSi)\le 0$ due to the optimality of $\hbLa$ with task covariance $\hB$.
Now we will bound $\risk(\hbLa, \B, \bSi)-\risk(\hbLa, \hB,\bSi)$ for arbitrary $\hbLa$, and it automatically works for $\risk(\bLa, \hB,\bSi)-\risk(\bLa, \B, \bSi)$. Note that in (<ref>) we know that

\begin{align}
&\quad \risk({\hbLa}',\Bp) = f(\btheta;\B, \bSi): = \sum_{i=1}^R\dfrac{\nfs(1 - \btheta_i)^2}{R(\nfs - \|\btheta\|^2)}\BbarRi + \frac{\nfs}{\nfs - \|\btheta\|^2}\sigma^2. \label{eq:risk f app}
\end{align}

This function is linear in $\B$ thus we know that

\begin{align}
|\risk(\bLa, \hB,\bSi)-\risk(\bLa, \B, \bSi)| \le \frac{\nfs}{d(\nfs-\|\btheta\|^2)}\calE.\label{eq:risk f app2}
\end{align}

Now we need to bound $\|\btheta\|^2$. With the constraint $\uth \le \btheta<1-\frac{R-\nfs}{\nfs}\uth$ and $\sum\btheta_i=\nfs$, we know that the maximum of $\|\btheta\|^2$ happens when $(R-\nfs)$ among $\btheta_i$ are $\uth$ and the others are $1-\frac{R-\nfs}{\nfs}\uth$. With this we have

\begin{align*}
\|\btheta\|^2 &\le (R-\nfs)\uth^2 + \nfs(1-\frac{R-\nfs}{\nfs}\uth)^2\\
&= (R-\nfs)\uth^2 + \nfs - 2(R-\nfs)\uth + \frac{(R-\nfs)^2}{\nfs}\uth^2 \\
&= \nfs - 2(R-\nfs)\uth + \frac{(R-\nfs)R}{\nfs}\uth^2
\end{align*}

\begin{align*}
\nfs - \|\btheta\|^2 \ge  (R-\nfs)\uth (2\nfs - R\uth).
\end{align*}

Plugging it into (<ref>) and (<ref>) leads to the theorem.

§.§ One step GD

What does this mean for gradient descent? If we just look at the GD analysis in <cit.>, we can replace the Frobenius norm by $\tr(UV^\top \bSi VU^\top)$ in their bound, but the sample complexity cannot be smaller (like $(d+t)r\rightarrow(r+t)r$) to claim that all global minimums are good. Because there's no regularization in the objective (note that $\frac{1}{2}\|U^\top U - V^\top V\|$ is not a regularizer. If $ex^\top$ has a nullspace, you can always make both $U,V$ big in the null space but the value is $0$), and the analysis is based on overdetermined LS. So you can imagine that if the sample size gets slightly smaller, we are approaching the spike in the double descent curve, and when it gets through to overparameterized regime, we have to take the least norm solution and the landscape looks different.

So we need to add two things: early stopping or regularization. Now we consider the simple problem

\begin{align*}
\min_{W\in\R^{t\times d}}\ \sum_{i=1}^n \|e_{t(i)}^\top W\x_i - y_i\|^2
\end{align*}

We start from $W = 0$ and follow the gradient. Let $\tilde \B\in\R^{t\times d}$ be the matrix whose rows are the task vectors $\bt_j$, and let $\tilde \B_j$ be the matrix whose $j$th row is $\bt_j$ and the other rows are $0$. The gradient is

\begin{align*}
-\nabla f(W) &= \sum_{i=1}^n y_ie_{t(i)}\x_i^\top\\
&= \sum_{i=1}^n e_{t(i)}\bt_{t(i)}^\top\x_i\x_i^\top + \eps_ie_{t(i)}\x_i^\top\\
&= \sum_{i=1}^n \tilde \B_{t(i)} \bSi + \eps_ie_{t(i)}\x_i^\top.
\end{align*}

\begin{align}
\frac{1}{n}\mtx{E}(-\nabla f(W)) = \frac{1}{t} \tilde \B \bSi. \label{eq:gd_mean}
\end{align}

Now the variance analysis is easier.

* Estimating the rows of $\tilde\B$. Fixing each $\bt_j$, we are just summing over $(\sum_i\ones(t(i) = j)\x_i\x_i^\top) \bt_j$. We know from Lemma <ref> the concentration of $\sum_i\x_i\x_i^\top$ so we know how many samples we need to estimate each $\bt_j$.

* Estimating $\tilde\B\bSi$. We use $\frac{1}{n}\mtx{E}(-\nabla f(W))$ in (<ref>).
The mean is $\frac{1}{t} \tilde\B\bSi$, the error decreases as

\begin{align*}
O(\sqrt{\frac{\tr(\bSi) \tr(\sbSi\B\sbSi) + 2\tr(\bSi\B\bSi)+ \sigma_\eps^2 \tr(\bSi)}{n}}).
\end{align*}

The proof is presented in latter sections.

Matrix factorization. Now we study the matrix factorization form

\begin{align*}
\min_{\Ub\in\R^{t\times r}, \Vb\in\R^{d\times r}}\ \sum_{i=1}^n \|\eb_{t(i)}^\top \Ub\Vb^\top\x_i - y_i\|^2 := f(\Ub, \Vb)
\end{align*}

Denote $\A_i = \eb_{t(i)}\x_i^\top$, $\delta_i = y_i- \eb_{t(i)}^\top \Ub\Vb^\top\x_i$. We consider the gradient flow where

\begin{align}
d\Ub / dt &= -\nabla_\Ub f(\Ub, \Vb)/dt = \sum_{i=1}^n \delta_i \A_i\Vb,\\
\ d\Vb / dt &= -\nabla_\Vb f(\Ub, \Vb)/dt = \sum_{i=1}^n \delta_i \A_i^\top\Ub.
\end{align}

Note that

\begin{align}
d\Ub^\top\Ub / dt =  d\Vb^\top\Vb / dt \label{eq:inbal}
\end{align}

\begin{align}
df/dt &= -2\sum_{i=1}^n \delta_i(\langle\A_i, \sum_{i=1}^j(\delta_j \A_j \Vb\Vb^\top + \delta_j\Ub\Ub^\top \A_j) \rangle)\\
&= -2( \langle\sum_{i=1}^n\delta_i\A_i, \sum_{i=1}^n\delta_i\A_i\Vb\Vb^\top\rangle + \langle\sum_{i=1}^n\delta_i\A_i,\Ub\Ub^\top \sum_{i=1}^n\delta_i\A_i\rangle)
\end{align}

Suppose the variables are initialized at $\Ub(0), \Vb(0)$ and follow the gradient flow, let the variables at time $\tau$ be $\Ub(\tau), \Vb(\tau)$. Then we know that $\sigma_i(\Ub(\tau)) - \sigma_i(\Vb(\tau)) = \sigma_i(\Ub(0)) - \sigma_i(\Vb(0))$ due to (<ref>).

\begin{align*}
S(\Ub, \Vb) = \sum_{i=1}^n\delta_i\A_i
\end{align*}

\begin{align*}
df/dt =-2( \langle S(\Ub, \Vb), S(\Ub, \Vb) \Vb\Vb^\top\rangle +  \langle S(\Ub, \Vb), \Ub\Ub^\top S(\Ub, \Vb)\rangle)
\end{align*}

\begin{align}
d\Ub\Vb^\top / dt = 2( S(\Ub, \Vb) \Vb\Vb^\top + \Ub\Ub^\top S(\Ub, \Vb))
\end{align}

What if we think about

\begin{align}
d\Ub / dt &= (\Ib - \Pc_\Ub)\SUV\Vb,\\
d\Vb / dt &= (\Ib - \Pc_\Vb)\SUV^\top\Ub
\end{align}

where $\Pc$ denotes the projection onto the column space. Then

\begin{align}
d\Ub^\top\Ub / dt &=  d\Vb^\top\Vb / dt = 0\\
df/dt &= -2( \langle\sum_{i=1}^n\delta_i\A_i, (\Ib - \Pc_\Ub)\sum_{i=1}^n\delta_i\A_i\Vb\Vb^\top\rangle + \langle\sum_{i=1}^n\delta_i\A_i,\Ub\Ub^\top \sum_{i=1}^n\delta_i\A_i (\Ib - \Pc_\Vb)\rangle)
\end{align}

In brief, we only need to study the tail bound and variance of $ye\x^\top$.

\begin{align}
ye\x^\top &= (\x^\top\bt_{t(i)} + \eps_i) e\x^\top\\
&= e\bt_{t(i)}^\top \x \x^\top + \eps_i e\x^\top
\end{align}

\begin{align}
\mtx{P}(|\bt^\top \x| \ge t ) &\le \exp(-ct/ \|\sqrt{\B}\x\|)\\
\mtx{P}(\|\sqrt{\B}\x\| - \tr(\sbSi\B\sbSi) \ge t) &\le \exp(c\min\{\frac{t^2}{\tr((\sbSi\B\sbSi)^2)}, \frac{t}{\|\sbSi\B\sbSi\|}\})\\
\mtx{P}(\|\x\|^2 - \tr(\bSi) \ge t) &\le \exp(-c\min\{\frac{t^2}{\tr^2(\bSi)}, \frac{t}{\|\bSi\|}\})\\
\mtx{P}(|\eps| \ge t) &\le \exp(-ct^2 / \sigma_\eps^2)
\end{align}

Now we bound $\|y\x\|$. With probability $1-\delta$, $\|x\|$ is bounded by (<ref>), and $|y| = |\bt^\top\x| \le O(\|\sbSi\bt\|(1 + \sqrt{\log(1/\delta)}))$. Then with probability $1-\delta$

\begin{align*}
\|y\x\| \le  (1 + \log(1/\delta))\tr(\bSi)\|\sbSi\bt\|.
\end{align*}

With probability $1-\delta$, we have

\begin{align*}
\|\sbSi\bt\| \le (1 + \sqrt{\log(1/\delta)})\sqrt{\tr(\sbSi\B\sbSi)}
\end{align*}

So with probability $1-\delta$, we have

\begin{align}
\|y\x\| \le  (1 + \log(1/\delta)^{3/2})\tr(\bSi) \sqrt{\tr(\sbSi\B\sbSi)}
\end{align}

\begin{align}
\mtx{E} \|y^2e\x^\top\x e^\top \| &= \mtx{E} (\x^\top\bt_{t(i)} + \eps_i)^2 \x^\top \x\\
&= \frac{1}{t} \mtx{E} (\x^\top \tilde\B^\top\tilde\B \x\x^\top\x) + \sigma_\eps^2 \tr(\bSi) \label{eq:onestepgd_exp}
\end{align}

To compare it with other two parts, we switch to the case where $\bt$ is generated from $\mathcal{N}(0, \B)$ but we can replace $\frac{1}{t}\tilde\B^\top\tilde\B = \B$.

We continue from (<ref>). Note that $\frac{1}{t}\mtx{E}(\tilde\B^\top\tilde\B) = \B$, let $\z\sim \mathcal{N}(0, I_d)$, so we have

\begin{align*}
\mtx{E} \|y^2e\x^\top\x e^\top \| &= \mtx{E} (\x^\top \B \x\x^\top\x) + \sigma_\eps^2 \tr(\bSi) \\
&= \mtx{E} (\z^\top \sbSi\B\sbSi \z\z^\top\bSi\z) + \sigma_\eps^2 \tr(\bSi).
\end{align*}

WLOG, we assume that $\bSi$ is diagonal (otherwise we can rotate by the eigenvectors of $\bSi$). Then

\begin{align*}
\mtx{E} (\z^\top \sbSi\B\sbSi \z\z^\top\bSi\z)
&= \sum_{i=1}^d \sigma_i \z^\top \sbSi\B\sbSi \z\z_i^2 \\
&= \sum_{i=1}^d \sigma_i (\sum_{j=1}^d\sigma_j\B_{jj} + 2\sigma_i\B_{ii}) .\\
&= \tr(\bSi) \tr(\sbSi\B\sbSi) + 2\tr(\bSi\B\bSi)
\end{align*}

The mean is $\frac{1}{t} \tilde\B\bSi$, the error decreases as

\begin{align*}
O(\sqrt{\frac{\tr(\bSi) \tr(\sbSi\B\sbSi) + 2\tr(\bSi\B\bSi) + + \sigma_\eps^2 \tr(\bSi)}{n}}).
\end{align*}

§.§ Robustness of Optimal Representation

For a diagonal $\B$ the function $\opt(\cdot):\R^d\rightarrow\R^d$ solves the optimal representation problem (<ref>) subject to $\max_i(\opt_{i,i}(\B))/\min_i(\opt_{i,i}(\B))\leq C$ (in case we need nice conditioning). For a general $\B$ with eigendecomposition $\Ub\Eb\Ub^T$, set $\opt(\B)=\Ub\opt(\Eb)\Ub^T$.

let's change this as we discussed: make an assumption of same eigenspace like the Wu and Xu paper, then finding optimal ref reduces to diagonal

Suppose $\bSi$ is identity. Let $\Bh$ be the estimate of the ground-truth task covariance $\B$. Let $\bLa$ and $\bLah$ be the optimal representations associated with $\B,\Bh$ respectively.

Given a few-shot task $\bt$ with associated training data $(\y,\X)$ define $\bth_\La$ and $\bth_\bLah$ obtained via (<ref>). Important: Note that we don't need to show closeness of $\La,\bLah$ (in terms of $\|\Bh-\B\|$). We just need to show that $\bLah$ achieves approximately same risk on the task distributions $\B,\Bh$ and $\La$ similarly achieves approximately same risk on $\B,\Bh$. By definition of optimality, this auto-implies $\bLah,\La$ achieves approximately the same risk and they are as good as each other!

As a result, we would like to argue that for any close pair $\B,\Bh$ the risk induced by (an arbitrary) representation $\La$ is close. To this aim, the first step is showing that $\bth_\La$ and $\bth_\bLah$ are close.

Following (<ref>), define the task averaged risk

\begin{align*}
\risk_{\B,\La}&=\E_{\bt\distas\Nn(0,\B)}[\risk_{\La}(\hat \bt_\La)]\\
&=\E_{\bt}[(\hat \bt-\bt)^\top \La^2(\hat \bt-\bt)].
\end{align*}

$\bullet$ TODO: Can we prove that?

\[
|\risk_{\Bh,\bLah}-\risk_{\B,\bLah}|\lesssim \tf{\B-\Bh}^2\quad\text{and similarly for $\La$}
\]

Possible considerations: This is essentially Lipschitzness of the model with respect to the task distribution. Randomness of $\X$ might be a problem, i.e. what if it is bad-conditioned, pseudo-inverse might not be stable etc. Maybe we can use very small ridge regression in (<ref>) to ensure $\bth_\La$ is well-behaved and small enough that the risk stays same.

Final argument: Combine this with the fact that $\risk_{\Bh,\La}\geq \risk_{\Bh,\bLah}$ (follows from the optimality of $\bLah$ for $\Bh$) implies the desired relation that $\bLah$ is good for distribution $\B$ i.e.

\[
\risk_{\B,\bLah}\leq \risk_{\Bh,\La}+C\cdot \tf{\B-\Bh}^2.
\]

Alternative discussion: If we can't prove the actual robustness, maybe prove/discuss it for representation with minimal bias which is simply $\La=\sqrt{\B}$ (from Xu/Wu paper).

Finally, we would like to use this robustness to prove an end-to-end learning bound with optimal representation (from representation learning to few-shot).

$\bullet$ TODO: State end-to-end bound. If necessary, *assume* the robustness of optimal rep OR use the rep with min bias…

Did we estimate $\B$ in meta-training sec?

Samet's approach. We discuss the robustness with respect to the shaping matrix here. For any shaping matrix $\bL\in\R^{d\times d}$. Let $\bt\in\Nn(0,\B)$, $\y = \X\bt+\vct{\eps}$ where every row of $\X$ is independently from $\Nn(0, \bSi)$, we define $\bt_\bL = \bL(\X \bL)^\dag \y$,

\begin{align*}
\risk(\bL, \B) &= \Eb_{\x,\bt,\X,\y}(\bt - \bt_\bL)^\top\bSi(\bt - \bt_\bL)  + \sigma_\eps^2
\end{align*}

\begin{align*}
\bLa = \arg\min_{\bL}\ \risk(\bL, \B),%\ \mbox{s.t.} \|\bL\|=1, \sigma_\min(\bL) > 1/\kappa.
\end{align*}

Thus $\bLa$ is the shaping that minimizes the risk. As we do not know $\B$, we use $\hat \bLa$ as the shaping matrix which is defined as

\begin{align*}
\hat \bLa = \arg\min_{\bL}\ \risk(\bL, \hB).
\end{align*}

From the optimality, we know that $\risk(\hbLa, \hB) \le \risk(\bLa, \B)$.

Now we consider $\risk(\hbLa, \hB) - \risk(\hbLa, \B)$.

\begin{align*}
&\quad \risk(\hbLa, \B)\\
&= \Eb(\bt - \hbLa(\X\hbLa)^\dag\y)^\top \bSi (\bt - \hbLa(\X\hbLa)^\dag\y)\\
&= \Eb(\bt^\top \bSi \bt) - 2 \Eb(\bt^\top\bSi \hbLa(\X\hbLa)^\dag(\X\bt+\vct{\eps})) \\
&\quad + \Eb\|\hbLa(\X\hbLa)^\dag(\X\bt+\vct{\eps})\|_{\bSi}^2\\
&= \tr(\B\bSi) - 2 \Eb(\tr(\B\bSi \hbLa(\X\hbLa)^\dag\X)) \\
&\quad +\Eb(\tr(\B[\hbLa(\X\hbLa)^\dag\X]^\top \bSi[\hbLa(\X\hbLa)^\dag\X] ) )\\
&\quad + \sigma_\eps^2\tr([\hbLa(\X\hbLa)^\dag]^\top[\hbLa(\X\hbLa)^\dag])
\end{align*}

where each row of $\X$ is $\Nn(0,\bSi)$. Similarly

\begin{align*}
&\quad \risk(\hbLa, \hB)\\
&= \tr(\hB\bSi) - 2 \Eb(\tr(\hB\bSi \hbLa(\X\hbLa)^\dag\X)) \\
&\quad +\Eb(\tr(\hB[\hbLa(\X\hbLa)^\dag\X]^\top \bSi[\hbLa(\X\hbLa)^\dag\X] ) )\\
&\quad + \sigma_\eps^2\tr([\hbLa(\X\hbLa)^\dag]^\top[\hbLa(\X\hbLa)^\dag])
\end{align*}

where each row of $\X$ is $\Nn(0,\bSi)$. We subtract them and get

\begin{align*}
&\quad |\risk(\hbLa, \hB) - \risk(\hbLa, \B)|\\
&=\Big| \tr((\hB-\B)\bSi) - 2 \Eb(\tr((\hB-\B)\bSi \hbLa(\X\hbLa)^\dag\X))\\
&\quad +\Eb(\tr((\hB-\B)[\hbLa(\X\hbLa)^\dag\X]^\top \bSi[\hbLa(\X\hbLa)^\dag\X] ) )\Big|\\
&\le \|\hB - \B\|_F \cdot\left(\|\bSi\|_F + 2\|\Eb( \bSi \hbLa(\X\hbLa)^\dag\X) \|_F\right.\\
&\quad\left. + \|\Eb([\hbLa(\X\hbLa)^\dag\X]^\top \bSi[\hbLa(\X\hbLa)^\dag\X])\|_F\right)
\end{align*}

$(\hbLa(\X\hbLa)^\dag\X)^2 = \hbLa(\X\hbLa)^\dag\X$, so that $eig(\hbLa(\X\hbLa)^\dag\X) = [{\bf 1}^m, {\bf 0}^{d-m}]$, not sure if this helps.

We define the last line as $F(\hbLa, \bSi)\|\hB - \B\|_F$. Similarly we know that $ |\risk(\bLa, \hB) - \risk(\bLa, \B)| \le F(\bLa, \bSi)\|\hB - \B\|_F$.

\begin{align*}
&\quad \risk(\hbLa, \B) - \risk(\bLa, \B)\\
&= \underbrace{\risk(\hbLa, \hB) - \risk(\bLa, \hB)}_{\le 0}\\
&\quad + \risk(\hbLa, \B)-\risk(\hbLa, \hB)\\
&\quad + \risk(\bLa, \hB)-\risk(\bLa, \B)\\
&\le (F(\bLa, \bSi)+F(\hbLa, \bSi))\|\hB - \B\|_F.
\end{align*}

Multiplicative/additive error

The risk in <ref> is of the form
$\risk(\B,\btheta) = \sum_{i=1}^d f_1(\btheta_i)\B_{ii} + f_2(\btheta)$. We need to make sure

1. If all $\max_i f_1(\btheta_i) = \bar f_1$, then if $\B$ is perturbed by $\delta$, $\risk$ changes at most $d\bar f_1$.

2. When $\B_{ii}$ is small, $f_1(\btheta_i)$ is not too big. Suppose

\begin{align*}
f_1(\btheta_i) &\le \bar f_1,\ \mathrm{if~}\B_{ii} \le \underbar{B}
\end{align*}

Then if $\B$ is perturbed by $\delta$,

\begin{align*}
\risk(\B+\delta,\btheta) - \risk(\B,\btheta) \le \frac{\delta}{\underbar{B}}\risk(\B,\btheta) + d\bar f_1
\end{align*}

My thought from here
Now we look at $\hat \bt = \bLa(\X\bLa)^\dag \y$. To show the robustness with respect to $\bLa$, we look at the matrix $\bLa(\X\bLa)^\dag$. Define a new matrix $\bL$ which is $\bLa$ after perturbation, and we will bound $ \bLa(\X\bLa)^\dag - \bL(\X\bL)^\dag$ by $\bLa - \bL$. $\X$ is a fat matrix and suppose it has full column rank, also suppose $\bLa, \bL$ be positive definite. Let $\X = \Ub[\bSi_\X,0][\Vb_1, \Vb_2]^\top$ be the SVD of $\X$. Let $\bLa = \Vb_1\bLa_1 + \Vb_2\bLa_2$, and $\bL = \Vb_1\bL_1 + \Vb_2\bL_2$.

First, $\X\bLa(\X\bLa)^\dag = \Ib = \X\bL(\X\bL)^\dag$ thus

\begin{align*}
\bLa(\X\bLa)^\dag - \bL(\X\bL)^\dag = \Vb_2\Vb_2^\top (\bLa(\X\bLa)^\dag - \bL(\X\bL)^\dag)
\end{align*}

Here $\Vb_2\Vb_2^\top$ is simply a projection matrix. Now we focus on $\Vb_2^\top (\bLa(\X\bLa)^\dag - \bL(\X\bL)^\dag)$.

\begin{align*}
&\quad \Vb_2^\top (\bLa(\X\bLa)^\dag - \bL(\X\bL)^\dag)\\
&= \bLa_2(\X\bLa)^\dag - \bL_2(\X\bL)^\dag\\
&= \bLa_2(\Ub\bSi_\X\bLa_1)^\dag -  \bL_2(\Ub\bSi_\X\bL_1)^\dag
\end{align*}

Now we need to bound this quantity by $\bLa_1 - \bL_1$ and $\bLa_2 - \bL_2$, which are the perturbation of $\bLa$ projected onto the row space of $\X$ and its orthogonal complement.

We continue with the expression and get

\begin{align*}
&\quad \| \Vb_2^\top (\bLa(\X\bLa)^\dag - \bL(\X\bL)^\dag) \|\\
&= \| \bLa_2(\Ub\bSi_\X\bLa_1)^\dag -  \bL_2(\Ub\bSi_\X\bL_1)^\dag \|\\
&\le \|(\bLa_2 - \bL_2)(\Ub\bSi_\X\bLa_1)^\dag\|\\
&\quad + \|\bL_2((\Ub\bSi_\X\bLa_1)^\dag - (\Ub\bSi_\X\bL_1)^\dag)\|\\
&\le  \|(\Ub\bSi_\X\bLa_1)^\dag\| \|\bLa_2 - \bL_2\| \\
&\quad + \|\bL_2\| \|(\Ub\bSi_\X\bLa_1)^\dag - (\Ub\bSi_\X\bL_1)^\dag\|\\
&\le \|(\Ub\bSi_\X\bLa_1)^\dag\| \|\bLa_2 - \bL_2\| \\
&\quad + \|\bLa_2\| \|\bSi_\X\| \| (\Ub\bSi_\X\bLa_1)^\dag \|^2 \|\bLa_1 - \bL_1 \| + o(\|\bLa - \bL\|)
\end{align*}

The last line uses the derivative of the pseudoinverse

\begin{align*}
&\quad d((\Ub\bSi_\X\bLa_1)^\dag)\\
&= (\Ub\bSi_\X\bLa_1)^\dag\Ub\bSi_\X (d\bLa_1)(\Ub\bSi_\X\bLa_1)^\dag\\
&\quad + (I - (\Ub\bSi_\X\bLa_1)^\dag(\Ub\bSi_\X\bLa_1))\cdot\\
&\qquad (d\bLa_1)^\top\bSi_\X\Ub^\top((\Ub\bSi_\X\bLa_1)^\dag)^\top(\Ub\bSi_\X\bLa_1)^\dag
\end{align*}

\begin{align*}
&\quad \| \Vb_2^\top (\bLa(\X\bLa)^\dag - \bL(\X\bL)^\dag) \|_F\\
&\le \|(\Ub\bSi_\X\bLa_1)^\dag\| \|\bLa_2 - \bL_2\|_F \\
&\quad + \|\bLa_2\|\|\bSi_\X\| \| (\Ub\bSi_\X\bLa_1)^\dag \|^2 \|\bLa_1 - \bL_1 \|_F
\end{align*}

================== DELETE ABOVE ==================