Subspace Recovery from Heterogeneous Data with Non-isotropic Noise
Abstract
Recovering linear subspaces from data is a fundamental and important task in statistics and machine learning.
Motivated by heterogeneity in Federated Learning settings, we study a basic formulation of this problem: the principal component analysis (PCA), with a focus on dealing with irregular noise.
Our data come from n users with user i contributing data samples from a d -dimensional distribution with mean µ i .
Our goal is to recover the linear subspace shared by µ 1 , . . . , µ n using the data points from all users, where every data point from user i is formed by adding an independent mean-zero noise vector to µ i .
If we only have one data point from every user, subspace recovery is information-theoretically impossible when the covariance matrices of the noise vectors can be non-spherical, necessitating additional restrictive assumptions in previous work.
We avoid these assumptions by leveraging at least two data points from each user, which allows us to design an efﬁciently-computable estimator under non-spherical and user-dependent noise.
We prove an upper bound for the estimation error of our estimator in general scenarios where the number of data points and amount of noise can vary across users, and prove an information-theoretic error lower bound that not only matches the upper bound up to a constant factor, but also holds even for spherical Gaussian noise.
This implies that our estimator does not introduce additional estimation error (up to a constant factor) due to irregularity in the noise.
We show additional results for a linear regression problem in a similar setup.
Introduction
We study the problem of learning low-dimensional structure amongst data distributions, given multiple samples from each distribution.
This problem arises naturally in settings such as federated learning, where we want to learn from data coming from a set of individuals, each of which has samples from their own distribution.
These distributions however are related to each other and in this work, we consider the setting when these distributions have means lying in a low-dimensional subspace.
The goal is to learn this subspace, even when the distributions may have different (and potentially non-spherical) variances.
This heterogeneity can manifest itself in practice as differing number of samples per user, or the variance differing across individuals, possibly depending on their mean.
Recovery of the subspace containing the means can in turn help better estimate individual means.
In other words, this can allow for learning good estimator for all individual means, by leveraging information from all the individuals.
The irregularity of the noise makes this task challenging even when we have sufﬁciently many individual distributions.
For example, suppose we have n individuals and for every i = 1 , . . .
, n , an unknown µ i ∈ R d .
For simplicity, suppose that µ 1 , . . . , µ n are distributed independently as
N (0 , σ 2 uu T ) for σ ∈ R ≥ 0 and an unknown unit vector u ∈ R d .
In this setting, our goal is to recover the one-dimensional subspace, equivalently the vector u .
For every i , we have a data point x i = µ i
+ z i where z i ∈ R
d is a mean-zero noise vector.
If z i is drawn independently from a spherical
Gaussian N (0 , α 2 I ) , we can recover the unknown subspace with arbitrary accuracy as n grows to inﬁnity because 1 n (cid:80) x
i x T i concentrates to E
[ x i x T i ]
= σ 2 uu T + α 2 I , whose top eigenvector is ± u .
However, if the noise z i is drawn from a non-spherical distribution, the top eigenvector of 1 n (cid:80) x i x T i can deviate from ± u signiﬁcantly, and to make things worse, if the noise z i is drawn independently from a non-spherical Gaussian N (0 , σ 2 ( I − uu T )
+ α 2 I ) , then our data points x i = µ
i
+ z i distribute independently as N (0 , ( σ 2 + α 2 ) I ) , giving no information about the vector u .
The information-theoretic impossibility in this example however disappears as soon as one has at least two samples from each distribution.
Indeed, given two data points x i 1 = µ i + z i
and x i
= µ i + z ifrom user
i , as long as the noise z i 1 , z i 2 are independent and have zero mean, we always have E
[ x i 1 x T
i 2 ]
= σ 2 uu T regardless of the speciﬁc distributions of z
i 1 and z i 2 .
This allows us to recover the subspace in this example, as long as we have sufﬁciently many users each contributing at least two examples.
As this is commonly the case in our motivating examples, we make this assumption of multiple data points per user, and show that this intuition extends well beyond this particular example.
We design efﬁciently computable estimators for this subspace recovery problem given samples from multiple heteroscedastic distributions (see Section 1.1 for details).
We prove upper bounds on the error of our estimator measured in the maximum principal angle (see Section 2 for deﬁnition).
We also prove an information-theoretic error lower bound, showing that our estimator achieves the optimal error up to a constant factor in general scenarios where the number of data points and the amount of noise can vary across users.
Somewhat surprisingly, our lower bound holds even when the noise distributes as spherical Gaussians.
Thus non-spherical noise in setting does not lead to increased error.
We then show that our techniques extend beyond the mean estimation problem to a linear regression setting where for each µ i , we get (at least two) samples ( x ij , x T ij µ i + z ij ) where z ij is zero-mean noise from some noise distribution that depends on i and x ij .
This turns out to be a model that was recently studied in the meta-learning literature under more restrictive assumptions (e.g. z ij is independent of x ij )
[Kong et al., 2020, Tripuraneni et al., 2021, Collins et al., 2021, Thekumparampil et al., 2021].
We show a simple estimator achieving an error upper bound matching the ones in prior work without making these restrictive assumptions.
1.1 Our contributions
PCA with heterogeneous and non-isotropic noise: Upper Bounds.
In the PCA setting, the data points from each user i are drawn from a user-speciﬁc distribution with mean µ i ∈ R d , and we assume that µ 1 , . . . , µ n lie in a shared k -dimensional subspace that we want to recover.
Speciﬁcally, we have m i data points x ij ∈ R d from user i for j = 1 , . . .
, m i , and each data point is determined by x ij = µ i + z ij where
z ij ∈ R
d is a noise vector drawn independently from a mean zero distribution.
We allow the distribution of z ij to be non-spherical and non-identical across different pairs ( i, j ) .
We use η i ∈ R ≥ 0 to quantify the amount of noise in user i ’s data points by assuming that z ij is an η i -sub-Gaussian random variable.
As mentioned earlier, if we only have a single data point from each user, it is information-theoretically impossible to recover the subspace.
Thus, we focus on the case where m i ≥ 2 for every i = 1 , . . . , n .
In this setting, for appropriate weights w 1 , . . .
, w n ∈ R ≥ 0 , we compute a matrix A :
[Equation]
where the inner summation is over all pairs j 1 , j 2 ∈ { 1 , . .
, m i } satisfying j 1 (cid:54) =
j 2 .
Our estimator is then deﬁned by the subspace spanned by the top- k eigenvectors of A .
Although the inner summation is over m i ( m i − 1) terms, the time complexity for computing it need not grow quadratically with m i because of the following equation:
[Equation]
The ﬂexibility in the weights w 1 , . . .
, w n allows us to deal with variations in m i and η i for different users i .
In the special case where η 1 = · · · = η n = η and m 1 = · · · = m n = m , we choose w 1 = · · · = w n = 1 /n
and we show that our estimator achieves the following error upper bound with success probability at least 1 − δ :
[Equation]
Here, θ is the maximum principal angle between our estimator and the true subspace shared by µ 1 , . . . , µ n , and σ 2 (cid:96) is the (cid:96) -th largest eigenvalue of 1 n (cid:80) ni
=
µ i µ T i .
Our error upper bound for general m i , η i , w i is given in Theorem 3.1.
We instantiate our error upper bound to the case where µ 1 , . . . , µ n are drawn iid from a Gaussian distribution N (0 , σ 2 UU T ) , where the columns of U ∈ R d × k form an orthonormal basis of the subspace containing µ 1 , . . . , µ n .
By choosing the weights w 1 , . . .
, w n according to m 1 , . . .
, m n and η 1 , . .
η n , our estimator achieves the error upper bound
[Equation]
under a mild assumption (Assumption 3.2), where γ (cid:48)
i is deﬁned in Deﬁnition 3.1 and often equals (cid:16)
[Equation]
PCA: Lower Bounds.
We show that the error upper bound (2) is optimal up to a constant factor by proving a matching information-theoretic lower bound (Theorem 3.7).
Our lower bound holds for general m i and η i that can vary among users i , and it holds even when the noise vectors z ij are drawn from spherical Gaussians, showing that our estimator essentially pays no additional cost in error or sample complexity due to non-isotropic noise.
We prove the lower bound using Fano’s method on a local packing over the Grassmannian manifold.
We carefully select a non-trivial hard distribution so that the strength of our lower bound is not affected by a group of fewer than k users each having a huge amount of data points with little noise.
Linear Models.
While the PCA setting is the main focus of our paper, we extend our research to a related linear models setting that has recently been well studied in the meta-learning and federated learning literature [Kong et al., 2020, Tripuraneni et al., 2021, Collins et al., 2021, Thekumparampil et al., 2021].
Here, the user-speciﬁc distribution of each user i is parameterized by β i ∈ R d , and we again assume that β 1 , . . .
, β n lie in a linear subspace that we want to recover.
From each user i we observe m i data points ( x ij , y ij )
∈ R
d × R for j = 1 , . . .
, m i drawn from the user-speciﬁc distribution satisfying y ij = x
T ij β
i + z ij for an O (1) -sub-Gaussian measurement vector x ij ∈ R d with zero mean and identity covariance and an η i -sub-Gaussian mean-zero noise term z ij ∈ R .
While it may seem that non-isotropic noise is less of a challenge in this setting since each noise term z ij is a scalar, our goal is to handle a challenging scenario where the variances of the noise terms z ij can depend on the realized measurements x ij , which is a more general and widely applicable setting compared to those in prior work.
Similarly to the PCA setting, our relaxed assumptions on the noise make it information-theoretically impossible to do subspace recovery if we only have one data point from each user (see Section 4), and thus we assume each user contributes at least two data points.
We use the subspace spanned by the top- k eigenvectors of the following matrix A as our estimator:
[Equation]
In the special case where η 1 = · · · = η n = η, m 1 = · · · = m n = m , and (cid:107) β i (cid:107) 2 ≤ r for all i , our estimator achieves the following error upper bound:
[Equation]
where σ 2 k is the k -th largest eigenvalue of 1 n (cid:80)
ni
=1 β i β T i (Corollary L.2).
Our error upper bound extends smoothly to more general cases where η i and m i vary among users (Theorem L.1).
Moreover, our upper bound matches the ones in prior work [e.g. Tripuraneni et al., 2021, Theorem 3] despite requiring less restrictive assumptions.
1.2 Related Work
Principal component analysis under non-isotropic noise has been studied by Vaswani and Narayanamurthy [2017], Zhang et al.
[2018] and Narayanamurthy and Vaswani [2020].
When translated to our setting, these papers focus on having only one data point from each user and thus they require additional assumptions—either the level of non-isotropy is low, or the noise is coordinate-wise independent and the subspace is incoherent.
The estimation error guarantees in these papers depend crucially on how well these additional assumptions are satisﬁed.
Zhu et al.
[2019] and Cai et al.
[2021] study PCA with noise and missing data, and Chen et al.
[2021] and Cheng et al.
[2021] study eigenvalue and eigenvector estimation under heteroscedastic noise.
These four papers all assume that the noise is coordinate-wise independent and the subspace/eigenspace is incoherent.
The linear models setting we consider has recently been studied as a basic setting of meta-learning and federated learning by Kong et al.
[2020], Tripuraneni et al.
[2021], Collins et al.
[2021], and
Thekumparampil et al.
These papers all make the assumption that the noise terms z ij are independent of the measurements x ij , an assumption that we relax in this paper.
Collins et al.
[2021] and Thekumparampil et al.
[2021] make improvements in sample complexity and error guarantees compared to earlier work by Kong et al.
[2020] and Tripuraneni et al.
[2021], but Collins et al.
[2021] focus on the noiseless setting ( z ij = 0 ) and Thekumparampil et al.
[2021] require at least Ω( k 2 ) examples per user.
Tripuraneni et al.
[2021] and Thekumparampil et al.
[2021] assume that the measurements x ij are drawn from the standard (multivariate) Gaussian distribution, where as
Kong et al.
[2020], Collins et al.
and our work make the relaxed assumption that x ij are sub-Gaussian with identity covariance, which, in particular, allows the fourth-order moments of x ij to be non-isotropic.
There is a large body of prior work on meta-learning beyond the linear setting [see e.g. Maurer et al., 2016, Tripuraneni et al., 2020, Du et al., 2020].
When collecting data from users, it is often important to ensure that private information about users is not revealed through the release of the learned estimator.
Many recent works proposed and analyzed estimators that achieve user-level differential privacy in settings including mean estimation [Levy et al., 2021, Esfandiari et al., 2021], meta-learning [Jain et al., 2021] and PAC learning [Ghazi et al., 2021].
Recently, Cummings et al.
[2021] study one-dimensional mean estimation in a setting similar to ours, under a differential privacy constraint.
The matrix A we deﬁne in (1) is a weighted sum of A i :=
[Equation]
over users i = 1 , . . . , n , and each A i has the form of a U -statistic
[Halmos, 1946, Hoeffding, 1948].
U -statistics have been applied to many statistical tasks including tensor completion [Xia and Yuan, 2019] and various testing problems [Zhong and Chen, 2011, He et al., 2021, Schrab et al., 2022].
In our deﬁnition of A i , we do not make the assumption that the distributions of x i 1 , . . .
, x im i are identical although the assumption is commonly used in applications of U -statistics.
The matrix A in (3) is also a weighted sum of U -statistics where we again do not make the assumption of identical distribution.
1.3 Paper Organization
In Section 2, we formally deﬁne the maximum principal angle and other notions we use throughout the paper.
Our results in the PCA setting and the linear models setting are presented in Sectionsand 4, respectively.
We defer most technical proofs to the appendices.
Preliminaries
We use (cid:107) A (cid:107) to denote the spectral norm of a matrix A , and use (cid:107) u (cid:107) 2 to denote the (cid:96) 2 norm of a vector u .
For positive integers k ≤ d , we use O d,k to denote the set of matrices A ∈ R d × k satisfying
A T A = I k , where I k is the k × k identity matrix.
We use O d to denote O d,d , which is the set of d × d orthogonal matrices.
We use col ( A ) to denote the linear subspace spanned by the columns of a matrix A .
We use the base- e logarithm throughout the paper.
Maximum Principal Angle.
Let U, ˆ U ∈
O d be two orthogonal matrices.
Suppose the columns of U and ˆ U are partitioned as U
=
[ U 1 U 2 ] , ˆ U =
[ ˆ U 1 ˆ U 2 ] where U 1 , ˆ U 1 ∈
O d,k for an integer k satisfying 0 < k < d .
Let Γ (resp. ˆΓ ) be the k -dimensional linear subspace spanned by the columns of U 1 (resp.
ˆ U 1 ).
Originating from [Jordan, 1875], the maximum principal angle θ ∈ [0 , π/ 2] between Γ and ˆΓ , denoted by ∠ (Γ , ˆΓ) or ∠ ( U 1 , ˆ U 1 ) , is deﬁned by sin θ = (cid:107) U 1 U T 1 − ˆ U
ˆ U T 1 (cid:107) =
(cid:107) U T 1 ˆ U 2 (cid:107) =
(cid:107) U T 2 ˆ U 1 (cid:107) .
It is not hard to see that the maximum principal angle depend only on the subspaces Γ , ˆΓ and not on the choices of U and ˆ U , and sin ∠ (Γ , ˆΓ) is a natural metric between k -dimensional subspaces (see Appendix A for more details where we discuss the deﬁnition of principal angles for any two subspaces with possibly different dimensions).
With the deﬁnition of the maximum principal angle, we can now state a variant of the Davis–Kahan sin θ theorem [Davis and Kahan, 1970] that will be useful in our analysis (see Appendix E for proof):
Theorem 2.1 (Variant of Davis–Kahan sin θ theorem) .
Let A, ˆ A ∈ R d
× d be symmetric matrices.
Let λ i denote the i -th largest eigenvalue of A .
For a positive integer k smaller than d , let θ denote the maximum principal angle between the subspaces spanned by the top- k eigenvectors of A and ˆ A .
Assuming λ k > λ k +1 ,
[Equation]
Sub-Gaussian and sub-exponential distributions.
We say a random variable x ∈ R with expectation E [ x ] ∈ R has sub-Gaussian constant b ∈ R ≥ 0 if E [ | x − E [ x ] | p ] 1 /p ≤ b √ p for every p ≥ 1 .
We say x has sub-exponential constant b ∈ R ≥
if E [ | x − E [ x ] | p ] 1 /p ≤ bp for every p ≥ 1 .
We say a random vector y ∈ R
d has sub-Gaussian (resp.
sub-exponential) constant b ∈ R ≥
if for every unit vector u ∈ R d (i.e., (cid:107) u (cid:107) 2 = 1 ), the random variable u T y ∈ R has sub-Gaussian (resp. subexponential) constant b .
We say y is b -sub-Gaussian (resp.
b -sub-exponential ) if it has sub-Gaussian (resp.
sub-exponential) constant b .
Principal Component Analysis
In the principal component analysis (PCA) setting, our goal is to recover the k -dimensional subspace Γ spanned by the user-speciﬁc means µ 1 , . . . , µ n ∈ R d of the n users.
From each user i , we have m i ≥ 2 data points
[Equation]
We assume the noise z ij ∈ R d is drawn independently from a mean zero distribution with subGaussian constant η i .
We do not assume that the variance of z ij is the same along every direction, nor do we assume that the distribution of z ij is the same for different ( i, j ) .
We ﬁrst show an error upper bound for our estimator when the user-speciﬁc means µ 1 , . . . , µ n are deterministic vectors (Section 3.1) and then apply this result to the case where µ 1 , . . . , µ n are drawn from a sub-Gaussian distribution (Section 3.2).
In Section 3.3 we prove an information-theoretic error lower bound matching our upper bound.
3.1 Fixed User-Speciﬁc Means
We ﬁrst focus on the case where µ 1 , . . . , µ n are deterministic vectors.
In this case, all the randomness in the data comes from the noise z ij .
Our estimator is the subspace ˆΓ spanned by the top- k eigenvectors of A deﬁned in (1).
For (cid:96) = 1 , . . .
, d , we deﬁne σ 2 (cid:96) to be the (cid:96) -th largest eigenvalue of
(cid:80) ni
=
w i µ i µ T i .
Since µ 1 , . . . , µ n share a k -dimensional subspace, σ (cid:96) = 0 for (cid:96) > k .
We prove the following general theorem on the error guarantee of our estimator:
Theorem 3.1.
Deﬁne ξ 2 = (cid:107) (cid:80) ni
=1 w 2 i µ i µ T i η 2 i /m i (cid:107) and let θ denote the maximum principal angle between our estimator ˆΓ and the true subspace Γ spanned by µ 1 , . . . , µ n .
For any δ ∈ (0 , 1 / 2) , with probability at least 1 − δ ,
[Equation]
We can simplify the bound in Theorem 3.1 by considering special cases:
Corollary 3.2.
Assume max { η 1 / √ m 1 , . . .
, η n / √ m n } = t
and we choose w 1 = · · · = w n = 1 /n .
For any δ ∈ (0 , 1 / 2) , with probability at least 1 − δ ,
[Equation]
In particular, when η 1 = · · · = η n = η , and m 1 = · · · = m n = m , error bound (7) becomes sin
[Equation]
We defer the complete proof of Theorem 3.1 and Corollary 3.2 to Appendices F and G. Our proof is based on the Davis-Kahan sin θ theorem (Theorem 2.1).
Since σ 2 k +
=
0 , Theorem 2.1 implies
[Equation]
This reduces our goal to proving an upper bound on the spectral norm of A − (cid:80) ni
=
w i µ i µ T i .
Since for distinct j 1 and j 2 in { 1 , . . .
, m i } we have E [ x ij 1 x T ij 2 ]
= µ i µ T i , our construction of A in (1) guarantees E [ A ] = (cid:80) ni
=
w i µ i µ T i .
Therefore, our goal becomes controlling the deviation of A from its expectation, and we achieve this goal using techniques for matrix concentration inequalities.
3.2 Sub-Gaussian User-Speciﬁc Means
We apply our error upper bound in Theorem 3.1 to the case where µ 1 , . . . , µ n ∈ R d are drawn iid from N (0 , σ 2 UU T ) for an unknown U ∈
O d,k .
We still assume that each data point x ij ∈ R d is generated by adding a noise vector z
ij ∈ R d to the user-speciﬁc mean µ i as in (5).
We do not assume that the noise vectors ( z ij ) 1 ≤ i ≤ n, 1 ≤ j ≤ m i are independent of the user-speciﬁc means ( µ i ) 1 ≤ i ≤ n , but we assume that when conditioned on ( µ i )≤ i ≤ n , every noise vector z ij independently follows a distribution with mean zero and sub-Gaussian constant η i .
We use the same estimator ˆΓ as before: ˆΓ is the subspace spanned by the top- k eigenvectors of A deﬁned in (1).
We determine the optimal weights w 1 , . . .
, w n in (1) as long as m 1 , . . .
, m n and η 1 , . .
, η n satisfy a mild assumption (Assumption 3.2), achieving an error upper bound in Theorem 3.4.
In the next subsection, we prove an error lower bound (Theorem 3.7) that matches our upper bound (Theorem 3.4) up to a constant factor, assuming d ≥ (1 + Ω(1))
k and δ = Θ(1) .
We prove our error upper bound in a slightly more general setting than µ 1 , . . . , µ n drawn iid from
N (0 , σ 2 UU T ) .
Speciﬁcally, we make the following assumption on the distribution of µ 1 , . . .
, µ n :
Assumption 3.1.
The user-speciﬁc means µ 1 , . .
. , µ n ∈ R d are mean-zero independent random vectors supported on an unknown k -dimensional subspace Γ .
Moreover, for a parameter σ > 0 , for every i = 1 , . . . , n , µ i has sub-Gaussian constant O ( σ ) , and the k -th largest eigenvalue of E [ µ i µ T i ] is at least σ 2 .
Under this assumption, we have the following lower bound on the σ 2 k in Theorem 3.1 (see Appendix H for proof):
Claim 3.3.
Under Assumption 3.1, let w 1 , . . .
, w n ∈ R ≥ 0 be user weights satisfying w 1 + · · · + w n = 1 and σ 2 k be the k -th largest eigenvalue of (cid:80) ni
=
w i µ i µ T i .
There exists an absolute constant C ∗ >such that for any δ ∈ (0 , 1 / 2) , as long as max 1 ≤ i ≤ n w i ≤ 1 /C ∗ ( k + log(1 /δ )) , then σ 2 k ≥ σ 2 / 2 with probability at least 1 − δ/ 2 .
The following deﬁnition is important for us to choose the weights w 1 , . . .
, w n in (1) optimally:
Deﬁnition 3.1.
Deﬁne
[Equation]
= (cid:16)
[Equation]
+
[Equation]
i σ 4 m 2 i (cid:17) − 1 and assume w.l.o.g. that γ 1 ≥ · · · ≥ γ n .
Deﬁne γ
[Equation]
i = γ i if i ≥ k , and γ (cid:48)
i = γ k
if i < k .
Intuitively, we can view γ i as measuring the “amount of information” provided by the data points from user i .
This is consistent with the fact that γ i increases as the number m i of data points from user i increases, and γ i decreases as the noise magnitude η i from user i increases.
With the users sorted so that γ 1 ≥ · · · ≥ γ n , the quantity γ (cid:48) i is then deﬁned to be γ k for the k most “informative” users i = 1 , . .
k , and γ (cid:48)
i = γ i for other users.
We make the following mild assumption on γ (cid:48)
i under which we achieve optimal estimation error:
Assumption
[Equation]
deﬁned in Claim 3.
[Equation]
By the deﬁnition of γ (cid:48)
i , it is easy to show that Assumption 3.2 is equivalent to (cid:80) n i =
k +1 γ i ≥
( C ∗ − 1)
k + C ∗ log(1 /δ ))
γ k .
Therefore, if we view γ i as the “amount of information” from user i ,
Assumption 3.2 intuitively requires that a signiﬁcant contribution to the total “information” comes from outside the k most “informative” users.
This assumption allows us to avoid the case where we only have exactly n = k users: in that case, we would have σ 2 k ≈ σ 2 /k 2 for uniform weights w 1 = · · · = w n (see [Rudelson and Vershynin, 2008] and references therein), as opposed to the desired σ 2 k ≥ σ 2 / 2 in Claim 3.3.
Assumption 3.2 is a mild assumption.
For example, when γ k = · · · = γ n , Assumption 3.2 holds as long as n ≥ C ∗ ( k +log(1 /δ )) .
Also, since γ (cid:48) 1 = · · · = γ (cid:48)
k ≥ γ (cid:48)
k +1 ≥ · · · ≥ γ (cid:48) n ≥ 0 , it trivially holds that (cid:80) ni
=1 γ (cid:48)
i ≥ kγ (cid:48) 1 .
Assumption 3.2 is relatively mild when compared to this trivial inequality.
Under Assumption 3.2, we show that it is optimal to choose the weights w 1 , . . .
, w n as
[Equation]
Speciﬁcally, if we plug (9) into Theorem 3.1 and bound ξ and σ k based on the distribution of µ 1 , . . . , µ n , we get the following error upper bound which matches our lower bound (Theorem 3.7) in Section 3.3.
We defer its proof to Appendix I.
Theorem 3.4.
Under Assumptions 3.1 and 3.2, if we choose w 1 , . . .
, w n as in (9) and deﬁne θ = ∠ (Γ , ˆΓ) , for δ ∈ (0 , 1 / 2) , with probability at least 1 − δ ,
[Equation]
For comparison, consider the setting when σ = η i = 1 for every i = 1 , . .
. , n .
The result then says that sin θ is bounded by approximately (cid:113) d
[Equation]
ni =
m i .
This is the same rate as we would get if we have (cid:80) ni
=
m i users each contributing a single independent data point with homogeneous spherical noise.
Thus as long as the data points are not too concentrated on fewer than k users, the heterogeneity comes at no additional cost.
3.3 Lower Bound
We prove a lower bound matching the upper bound in Theorem 3.4 up to constant in the setting where δ = Θ(1) , d ≥ (1 + Ω(1)) k .
For every positive integer d , there is a natural “uniform” distribution over O d given by Haar’s theorem [Haar, 1933]
(see e.g. [Diestel and Spalsbury, 2014] for a textbook).
We denote this distribution by
Haar ( O d ) .
A random matrix A drawn from Haar ( O d ) has the following invariance property: for any deterministic matrix B ∈ O d , the random matrices A, AB and BA all have the same distribution.
For an integer k ≤ d , we can construct a random matrix
A 1 ∈ O d,k by ﬁrst drawing A ∈ R d × d from
Haar ( O d ) and then take the ﬁrst k columns of A .
We denote the distribution of A 1 by Haar ( O d,k ) .
The invariance property of Haar ( O d ) immediately implies the following claims:
Claim 3.5.
Let A ∈ O d be a random matrix drawn from Haar ( O d ) and let B ∈ O d,k be a ﬁxed matrix.
Then AB distributes as Haar ( O d,k ) .
Proof.
The matrix B can be written as the ﬁrst k columns of a matrix C ∈ O d .
Now AB is the ﬁrst k columns of AC , where AC distributes as Haar ( O d ) by the invariance property.
This implies that
AB distributes as Haar ( O d,k ) .
Claim 3.6.
Let B ∈ O d,k be a random matrix.
Assume for every ﬁxed matrix A ∈ O d , the random matrices B and AB have the same distribution.
Then B ∼ Haar ( O d,k ) .
Proof.
If we draw A independently from Haar ( O d ) , the random matrices B and AB still have the same distribution.
By Claim 3.5, AB distributes as Haar ( O d,k ) , so B must also distribute as
Haar ( O d,k ) .
With the deﬁnition of Haar ( O d,k ) , we state our lower bound in the following theorem:
Theorem 3.7.
Let k, d, n be positive integers satisfying k < d and k ≤ n .
Let m 1 , . . .
, m n be positive integers and σ, η 1 , . . .
, η n be positive real numbers.
Suppose we draw U ∈
O d,k from
Haar ( O d,k ) and then draw µ 1 , . .
. , µ n independently from N (0 , σ 2 UU T ) .
For every i = 1 , . . . , n , we draw m i data points x ij for j = 1 , . . .
, m i as x ij =
µ
i + z ij , where each z ij is drawn independently from the spherical Gaussian N (0 , η 2 i I ) .
Let ˆΓ be any estimator mapping ( x ij ) 1 ≤ i ≤ n, 1 ≤ j ≤ m i to a (possibly randomized) k -dimensional subspace of R d .
Let θ denote the maximum principal angle between ˆΓ(( x ij ) 1 ≤ i ≤ n, 1 ≤ j ≤ m i ) and the true subspace Γ = col ( U ) .
If real numbers t ≥ 0 and δ ∈ [0 , 1 / 2) satisfy Pr[sin θ ≤ t ] ≥ 1 − δ, then
[Equation]
where γ 1 , . . . , γ n are deﬁned in Deﬁnition 3.1.
Note that γ (cid:48) i
= γ i for i ≥ k , so our upper bound in (10) matches the lower bound (11) up to a constant factor assuming δ = Θ(1) and d ≥ (1 + Ω(1)) k .
We use the local Fano method to prove the lower bound using the technical lemmas in Appendix D.
In particular, we reduce our goal to proving an upper bound on the KL divergence between Gaussian distributions whose covariance matrices are deﬁned based on matrices U, ˆ
U ∈
O d,k with (cid:107) UU T − ˆ U ˆ U T (cid:107) F bounded.
We prove the following lemma in Appendix J that upper bounds the KL divergence using (cid:107) UU T − ˆ U ˆ U T (cid:107) F :
Lemma 3.8.
For σ ∈ R ≥ 0 , η ∈ R > 0 , U, ˆ U ∈
O d,k , deﬁne Σ
= σ 2 UU T + η 2 I and ˆΣ =
σ 2 ˆ U
ˆ U T + η
I .
Then,
[Equation]
Lemma 3.8 and the results in Appendix D allow us to prove a version of (11) in which the sum in the demoninator is over i = 1 , . . . , n .
This, however, is weaker and less useful than (11) in which the sum in the denominator is over i =
k, k + 1 , . .
. , n .
To prove Theorem 3.7, we extract a hard distribution in which the data points from users 1 , . .
, k − 1 are “useless” in terms of subspace recovery.
Let Γ 1 be the ( k − 1) -dimensional subspace spanned by µ 1 , . . .
, µ k − 1 .
We let v 1 , . .
, v k − 1 be a random orthonormal basis of Γ 1 , and we append another vector v k ∈ Γ to form an orthonormal basis v 1 , . . .
, v k of Γ .
We deﬁne V 1 =
[ v 1 · · · v k − 1 ] ∈ O d,k − 1 and V =
[ v 1 · · · v k ] ∈ O d,k .
In Figure 1 we show a graphical model demonstrating the dependency among the random objects we deﬁned.
Let us focus on the joint distribution of ( V 1 , V, ( µ 1 , . . .
, µ k − 1 )) .
By the invariance property, for any matrices ˜ V 1 ∈ O d,k − 1 , ˜ V ∈ O d,k , measurable set S ⊆
( R d ) k − 1 , and orthogonal matrix G ∈
O d ,
Pr[( µ 1 , . . . , µ k − 1 ) ∈ S | V =
˜
V , V 1 = ˜ V 1 ] =
Pr[( µ 1 , . .
. , µ k − 1 ) ∈ S G | V = G ˜ V , V 1 = G ˜ V 1 ] ,
[Table]
+OF4ceCrERnsH92x/8M/57O9lp5mpeslqAsE9SYCYkrO3VUW84E+HZaG6gou6aXMAlSUQlm6hYNe3wQPDkuSh2OsnjhvZ/hqDSmkVkgJbVXZj02d/4vNqlt0Z86rqragmLLh4paYFvi+e9xzjUwK5ogKNM89IrZFdWU2TCjdjvNoQiDXPTjZJOJGrz7ejL0Lj7ChPTD9bHv1zEN+YrqJcdHOPkQsCRQCn6wMoxX5S698c6lD+re+DWkuY/Mazbeh72Q9S08Fue9Lkm6vTPSGQzR0rbQW/QOHSKCjtEAfUGnaIQYEugn+oV+RxvR+4hEyRKNWqucN+iBRZ/+AfWEx+U
[Figure]
> where S G = { ( G ˜ µ 1 , . . .
, G ˜ µ k − 1 ) :
(˜ µ 1 , . . . , ˜ µ k − 1 ) ∈ S } .
For any ˜ V , ˜ V (cid:48) ∈
O d,k whose ﬁrst k − 1 columns are both ˜ V 1 , there exists G ∈ O d such that ˜ V (cid:48)
=
G ˜ V and thus ˜ V 1 =
G ˜ V 1 .
This implies that for any ˜ µ ∈ col ( ˜ V 1 ) , we have G ˜ µ = ˜ µ , and thus ( S ∩ col ( ˜ V 1 ) k − 1 ) G = S ∩ col ( ˜ V 1 ) k −for any measurable S ⊆
( R d ) k − 1 .
Here, col ( ˜ V 1 ) k − 1 = { (˜ µ 1 , . . .
, ˜ µ k − 1 ) :
˜ µ i ∈ col ( ˜ V 1 ) for i = 1 , . .
, k − 1 } ⊆
( R d ) k − 1 .
When conditioned on V 1 = ˜ V 1 , for every i = 1 , . .
, k − 1 we have µ i ∈ Γ 1 = col
( V 1 )
=
col ( ˜ V 1 ) , which implies that ( µ 1 , . .
. , µ k − 1 ) ∈ col
( ˜ V 1 ) k − 1 .
Therefore,
[Equation]
This implies that ( µ 1 , . . . , µ k − 1 ) and V are conditionally independent given V 1 .
Therefore, the joint distribution of ( V 1 , V, ( µ 1 , . .
. , µ k − 1 )) can be formed by ﬁrst drawing V and V 1 , and then drawing µ 1 , . . .
, µ k − 1 based only on V 1 and not on V .
Since µ k , . . .
, µ n are drawn iid from N (0 , σ 2 UU T ) =
N (0 , σ 2 V V T ) , we have the graphical model shown in Figure 2.
[Equation]
[Table]
[Figure]
By Claim 3.6, the marginal distribution of V is Haar ( O d,k ) .
By Claim 3.
[Figure]
we can implement this distribution by ﬁrst drawing W ∼ Haar ( O d ) and then drawing E independently from any distribution over O d,k and let V = WE .
We choose the distribution of E later, where we ensure that the ﬁrst k − 1 columms of E is always
[Equation]
This guarantees that the ﬁrst k − 1 columns of W and V are the same, and thus V 1 is exactly the ﬁrst k − 1 columns of W , resulting in the graphical model shown in Figure 3.
[Equation]
[Table]
[Figure]
[Table]
[Figure]
Note that in Figure 3 there is no directed path from E to ( µ 1 , . . . , µ k − 1 ) .
Intuitively, this means that knowing ( µ 1 , . . . , µ k − 1 ) gives us no information about E .
Now by choosing the distribution of E appropriately, we can prove (11) in which the denominator does not contain γ 1 , . . .
, γ k − 1 .
We defer the complete proof of Theorem 3.7 to Appendix K.
Linear Models
In the linear models setting, the data distribution of user i is parameterized by an unknown vector β i ∈ R d .
As before, we assume that the vectors β 1 , . . .
, β n from the n users lie in an unknown k -dimensional subspace Γ .
Our goal is to recover the subspace using the following data.
For every i = 1 , . . . , n , we have m i data points from user i : ( x i 1 , y i 1 ) , . . .
, ( x im i , y im i ) ∈ R d × R .
For every j = 1 , . . .
, m i , we assume the measurement x ij ∈ R d is a random vector drawn independently from an O (1) -sub
-Gaussian distribution with zero mean and identity covariance matrix.
The measurement outcome y ij is determined by y ij =
x T ij β
i + z ij , where the random noise z ij ∈ R can depend on the measurements x i 1 , . . .
, x im i .
When conditioned on x i 1 , . .
, x im i , we assume every z ij for j = 1 , . . .
, m i is independently drawn from an η i -sub-Gaussian distribution with zero mean, but we do not assume that the conditional distribution of z ij is the same for every j = 1 , . . .
, m i .
The (in)dependence among x ij and z ij for i = 1 , . .
. , n and j = 1 , . .
, m i can be summarized by the example graphical model in Figure 4.
[Table]
/sPWw++jxk6fPtneeT01ZawYTVopSH2fUgOAKJpZbAceVBiozAUfZ+cc2fnQB2vBSfbVNBXNJTxQvOKM2uKbfFi4Z+MV2L+7HS8O3BVmL3ujl4V/+a/z7YLHTydO8ZLUEZZmgxsxIXNm5o9pyJsB309pARdk5PYFZkIpKMHO3bNfjN8GT46LU4SiLl97rGY5KYxqZBVJSe2o2Y63zf7FZbYvh3HFV1RYUWz1U1ALbErd/xznXwKxogqBM89ArZqdUU2bDhLrdNIcijHHZj5NNJmrw7sunsXfxHiZkGK73Q7+JacjX1CDZ38PJu4AlgVJwyUopqcpdeuWdS2/UvfIbSHMdaWs2vt0L2dzCbTEd9EnSHxyS3miMVraFXqHXaBcRtI9G6DM6QBPE0Bn6jn6gn1EU7UZvI7JCo8465wW6YdGHf2Zqxzg=</latexit>
Since we allow the noise z ij to depend on the measurements x ij , it is information-theoretically impossible to recover the subspace if we only have one data point from every user.
Consider the scenario where every β i is drawn independently from N (0 , σ 2 uu T ) for an unknown unit vector u ∈ R d and every x ij is drawn independently and uniformly from {− 1 , 1 } d .
If we set z ij to be z ij = x T ij ν ij where ν ij is independently drawn from N (0 , σ 2 ( I − uu T )) , then every y ij satisﬁes y ij = x
T ij ( β i + ν ij ) where β i + ν ij distributes as N (0 , σ 2 I ) independently from x ij .
This implies that the joint distribution of (( x i 1 , y i 1 ))
i
=1 ,...,n does not change with u , i.e., we get no information about u from one data point per user.
Thus, we assume m i ≥ 2 for every user i .
In this case, we achieve error upper bounds that match the ones in [Tripuraneni et al., 2021] despite our relaxed assumptions on the noise.
Our estimator is the subspace ˆΓ spanned by the top- k eigenvectors of A deﬁned in (3).
We defer the analysis of our estimator to Appendix L.
Acknowledgments and Disclosure of Funding
Part of this work was performed while LH was interning at Apple.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[N/A]
Our work aims to develop better subspace recovery methods for a diverse population of users with heterogeneous data distributions.
Researchers and practitioners working on transfer learning, federated learning, data analysis, and statistics will beneﬁt from our work.
We are not aware of any negative societal impacts of our work.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[N/A] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[N/A] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[N/A] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[N/A] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[N/A] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Basics about Principal Angles
We provide a formal deﬁnition of the principal angles introduced by Jordan [1875] and brieﬂy discuss some basic properties of this notion (see e.g. [Stewart and Sun, 1990] for a textbook).
Let U, ˆ U ∈
O d be two orthogonal matrices.
Suppose the columns of U and ˆ U are partitioned as U
=
[ U 1 U 2 ] , ˆ U =
[ ˆ U 1 ˆ U 2 ] where U 1 ∈ O d,k
and ˆ U 1 ∈ O d, ˆ k for integers k, ˆ k satisfying 0 < k ≤ ˆ k < d .
Starting from a singular value decomposition of U T 1 ˆ U 1 as R T 1 Σ ˆ R 1 , one can ﬁnd an integer (cid:96)
∈
Z ≥ 0 and angles 0 < θ 1 ≤ · · · ≤ θ (cid:96) ≤ π/ 2 such that there exist orthogonal matrices
P, ˆ P ∈
O d , R 1 ∈
O k , R 2 ∈ O d − k , ˆ R 1 ∈
O ˆ k , ˆ R 2 ∈
O d − ˆ k satisfying
U = P (cid:20) R
R(cid:21) , ˆ U =
ˆ P (cid:20)
ˆ R  ˆ R(cid:21) , and (12)
P
[Equation] [Equation]
=  
I
[Equation]
(cid:96)     cos Θ sin Θ
I ˆ k − k   − sin Θ cos Θ
I d − ˆ k
[Equation]
(cid:96)   , (13) where cos Θ = diag(cos θ 1 , . . .
, cos θ (cid:96) ) and sin Θ = diag(sin θ 1 , . . . , sin θ (cid:96) ) .
It is easy to see that cos θ 1 , . . .
, cos θ (cid:96) are exactly the singular values of U T
1 ˆ U 1 that are not equal to 1 , so (cid:96), θ 1 , . . .
, θ (cid:96) are unique.
Let Γ (resp. ˆΓ ) be the k -dimensional (resp.
ˆ k -dimensional) linear subspace spanned by the columns of U 1 (resp.
ˆ U 1 ).
The (non-zero) principal angles between Γ and ˆΓ are deﬁned to be θ 1 , . .
, θ (cid:96) .
It is not hard to see that the principal angles depend only on the subspaces Γ , ˆΓ and not on the choices of U and ˆ U .
The maximum principal angle between Γ and ˆΓ , denoted by ∠ (Γ , ˆΓ) or ∠ ( U 1 , ˆ U 1 ) , is deﬁned to be θ (cid:96) (or 0 when (cid:96) = 0 ).
When k = ˆ k , using (12) and (13), it is not hard to show that the non-zero eigenvalues of U 1 U T 1 − ˆ U
1 ˆ U T 1 are ± sin θ 1 , . .
. , ± sin θ (cid:96) .
This implies (cid:107) U 1 U T 1 − ˆ U
ˆ U T 1 (cid:107) =
sin ∠ (Γ , ˆΓ) and (cid:107) U 1 U T 1 − ˆ U
1 ˆ U T 1 (cid:107) F = (cid:113) 2 (cid:80) (cid:96)i =1 sin 2 θ i .
In particular, sin ∠ (Γ , ˆΓ) is a natural metric between k -dimensional subspaces.
B Basic Facts and Concentration Inequalities
Lemma B.1 ([See Vershynin, 2018, Proposition 2.6.1]) .
If X 1 , . .
. , X n ∈ R are independent random variables with sub-Gaussian constants b 1 , . . .
, b n ∈ R ≥ 0 respectively, then (cid:80) ni
=1 X i has subGaussian constant O (cid:16)(cid:112)(cid:80) ni
=1 b 2 i (cid:17) .
Lemma B.2 ([See Vershynin, 2018, Lemma 2.7.6]) .
If X ∈ R is a mean-zero random variable with sub-Gaussian constant b ∈ R ≥ 0 , then X 2 has sub-exponential constant O ( b 2 ) .
Lemma B.3 ([See Vershynin, 2018, Theorem 2.6.2]) .
Let x 1 , . . .
, x n ∈ R be independent random variables.
For every i = 1 , . . . , n , assume x i has mean zero and sub-Gaussian constant b i ∈ R ≥ 0 .
Then for every δ ∈ (0 , 1 / 2) , with probability at least 1 − δ ,
[Equation]
n (cid:88)
[Equation]
Lemma B.4 ([See Vershynin, 2018, Theorem 2.8.1]) .
Let x 1 , . . .
, x n ∈ R be independent random variables.
For every i = 1 , . . . , n , assume x i has mean zero and sub-exponential constant b i ∈ R ≥ 0 .
Then for every δ ∈ (0 , 1 / 2) , with probability at least 1 − δ , (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
i
[Equation]
O 
[Equation]
n
[Equation]
i
[Equation]

 .
C Vector and Matrix Concentration Inequalities
Lemma C.1 ([See Vershynin, 2018, Corollary 4.2.13 and Exercise 4.4.3]) .
There exists an absolute constant C > 1 and a set O (cid:48)
d ⊆ O d, 1 for every positive integer d such that
| O (cid:48)
d | ≤ 2 Cd for every d ∈ Z > 0 . 2.
for every d ∈ Z > 0 and every symmetric matrix A ∈ R d × d , (cid:107) A (cid:107) ≤ C sup u ∈ O (cid:48) | u T Au | .
for every d 1 , d 2 ∈ Z > 0 and every matrix A ∈ R d 1 × d 2 ,
[Equation]
A (cid:107)
[Equation]
sup u
[Equation]
d
[Equation]
Lemma C.2.
Suppose x 1 , . . .
, x n ∈ R
d are independent random vectors, and each x i is mean-zero and 1 -sub-Gaussian.
Suppose w 1 , . . .
, w n ∈ R are ﬁxed real numbers.
Deﬁne E = (cid:80) ni
=
w i x i x T i and ¯ E
= E
[ E ] .
Then for any δ ∈ (0 , 1 / 2) , with probability at least 1 − δ ,
[Equation]
≤ O
[Equation]
Proof.
Let O (cid:48) denote the set O (cid:48)
d ⊆ O d, 1 guaranteed by Lemma C.1.
For any ﬁxed u ∈ O (cid:48) , by
Lemma B.2, ( x T i u ) 2 has sub-exponential constant O (1) , and thus w i ( x T i u ) 2 has sub-exponential constant | w i | .
Deﬁne δ (cid:48) :
= δ/ | O (cid:48) | .
By Lemma B.4, with probability at least 1 − δ (cid:48) , (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
By a union bound, with probability at least 1 − δ , the above inequality holds for every u ∈ O (cid:48) .
By the deﬁnition of O (cid:48) ,
[Equation]
Combining (14) and (15) and noting that log(1 /δ (cid:48) )
=
O ( d + log(1 /δ )) completes the proof.
Lemma C.3.
Suppose x 1 , . . .
, x n ∈ R
d are independent random vectors, and each x i is mean-zero and 1 -sub-Gaussian.
Suppose b 1 , . . .
, b n ∈ R are ﬁxed real numbers.
Then for any δ ∈ (0 , 1 / 2) , with probability at least 1 − δ ,
[Equation]
[ x· · · x n ] diag( b 1 , . . . , b n ) (cid:107) ≤ O   (cid:118)(cid:117)(cid:117)(cid:116) n (cid:88) i =
b 2 i + ( d + log(1 /δ )) max i b 2 i 
 .
Proof.
Deﬁne F =
[ x· · · x n ] diag( b 1 , . . . , b n ) and E = FF T = (cid:80) ni
=
b 2 i x i x T i .
By
Lemma C.2, with probability at least 1 − δ , (cid:107) E − E
[ E ] (cid:107) ≤ O   (cid:118)(cid:117)(cid:117)(cid:116)
( d + log(1 /δ ))
n (cid:88) i =
b 4 i + ( d + log(1 /δ )) max i b 2 i 
 .
For every unit vector u ∈ R d , x T i u is 1 -sub-Gaussian, and thus E [( x T i u ) 2 ] ≤ O (1) .
This implies that u E [ x i x T i ] u ≤ O (1) for every unit vector u ∈ R d and thus (cid:107) E
[ x i x T i ] (cid:107) ≤ O (1) .
Now we have (cid:107) E
[ E ] (cid:107) ≤ n (cid:88)
i =
b 2 i (cid:107) x
i x T i (cid:107) ≤ O (cid:32) n (cid:88) i =
b 2 i (cid:33) .
Combining (16) and (17) and using the fact that (cid:80) ni =1 b 4 i ≤ (cid:0)(cid:80) ni
=1 b 2 i (cid:1) max i b 2 i , with probability at least 1 − δ , (cid:107) E (cid:107) ≤ (cid:107) E − E
[ E ] (cid:107) + (cid:107) E
[ E ] (cid:107) ≤ O (cid:32) n (cid:88) i
=
b 2 i + ( d + log(1 /δ )) max i b
i (cid:33) .
The lemma is proved by (cid:107) F (cid:107) =
(cid:112) (cid:107) E (cid:107) .
Lemma C.4.
Let X ∈ R
d
× m be a ﬁxed matrix and let v ∈ R
m be a mean zero random vector with sub-Gaussian constant 1 .
For δ ∈ (0 , 1 / 2) , with probability at least 1 − δ ,
[Equation]
where r is the rank of X .
Proof.
By the singular value decomposition, there exists P ∈
O d such that all but the ﬁrst r rows of
PX are zeros.
Let Y ∈ R r × m denote the ﬁrst r rows of PX .
We have
[Equation]
Let O (cid:48) be the O (cid:48) r in Lemma C.1.
For every u ∈ O (cid:48) , u T Y v has sub-Gaussian constant O ( (cid:107) u T Y (cid:107) 2 ) .
By Lemma B.3, for δ (cid:48)
=
δ/ | O (cid:48) | , with probability at least 1 − δ (cid:48) ,
[Equation]
By a union bound, with probability at least 1 − δ ,
[Equation] [Equation]
By the deﬁnition of O (cid:48) ,
[Equation]
The lemma is proved by combining (18), (19) and (20) and noting that (cid:107) Y (cid:107) =
(cid:107) X (cid:107) and log(1 /δ
(cid:48) )
=
O ( r + log(1 /δ )) .
Lemma C.5.
Let Z 1 , . . .
Z n ∈ R d 1 × d 2 be mean-zero independent random matrices.
Suppose for real numbers R, p 1 , . .
. , p n ∈ R ≥ 0 , we have Pr[ (cid:107)
Z i (cid:107) ≥ R ] ≤ p i .
Moreover,
[Equation]
Then with probability at least 1 − 2( d 1 + d 2 ) e − t − (cid:80) ni
=1 p i , (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) n (cid:88) i =
[Equation]
≤ O 
[Equation]
n (cid:88) i =
[Equation]
2 i
[Equation]
+ Rt   .
Proof.
Since probabilities are always nonnegative, the lemma is trivial if (cid:80) ni
=1 p i ≥ 1 or if t ≤ c for a sufﬁciently small positive constant c .
We thus assume (cid:80) ni
=1 p i ≤ 1 and t ≥ Ω(1) .
Deﬁne Z (cid:48)
i
=
Z i 1 ( (cid:107) Z i (cid:107) ≤ R ) .
For every unit vector u ∈ R d , by Cauchy-Schwarz,
[Equation]
p
[Equation]
By Jenson’s inequality,
[Equation]
Combining,
[Equation]
for every unit vector u ∈ R d 2 , which implies that (cid:107) E
[ Z (cid:48)
i ] (cid:107) =
(cid:107) E
[ Z i ] − E
[ Z (cid:48)
i ] (cid:107) =
(cid:107) E
[ Z i 1 ( (cid:107) Z i (cid:107) > R )] (cid:107) ≤ σ i √ p i .
We apply the matrix Bernstein inequality [see Vershynin, 2018, Exercise 5.4.15] to Z (cid:48)
i .
For every i , (cid:107) Z
[Equation]
i − E
[ Z (cid:48)
[Equation]
] (cid:107) ≤ (cid:107) Z (cid:48)
i (cid:107) + (cid:107) E
[ Z (cid:48)
[Equation]
] (cid:107) ≤ (cid:107) Z (cid:48)
i (cid:107) + E [ (cid:107) Z (cid:48)
i (cid:107) ] ≤ 2 R , and
E [( Z (cid:48)
i − E
[ Z (cid:48)
i ])
( Z (cid:48)
i − E
[ Z (cid:48)
i ]) T ] =
E [ Z (cid:48)
i ( Z (cid:48) i ) T ] − E
[ Z (cid:48)
i ]
E [ Z (cid:48)
i ] T (cid:22) E
[ Z (cid:48)
i ( Z (cid:48) i ) T ] (cid:22) σ 2 i I.
The matrix Bernstein inequality implies that with probability at least 1 − 2( d 1 + d 2 ) e − t , (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) n (cid:88) i =
Z (cid:48)
i − n (cid:88) i =
E [ Z (cid:48)
i ] (cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ≤ O   (cid:118)(cid:117)(cid:117)(cid:116) n (cid:88) i =
σ 2 i t + Rt   .
By the union bound
[Table]
probability
[Table]
least
[Table]
n (cid:88) i =
σ 2 i
[Equation]
(by Cauchy-Schwarz)
The lemma is proved by our assumption that (cid:80) n i =1 p i ≤ 1 and t ≥ Ω(1) .
D Fano’s Inequality and Local Packing Numbers
For two probability distributions p 1 , p 2 over the same set X with probability densities p 1 ( x ) and p 2 ( x ) , their KL divergence is deﬁned to be
[Equation]
For a joint distribution p over random variables X
and Y , we deﬁne the mutual information between
X and Y to be
[Equation]
where p X (resp.
p Y ) is the marginal distribution of X (resp.
Y ), and p X × p Y is the product distribution of p X and p Y .
The following claim is standard:
Claim D.1.
For random variables X, Y, Z , if Y and Z are conditionally independent given X , then
[Equation]
Theorem D.2 (Fano’s Inequality) .
Let X, Y, ˆ
X
be random objects.
Assume their joint distribution forms a Markov chain X
→ Y → ˆ X .
Assume that the marginal distribution of X is uniform over a ﬁnite set X .
Then,
[Equation]
Lemma D.3.
In the setting of Theorem D.2, let P x denote the conditional distribution of Y given
X = x .
Then,
[Equation]
Proof.
Let
P x ( y ) denote Pr
[ Y = y
| X = x ] .
Since X distributes uniformly over X , the marginal distribution of Y is given by Pr[ Y = y ]
= (cid:80) x ∈X P x ( y ) / |X| .
By the deﬁnition of I ( X ; Y ) ,
I ( X ; Y
[Equation]
Jensen’s Inequality
[Equation]
= 1 |X| (cid:88) x (cid:48) ∈X |X| (cid:88)
x ∈X
E Y ∼ P x (cid:48) (cid:20) log P x (cid:48)
( Y )
P x ( Y ) (cid:21) = |X|(cid:88)
x ∈X (cid:88) x (cid:48) ∈X
D kl (
[Equation]
x (cid:48) (cid:107) P x )
[Equation]
max
[Equation]
Lemma D.4 ([see Cai et al., 2013, Lemma 1]) .
There exists an absolute constant C > 0 with the following property.
For any positive integers k ≤ d and real number t ∈ (0 , 1 /C ] , there exists a subset O (cid:48) ⊆
O d,k with size at least 10 k ( d − k ) such that
For distinct U, V ∈ O (cid:48) , (cid:107) UU T − V V T (cid:107) F > t . 2.
For any U, V ∈ O (cid:48) , (cid:107) UU T − V V T (cid:107) F ≤ Ct .
E Proof of Theorem 2.
Theorem 2.1 (Variant of Davis–Kahan sin θ theorem) .
Let A, ˆ A ∈ R d
× d be symmetric matrices.
Let λ i denote the i -th largest eigenvalue of A .
For a positive integer k smaller than d , let θ denote the maximum principal angle between the subspaces spanned by the top- k eigenvectors of A and ˆ A .
Assuming λ k > λ k +1 ,
[Equation]
Proof.
We assume 2 (cid:107) A − ˆ A (cid:107) λ k − λ k
+
< 1 because otherwise the theorem is trivial.
Let ˆ λ
i denote the i -th largest eigenvalue of ˆ A .
By Weyl’s inequality,
[Equation]
where we used our assumption 2 (cid:107) A − ˆ A (cid:107) λ k − λ k +1 < 1 in the last inequality.
This implies
[Equation]
By the Davis-Kahan sin θ theorem [Davis and Kahan, 1970],
[Equation]
[Equation]
Combining (23) and (24) completes the proof.
F Proof of Theorem 3.
We recall:
Theorem 3.1.
Deﬁne ξ 2 = (cid:107) (cid:80) ni
=1 w 2 i µ i µ T i η 2 i /m i (cid:107) and let θ denote the maximum principal angle between our estimator ˆΓ and the true subspace Γ spanned by µ 1 , . . . , µ n .
For any δ ∈ (0 , 1 / 2) , with probability at least 1 − δ ,
[Equation]
By (8), it sufﬁces to upper bound the spectral norm of A − (cid:80) n i =
1 w i µ i µ T
i .
We achieve this goal by proving Lemmas F.1 to F.3 in which we bound the spectral norm of each term on the right-hand-side of the following decomposition:
[Equation]
where ¯ z
[Equation]
Lemma F.1.
In the setting of Theorem 3.1, for any δ ∈ (0 , 1 / 2) , with probability at least 1 − δ ,
[Equation]
Proof.
Let O (cid:48) denote the set O (cid:48)
d ⊆ O d, 1 guaranteed by Lemma C.1.
For any ﬁxed u, v ∈ O (cid:48) , by Lemma B.1, ¯ z T i v is η i / √ m i -sub-Gaussian, and thus w i ( µ T i u )(¯ z T
i v )
is w i | µ T i u | η i / √ m i -subGaussian.
By Lemma B.3, for any δ (cid:48) ∈ (0 , 1 / 2) , with probability at least 1 − δ (cid:48) ,
[Equation]
By the deﬁnition of ξ , n
[Equation]
Plugging this into (26), and setting δ (cid:48)
= δ/ | O (cid:48) | 2 , with probability at least 1 − δ/ | O (cid:48) | 2 ,
[Equation]
where we used the fact that log | O (cid:48) |
= O ( d ) to obtain the last inequality.
By a union bound over u, v ∈ O (cid:48) , with probability at least 1 − δ , sup u ∈ O (cid:48) sup v ∈ O
[Equation]
| u
[Equation]
ni
[Equation]
w
[Equation]
µ
[Equation]
¯ z
[Equation]
(cid:1) v | ≤
O ( ξ (cid:112) d + log(1 /δ ) .
The lemma is proved by combining this with (cid:13)(cid:13)(cid:80) ni
=
w i µ i ¯ z T i (cid:13)(cid:13) ≤ sup u ∈ O (cid:48) sup v ∈ O (cid:48)
| u T (cid:0)(cid:80) ni
=
w i µ i ¯ z T i (cid:1) v | by Lemma C.1.
Lemma F.2.
In the setting of Theorem 3.1, deﬁne Z = (cid:80) ni
= w
[Equation]
For any δ ∈ (0 , 1 / 2) , with probability at least 1 − δ ,
[Equation]
Proof.
By Lemma B.1, (cid:80) m i j =1 z ij has sub-Gaussian constant b i for some b i ∈ R ≥ 0 satisfying b i
= O ( √ m i η i ) .
Deﬁne ˆ
z
i =
b i (cid:80) m i j =1 z ij
and now ˆ z i has sub-Gaussian constant 1 .
We write Z in terms of ˆ z i :
Z =
[Equation]
Applying Lemma C.2 completes the proof.
Lemma F.3.
In the setting of Theorem 3.1, deﬁne Z = (cid:80) ni =
[Equation]
(cid:80) m i j =1 z ij z T ij .
For any δ ∈ (0 , 1 / 2) , with probability at least 1 − δ ,
[Equation]
Proof.
Deﬁne ˆ
z ij = 1 η i z ij and now ˆ z ij has sub-Gaussian constant 1 .
We write Z in terms of ˆ z ij :
[Equation]
Applying Lemma C.2 completes the proof.
We are now ready to ﬁnish the proof of Theorem 3.1.
Proof of Theorem 3.1.
We start with the last term in (25):
[Equation]
It is clear that E (cid:104)(cid:80)
[Equation]
= 0 , so combining Lemmas F.2 and F.3, with probability at least 1 − δ ,
[Equation]
where we used the fact that the right-hand-side of (27) is at most the right-hand-side of (28) (up to a constant factor).
By Lemma F.1,
[Equation]
Plugging (29) and (30) into (25) and then into (8) proves the theorem.
G Proof of Corollary 3.
Proof.
By the deﬁnition of ξ 2 in Theorem 3.1, ξ 2 = (cid:107) (cid:80) ni =1 µ i µ T i t 2 n 2 (cid:107) ≤ t 2 σ 21 n .
Plugging this into (6), we know that the following inequality holds with probability at least 1 − δ :
[Equation]
Since sin θ ≤ 1 always holds, inequality (31) implies
[Equation]
Since min { 1 , y } ≤ √ y for every y ∈ R ≥ 0 , we have
[Equation]
Plugging (33) into (32) proves the corollary.
H Proof of Claim 3.
Proof.
Without loss of generality, we can assume that the subspace Γ in Assumption 3.1 is the subspace containing all vectors with all but the ﬁrst k coordinates being zeros.
We let ˆ µ i ∈ R k denote the ﬁrst k coordinates of µ i .
By Lemma C.2, with probability at least 1 − δ/ 2 ,
[Equation]
≤ O 
[Equation]
where we used the fact that (cid:80) ni
=
w 2 i ≤ max i w i ≤ 1 /C ∗ ( k +log(1 /δ )) .
Choosing C ∗ large enough, (34) implies
[Equation] [Equation]
By Assumption 3.1, the k -th largest eigenvalue of E [ µ i µ T i ] is at least σ 2 .
This implies that the k -th largest eigenvalue of E [ (cid:80) ni =1 w
i ˆ µ i ˆ µ T i ] is at least σ 2 .
The claim is proved by combining this with (35) and noting that (cid:80) ni
=1 w i ˆ µ i ˆ µ T i and (cid:80) ni
=
w i µ i µ T i have the same k -th largest eigenvalue.
I Proof of Theorem 3.
Theorem 3.4.
Under Assumptions 3.1 and 3.2, if we choose w 1 , . . .
, w n as in (9) and deﬁne θ = ∠ (Γ , ˆΓ) , for δ ∈ (0 , 1 / 2) , with probability at least 1 − δ , sin θ ≤ O (cid:32)(cid:115)
[Equation]
Proof.
By Assumption 3.2 and our choice of w i , we have w i ≤ 1 /C ∗ ( k + log(1 /δ )) .
By the deﬁnition of C ∗ , with probability at least 1 − δ/ 2 , σ
[Equation]
[Equation] [Equation]
By Lemma C.3, with probability at least 1 − δ/ 4 ,
[Equation]
Replacing δ in Theorem 3.1 by δ/ 4 and plugging (36) and (37) into (6), by the union bound, with probability at least 1 − δ ,
[Equation] [Equation]
By deﬁnition, η 2 i σ 2 m i + η 4 i σ 4 m 2 i = 1 /γ i ≤ 1 /γ (cid:48) i .
Therefore, inequality (38) implies sin θ ≤ O
[Equation]
Plugging w i = γ (cid:48)
i (cid:80) n(cid:96)
=1 γ (cid:48) (cid:96) into the inequality above, sin
[Equation]
Since sin θ ≤ 1 always holds, inequality (39) implies sin θ ≤ O
[Equation]
The theorem is proved by the following facts:
[Equation]
J Proof of Lemma 3.
Proof.
The KL divergence between mean-zero Gaussians can be computed by the following formula:
[Equation]
It is clear that det Σ =
[Equation]
It remains to compute tr(Σ − 1 ˆΣ) − d .
Deﬁne J = (cid:20)
I k
(cid:21) ∈ R
d × d .
By the deﬁnition of principal angles in Section 2, there exists
P, ˆ P ∈
O d such that UU T =
PJP T ,
ˆ U
ˆ U T
= ˆ PJ ˆ P T and
P T ˆ P =  
I k − (cid:96)    cos Θ sin Θ  − sin Θ cos Θ
I d − ˆ k − (cid:96) 
 , where θ 1 , . . .
, θ (cid:96) are the principal angles between col ( U ) and col ( ˆ U ) , cos Θ = diag(cos θ 1 , . . .
, cos θ (cid:96) ) , and sin Θ = diag(sin θ 1 , . . . , sin θ (cid:96) ) .
Now we have ˆ U ˆ U T
= ˆ PJ ˆ P T =
P ( P T ˆ P ) J ( P T ˆ P ) T P T =
P ˆ JP T where
[Equation] [Equation]
− cos
[Equation]
sin Θ (sin Θ
[Equation]

 .
Therefore, Σ = P ( σ 2 J + η 2 I ) P T and ˆΣ =
P ( σ 2 ˆ J + η 2 I ) P T , which implies Σ − 1 ˆΣ = P ( σ 2 J + η 2 I ) − 1 ( σ 2 ˆ J + η 2 I ) P T .
Since σ 2 J + η 2 I = diag( σ 2 + η 2 , . . .
, σ 2 + η(cid:124) (cid:123)(cid:122) (cid:125) k , η 2 , . . .
, η(cid:124) (cid:123)(cid:122) (cid:125) d − k ) and the diagonal entries of σ 2 ˆ J + η 2 I are
[Equation]
we have
[Equation] [Equation]
Plugging (41) and (42) into (40) completes the proof.
K Proof of Theorem 3.
We recall Theorem 3.7:
Theorem 3.7.
Let k, d, n be positive integers satisfying k < d and k ≤ n .
Let m 1 , . . .
, m n be positive integers and σ, η 1 , . . .
, η n be positive real numbers.
Suppose we draw U ∈
O d,k from
Haar ( O d,k ) and then draw µ 1 , . .
. , µ n independently from N (0 , σ 2 UU T ) .
For every i = 1 , . . . , n , we draw m i data points x ij for j = 1 , . . .
, m i as x ij =
µ
i + z ij , where each z ij is drawn independently from the spherical Gaussian N (0 , η 2 i I ) .
Let ˆΓ be any estimator mapping ( x ij ) 1 ≤ i ≤ n, 1 ≤ j ≤ m i to a (possibly randomized) k -dimensional subspace of R d .
Let θ denote the maximum principal angle between ˆΓ(( x ij ) 1 ≤ i ≤ n, 1 ≤ j ≤ m i ) and the true subspace Γ = col ( U ) .
If real numbers t ≥ 0 and δ ∈ [0 , 1 / 2) satisfy Pr[sin θ ≤ t ] ≥ 1 − δ, then
[Equation]
where γ 1 , . . . , γ n are deﬁned in Deﬁnition 3.1.
Proof.
The theorem is trivial if t is lower bounded by a positive absolute constant, so without loss of generality, we can assume that t ≤ 1 / (2 √ 2 C ) for the constant C in Lemma D.4.
Now we describe the distribution of E on which we prove the error lower bound (11) that does not depend on γ 1 , . .
, γ k − 1 .
As we mentioned in Section 3.3, the ﬁrst k − 1 columns of E are ﬁxed to be
[Equation]
I
[Equation]
− 1 0 (cid:21) , so we only need to describe the distribution of the k -th column of E .
By Lemma D.4, there exists O (cid:48) ⊆
O d − k +1 , 1 with size at least 10 d − k
such that
1. for distinct u, v ∈ O (cid:48) , (cid:107) uu T − vv T (cid:107) F > 2 √ 2 t ; 2. for any u, v ∈ O (cid:48) , (cid:107) uu T − vv T (cid:107) F ≤ O ( t ) .
The k -th column of E is chosen by ﬁrst drawing a uniform random u ∈ O (cid:48) and then prepend zeros to it.
That is
[Equation]
We let E denote the support of the E so that E is chosen uniformly at random from E .
By the law of total expectation, conditioned on W, V 1 , ( µ 1 , . . .
, µ k − 1 ) and ( x ij ) 1 ≤ i<k, 1 ≤ j ≤ m i being ﬁxed to some speciﬁc value, we still have Pr[sin θ ≥ t ] ≤ δ , where now Pr[ · ] represents the conditional probability.
We will abuse notation and omit explicitly writing out the conditioning throughout the proof.
After the conditioning, the randomness in V comes only from the randomness in E ∈ E , and the distribution of V is the uniform distribution over the set V :
= { WE : E ∈ E} .
For any two matrices
V (cid:48) , V (cid:48)(cid:48) ∈ V , there exists E (cid:48) , E
(cid:48)(cid:48) ∈
E such that V (cid:48)
=
WE (cid:48) and V (cid:48)(cid:48)
=
WE (cid:48)(cid:48) .
Moreover, there exists u, v ∈ O (cid:48) such that E (cid:48) =
(cid:20)
I
[Equation]
u
[Equation]
and E (cid:48)(cid:48)
= (cid:20)
I k
[Equation]
v
[Equation]
This implies that
[Equation]
Moreover, when V (cid:48) (cid:54) =
V
[Equation]
, we have u (cid:54) = v and thus
[Equation]
where we used the fact that (cid:107) uu T − vv T (cid:107)
= sin ∠ ( u, v ) and 2 √ 2 t < (cid:107) uu
[Equation]
− vv
[Equation]
(cid:107) F = (cid:113) 2 sin 2 ∠ ( u, v ) .
We deﬁne ˆ V ∈ V so that the maximum principal angle between col ( ˆ V ) and ˆΓ is minimized.
When sin θ ≤ t (i.e., sin ∠ ( col ( V ) , Γ) ≤ t ), for any V (cid:48) ∈ V different from V ,
[Equation]
=
[Equation]
sin ∠
[Equation]
col
[Equation]
V
[Equation]
Γ) , where we used sin ∠ ( V, V (cid:48) )
= (cid:107) V V T − V
[Equation]
( V
[Equation]
) T (cid:107) > 2 t by (45).
Therefore, sin θ ≤ t implies ˆ V = V and thus Pr[sin θ > t ] ≥ Pr[ ˆ V (cid:54) = V ] .
Now we apply Theorem D.2 to the following Markov chain
[Equation]
to get
[Equation]
Now we bound I ( V ; ( x ij ) k ≤ i ≤ n, 1 ≤ j ≤ m i ) .
According to the graphical model in Figure 3, µ k , . . . , µ n are conditionally independent given V , and thus ( x ij ) 1 ≤ j ≤ m i are conditionally independent for i
=
k, . .
. , n given V .
By (21),
[Equation]
Let ¯ x i denote 1 m i (cid:80) m i j =1 x ij .
Since x ij for j = 1 , . . .
, m i are drawn iid from N ( µ i , η 2 i I ) , it is a standard fact that ¯ x i is a sufﬁcient statistic for µ i , i.e., ( x ij ) 1 ≤ j ≤ m i and µ i are conditionally independent given ¯ x i .
Therefore, for any measurable set S ∈ ( R d ) m i ,
[Equation]
Since ( x ij ) j =1 ,...,m i and V are conditionally independent given µ i ,
[Equation]
ij
[Equation]
V,
[Equation]
i , ¯ x i ] =
Pr[
[Equation]
ij
[Equation]
Combining (48) and (49), ( x ij ) 1 ≤ j ≤ m i and V are conditionally independent given ¯ x i .
By (22),
[Equation] [Equation]
Since ¯ x i = µ i + 1 m i (cid:80) m i j
=1 z j , the conditional distribution of ¯ x i given V is N (0 , σ 2 V V T + η
i m i I ) .
By Lemma 3.8 and inequality (44),
I ( V
[Equation]
¯ x
i
[Equation] [Equation]
Combining (47), (50), and (51),
[Equation]
The theorem is proved by plugging this into (46) and noting that log |V| ≥ 2( d − k ) .
L Our Results in the Linear Models Setting
In the linear models setting (see Section 4), our estimator is the subspace ˆΓ spanned by the top- k eigenvectors of A deﬁned in (3).
We use θ to denote the maximum principal angle between our estimator ˆΓ and the true subspace Γ .
To describe our guarantee on θ , it is convenient to ﬁrst deﬁne the following quantities for every i = 1 , . . .
, n :
[Equation]
We prove the following guarantees for our estimator in Appendices L.1 and L.2:
Theorem L.1.
There exists an absolute constant C > 0 such that for any δ ∈ (0 , 1 / 2) , with probability at least 1 − δ ,
[Equation]
Corollary L.2.
Suppose m ∈ Z > 0 and η ∈ R > 0 satisfy 2 ≤ m ≤ min 1 ≤ i ≤ n m i and max 1 ≤ i ≤ n η i ≤ η .
Choosing w 1 = · · · = w n = 1 /n , for any δ ∈ (0 , 1 / 2) , with probability at least 1 − δ , sin θ ≤ O
[Equation]
If we further assume that max 1 ≤ i ≤ n (cid:107) β i (cid:107) 2 ≤ r
[Equation]
then (52) implies
[Equation] [Equation]
In [Tripuraneni et al., 2021], it is further assumed that η = Θ(1) , r = Θ(1) , (cid:107) β
i (cid:107) 2 = Θ(1) , δ = ( mn ) − 100 .
Deﬁning κ = 1 n (cid:80) ni =1 (cid:107) β
i (cid:107) 22 /kσ 2 k = Θ(1 /kσ 2 k ) as in [Tripuraneni et al., 2021] and using the fact (cid:80) ni
=1 (cid:107) β
i (cid:107) 42 ≤ (cid:0) max i (cid:107) β i (cid:107) 22 (cid:1) (cid:80) ni
=1 (cid:107) β
i (cid:107) 22 , our bound becomes
[Equation]
matching the bound in [Tripuraneni et al., 2021, Theorem 3].
As in the PCA setting, our proof of Theorem L.1 is based on the Davis-Kahan sin θ theorem (Theorem 2.1) and a bound on the spectral norm of the difference between the matrix A and its expectation.
We prove and make crucial use of a generalization of the matrix Bernstein inequality (Lemma C.5) which turns out to be slightly more convenient than a similar inequality used in [Tripuraneni et al., 2021].
L.1 Proof of Theorem L.
Plugging y ij = x T ij β
i + z ij into the deﬁnition of A ,
[Equation]
When j 1 (cid:54) =
j 2 , conditioned on x ij 1 and x ij 2 , the noise terms z ij 1 and z ij 2 are independent and have zero mean.
Therefore,
[Equation]
where we used the fact that x ij 1 and x ij 2 are independent and have zero mean and identity covariance matrix.
Thus, the expectation of A is ¯ A = (cid:80) ni
=1 w i β i
β T
i , and as in the PCA setting, our goal is to bound the spectral norm of the difference A − ¯ A .
We decompose A as
[Equation]
T
[Equation]
G, where
[Equation]
G = n (cid:88) i =
w i m i ( m i − 1) (cid:88)
j 1 (cid:54) =
jx
[Equation]
1 z
[Equation]
z ij 2 x T ij 2 .
For every i = 1 , . . . , n and j = 1 , . . .
, m i , deﬁne b ij = z ij x ij and h ij = x ij x T ij β i .
Now we can rewrite E, F, G as E
= E 1 − E 2 , F = F 1 − F 2 , G
= G 1 − G 2 , where
E 1 = n (cid:88) i =
w i m i ( m i − 1)   m i (cid:88)
j = h ij 
   m i (cid:88)
j = h ij 

T ,
E 2 = n (cid:88) i =
w i m i ( m i − 1) m i (cid:88)
j = h ij h T ij ,
F 1 = n (cid:88) i =
w i m i ( m i − 1)   m i (cid:88)
j = h ij 
   m i (cid:88)
j = b ij 

T ,
F 2 = n (cid:88) i =
w i m i ( m i − 1) m i (cid:88)
j = h ij b T ij ,
G 1 = n (cid:88) i =
w i m i ( m i − 1)   m i (cid:88)
j = b ij 
   m i (cid:88)
j = b ij 

T ,
G 2 = n (cid:88) i =
w i m i ( m i − 1)
[Equation]
Deﬁne
[Equation]
The following lemma controls the deviation of E 1 , E 2 , F 1 , F 2 , G 1 , G 2 from their expectations ¯ E 1 , ¯
E 2 , ¯ F 1 , ¯ F 2 , ¯ G 1 , ¯ G 2 in spectral norm:
Lemma L.3.
There exists an absolute constant C > 0 such that for any δ ∈ (0 , 1 / 2) , each of the following inequality holds with probability at least 1 − δ :
[Equation]
Before we prove Lemma L
[Equation]
3, we ﬁrst use it to prove Theorem L
[Equation]
Proof of Theorem L.1.
By (54),
[Equation]
Setting δ in Lemma L.3 to be δ/ 6 and using the union bound, for an absolute constant C > 0 , with probability at least 1 − δ ,
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Equation]
where we used the fact that the right-hand-sides of (56), (58) and (60) are upper bounded by a constant times the right-hand-sides of (55), (57), (59), respectively.
Since ¯ A = (cid:80) ni =1 w i β i β T i , the theorem is proved by
[Equation]
due to Theorem 2.1.
In the lemmas below, we prove helper inequalities that we use to prove Lemma L.3.
In these lemmas, we focus on a single user i and thus omit the subscript i .
Lemma L.4.
Let x 1 , . . .
, x m ∈ R
d are independent random vectors.
Assume every x i has zero mean and is O (1) -sub-Gaussian.
Conditioned on x 1 , . . .
, x m , let z 1 , . . .
, z m ∈ R be η -sub-Gaussian independent random variables with zero mean.
Deﬁne b = (cid:80) mj
=
z j x j .
For any δ ∈ (0 , 1 / 2) , with probability at least 1 − δ ,
[Equation]
Moreover,
[Equation]
Proof.
Deﬁne X
= [ x 1 · · · x m ] ∈ R d
× m and z =
[ z 1 · · · z m ] T ∈ R m .
Now b = Xz .
By
Lemma C.3, with probability at least 1 − δ
[Equation]
Let r denote the rank of X .
By Lemma C.4, with probability at least 1 − δ/ 2 ,
[Equation]
Combining (64) and (65) using the union bound, with probability at least 1 − δ ,
[Equation]
This proves (61).
For every ﬁxed unit vector u ∈ R d ,
[Equation]
This proves (62).
[Equation]
Since z 1 , . . .
, z m are independent and have zero mean when conditioned on x 1 , . . .
, x m , the term E [ u T z j 1 x j 1 x T j 2 z j 2 z j 3 x j 3 x T j 4 z
j 4 u ] is zero unless ( j 1 , j 2 , j 3 , j 4 ) belongs to
S :
= { ( j 1 , j 2 , j 3 , j 4 ) ∈ { 1 , . . .
, m } 4 : ( j 1 =
j 2 ∧ j 3 =
j 4 ) ∨ ( j 1 =
j 3 ∧ j 2 = j 4 ) ∨ ( j 1 =
j 4 ∧ j 2 = j 3 ) } .
Letting e 1 , . .
, e d be an orthonormal bases of R d , we have
[Equation]
≤ (cid:88) ( j 1 ,j 2 ,j 3 ,j 4 ) ∈ S d (cid:88) (cid:96) =
η8 E [( z j 1 /η ) 8 + ( z j 2 /η ) 8 + ( z j 3 /η ) 8 + ( z j 4 /η )
[Equation]
u T
[Equation] [Equation]
2 e
[Equation] [Equation]
e
[Equation]
(cid:96)
[Equation] [Equation]
This proves (63).
Lemma L.5.
In the same setting as the lemma above, deﬁne h
= (cid:80) mj
=1 x j x T j β for a vector β ∈ R d .
For any δ ∈ (0 , 1 / 2) , with probability at least 1 − δ ,
[Equation]
Also, with probability at least 1 − δ ,
[Equation]
Moreover,
[Equation]
Proof.
Deﬁne X
= [ x 1 · · · x m ] =
R d × m .
We have h =
XX T β .
Since each of the m entries in X T β has sub-Gaussian constant O ( (cid:107) β (cid:107) 2 ) , by Lemma C.3, with probability at least 1 − δ/ 2 , (cid:107) X T β (cid:107) 2 ≤ O ( (cid:107) β (cid:107) 2 (cid:112) m
+ log(1 /δ )) .
By Lemma C.3, with probability at least 1 − δ/ 2 , (cid:107) X (cid:107) ≤
O ( (cid:112) d + m + log(1 /δ )) .
This proves (66) by a union bound.
We show that with probability at least 1 − δ ,
[Equation]
If d ≥ m , this follows from (66) and
[Equation]
If d ≤ m , with probability at least 1 − δ ,
[Equation]
Therefore, inequality (69) holds with probability at least 1 − δ .
Now we show that (cid:107) E
[ h ] E [ h ] T − E
[ hh T ] (cid:107) ≤ O ( m (cid:107) β (cid:107) 22 ) .
Indeed,
[Equation]
Since x j 1 and x j 2 are independent when j 1 (cid:54) =
j 2 ,
[Equation]
It is clear that E [ x j x T j β ] = β , and for every unit vector u ∈ R d ,
[Equation] [Equation]
This implies that (cid:107) E
[( x
[Equation]
j
[Equation] [Equation]
β ) T ] (cid:107) ≤ O ( (cid:107) β (cid:107) 22 ) .
Now it is clear from (71) that (cid:107) E
[ h ] E [ h ] T −
E [ hh T ] (cid:107)
[Equation]
22 ) , which implies
[Equation]
Plugging (69), (66)
[Equation]
and (70) into the inequality above, we get (67).
Finally,
[Equation]
where
[Equation]
When j 1 , j 2 , j 3 , j 4 are distinct, H j 1 j 2 j 3 j 4 is zero.
When j 1 , j 2 , j 3 , j 4 are not distinct
[Equation]
letting e 1 , . .
, e d be an orthonormal basis for R d
[Equation]
for every ﬁxed unit vector u ∈ R d
[Equation]
we have
[Equation]
This implies
[Equation]
j 3 x T j 3 ββ T x
[Equation]
4 x T j 4 ] (cid:107) =
O ( (cid:107) β (cid:107) 42 d ) .
Similarly, (cid:107) E [ x j 1 x T j 1 ββ T x j 2 x T j 2 ] (cid:107) =
O ( (cid:107) β (cid:107) 22 )
because u T E
[ x j 1 x T j 1 ββ T x
j 2 x T j 2 ] u = E
[Equation] [Equation]
≤
[Equation] [Equation]
(cid:107) 22 ) .
Therefore, (cid:107) H j 1 j 2 j 3 j 4 (cid:107) ≤
O ( (cid:107) β (cid:107) 4 2 d ) when j 1 , j 2 , j 3 , j 4 are not distinct.
Now (68) follows from (72).
Lemma L.6.
In the lemma above,
[Equation]
bh
[Equation] [Equation]
≤ O ( η 2 (cid:107) β (cid:107) 22 m 3 d ) , and (cid:107) E [ bh T hb T ] (cid:107) ≤ O ( η 2 (cid:107) β (cid:107) 22 m 3 d ) .
Proof.
For every unit vector u ,
[Equation]
If j 2 (cid:54) =
j 3 , we have E [ z j 2 z j 3 | x j 1 , x j 2 , x j 3 , x j 4 ] = 0 and thus E [ u T x j 1 x T j 1 βz j 2 x T j 2 x j 3 z j 3 βx j 4 x T j 4 u ] =
[Equation]
[Equation]
Plugging this into (74) and noting that the size of { ( j 1 , j 2 , j 3 , j 4 ) ∈ { 1 , . .
, m } 4 : j
2 =
j 3 } is
O ( m 3 ) , we get (cid:107) E [ hb
[Equation]
bh
[Equation]
] (cid:107) ≤ O (
[Equation]
(cid:107)
[Equation]
(cid:107) 22 m
[Equation]
d ) .
Similarly,
[Equation]
If j 1 (cid:54) =
j 4 , then E [ u T z
j 1 x
[Equation]
β
[Equation]
x
[Equation]
βz
[Equation]
x
[Equation]
u ] = 0 .
When j 1 = j 4 ,
[Equation]
β (cid:107) d (cid:88) (cid:96) =
E [ z 2 j 1 (( u T x j 1 ) 6 + (( β/ (cid:107) β (cid:107) 2 ) T x j 2 ) 6 + ( x T j 2 e (cid:96) )
+ ( e T (cid:96) x
j 3 ) 6 + ( x T j 3 ( β/ (cid:107) β (cid:107) 2 ))
6 + ( x T j 4 u ) 6 )]
= O ( dη 2 (cid:107) β (cid:107) 22 ) .
Plugging this into (75) proves (cid:107) E [ bh T hb T ] (cid:107) ≤
O ( η 2 (cid:107) β (cid:107) 22 m 3 d ) .
We can now ﬁnish the proof of Lemma L.3.
Proof of Lemma L.3.
To bound (cid:107) E 1 − ¯ E 1 (cid:107) , we set Z (cid:48) i =
w i m i ( m i − 1) (cid:16)(cid:80) m i j =1 h ij (cid:17) (cid:16)(cid:80) m i j =1 h ij (cid:17) T and
Z i
= Z (cid:48)
i − E
[ Z (cid:48)
i ] .
By (67), with probability at least 1 − δ/ (2 n ) , (cid:107) Z i (cid:107) ≤ O (cid:0) w i (cid:107) β
i (cid:107) 22 q i (cid:1) .
Also, by (68), (cid:107) E
[ Z i Z T i ] (cid:107) =
(cid:107) E
[ Z T i Z i ] (cid:107) =
O (cid:18) w 2 i (cid:107) β i (cid:107)
d m i (cid:19) .
Lemma C.5 proves that (55) holds with probability at least 1 − δ .
To bound (cid:107) E 2 − ¯ E 2 (cid:107) , we set Z (cid:48) ij =
w i
m i ( m i − 1)
h ij h T ij and Z ij =
Z (cid:48)
ij − E
[ Z (cid:48) ij ] .
By (67), with probability at least 1 − δ/ (2 nm i ) , (cid:107) Z ij (cid:107) ≤ O ( w i (cid:107) β i (cid:107) 22 s i ) .
Also, by (68), (cid:107) E
[ Z ij Z T ij ]
[Equation]
O
[Equation]
d m
[Equation]
Lemma C.5 proves that (56) holds with probability at least 1 − δ .
To bound (cid:107) F 1 − ¯ F 1 (cid:107) , we set Z i =
w i m i ( m i − 1) (cid:16)(cid:80) m i j =1 h ij (cid:17) (cid:16)(cid:80) m i j =1 b ij (cid:17) T .
Note that E [ Z i ] = 0 .
By (61) and (66), with probability at least 1 − δ/ (2 n ) , (cid:107) Z i (cid:107) ≤ O ( w i η i (cid:107) β i (cid:107) 2 q i ) .
Also, by (73), max (cid:8) (cid:107) E
[ Z i Z T i ] (cid:107) ,
(cid:107) E
[ Z T i Z i ] (cid:107) (cid:9) =
O (cid:18) w 2 i η 2 i (cid:107) β i (cid:107)
d m i (cid:19) .
Lemma C.5 proves that (57) holds with probability at least 1 − δ .
To bound (cid:107) F 2 − ¯ F 2 (cid:107) , we set Z ij = w i
m i ( m i − 1)
h ij b T ij .
Note that E [ Z ij ] = 0 .
By (61) and (66), with probability at least 1 − δ/ (2 nm i ) ,
[Equation]
Z ij
[Equation]
Also, by (73), max
[Equation] [Equation]
Lemma C.5 proves that (58) holds with probability at least 1 − δ .
To bound (cid:107) G 1 − ¯ G 1 (cid:107) , we set Z (cid:48) i =
[Equation]
and Z
[Equation]
Z
[Equation]
By
[Equation]
61) and (62), with probability at least 1 − δ/ (2 n ) ,
[Equation]
Also, by (63),
[Equation]
(cid:19) .
Lemma C.5 proves that (59) holds with probability at least 1 − δ .
To bound (cid:107) G 2 − ¯ G 2 (cid:107) , we set Z (cid:48) ij = w i m i
[Equation]
m i − 1) b ij b T ij and Z ij = Z (cid:48)
ij − E
[ Z (cid:48) ij ] .
By (61) and (62), with probability at least 1 − δ/ (2 nm i ) , (cid:107) Z ij
[Equation]
O [Equation] w i
[Equation]
2 i r i ) .
Also, by (63),
[Equation]
Lemma C.5 proves that (60) holds with probability at least 1 − δ .
L.2 Proof of Corollary L.
Proof.
We note that
[Equation]
Therefore, by Theorem L.1, with probability at least 1 − δ ,
[Equation]
Since sin θ ≤ 1 always holds, the inequality above implies sin
[Equation]
O
[Equation]
Using 1 n (cid:80) ni =1 (cid:107) β
i (cid:107) 22 = tr( 1 n (cid:80) ni
=1 β i β T i ) =
tr( (cid:80) ni
=1 w i β i β T i ) ≥ (cid:80) k(cid:96)
=1 σ 2 (cid:96) ≥ kσ 2 k , we have
[Equation]
mnσ
[Equation]
[Equation]
Plugging this into (76) proves (52).
When (cid:107) β
i (cid:107) 2 ≤ r for every i = 1 , . . . , n , (52) implies sin θ ≤ O
[Equation]
Since sin θ ≤ 1 , the inequality above implies
[Equation]
Using r 2 ≥ 1 n (cid:80) ni =1 (cid:107) β i (cid:107) 22 ≥ kσ 2 k ,
[Equation]
Plugging this into (77) proves (53).