The Implicit Bias of Heterogeneity towards Invariance: A Study of Multi-Environment Matrix Sensing
Authors: Yang Xu
Yihong Gu
Cong Fang
Venue: NeurIPS 2024 poster (2024)
Keywords: implicit bias
matrix sensing
invariance learning
non-convex optimization
Primary area: learning_theory
TL;DR: The implicit bias of heterogeneity helps model learn invariance.

Abstract: Models are expected to engage in invariance learning, which involves distinguishing the core relations that remain consistent across varying environments to ensure the predictions are safe, robust and fair. While existing works consider specific algorithms to realize invariance learning, we show that model has the potential to learn invariance through standard training procedures. In other words, this paper studies the implicit bias of Stochastic Gradient Descent (SGD) over heterogeneous data and shows that the implicit bias drives the model learning towards an invariant solution. We call the phenomenon the implicit invariance learning. Specifically, we theoretically investigate the multi-environment low-rank matrix sensing problem where in each environment, the signal comprises (i) a lower-rank invariant part shared across all environments; and (ii) a significantly varying environment-dependent spurious component. The key insight is, through simply employing the large step size large-batch SGD sequentially in each environment without any explicit regularization, the oscillation caused by heterogeneity can provably prevent model learning spurious signals.  The model reaches the invariant solution after certain iterations. In contrast, model learned using pooled SGD over all data would simultaneously learn both the invariant and spurious signals. Overall, we unveil another implicit bias that is a result of the symbiosis between the heterogeneity of data and modern algorithms, which is, to the best of our knowledge, first in the literature.

The Implicit Bias of Heterogeneity towards Invariance: A Study of Multi-Environment Matrix Sensing
Yang Xu"
School of Mathematical Sciences
Peking University xuyang1 1oyyang0014@kk edu.
Yihong Gu
Department of Operations Research and Financial Engineering Princeton University
yähongg@princeton.. edu
Cong Fang"
School of Intelligence Science and Technology Peking University
fangcong@pku. edu.
Abstract
Models expected engage in invariance learning, which involves distinguish- ing the core relations that remain consistent across varying environments to ensure the predictions safe, robust and fair. While existing works consider specific algorithms to realize invariance learning, we show that model has the potential to learn invariance through standard training procedures. In other words, this paper studies the implicit bias of Stochastic Gradient Descent (SGD) over heterogencous data and shows that the implicit bias drives the model learning towards invariant solution. We call the phenomenon the implicit invariance learning. Specifically, we theoretically investigate the multi-envine environment low-rank matrix sensing prob- lem where in each environment, the signal comprises (i) lower-rank invariant part shared across all environments; and (ii) significantly varying environment- dependent spurious component. 'The key insight is, through simply employing the large step size large-batch SGDD sequentially in each environment without any explicit regularization, the oscillation caused by heterogencity can provably prevent model learning spurious signals. The model reaches the invariant solution after certain iterations. In contrast, model learned using pooled SGD over all data would simultaneously learn both the invariant and spurious signals. Overall, we unveil another implicit bias that is result of the symbiosis between the heterogeneity of data and modern algorithms, which is, to the best of our knowledge, first in the literature.
Introduction
In real applications, the machine learning models are often heavily over over-paranin parameterized, which means that the number of parameters exceeds the number of data. For peramaarmeteri models,
"Equal Contribution.
t "orresponding author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Algorithm PooledSGD
Algorithm 2 HeteroSCD
Parameter: 0
Parameter:
fort= 1.... do
for 1,... do
Batch from entire dataset Update 0 t (0; B)
Batch B from environment e ~
Update t- 0 ((0; B)
end for
end for
Core Signal
Core Signal
Spurious Signal
Spurious Signal
Value
Value
Iters
Iters
Figure 1: An illustration comparing training from aggregated data versus from heterogeneous data. The left example resembles the case where the model is trained on complete datasets, resulting in a stable spurious signal that the model tends fit. The right example simulates two-environment case where the spurious signal changes at each step. This oscillation creates a contraction effect, preventing the model from fitting the spurious signal.
the generalization the general case becomes ill-posed. One key insight to generalize well the implicit preference of the optimization algorithm which plays the role of regularization/bias [45] [22]. Nowadays, there are several kinds of implicit bias discovered from optimization algorithms under different models and settings. One common feature of the bias is the simplicity which concludes that (stochastic) gradient-based algorithms perform the incremental learning with the model complexity gradually increasing. Therefore, benign generalization is possible even when the number of training data is limited. For example, Li et al. (31) Gunasekar et al. [8] show that unregularized gradient descent models. Kalimeris et al. (27). Gissin et al. [16], Jiang et al. [24]. and Jin et al. [25] further show that (Stochastic) Gradient Descent ((S)DD learn models from simple ones complex ones. Most of the existing works study the implicit bias of algorithms over single distributional environment data.
find the low-rank solution efficiently for matrix sensing
However, data in modern practice often collected from multiple sources, thus exhibiting certain hetcrogencity. For example, medical data may come from multiple hospitals, and training sets for large language models consist of numerous corpus from the Internet []]. So what is the impact of implicit bias for standard training algorithmss over heterogeneous data?
This paper initializes the study and shows that implicit bias of SGD on an over- over-parane parameterized model heterogeneous data and shows that the implicit bias can not only save the number of training data but also, more importantly, drive the model learning the invariant relation across diverse environments.
using
Learning the invariant relation that remains consistent across varying environments 43 has garnered significant attention in recent years. Though the association- standard machine learning pipelines can achieve a good performance with identical data distributions, higher requirement is to make predictions robustly generalize over diverse downstream environments. Learning invariance produces reliable, fair, robust predictions against strong structural mechanism perturbation. More importantly. opens the door pursue causality blind any prior knowledge and can unveil direct causes when the heterogeneity among environments is sufficient [17] 43 While existing works consider specific algorithms to realize invariance learning, this work shows that implicit bias of algorithms over heterogeneous data has the potential to automatically learn the invaraince. We call the phenomenon the implicit invariance learning, partially explains why active invariance learning may not be necessary in practice [42. Our key insight is:

---

The heterogeneiry of the data, and the large step size adopted in the optimization algorithm jointly provide strong multiplicative oscillations in the spurious signal space, which prevents the model from moving in the direction of unstable and spurious solutions, thus resulting an implicit bias to the invariant solution.
We illustrate it rigorously through a simple, canonical but insightful model multi-environment matrix sensing, where in each environment the signal consists of two parts: an invariant low-rank matrix A' E RRdd and an environment-varyyng spurious low-rank matrix (e) E Rdxd where environment e E E, the set of environments. For each environment E &, the joint distribution of (X(+), y satisfies y/*) (x(s), A*) + (x(e), A (e)) with matrix inner product (A, B) = Trace A). Here the response. We consider the case that association does not coincide invariance (or causality), where averaging over all the environments, the best prediction y given X
( E Rdxd a random linear measurement and y
f (X): = (X)
X,EA(A
with
Eg A (+) 00
invoriont part
spurions part
In this case, it is not surprising that given enough data, the standard empirical risk minimizer algorithm, for example, running SGD on pooled data, will return a solution that converges to f*. which diverges from the invariant solution. In this paper, we will show that surprisingly. if each batch sampled from data in one environment rather than data in all the environments, the heterogeneity in the environments together with the implicit regularization effects in the SGD algorithm can drive it towards the invariant solution. This can be stated informally as follows.
Theorem (Main result, informal). Under sufficient heterogeneity condition and some regularity conditions in matrix sensing, if we adopt an wver-parameterizer model and runs stochastic gradient descent where batches are sampled from environment, i.e., HeteroSGD (Algorithm) 3) then
|1 A A+ or(1).
Instead, the standardd approach, ie., PooledSGD, will return solution PooolddecD satisfying
|0poo1 A A (1)
thus
||0pooleascp
A F = 9p((1.
An illustration of our result is shown in Figure I Our result demonstrates that implicit bias of commonly used algorithms over heterogeneous data has the potential to drive the model to learn the invariant relation. Such result thereby provides an explanation for why models may attain some robust and even causal prediction after SGD training.
We emphasize that the previous implicit bias studies are restricted to the same data distribution generalization, under which the population-level minimizer f* minimizing the loss with infinite data the target in pursuit. However, both the population-level minimizer and those "good" solutions under previous studies diverge from an invariant solution general and are no longer benign this context, this is termed as "curse of endogeneity" 22 13]
Notations. We use the conventional notations 0()), to ignore the absolute constants, Ö(-), 6(-), 22-) to further ignore the polynomial logarithmic factors. Similarly, a S means that there exists absolute constant C. > 0 such that S Cb. We also denote it a = o(b). Unless otherwise specified, we use lowercase bold letters such as v to represent vectors, and use ||v|| to denote its Euclidean norm. We use uppercase bold letters such as X to represent matrices and use X|| || | X|| to denote its operator norm, Frobenius norm and nuclear norm, respectively. We x xX) to denote the condition number, which is F (X)/ (X). We define
< b, b >> if
Z = (1) if the random variable Z satisfies Z 4 0.
2 Related Works
Implicit Regularization. believed that implicit bias key factor in why over-parameterized models can generalize well. Through the analysis of certain settings, existing results suggest that
GD/SGD prefers solutions with specific properties [45] 19 For the matrix sensing problem, several works
3X 23 or specific local landscapes
3 9 32 38).
18
46 52, 24 25 analyze
3

---

the (S)GD dynamics to show how (S)GD recovers the ground truth low-rank matrix. Recently, the effects of large step size have aroused much attention, particularly the -lge-lf-stabilityy phenomenon []. Lu et al. 37 investigates the phenomenon "benign oscillationn, which suggests that SGD with a large learning rate can effectively help neural networks learn weak features thereby benefiting generalization. Several works (20] 48 show that label noise with large step size has sparcifying effect for sparse linear regression. This paper instead studies eetiromenii scenarios and fills in the understanding of the impact of randomness matrix sensing problems.
Federated Learning. Federated learning 39 26 is a machine learning paradigm where data is
stored separately and locally on multiple clients and not exchanged, and clients collaboratively train model. Extensive work has focused on designing effective decentralized algorithms (e.g. [39 (29])
while preserving privacy (e.g. [TO, (7).. The importance of fairmess in federated learning has also One important issue federated learning to handle the heterogeneity across the data and hardware. Our work shows that by training with certain stochastic gradient descent methods, the system can automatically remove the bias from the individual environment and thus learn the invariant features. Our work provides insights into discovering the implicit regularization effects of standard decentralized algorithms.
garnered attention 30 33
since invariant covariates correspond to direct cause. From theoretic aspects, Fan et 13 proposes the EILLS method that provably achieves invariant variable selections under mild conditions for linear models. Invariance learning has raised much attention in machine learning since Arjovsky et al. proposes the structural-agnoo
Invariance Learning. This research line initiates from causal inference literature
43
40
15
framework IRM. Subsequent works analyze its limitations [44] 28 or
as regularization and reweighting. About the failure of classical methods, Wald et al. 49 construct hard problem and show that interpolation-bast -asse methods fail to learn invariance.
propose variant methods [50 36
34
35
21
51
To the best of our knowledge, all the existing works consider specific algorithms to realize invariance learning or constructing hard cases that classical methods fail. In contrast, this paper studies commonly used training algorithms and aims to understand how the algorithms can go beyond learning associations achieve invariance learning certain scenarios.
3
Main Results
3.1 Problem Formulation
Data Generating Process. Suppose we observe data from a set of environments E sequentially. Let
D be some distribution on & At each time t 0, 1, l...., we receive m samples ((x (ee) (oe)
.y ) C
Ddxd IR from environment e D satisfying
y = (X! (e) A+) + (e)) A (00), = m,
(1)
where A is an unknown rank 7i d x d symmetric and positive definite matrix that represents the true signal invariant across different environments, (+)) is an unknown d x d symmetric matrix with rank most 12 that represents the spurious signal that may vary. Here (A, B) = tracc(B (B A). We aim to estimate the A using data from heterogeneous environments.
Algorithm. We consider running batch gradient descent where at each step t one gradient update is performed using the data from environment es. To be specific, we parameterize our fitted model as = (A, UUT) with d d matrix U for the sake of simplicity. One can generally use the parameterization X UUT VVT by the same technique of HaoChen et al. [20], Fan et al. [14]. We initialize U as U, = al, for some small enough > O. At timestep we run one-step gradient descent on the standard least squares loss using {(x) (ex) (e+) )
over-parametriization of the model,
y
L(U)-
1
(e))
(x(1) UUT))²
2m
That is, Un = al. and
U = U, U,-,V((U,,) = [4-
x U,U.T) 1 Σ
-/ -
U.
(3)
i=1

---

Algorithm 3 HeteroSGD
Set Uo = ala, where a is a small positive constant to be determined later. Set large step size 11 (1). for / /orl=10 1,.., T do
Receive samples {(X; (et) (ex)
M// Jm
from current environment ez.
Gradient Descent U1+1 U, -
m
((x(*) U,U,) y/ (e+) (())) U..
end for Output: Ur-
for O,...,T T 1. See a complete presentation in Algorithm 3
The algorithm adopts constant level step size yj and log(a-1) level number of iterations T, i.e. (1) and T = "(log(a-1)) and use UT as our estimate of A*
Standard Method: Pooled Stochastic Gradient Descent. As comparison, consider the standard approach where data in each batch come from different environments and the weights follow from D. To be specific, the pooled stochastic gradient descent over all environments adopted the update rule
U U-U- VZ((), where L(U)=
Σ
(v)
(X
)
e, D. (4)
2m7
3.2 Assumptions
We first impose some standard assumptions used in matrix sensing. Since we are dealing with learning true invariant signals from heterogeneous environments, several conditions on the structure of the invariant signal A' and the spurious signals (e) should be imposed.
Assumption 1 (Invariant and Spurious Space). There exists U 12dxr and V IRdxra both with orthogonal columns, i.e., (U*) U = Ir, and (V*) V = Ir, such that
(a). Colog Sn A 22 and d 2 (r1+ + r2) for some large absolute constant C.
(b). A U*(U*)T.
(c). A (c) V-E (V*) with some symmetric 12 X matrix E(e) for any E.
(d). |I0*) ei for some small quantity e 2
In Condition (b), we assume that the singular values of the true signal A are the same simplify the presentation since our main focus is to reduce the spurious signals. It holds for the basic case when there is only one invariant signal, i.e. r1
The analysis for varying singular values using the technique of Li al. 31 is deferred to Section D in Appendix. Other assumptions
usual and easy to achieve. Condition (a) requires that the total dimension of invariant signals and spurious signals are small relative to the ambient dimension d. Condition (c) resembles the RIP condition [6] in sparse feature selection [3] Condition (d) says the overlap of invariant subspace and spurious subspace should be small. Such condition can be easily satisfied for random projections in high dimensions where P1 r2 < d, under which have e = =(v((r r2)/d), Proposition 1 below.
Proposition 1. Ler M1 E Rdxes and M2 Rdxraa be two mutually independent random matrix with i..dd. (O, entries. Denote their QR decompositions as M, U, and M2 U²R2 respectively. Then there exists universal constant C₁ > O such that
U S /// d [())]
+ 72
(5)
with probability at least 1 4exp (dC - 2exp (-C (ri +2)
5

---

Assumption 2 (Regularity on Spurious Signal E()),, There exists some constant-level quantity M1, M2 such that
Var D
sup xE(r)
(e)
< Mi
and
min (E(ra) 1 1 BenD (c)) > M2.
(6)
where M1 < M2 for some universal constant Co 0. Moreover, E(e) is strongly diagonal dominant for any e E &, i.e.
sup max 2 3E rE(ral 37
Co
M2 L.5
where > 0 is some universal constant.
The first inequality in 66) requires that all the spurious signals have a uniform bound, under which a fixed step size can be adopted. The second inequality in (6) requires that the heterogeneity of the spurious signals large compared to the bias of the spurious signals. For example, some variables receive different interventions in different environments. The condition (7) is imposed to prevent the explosion of spurious signals during training. When the diagonal and off off-diagonal elements are the same order, empirical studies and theoretical analyses in some toy examples illustrate the failure of recovering A Condition (d) in Assumption [ and (6) resemble the RIP condition in sparse feature selection. Example D can fulfill all our conditions.
Finally, we impose assumptions on measurements. Recall the RIP condition 6
Definition (RIP for Matrices (6]). A set of linear measurements X1....,
X.m satisfy the restricted
isometry property (RIP) with parameter (s, d) if the following inequality
(1 (1-M)
Σ Xx,
M) < (1 (1+ M
(8)
holds for any d x d matrix M with rank at most 8.
Assumption 3 (RIP Condition for Linear Measurements), x ((a) parameter 4(ri + r2) and 8 S [Mglog(d) V/n for all e E E.
satisfies the RIP with
It is known from Candès and Plan 6 that for symmetric Gaussian measurements, sample complexity [(ds6 M2) d poly(r, log(d)) < d² suffices.
3.3 Convergence Analysis
The main conceptual challenge in the problem is that any with UUT A no longer local minimum since Ee-D (x]) is non-zero and could even be comparable A*. This further implies that running stochastic gradient descent on pooled data will fail to recover A+. However, our main result below shows that simply adopting online gradient descent with "heterogeneous batches" can successfully recover the true, invariant signal from heterogeneous environments.
Theorem (Main Theorem). Under Assumption 3 suppose further that e < < Define M2 log(d)) 1.5 8r2 /ri + r2 for some absolute constant Cu. If we choose 1) E (24M 64 M, and a E (1/d*, 1/d"), then running Algorithm 3 I = O(log(a-1)) steps, the algorithm outputs UT that satisfies
I
UU AA'IF S C' C'ax(8*22 V M2 log'
(9)
for some absolute constant C, with probability over 0.99.
Consider the case where r2, M₁ are sufficiently large but is regarded at constant level, and the batch size m, ambient dimension satisfy dl log (d) < m. It follows from the RIP result [ with 5 = @ d/m) and Theorem 2 that one can adopt = @(d-1), = O(1, and early stop T = O(log d) such that
P |0-7 - A S Ci (og(d) 21 1- (dloo[(()/m))
(10)

---

provided 61 I CT /d/m with some large enough constant C1 > O. In this case, it follows from (TO) that one can distinguish the true invariant signals from those spurious heterogeneous ones since
max { UU - 12 .l.(=l((v) U
or()).
(11)
The underlying reason why the online gradient descent can recover A is that the heterogeneity of A (e) and the randomness in the SGD algorithm jointly prevent it from moving in the direction of spurious signals. At the same time, the standard RIP conditions and the almost orthogonality between ensure steady movement towards the invariant signals.
U and V in Condition T
Conversely, running pooled stochastic gradient descent using all data will result in a biased solution:
Theorem 3 (Negative Result for Pooled SGD). Under the assumptions of Theorem and some mild conditions. for the certain case where U T V and EEED E(e) E = Irg. if we perform SGD over all samples with batch size = ②2d poly(ra + r2, M2, log(d))) and ends with T (log d), then
Ut keeps approaching U T + VVV* in the sense that
U U V.V
|| S s0(),
(12)
during which for all t = 0,1....,,
|..U A' l &
(13)
The convergence [2) similar 99 in derivation. To see this, since each update uses batch from the whole data, the update in effect degenerates to the case for one environment with no heterogeneity.
Now the one-environment invariant solution A in (9) is exactly equal to U* + V*V* in ([2) One can also show that for sufficiently large f. U, is sufficiently away from A indicating that the biased estimation is not attributed to early stopping.
Our framework can be applied to learning the invariant features for a two-layer neural network with quadratic activation functions, by recognizing the fact that [11]
1 LT (Ux) (xx") UU
(14)
where q(-) is the element-wise quadratic function. The following example shows that Theorem implies success of invariant feature learning for 2-layer NN when the ground truth invariant and variant features are independent random vectors sampled from normal distribution.
Example (Two-Layer NN with Quadratic Activation). Let aj, ,aa, E IRd be random vectors sampled from normal distribution N NO, L For environment E E, suppose the target function is
determined by Tı invariant features and r2 variant admits that for each sample (x.² ())")))
Σ
[(axx(()))
(0) Σ aj
q(aTx()))
(c)
Σ
aj (c) [aa] ) (15)
yi
a,a
j=
j=
which is equivalent to matrix sensing problem with
A
-a 3=1
a
A (6) = Σ j=rit1
aja and x((e) = =x (e))
(16)
And our goal is to train a two-layer NN to capture the invariant features (ai...., ... ar,). In this example, the invariant component and the spurious component have more intuitive characterization: they are two disjoint groups of neurons. Moreover, it can be shown that the invariant and variant features are
nearly orthogonal (Proposition 11. Then if {a)") Hie satisfies sup (la. [[) [] max, (1+|B_0 (3)
co for
min, [a (
some absolute constant Co, the variant version of Algorithm 3 returns a solution that only significantly
selects invariant features with probability over 0.99. See Section C and Theorem for details.
4
Proof Sketch
We define the invariant part R, Bdxr, spurious part Q; E Rdxra in U, as
R, = U, U+
and Q = U. V
(17)
7

---

and let the residual be the enror part, that is,
E, :=: U, (U'R.T V'Q = (I- ''
(18)
It is worth noticing that Idy. U and Idy. V.V+ are both orthogonal projections, and
Idyes := I - Idu. - Idv. is not.
It follows from the model ① and the gradient update that
m
U.U. - - -A- (e))) X U..
U2+1 = U. - "I
(19)
We use operator E,, (M) E Rdxd to denote the RIP error of the batch at time step t for some d x d matrix M, i.c.,
1 m
Es, (M)
m Σ
M)X M.
(20)
We also write E, o (M) = E. (M) when there no ambiguity, and we simply denote matrix E, E, (U.UT - V+T,V+T) with E, E(e)), Then the gradient update of U,
can be written as
Ul41 U, 7 (UuU UU
U. - -
VE.V
E.U, RIP Eror
(21)
Combining our definition 17 with 21
we obtain
R = (U, (U.U, UU V'EVV' U, nE.U,) T
U²
(22)
= -I- U. IJR V' E.
Dominatin Dynamics
Interaction Error
RIP Error
Q41 = U, -U.U UUU V'ERV' - J. nE,U,) V
= Q: U. + QE to V. - U.EVV.
(23)
Fluctuation Dynamics
Interaction Error
RIP Error
For the error part, combining 18 with 21
yields
E = Idres (U, - (U--,(UU,) - T V'E J. - -NE,U,
E, u,U U,) +n Idres (UU + V- U. Edre,E,UUU
(24)
Shrinkage Dynamics
Interaction Frror
Eror
For the invariant part R2, though different singular values of R, will grow at different speeds because of the randomness from RIP error and et, claim that all the singular values of R, close to R, during the training process, where the scalar sequence R, is defined recursively as
Re+1 (1 nR? + R) Ro
(25)
The dynamics of Q are very complicated because of the randomness of e, and the RIP error. Such dynamic will also impact that of R, and E, through the complicated dependencies between these three parts, which will also make it difficult to utilize probability inequalities applicable under independence. Instead, we claim that such "fluctuation dynamics" of Q2 can be controlled as
. 0 log(ri + r2))
||QuI| poly(log(d), r,
Mi)L, r,,)),,,,,LLL... with L, fooo +r,,,) t 0 log(r1 r2))
(26)

---

We now offer an informal illustration for how oscilla- tions "shrink" spurious signal. We simply omit error
Satte Signal Cocl Cocilaorg Sgra
Sigral asci Bignal
terms. When matrix ( is diagonal, from (23). the i-th column q of Q₁ should satisfies: || _ (1 | )
n
Let 5 def >((+1) and assume M [<] a.s..
Introduce a concave function [20] o(x) = x (0, 1). When r7M < T8 do second-order Taylor's expansion at
50 Iters
100
50 Iters
75
100
((1) in (a) below:
Figure 2: The left figure shows E Blqllll and the right figure shows |/ |1
E [((l
I (0( + nE) <(llqel)
+ m - Var(E)
where E[. is w.r.t. 5. So spurious signal keeps small when 16.
(1-7) Varlg] See the figure for illustration in Figure 2 While [ (shown in the left figure) increases since the signals have positive expectations, E
4
ET q: [[.1]] (shown in the right figure) decreases. Note that the above intuition is informal and the formal argument is deferred in Lemma 6 and Lemma in Appendix.
The entire training process
be divided into two phases. In Phase the invariant signals R. increase rapidly while the spurious signals Q, fluctuate but remain at low level. Phase ends in O (og( ()) steps when R2 attains (l)dorr (see Theorem 4). In Phase the magnitudes of Q and E, stay low, while all the singular values of R₃ approach 1 (See Theorem 5p. We defer the details to Appendix.
5 Simulations
In this section, we present our simulations. We design three sets of experiments. In the first set of experiments, we show with the growth of enviromment heterogeneity, invariance learning is achievable. For the second set of experiments, we show that given heterogeneous data, invariance learning is achievable with the growth of step size For the third set of experiments, we compare HeteroSCD (Algorithm 3 and Pooled SGD. Section B.2 we also perform simulations for Pooled SGD with small batch size.
In below two sets of experiments, set the scale of initialization a 10-3 problem dimension d = 100, r1 = 1 and 122 = 1. Let the true signal be A' uu Denote the heterogeneity parameter by M. The environment is generated by A (e) = A + vv where s(e) ~ Unif (1 - M, + M). and the default of 77 is 0.05. The numher of linear measurements set be m 8000 with elements
following from i.i.d (0, 1). For the third sets, we set (r1, r2, d. (E(te) = (3, 2, 40, 0.5), m 2800 5600 without replacement for Pooled SGD, The plots show signal recovery proportion, 1.0 indicates fully recovery.
for HeteroSGD and
1.0
True Signal M=6.0
1.0
0.8
Spurious M=6.0
Jqnal Value
0.8
Value
0.6
Error M=6.0
0.6
True Signal 1=1.0
Spurious 1=1.0
0.4
0.2
Error M=1.0
S
0.2
True Signal
Spurious Signal
0.0
50
100
150
200
4
6
8
10
Iters
M
Figure 3: The left figure shows that the heterogeneity facilitates us to eliminate the spurious signal and learn the invariance. The right figure shows that both true and spurious signals flow up when M small, the "phase transition" happens around M
A smaller step size can reduce the noise arising from heterogeneity, making the dynamics more similar to those of Gradient Descent.

---

1.0
True Signal 0.05
1.4
True Sighal Spurious Signal
0.8
Spurious n 0.05
Valu 18
1.2
1.0
Value
0.6
Error 0.05
0.8
True Signal 0.01
0.4
0.6
Spurious n = 0.01
0.4
0.2
Error n 0.01
0.2
0.0
0.0
50
100
150
0.01
0.02 0.03 0.04 0.05 0.08 0.07 0.08 0.09
(M = 5.0)
Figure 4: The left figure shows that the large step size helps eliminate the spurious signal. The right figure shows that both true and spurious signals flow up when small, and when 7) 2 0.05, the spurious signal is eliminated.
1.0
1.0
True Signal
True Signal
0.8
Spurious
0.8
Spurious
Value
0.6
Error
value
0.6
Error
0.4
0.4
0.2
0.2
0.0
0.0
25
50
100
125
150
175
100 Iters
125
150
175
Iters
Figure 5: The left figure shows that heterogeneity helps eliminate the spurious signal. The right figure shows that Pooled SGD fits invariant signal and spurious signal simultaneously without distinction.
6 Conclusions
This paper explains that implicit bias of heterogeneity leads the model learning towards invariance and causality. We show that under heterogeneous environments, online gradient descent with large step sizes can select out the invariant matrix in the over- perepparaettriieed matrix sensing models. We conjecture that both heterogeneity and stochasticity are indispensable. Over-parameterizzaion may not be. We leave future studies to understand the necessity of the three factors.
7 Acknowledgement
C. Fang was supported by National Key R&D Program of China (0222D0114900)) the NSF China .No.62376008).
10

---

References
[1] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
[2] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization. arXiv preprint arXiv: /907.02997. 2019.
[3] Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process. Conference on Learning Theory, pages 483-513. PMLR, 2020.
[4] Emmanuel J Candes. The restricted isometry property and its implications for compressed sensing. Comptes rendus. Mathematique, 346(9-10):589 -9999959922200 2008.
[5] Emmanuel Candes and Terence Tao. Decoding by linear programming, IEEE transactions on information theory, 51(12):4203-422 2005.
[6] Emmanuel J. Candès and Yaniv Plan. Tight oracle inequalities for low-rank matrix recovery from minimal number of noisy random measurements. IEEE Transactions on Information Theory, 57((4)2342
2011. doi: 10. 10.1109/T112201121111177
[7] Wei-Ting Chang and Ravi Tandon. On the upload versus download cost for secure and private matrix multiplication. In 2019 IEEE Information Theory Workshop (ITW), pages 1-5, IEEE, 2019.
[8] Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 2020.
[9] Alex Damian, Tengyu Ma, and Jason D Lee. Label noise SGD provably prefers flat global minimizers. Advances in Neural Information Processing Systems, 34, 2021.
[10] Cynthia Dwork, Krishnaram Kenthapadi, Frank McSherry, Ilya Mironov, and Moni Naor. Our data, ourselves: Privacy via distributed noise generation. In Advances in Cryptology- EUROCRYPT 2006: 24th Annual International Conference on the Theory and Applications of Cryptogruphic Techniques, St. Petersburg, Russia, May 28-June 1, 2006. Proceedings 25, pages 486-503. Springer, 2006.
[11] Mathieu Even, Scott Pesme, Suriya Gunasekar, and Nicolas Flammarion. (s) gd over diagonal linear networks:
Implicit bias, large stepsizes and edge of stability. Advances in Neural Information Processing Systems, 36, 2024.
[12] Jianqing Fan and Yuan Liao. Endogeneity in high dimensions. Annals of Statistics, 42(3):872, 2014.
[13] Jianqing Fan, Cong Fang, Yihong Gu, and Tong Zhang. Environment invariant linear least squares. arXiv preprint arXiv:2303.03009. 2023.
[14] Jianqing Fan, Zhuoran Yang, and Mengxin Yu. Understanding implicit regularization in over- parameterized single index model. Journal of the American Statistical Association, 118(544): 2315 2328, 2023.
[15] AmirEmad Ghassami, Saber Salehkaleybar, Negar Kiyavash, and Kun Zhang. Learning causal structures using regression invariance. Advances in Neural Information Processing Systems, 30, 2017.
[16] Daniel Gissin, Shai Shalev-Shwartz, and Amit Daniely. The implicit bias of depth: How incre- mental learning drives generalization. In International Conference on Learning Representations, 2019.
[17] Yihong Gu, Cong Fang, Peter Bühlmann, and Jianqing Fan. Causality pursuit from heteroge- neous environments via neural adversarial invariance learning. arXiv preprint N.Xx:2445.0041555 2024.
11

---

[18] Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. Advances in Neural Information Processing Systems, 30, 2017.
[19] Suriya Gunasekar, Jason D Lee, Daniel Soudry, and Nati Srebro. Implicit bias of gradient descent on linear convolutional networks. Advances in Neural Information Processing Systems,
3 31, 2018.
[20] Jeff Z. HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. In Conference Learning Theory, pages 2315-2357. PMLR, 2021.
[21] Badr Youbi Idrissi, Martin Arjovsky, Mohammad Pezeshki, and David Lopez-Paz. Simple data balancing achieves competitive ceatt-rroupuuccucc In Conference Causal Learning and Reasoning, pages 336-351. PMLR, 2022.
[22] Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. In International Conference on Learning Representutions, 2019.
[23] Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. Ad- vances in Neural Information Processing Systems, 33, 2020.
[24] Liwei Jiang, Yudong Chen, and Lijun Ding. Algorithmic regularization in model-free over- parametrized asymmetric matrix factorization. SIAM Journal on Mathematics of Data Science, 5(3):723-744, 2023.
[25] Jikai Jin, Zhiyuan Li, Kaifeng Lyu, Simon Shaolei Du, and Jason D Lee. Understanding incre- mental learning of gradient descent: A fine-grained analysis matrix sensing. International Conference Machine Learning, pages 15200-15238. PMLR, 2023.
[26] Peter Kairouz, H Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Ar- jun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, et al. Advances and open problems in federated learning. Foundations and trends@) in machine learning, 14(1-2):1-21) 2021.
[27] Dimitris Kalimeris, Gal Kaplun, Preetum Nakkiran, Benjamin Edelman, Tristan Yang, Boaz Barak, and Haofeng Zhang. Sgd on neural networks learns functions of increasing complexity. Advances in Neural Information Processing Systems, 32, 2019.
(28) Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro. Does invariant risk minimization capture invariance? In International Conference on Artificial Intelligncce and Statistics, pages 4069-4077. PMLR, 2021.
[29] Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference Machine Learning, pages 5132 5143. PMLR, 2020.
[30] Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization. In International Conference on Machine Learning, pages 6357-6368. PMLR, 2021.
[31] Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over parameterized matrix sensing and neural networks with quadratic activations. In Conference on Learning Theory, pages 2-47. PMLR, 2018.
[32] Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after sgd reaches zero loss?-a mathematical framework. In International Conference on Learning Representations, 2021.
[33] Shiyun Lin, Yuze Han, Xiang Li, and Zhihua Zhang. Personalized federated learning towards communication efficiency, robustness and fairness. Advances in Neural Information Processing Systems, 555300711-3048 2022.
[34] Yong Lin, Hanze Dong, Hao Wang, and Tong Zhang. Bayesian invariant risk minimization. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 16021-16030, 2022.
12

---

[35] Yong Lin, Shengyu Zhu, Lu Tan, and Peng Cui. Zin: When and how to learn invariance without environment partition? Advances Neural Information Processing Systems, 35, 2022.
[36] Chaochao Lu, Yuhuai Jose invariant risk minimization: causal approach. arXiv preprint arXiv:2 12353, 2021.
Jose Miguel Hernández-Lobato, and Bernhard Schölkopf. Nonlinear
[37] Miao Lu, Beining Wu, Xiaodong Yang, and Difan Zou. Benign oscillation of stochastic gradient descent with large learning rate. In International Conference on Learning Representations, 2023.
[38] Kaifeng Lyu, Zhiyuan Li, and Sanjeev Arora. Understanding the generalization benefit of normalization layers: Sharpness reduction. Advances in Neural Information Processing Systems, 35, 2022.
[39] Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication--fficee learning of deep networks from decentralized data. In Artificial intelligence and statistics, pages 1273-1282. PMLR, 2017.
[40] Nicolai Meinshausen, Alain Hauser, Joris M Mooij, Jonas Peters, Philip Versteeg, and Peter Bühlmann. Methods for causal inference from gene perturbation experiments and validation. Proceedings of the National Academy of Sciences, 1 113(27)7761--366, 2016.
[41] Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent on separable data. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 3420-3428. PMLR, 2019.
[42] Vivian Y Nastl and Moritz Hardt. Predictors from causal features do not generalize better to new domains. arXiv preprint arXiv:2402.0989/ 2024.
[43] Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant prediction: identification and confidence intervals. Journal of the Royal Statistical Society Series B: Statistical Methodology, 78(5):947-1012, 2016.
[44] Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. The risks of invariant risk minimiza- tion. In International Conference on Learning Representations, volume 2021.
[45] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19 (1):2822-2878, 2018
[46] Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruction. Advances in Neural Information Processing Systems, 34, 2021.
[47] Roman Vershynin. High-dimennsional probability: An introduction with applications in data science, volume 47, Cambridge university press, 2018
[48] Loucas Pillaud Vivien, Julien Reygner, and Nicolas Flammarion. Label noise (stochastic) gradient descent implicitly solves the lasso for quadratic parametrisation. In Conference on Learning Theory, pages 2127-2159. PMLR, 2022.
[49] Yoav Wald, Gal Yona, Uri Shalit, and Yair Carmon. Malign overfitting: Interpolation and invariance are fundamentally at odds. In International Conference on Learning Representations, 2023.
[50] Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. In International Conference Machine Learning, pages 11214-11224, PMLR, 2020.
[51] Cheng Zhang, Stefan Bauer, Paul Bennett, Jiangfeng Gao, Wenbo Gong, Agrin Hilmkil, Joel Jennings, Chao Ma, Tom Minka, Nick Pawlowski, et al. Understanding causality with large language models: Feasibility and opportunities. arXiv preprint arXiv:2304.05524, 2023.
[52] Jiacheng Zhuo, Jeongyeol Kwon, Nhat Ho, and Constantine Caramanis. On the compu- tational and statistical complexity of over-parameterized matrix sensing. arXiv preprint 202.22200..0275. 2021.
13

---

Contents
Introduction
2
Related Works
3
Main Results
3.1
Problem Formulation
3.2
Assumptions
3.3
Convergence Analysis
Proof Sketch
5
Simulations
6
Conclusions
10
Acknowledgement
10
Deferred Proofs in Theorem 2
14
Restricted Isometry Properties
14
A.2
Additional Auxiliary Sequences
15
A.3
Useful Lemmas
18
Bounds of Q
20
A.5
Bounds of E,
20
A.6 Analysis for Phase
23
Analysis for Phase
24
B
Deferred Proofs
26
B. Proof of Proposition]]
26
B.2 The Failure of Pooled Stochastic Gradient Descent
28
Neural Networks with Quadratic Activations
30
D The x(A))> >
Case
32
A Deferred Proofs in Theorem
This section is organized follows: In Section A.1
state some useful properties from the
definition of RIP. In Section A.2 and A.3 we formally define the auxiliary sequences we use to
control the dynamics and develop several useful lemmas frequently use. In Section S.ction|A.4| and and A.5
we bound Q and E, respectively. In Section A.6 and A.7
we prove Theorem 4 and Theorem
A.1 Restricted Isometry Properties
In this section, we list some useful implications of the definition of RIP property. Below we assume the set of linear measurements A (ex)
..... A Ale+) e Ixdxd satisfy the RIP property with parameter
(ex)
14

---

(r, and denote E, (M) (M)== (X!_)) Some lemmas are direct corollaries and some lemmas serve as extensions to rank above case. The proof of these lemmas can be found in Li al. [3]]
(c))
M)x M for some symmetric d x d matrix M.
Lemma 1. Under the assumplion of this subsection, if X, Y are d x d matrices with rank at most then
KE o (X), Y) 8 F|||F
(27)
Lemma 2. Under the ussumption of this subsection, if id d matrix with rank at most r and Z is a d. d' matrix, then
||E 9 (x)Z - 6 X |r|
(28)
Lemma 3. Under the assumption of this subsection, if X, Y are d x d matrices and Y has rank at most r, then
(E (X), Y) 8 X|| |Y[[-
(29)
Lemma 4. Under the assumption of this subsection, if X id d x matrix and Z is a x matrix, then
| (X)Z| - | |
(30)
Lemma D is from Candes [4] The other three lemma can be derived from Lemma thhouugh selecting Z or decomposing X into a series of rank-1 matrices 31
A.2 Additional Auxiliary Sequences
In this section, we additionally define some auxiliary sequences. Some for calibrating the dynamics, that is, describe how the dynamic progresses without error or randomness and track the trajectories with the accumulation of error. Some are used for characterizing the impact of randomness the dynamic.
The next two deterministic sequences help to track the dynamic of singular values of R, when it accumulates errors each step.
Definition 2. We define the following two deterministic sequences:
Re+1 = (1 11-nn R
17)
R4
+
log
R
Rg = a
32
(31)
R = (1- R + R
7
log
R
R = a.
32
The next lemma shows that the deviation between R, and R2 can be bounded.
Lemma (Bounded Deviation between R2 and Ry). Let the sequence Re be defined as be the first time R2 enters the region ( 1/5 1)) we have
25 Let Ti
R₂ S (1 + 1/6)
(32)
R, 2 (1 1/6)
for any 1 - 00,...,,
Proof. Fist, for RS have that
Re+1
(1 nR + 1) R4 + 21 log
(2) R
(1+ 1
1
R
32
<
log
(33)
Re+1
nR + r))R,
32
R₂
takes Ti < 1 log (t) steps for R2 to reach ( - 17, t) We can conclude that
R
log - )1))² < exp(1/8) < 1
< 1 + 32
(34)
RT.
where we use - x 2 exp( - 2x), ; 1 + N exp( ( for æ E 1/2]. Similarly, For R,, we have
17 log - R > (1 - RR
17 log Re, (35)
R,41 = (1 (1- + R4
32
32 6
15

---

and
R R+
(1 (1-1/h vR + 11)
32 B
2
7
log
-1
)
R₁
R
1 log
(12
(
(36)
R++1
(1
nR + R,
R
32
6
a
which implies
Rr:
7
21
> 1
(37)
R7;
48
6
One can also see that, R2 < +1+ 1/6) and R 2 (1 - 1/ holds for any / s T, which completes proof.
Next, we formally define the calibration line L4. In later parts, we can show that the norm of each column of Q₁ behaves like biased random walk with reflecting barrier Lt.
Definition 3. Let az, be defined as above. For 1,.,,, define the calibration line:
L V 40M8 rr2
(38)
Next, we define a stochastic process of based on Σ, The reason why we define this sequence is that though the randomness only directly affects Qt the dynamic of E, and R. also shares the randomness, therefore the dynamics become difficult to reason about since they are deeply coupled. Therefore, we define this "external" random sequence to dominate them.
Definition 4 (Controller Sequence). We fix the violation probability p= /./ log(d)) for some small absolute constant Cu. For each fixed i, we define stochastic process q. for = 0, 0, 2.... with d. = at, and
at = Sot (1+ 1 + 2n)d V L4+1
-111+
if there exists T S such that of 2 2p pp-11p1 1.5 1-1.
L.p
(1+1)
otherwise
qf is used for providing an upper bound the norm of columns of Q,. Before of hits the upper absorbing boundary p 1.5,115 Lt, can be considered as a "reflection and absorbing" process, with reflection barrier T, and absorbing barrier p-1.5 for hi:
1.55 L., The following lemma gives an upper bound
Lemma 6 (Upper bound for 9;). With probability over 0.995 over the randomness of the E(e), for
all = 1,2, 12 and 0, .., T2, we have
at < x 1.5 1.5
Lt-
(39)
To prove this, define family of random sequences X k
Definition For each 1,...,?3 r2, we construct a family of non-negative stochastic processes {X o for O..... , T2 as follows:
L,
0 0<tsk<Iz
(1+ (1
(ec) + 2n)
0 0sk<lsT
(40)
(X (XE,)) can expressed the following form:
Lo
(1 +
m() + Lo
(Ie (1 + m + +m)) Lo (ea) 21)
TT n + 2n)
Lo
Lı
L
[T=, nD((e)) 21))
L1
(Xixd)
La
L2
L
[I' (p(()) + 2n)
L2
:
Ltz
Lra
Lr
Laz
16

---

It can be noticed that X and qf have close relations. At the beginning we have of = X t 0, 1, ... progress along the 0-th row. If of gets lower than the calibration line Ly at some timestep = to. switches to the the fy-th row Xt = to.... until the next time it gets lower than calibration line Lt, and so on. We can see that of always progress along a certain row. Thus
P (3k such that of = X 11, eie [ra] and [Ta].
(41)
Theerfore, any uniform bound of Xi, can also be a bound for of- Later in the context, we analyze Xi, for each so we omit the argument in Xi,, for convenient notation.
We define o-field F. = (<)))) 1)) for f = 1... T2 and Fo = o Then we have Fo F C... Fr, form filtration. The next lemma shows that certain power of (XK,e)e is non-negative supermartingale w.r.t. Ft-
Lemma 7. For each = I,..., 12 and k o,...,, if the learning rate n satisfies 1j E My a4M))) [x2/3T7, then the process k,t It a non-negative supermartingale with respect to Fr.
Proof. First, its easy to verify the adaptiveness 22/3 E Fi since ! E Fe for all = 0, 0,1,.... T2. Next, note that
x2/3
st+1sk
E
X 1
2/3 FF
2/3
(42)
(1+ + 2n)
x2
,t 2 k
So suffices to prove that
EanD (1+ + +n2))/3 < 1.
For any 7 E (0, 1) and x - 1| 16 from Taylor's expansion, we have
< + (1 (1-))( - 1) 1 - 1111--1x---00 1)
(43)
Therefore,
Ee; (1 + n/(( + 2n)) 2-1
< 72 E - - -7)y Vare, [
(44)
Hence, it suffices to choose 77 7 such that
1 I 4 m Var Varl=/oa)1 n < 64M1
1
(2 E
(45)
When 7 = 112 E
( 24 Ma 64M,
suffices. Hence we prove
[1 + n + 2//)//|
< and we can
conclude that
s x2.41112 s x2/3.
(46)
Now we are ready to prove Lemma 6
Proof of Lemma 6 From the above observations, before of hits the upper absorbing boundary p- p-1.6,1.5 L,, there always exists some & such that of Xx.-- Therefore, of hits p 1,5,1 L, implies there exists some that X hits p- -1.5,115 L, So it suffices to bound Xg,1- For any fixed Is ..., T2, we denote two stopping times:
der Te Tk
min kS1ST:
2x2/33 (Lz) 2/33:
(47)
Tk Te ^
min
2x2/3 2 (p 1.5 115 1.5
L,) L2) 2/33.
17

---

One gets that
(a)
( I
P (
< kk) S
EX2/3
E.2-/33
(1-155125 5. 12)2/33 ASTEATR
(p- (p-15 LL )2/3
k.U
(48)
(2x)
(p 1. L,)2/3 (p (p-11511-5--
Where the inequality (a) is from Markov's inequality. Inequality (b) is from the optional stopping time theorem for supermartingales and inequality (c) is from the fact that L is non-decreasing. Therefore, we can conclude that:
A( r2; t S 72 such that of S pp-1._pp1000 1.5 T,))
r2k such that T < T2)
T--1
Σ
( <
(49)
r2
k=0
(p
1.5 11155-
T2p
where the first inequality is simply union bound over i == 1,,2, 1.,,..., Then
T2p O M2 log(d)
< ((M2 log(d))
- 0.01,
(50)
M2 log(d)
where the constant hidden in (-) only depends on the choice of a. Since log(1/a) < log(d), the constant hidden in O(-) is absolute. Therefore, the last inequality holds with sufficiently small C,, which does not depend on other parameters.
A.3
Useful Lemmas
In this part we bound some quantities that frequently encounter the error terms. These lemmas will simplify our proofs in later parts.
The next lemma helps to bound the "interaction error" arose from the non-orthogonality of V and U*
Lemma 8. Let Re, Qu E, and be defined above. We have
| U, - RR,R, + QQQ. + E, E.)l s 6e, |
(51)
Proof.
From the definition of R,, Qu, E,, we have
U.T.=(UMM = -U-RR V*Q." T + +E,) (URR.T + V.Q. T +E,)
= -RRR + Q.Q. + +E, E, + U. (Idrg- + Idles
(52)
+ Idy + Idhedddy. + +dy- + Idv.Idu.) )
Note that
Ild Idv. - VVV+ Sel
(53)
and
IIdre || = ||(I Idu- Idv-) Idu-l = II- Idy. S c1.
(54)
Similarly, for the other terms, we can prove that all the six terms the bracket in the last line o2 have operator norm 61. This completes the proof.
The next lemma helps to bound the RIP error in the dynamic of U.. Lemma 9 (Upper Bound for E,). Under the assumption of Theorem if E. Qel. R
1.1
and Eg 1 < 1, we have that:
U, T, (U.U, UTUT. V.E. -----V-V
U.I
(55)
18

---

Proof.
o U,U," -UU,--U'UT - V'EV
E (V'EV--T) + +
(,,E, || ||| F o (U.U,T - -U- E, E,E, -EEEEE,EE ||
(U.U,T-U
(1Eill + |
(56)
+ + R. - 11 + Q Q /- + Bxll ( + (Qills) + +2 (Rol
- 6( + +1+ 3/1 4/12 + 8(v11 + vr2) + +v/))
v 2M + r2,
where in (a) we use the linearity of E, o (.) and the triangle inequality. (6) we use Lemma 2 for the first term, Lemma forr the second term, and the expansion:
U.U,T - -U - EE,E (RR R, TJU V'Q. QQ.
+ E,, + +Q. (V +UR E,, (57)
+VQQ TRUTT U'R, TQ.
for the third term which shows that U,U," T U. - E,,, has rank no more than 2(r, + 72). Hence conclude that:
U, T, (U.U,- -
V'E.V+ - $ 22M10 + |
(58)
The following lemma tells how to bound the interaction error and RIP error using the auxiliary sequences Rg and Le we have already defined:
Lemma 10 (Bound Using calibration Line), Under the assumptions of Theorem 2 if E
<
R1 S
min(4R,, 1.1} and Q&lll TT2P" 1.5.1.5 Li,, we have
(1) R
(Ma + 2M18v/1 + 12) | S Ly
log
(59)
576
Proof. From triangle inequality and the condition of this lemma (2e1 S 8), we have that:
(M1a +2M88n + r2) |.|| S
Míōri r2( Rall
|Qll [E
(60)
M, M rr +22 + /r22 1.5 1.5 L.).
Then it suffices to check:
(4)
20M48 P2R, I L
(b)
[18 T2P
1.5
vr r2L, S L
(0)
20M r1 +
T2R,
log
(2) R
1152
(d)
3M T2p 1.5,1. 1.5 /ri + rايي 1152
5 log (2) Re:
where (a) is from the definition of Lz, (b) and (c) are from the assumption on à in Theorem 2 and (d) from the assumption on 8 (the absolute constant e the condition for H) and the fact that T., S R Hence the proof is completed.
0
19

---

A.4 Bounds of Qu
For evaluating the magnitude of || Q we consider its columns. denote each column of Q2 as q
(8) for i = 1,2, ..... 12. And use 9f defined above to upper bound them. Once we provide uniform bound for all q; we can also bound |Qull-
Lemma 11. Under the assumption of Theorem 2 under the event of q < p 1.. L, with p 2 6 for all i 1,2,.... ................................ . and t 0, ,, T,
2/3
E, - |R, 1.1. Eell?
and ||9 - 9j for
all j = 1..... T2, then we have
q
(++1 || S qi
++ <p 1-1 Let1
(61)
for all i 1.2,...,,7?
Proof. From the dynamic of Q₁:
Q++1 Q: - n U,Q, U0Q0+++02222 - -n((1 + 2M |0.1]
that for each column q. ( of Q
| 1 || $ (1I (T- ., U. + 1
+ 1 Σ
b 1 || + (e 22200 + ||
< (1 (1+ || (1) 1|2 +
Σ
1 ||
(e)) fス
||
nL..
where we use Lemma
10 For the second term we have:
7
Σ j
11)
"l
Co
1.5 1
(b) 5 S nlo
1.5, log 1.5 d) S 1,
(c)
(62)
where in (a) we use Assumption (c) and induction hypothesis that [|q < 1-pp--p0 1.5 11 1.5 Lr, in (b)
<p-11.p
we use the definition of p (Definition 4 and in (c) we use Assumption (a) and Assumption 2 with sufficiently small c, (which depends solely another universal constant c.)). Hence we have
1q (++1) < (1
"aall12 + 21L4.
(63)
There two probable cases:
q 19 S T.,, then |q (t-
S (1 n 22) S (1 + M + 2n)g S 9. 41,
If |
Ly then |
(++1)
S (1 + E
2n) |ل S (1 + + + 2n) S af
Both lead to the results we desire.
With the above lemma, we also give bound for |Qil:: Corollary 1. Under the condition of Lemma 77 have that
QQ+1ll & 1|Qut1llr S
(gtt < x p-16 Vialeti-
(64)
A.5
Bounds of E,
In this section, we bound the increments of both the operator norm and the Frobenius norm of E, The next lemma provide an upper bound for |Eal]- Lemma 12 (Increment of Spectral Norm of E,). Under the assumption of Lemma we have
E S E, + nT.
(65)
20

---

Proof. From the dynamic of E, 24 we can derive that
Ec (I 1f U.) + (|Idree (U-U* V V.TV| + Eell) ||U,
EEtil <
IVE IA3 IVg
E + 17 ((c + M1 + 2M + |
| + 15M + R + Q+ll)
| + ,
where (a) from Lemma and the fact that [[dres Uell. ||ddres U.|l S c1 (b) and (c) are derived
similarly as Lemma 10
The next lemma bounds the F-norm of error component E.,
Lemma 13 (Increment of the F-norm of Error Dynamic), Under the assumption of Lemma 17 and we further assume that |Ee] & 8M1v/1 + r2 (1/)). then the Frobenius norm of Es41 can be bounded by
E S (1 000000mm011 (1)
E. + m0066 + r2) 1.5 log(1/a)),
(66)
which immediately implies,
|E x (1++0m)M) -1 (M
(67)
Sno + r2) log (1/a).
Proof. We expand |E++1]] ppp from the dynamic of E,
E ||2- = E, (I - 17 U, U.) + r/Idu (U'U+ + V.E.V.T) U, -,
EI U, E
2n Er n,, U,), Idess,,
+ yIdres (U.U*T= + V.E.V' U
(68)
+ EE (I - --U alddes (U. -VE.V U.) )U,)
+ IIdres (U.U+ T V*E.V+ U, nId E,U,)
(1) + (2) + (3) + (4) + (5) + (6).
Now we bound the six parts separately. For the first part, since 0 N u U, = I, we have
(I-nU, UNI < E.
(69)
For the second part:
(2) 1² ( Er, Idres (,,UUu, (o)
< 1/0 ||U,U,T - T -E,E, TI. |,EEE, I]. VIV
(|E:(U.U,T E,E, 7) 1 |,E.E.E 11.)
s 178 (IU.U, -UU - E, || + E,,,, ||. [1++ (70)
||E||] (|U.U. -E,,, 1| + +|E II.)
(0)
2 /8Miv + r22 5M, vrn -2
r2 (0(vri - r2) + [E.
< 228M((r +f(ri+r22) 1.5
21

---

In (a) we use similar technique as in Lemma 9 to divide U,U, - U*U*T V V* into three parts so that we can use Lemma [ and Lemma (b) we use Lemma
9 to bound the first two terms and (57 to bound the third term with the assumption that [Rel], QQ: E, < 2. For the third part:
(3) (a) 2n E Id E, (I U,, U,),,
-22 EE, E, U, (.) T + R.U Q. Q.V.T))
(
x 210 (IU,U, - U.- E, F + E,E, T ||. + |211
(I,,(I- U,) |- + EI U R F + -I. Q.)1)))))))))))))
2 2n00 / +2) [ || || R R. E Qulle)
20 +2) E / (2 +
E.
2 nOMivii + E. 3 + 2 M (ri +2) 1.5 log 1og(1/a);
(71)
where in (a) we use the fact that
A C). In (b) we separate E, as in (0). In
(c). for the first term use ||A |||| (and similarly for nuclear norm) and | n)U U. 1.
the upper bound for E, appeared in Lemma 9 for the second term
For the fourth part:
(4) S 2 R +
8n 1
r2),
(72)
where the first inequality is from Cauchy's inequality and the fact that [[Idres U [1, IIdres V+I S el
For the fifth part:
(5) - E. (I - nu. nUU U.) Indde.s (U--T++ V-E. U,
I | (ri r2)
(73)
S (5 / rogg(1 61 Mri +
s n Mf(M r.5 log(1/a))
where the first inequality is from the norm inequality (X, Y) S |X|| and in the last inequality we use the fact that 61 < 5.
For the sixth part:
(6) n7² (Er,, (UUU++VV + +VEV U.u.) ),,T),),)))))
(0)
>/||E<| [I (UU + V'EVV U.U.ll.
51 ||xy] U. |Idres ? (U+ + VVE."") I-
(74)
V 1/ /8Mi + | Mir2)
(1)
2 Mf +.2)
In (a) we use the norm inequality X Y) < [Idres U.ll. [Id U,l
|||||||| and in (b) we use [Idres S and
< 61.
Now combining Equation 69 (79. along with the fact that Ep 1 < da² dd-
82y1.5 can
derive the result we desire.
22

---

A.6 Analysis for Phase
In this section, we give rigorous analysis for phase 1:
Theorem 4 (Phase 1 analysis). Under the assumptions of Theorem 2
During the first Ti =
O(, log( :)) steps, with probability at least 0.995, the following holds for any t E [0, Ti), that
( > (1+ m/3)a, (R,) for all j E (ri):
[Q:l1e Sp pp-1..12 1.. T2 L S 8 0.01, where Le is formally defined in (38)
|E S 6 11 + r2 R and |E z 8 Mf + r2) log(1/a)?
Finally, we have R ) σr R E (1: l.).
In the below contexts, unless otherwise specified, we abbreviate the largest and smallest singular values of R, as o and (t)
The next lemma tells that, if ||Qul and |E,|| are both small, then R, increases stcadily and the deviation between its singular values is small.
Lemma 14 (Dynamic of Singular Values of R, in Phase 1). For some t S T 1, under the
assumptions of Lemma 11 if 11. Qell < 96 log -1 (L/a)) and R2 S (() S 01 < Rr, then
(1)
R++1 S or
1)+ S (t+1) S Ret1-
(75)
and
o (++1 2 (1+ 11+1 (t)
(76)
(+1) 2 (1+ n/3o/
o
Proof. From the dynamic of R,:
R₁+1 = (I (I---10 n U, + IIR + U, V' U UU
= II.,, + IIR, - QQ.Q. T + E, E, + nUT [MM + (2
We use of and - to denote the largest/smallest singular value of R. To control the dynamic of (1) o and dr we need to bound the magnitude of the error term, that is
(t)
Q + E, E, R. + U, [2.55M /ri +
||0/1 2 + E 1 + log (1/a) 0 (t)
1
S 17
32
(77)
$ 17
+ 96
+ 96
log (1/a)o(
96
32 log
(1/a)a1
where in the first inequality we use the assumption for Q+Il and [E and 8. Therefore, from Weyl's inequality, have that
(++1) <(1 no n) (t)
(
32 log (1/a) ((1)
(78)
Tr: 2 (1 Njon + +1)00
(++1)
(t) (t)
3 (lg/) (t)
Using the assumption that
(4) S R₁ 01
a 2 R, 1= 1..,,...,,77
(79)
And Lemma 5 we can conclude that
(1 1/6)R(+1 < R++1 5 (++))
< 0
(t+1) R441 S (1 + 116))
(80)
For the increasing speed of on note that
o0 < 20 therefore
((+1) 2 (1 1n + 1 1)) ()
(81)
on 2 (1 1n 2 (8)
(++1)
This proves the desired result.
23

---

Now that the supporting lemmas are prepared, we can begin the proof of Theorem
Proof of Theorem The initial value of Ug implies that
|Uoll |Ro|| ||Qol || = Eoll = a, | F S a²dd.
(82)
17, 1/3). We have O,..., Ti happens with probability that the event of of < 1pp16,6 r 5 I, for all i ...... and 14 to inductively prove:
Recall that the time T - _ logll111 1/)) ) the first time R, enters the region
over 0.995. In this event, we can use Lemma
For the operator norm of E,, we have that for all t - Ti
E2
a
L
t=
240
a(1 + 17-
log
n
1-40M88T r2
11
(1+ (1+ (1++3+133)+ (1 n/3) 2....) (1+)))+++...))
3
(83)
250 log (2)
+ 40M + r2(1 + t//3)
S 40M1&8Ti+
log (1/a).
96
where the second inequality we use Lemma [4 that R& increases with rate not less than (1 + 3/n).
For the Frobenius norm of E,:
EE, < Tino Mf + r2) log(1/a) S 5f + r2)
log (1/a) <1
(84)
For |Q: we Corollary
1-1.5 1.5 Vr2Ln, S p 1-1 VrzdMIVTi +r2
log (1/aa).
(85)
96
For R4, we have for 1 S Ti:
(1 1/61/) - R, < o <o
((+1) (++1)
R S (1+ 1/6
(86)
For the condition ||E||| < [R,||
< " |R.| < [R |||
|E+i|l
||E|| S n2 <
R2
[Rell
(87)
10
Hence the proof is completed.
A.7 Analysis for Phase
In phase 1, the signal component R. grows at a stable speed from a to O(1) while the spurious component Q, and the error component E, are kept at low levels. In phase 2, we will characterize how R, approach and how to continually keep Q2 and E,
Lemma 15 (Stability of R2). If there exists some real number g satisfying
0.01 > 9 1Q111 + ||E,]| + U, Ex
(88)
for all t = T + + 1,.., Ti + T 1, then we have
1 1 5g < a! (1) < o0 < 11+9.
(89)
for all 1 -= Ti + 0 log (1) )....... +T
24

---

Proof. First we consider the upper bound for of" Similar to Equation 78
we have
o (++ (1- noi (t) + - =
(90)
Note that Equation (90) is equivalent to
(1+1) 2 of ((0 + "ell:/++++///i /1+
(91)
With of (1) (Ti) 2. one can see that 01 never goes above vI S 11 + 9.
Now we consider or, After phase I have o. (1/3 1. If o S 50!! similarly we have:
F11 2 (1 - 1101 (1)
(1+1)
+ -リ - !!
(92)
Which implies that, if dr
(t)
- 5g.
VI- 5g 0r
(1+1)
< (v 1 - 5g 0 )1 - noo!! + v - -
(93)
$ (1 7 n) /V1 - 5g - o!!!)
Therefore, o will get larger than -T 5g g² 2 5g at some time S T1 + log 5 Also
that, drl keeps increasing before it gets larger than / 5g- And once it surpasses VI 5g. it never falls below than /I vI-5g 5g again. Therefore, 0 I 50 is
from Equation 92
(t)
satisfied and the proof proceeds.
Now we can state and prove:
Theorem 5 (Phase 2 Analysis). Under the assumptions of Theorem Let T3 = Ti 0 log((ri + r2)/8) s ( log( Then with probability at least 0.995, we have
σ1 [Rr,) (R o, E (1 (1- MF 8Mivr1 + +22), + M 8 +22)). (94)
for = =Ti + I,.... ......TT, T2, we have
1Qillf S 1p51. VrzLe Sp 15 /
40M
| S 8Mi + 12 log(1/a)) and |E S 82 M2 + +2) log (la)
Proof of Theorem 5 The error g in Lemma 15 is no less than s2 MT (ri >> order, therefore
Th- = Ti log (I/a) suffices for of to reach 5g. Then similar to the induction the proof of Theorem we can derive(in the same high probability event):
| = 40m5M + 4076 / + r2(1 TT) < 808 + loglo//)))) <
0.01:
|E,1 S tr M2 (7i + T2) 1.5 log([//)) 62 ( (og
|Q.1l S 5-11 Vr2l. S 1-1 4, -15 < 0.01;
Hence the assumption in Lemma 15 is satisfied, with
g gSp - 38 M V 8M /ri +
(95)
Therefore,
||| Rr,ll
1|| S
382 33
8M, vri +
(96)
25

---

Proof of Theorem
Using Theorem and Theorem 5 with T = T2, we have
(a)
TT, U A A
ET, P T2 Rr² [| + 10 Qralr
Er, Qr E| Re F + R| Qr lI
11,,1 8 M v 8M1 /r + r2) vni Q
- E7, QT, 2 BT, R. | + 2| QT, || R
V 2 8M + 1.5 O w2)/ri
- M1 + +M 0
8M
(8 Mi V Mz
where in (a) decompose UU, (Sec (57 and triangle inequality. In (b) and (c) use Theorem
and repeatedly use the fact |AB|] F & B F. This completes the proof.
B Deferred Proofs
B.1 Proof of Proposition 1
In the below contexts, notations such as C, c, C1, C1 always denote some positive absolute constants. Such notation is widely adopted in the field of non-asymptotic theory.
We first state some useful definitions and lemmas:
Definition 6 (c-Net and Covering Numbers). Let (T, d) be metric space. Let 0. For subset K C T, a subset M C is called an e-net of K if every point in K is within distance of some point in M. We define the covering number of K to be the smallest possible cardinality of such M, denoted as N(K,)).
Lemma 16 (Covering Number of the Euclidean Ball). Let Sn-1 denote the unit Euclidean sphere in R" The following result satisfies for any 0:
N(Sa-1 160= 2 (2+1)"
(97)
Lemma 17 (Two-sided Bound on Gaussian Matrices). Let A be an d x matrix whose elements Aij are independent N N(0, 1) random variables. Then for any &. 2 we have
Vd C(Vr + t) (0, S (A) S Vd C(VF+t)
(98)
with probability at least eexp(-12)
Lemma 18 (Approximating Operator Norm Using e-nts)). Let A an m x n matrix and E [0, 1/2). For any e-net Mi for the sphere - and any e-net M2 of the sphere Sm-1, we have
1
sup
(Ax, S |A|| S 1 2e NEMi
sup
yAx
(99)
XEM
Moreover: if n. then we have
1 2c e²
sup (Ax,x) < |A|| S
sup ZEM.NEM.
((x,
(100)
XEMi
Lemma 19 (Concentration Inequality for Product of Gaussian Random Varables), Suppose X and Y are independent ((0, 1) random variables. Then (X,Y) is a sub-exponential random variable. Therefore for (X1,...; Xm, Yı,--., Y,,) N(0, Iam), the following holds for any &. 0:
m
Xi, Y2)
>6
exxp (-emin((?,, (1))))))))))))))))))))m))))))( -m).
(101)
m
26

---

Proof. Note that
(x 1
1 V/2 y
+
X
(102)
2
V2
2
/22
The two terms are independent and following Gamma distribution I (G. 1). Since Gamma distribution
random variables sub-exponential, (X, Y) is sub-exponential too. The concentration inequality
follows from Bernstein's inequality. (See Theorem 2.8.2 of Vershynin 47
Now we prove Proposition
First we provide a bound for [M, M2]1. We fix e = 1/4, and we can find an e-net M1 of the sphere S and e-net M2 of the sphere S*²-1 with
Proof of Proposition Z
MMl S 9r MM2 9
(103)
For each x E M, and y E M2, we have for 0 <u<l,
( x M M May
u =P 1 (M M2y) > u
(104)
< 2 -ccpp-
where use the fact that M, x and M2y are independent N (0, La) random vectors and application
of Lemma 19 We let / / - for t
ritr2" have:
d
( M M2 2
11 + r2
5 P
max d NEM.
x M M2y 2 2V
ri +22 .t
d
(b)
< tr 2exp (-c2(ri + r2)
2exp (-(ri + (22 log ())g
(105)
where in (a) we use Lemma 18 in (b) we apply union bound over all M1 and y e M2.
Next, we bound |R, and [R² I|. Recall the QR-decompositions of M, and M2:
M R and M2 22
(106)
which implies M M R R2 and M. M2 R2 R2, and consequently |R or, (Mi)-
and |R5 = or (M2) -1 From Lemma 17
2
Vd (M)
P R
= or
<2exp(-cdd).
(107)
Vdl
2
And similarly for [Rg"" Finally, for / V
1 U 24/
ri + r2
P
d
=P
|R MiiTMMaa MrTM,Ri'' || > 4t
Ti +
d
(108)
< P |R 2 Vd +P ||55||| 2 [/l
2
2
+ PP
M
|
ri 11+12
< 4 exp (-c + 2 exp (-22 + r2)
This completes the proof.
27

---

B.2 The Failure of Pooled Stochastic Gradient Descent
we have a separation that Pooled Gradient Descent fails to select out the invariant signal, whereas the HeteroSGD can succeed. This isolates the implicit bias of online algorithms over heterogeneous data towards invariance and causality.
From Theorems and 3 for the hard case in Theorem 3
In this section, we give a rigorous proof for Theorem We first demonstrate the failure of PooledGD Theorem 6 (Negative Result for Pooled Gradient Descent). Under the assumptions of Theorem 2 for the certain case where U V and BeED ()) r we perform GD over all samples from all environments and ends with T Ollog d), then U, keeps approaching U* + +VV+ in the sense that
- - U - V'V T| 20 + o(1,
(109)
during which for all t 1. T:
U.U. - A' F Arz.
(110)
Proof of Theorem 6 Firstly, we emphasis that Theorem 2 also applies the case where there is only one environment and no spurious signals, the m samples are generated as: (We use underlined notations to distinguish this setting from others)
X
m
(111)
In such cases, there is no randomness and U, U, deterministically learns A and all the singular values of R, grow at similar speeds.
Under the conditions in Theorem 6 we first construct single-enviromment case. Let U and V be defined in Theorem 6 we let the invariant signal A U. + V*V+T and there is no spurious signal. Then the updating rule is:
1 YX. - A U,
U -
(112)
U U.UUU
A A -. E (U.U. A A
Using Theorem 2 we can prove that U,U, continuously approaches A T U + V in phase & 2, during which:
In phase 1, R < 1/2 therefore |,,,, -. VTI.
In phase 2, all the singular values of |R, get larger than 1/6, from Weyl's inequality, we have that the top singular values of U, U.U T are all larger than 1/6. Hence
U.U. U 1 V22
Therefore, U.U. U. IF Z Vri Ar for all I. t= O,..., T.
Now prove Theorem 6
The updating rule can be written as
U.U. - A - A x() U,
U411 = U, - Be
m 1=1
= - (U.U. (..-n(UUU,
A² E (6) U.-nBan E (U.UT - A - A A ] U,
= U, - n (U.U - (VO-. + V.V U, BENDD [Ec o U.UUT - A - A U
We compare this updating rule with (12)). The only difference is the RIP error term. However,
the upper bounds for E (U,U. A A!o) used in the proof also apply for the expectation
EgnD [E (U.U.- A A So we can derive the same conclusion that
1 - EENDA (e) s o(1)
(113)
28

---

for T = log(1/a)), during which we have that for all t = 0, 0.,,...,, T:
|U,U, A V
(114)
Now we are ready to prove Theorem 3 Assume that at each time = 0,...,1, we receive m samples [X] (1) (!)
Di Ji=i each sample independently sampled from environment εx, D, satisfying
X A A (ew,))
(115)
and imply the Stochastic Gradient Descent
m U+1 Ia
1 17 Σ (X) m
U.U,)
a xx U..
(116)
For technical convenience, assume that X is the symmetric Gaussian matrix with diagonal elements from N(0, 1) and off-diagonal elements from ((0, 1/2). We further assume x! independent of C,,s. This corresponds the cases where each environment has infinitely many samples and the linear mcasurements from different environments share the same distribution.
Proof of Theurem
Denote = EE A (e), Then we have
m
U,U, A - -A)1x(0)) U.
U U241
[a 11 m
122
(0) A (aw) -
A)x(") U.
+ 1
The first term is the dynamic of single environment matrix sensing problem, and the second term a zero-mean noise arising from SGD. Once we can prove that the second term is small with high probability, then the dynamic will be similar to the dynamic of single environment matrix sensing problem, thereby we can get high-probability version of the result of Theorem 3
Now we control the SGD noise term. Let M3 be 1 -eet of the sphere S with IMall S 9 Then for any d x d matrix M, have |M] < 4maxx.., xx Mxx.. For any fixed x MMs one can see that (4) A (er.s) A)x Mx has zero mean and is the product of two uub-Gaussian random variable with Sab-Gaussian parameter more than 2M (r1 + r2) and 2. Therefore, it is sub-exponential random variable with parameter no more than CM, (r1 + r2) for some universal constant C 1, Then applying the Bernstein's Inequality (47 and taking the union bound over Ms. we can obtain that
[ (e..) A)x Mx > CMi(ri + r2)(( r2)l// (
P sup M
m
< 2 2- exp(-t) ex)t)) (117)
Setting t = 10d and m = dpoly(ri + r2, Mi + M2, log d), we can obtain that with probability over 1 exp(epd d),
Σ m
m
A (era)
A|x
1
(118)
poly(r: + r2, M1 + M2, log
Therefore, in this case the SGD emorr be upper bounded in the same way the RIP error the level of o(1/ poly(r1 + M2, + M2, log d)). This implies that the SGD error will not significantly affect the dynamic with probability over T exp( d). Therefore (113) and (14) hold with probability over 0.99.
Theorem 3 and Theorem indicate that the failure is because the signal is averaged when calculating gradients when perform GD SGD over pooled datasets. To the best of our knowledge, it is
29

---

intrinsically hard to provide a rigorous statement when the batch size is small. We would like to leave the theoretical analysis as a future work. In the following simulation, we aim demonstrate empirically that Pooled SGD fails to learn invariance with small batch size. We consider the
| = 2 case and the environments are generated by A U (s + M)VVV. and A U + (s V)V*V* where (U*, V*) is column orthonormal. Then the invariant solution is A' U.U and the spurious solution is A' + A U + VV We set (a, d, T1, r2, 8, M, m) = (10-3 30, 5, 5, 0.5, 4, 80), use Gaussian measurements as Section 5 and let be sufficiently large. The following shows the F-norm between U,U, and A or A + A.
PooledSGD with Batch Size=80
PooledSGD with Batch Size=80
4.0
0.01
4.0
0,01
3.5
0.05
3.5
n 0.05
n= 0.1
n 0.1
3.0
3.0
M
2.5
< 2.5
2.0
2.0
1.5
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
t/T
t/T
Figure 6: These figures shows that when the batch size is small, the trajectory will be far away from A and A + A, suggesting that the algorithm is not stable in this regime.
C Neural Networks with Quadratic Activations
In this section diseuss how apply our results nerral networks with quadratic activations. In particular, Example I
As discussed above.
y
Σ
c(a
Σ
10 ((aTx(() =
aja
Σ
(e)
a,a
j=ri+1
i=1
jmr+
(119)
and it is equivalent to matrix sensing problem with
F1 A' aaa (0
Σ jm1+1
(c) aja and x =x
x(0 T
(120)
The main difference is that, when the samples x; are i.i.d. N Ia), the set of linear measurements [xixI,... ...,, Xm. x } no longer satisfies the RIP property. However, the following lemma tells that, with proper truncation, the set of measurements enjoys similar properties.
Lemma 20 (Lemma 5.1 of Li et al. (31). Let (X....., X.m) {xxxx] T T.... Xmx where x;'s are i.i.d. (0, I). Let R log (+). Then, for every q, [0, 0.01] and m 2 d a /8/82, with probability least 9. we have that for every symmetric matrix A:
1
AXX - 2A tr(A)I S S/|A|]
(121)
If has rank at most r and operator norn at most 1, we have:
1
AXX. 2A tr(A)I
< 18.
(122)
To accommodate this difference, we adopt the modified version of loss function and algorithm from Li et al. (3I].
30

---

Algorithm 4 Modified Algorithm For Neural Network with Quadratic Activations
Set Un = ala, where is a small positive constant. Set step size 1)- for I,..., T 1 do
Receive samples (x(1) Mi
(ex) (cd) from current environment et-
Calculate gf = 1 T {
(e)) ) 1.2.....1 m.
Calculate modified loss function (U, +. - ( - -W; v/(4))
(er)
Gradient Descent U, = U. v/ (U,).
Let 72 = ||A + A (
Shrinkage U++1 T- U,
end for Output: Ur.
Remark 1. Here we encounter the same caveat that Algorithm 4 requires our knowledge on Ti- As discussed in Li et al. (37), the algorithm is likely to be robust if Tt is replaced by its moment estimation.
Now outline the proof sketch of Example 1
Theorem 7 (Two-Layer NN with Quadratic Activation). Let a..... R be independent random vectors sampled from normal distribution NO, d Ia). For environment e E, suppose the target function is determined by rı invariant features and 12 variant admits that for each sample
(x. Mi
(e)
)
Σ
x(axxx x
Σ af [(a
x
(c) x
(0)
Σ j=1
aaa
Σ j=ri+1
aj d aja
(123)
jmri+1
Suppose we train the following two-layer NN:
f(x) Σ q(,x)),
(124)
and the initialization of parameters {uj} satisfies E/=1 nnuj = al. If (a)"))
his satisfies
supe, (le ma-max, (1+ 1"
for some absolute constant co, sample complexity m d poly(r, log(d), suped fla (). a (d- 1,d-1) and n
miny ( [a 10
max; (I (ar. then Algorithm re-
turns solution that satisfies
4 - Al1e<o00
(125)
with probability over 0.99.
Proof. similar to the proof of Theorem 1.2 al. (31) the modified algorithm is in fact equivalent 21 with RIP parameter (r; when 2-d22 Hence it is fully reduced the matrix sensing problem.
Now we verify the conditions for A Σラ a and ( Ej=++1 aj / u, = l,...,r are independently and uniformly sampled from sphere, we have that
aja Since
With high probability over the randomness of (a:ti the eigenvalues of A lie within
(1 1
(See Theorem 4.6.1 of Vershynin 47
The angle between Col(A*) and Col(A (e)) is of O(vr, + + order.
31

---

Therefore we can construct two column orthogonal matrix U and V such that U T V = 0 and (U""),
sincoll co (ol)) S vri +r/vvd. Hence we can apply The- orem on A U and (c) := diag(( (0) v+ Such approximation only raises ( 22/ Vd multiplicative error, which is negligible. And we can easily verify À (6) satisfies
Assumption Then this result follows from the proof of Theorem 9
D The k(A)) > Case
In this section we show how to generalize our results to the x(A*) > 1 case by leveraging the adaptive subspace technique proposed by Li et al. [31] for single environment setting. This framework mainly consists of the following steps:
First, instead of using the fixed subspace col(U*), we use adaptive one S,, where So col(U*) and S1+ (I-r); where M, 1 >i (X,, U.U. X, And we denote Z, Ids, U and H = (Id - Ids,) U, Which makes the updating of H. substantially disentangled from Zμ
Second, we reason about the updating rule of Z, Since the subspace is updated at each step, the updating rule of Z, becomes indirect. We introduce Z, (Id ,HH T "Zz - 2n + "Iss, M,H so that Z,+1 ~ Z, (Z,) It can be shown min (Z,) continually increases until it gets larger than 2/R¹
During this iteration, we can keep Z, near Ze for each and the principal angle Og between S, and CU)(U satisfies sin(O:) S npt where = θ (frr)
Finally, when min (Z,) is sufficiently large and principal angle is small, we use the local restricted strongly convex property of C around A to prove U.l converges with rate 1 O(n/x).
For the multi-environment setting, we have the following result under a slightly stronger assumption on the heterogeneity:
Theorem 8 (General Theorem). Under Assumption 7 andd? 2 suppose the heterogeneity parameter
M2 Z 12, < 5. and the RIP parameter S poly(r.logdd)MM +
We choose the E
1.--1.1M, (24M₂1, 64 M, A ) and a (1/d², 1/d"), then running Algorithm steps, the algorithm outputs Up that satisfies
in T = O(log(a
|U =11ls0(11
(126)
with probability over 0.99.
Since the full proof of the adaptive subspace technique is involved, for clear representation, we point the main differences from the single-enviromment case. We need to address the following three issues: (1) How to introduce the spurious component Q₁ into the original framework; (2) Whether the spurious signal (e) significantly perturbs the dynamic of Z; and (3) How to give a phase analysis when there is no local restricted strongly convexity around At?
We first cope with (1). With abuse of notation, we adopt the M,, Z,, H, and additionally define
v (ada) Idy H. and E ((ada) = (Id Idy. H, We can prove that
v ((dda)
Idy- + (A(e)))
(ee) + OOnDVrM + (1 + m)) sin((e))) (v(ada) +E E(oda))
t+1
Idv. + 1 (ox) vv(eda) + small terms,
(127)
E (ada)
(Idres + OOnDVrM + (1 + mM, sin():))) (v(ada) E. p(ada))
E
(ada) + small terms.
If we can ensure sin( S 8 poly(r, M1 + M2, (d)), we can get similar dynamics as 233 and 24 then apply similar techniques in Section A.4 and Section A.5 to ensure V, and E, are no more than ö poly(r, Mi + M2, log(d)). w.h.p. Moreover, the dynamics in (127) is multiplicative, which means if we decrease a by comparing to Theorem V, and E, can further upper bounded by 1 & Mi + M2, log(d) in phase 1.
32

---

For issue (2), the spurious signal A (6) brings error about 1 1+0(((+()))) multiplicative factor, which can be absorbed by the inherent RIP error of A". Another difference that, at the beginning ||&lll or E may be substantially larger than Z, due to the oscillation. We emphasis that such interference happens in RIP error term or non-orthogonal error term, multiplied by 6, (i,()) or 61. We can ensure such interference is negligible when 5 X poly(r.log(d),.M MM2.K) Therefore, the dynamic (r, Finally for issue (3), when min (z)) Z 2V and sin( = ò poly(r; ldg(d), M1 + M2, < 1. We get back the original subspace col(U*) and col(V*). We have U, = Z. O(sin(0.)) V. v-(ada) + ((snn(ft)) E, = E(ada) + Osin (fe)) and E E. E( ||r I vfsin(0.). Then use the technique from phase 2 analysis (Theorem (5) complete the proof. We leave the extension of this theorem for the case where MI, M2 are constant level for future studies.
of 7 is benign, and the principal angle can be bounded by (ogg)d M1 + M2, k) <
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
See conclusion for what we leave for future studies.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
33

---

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Question:
Answer: [Yes]
Justification: We provide assumptions in main text, proof details in appendix, and references for certain tools.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: See simulation section.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
34

---

(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer:
[No]
Justification:
The simulations are a simple validation of our theory.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
gablic/guides
sondeSuummsssioPooiioo for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
//ipsscccppulee
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
See simulation section.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
35

---

Answer: [No]
Justification:
Our theoretical results provide success probability.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enoor bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [No]
Justification: Computing environment is not important for our problem settings.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
All the authors have reviewed the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
36

---

Question:
Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
We investigate theoretical models.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: Our paper does not pose such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification:
The paper does not use existing assets.
Guidelines:
The answer NA means that the paper does not existing assets.
37

---

The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the datasets..paerssttt paperswi thcode. com/datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
38

---

Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
39


Reviews:
[{"review_id": "jp8rwBn6Qc", "reviewer": "Reviewer_uFbc", "summary": "This paper studies the difference between the solutions to multi-environment matrix-sensing obtained by gradient descent and \"heterogenous\" stochastic gradient descent. The authors show through analytical results and simulations that HeteroSGD helps discover invariant solutions while gradient descent converges to solutions that contain both invariant and spurious components. The authors describe multi-environment matrix sensing as a mathematical model where the matrices being sensed consist of an invariant component and a spurious component. The solutions of two optimization problems are considered: a) gradient descent on a pooled batch of data from all environments and b) heteroSGD where each iteration samples a batch of data from a different environment. The authors introduce assumptions on the near-orthogonality between the invariant and spurious subspaces, the heterogeneity of the environments, and that the measurements satisfy RIP. Under these conditions, they show that heteroSGD discovers the invariant signal while gradient descent converges to solutions that contain both invariant and spurious components. The authors also run simulations to support their claims.", "questions": "Some questions in weaknesses. Others below:\n\n1. Did the authors try searching for different learning rates for GD vs SGD? Prior results on the interaction between batch size and learning rate indicate that the learning rate for GD and SGD should likely be tuned separately. Does the spurious solution still exist when this is taken into account?\n\n2. Does the separation between pooled data vs environment specific data hold when SGD is considered (rather than GD)? Is the invariant learning captured by stochastic gradients or the fact that each update only uses data from one environment? Hopefully this can be answered theoretically and through simulations.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and easy to follow. All assumptions are stated clearly and the authors provide a baseline (gradient descent) to compare their proposed algorithm against.", "weaknesses": "1. The problem of multi-environment matrix sensing seems under-motivated. While the authors introduce a clean mathematical model for this problem, they do not provide instances of problems that satisfy this model. Is it reasonable to model environment specific signals as being incoherent with the invariant signal? Is it reasonable to assume heterogenous environments? In what situations are these assumptions satisfied? While I understand the focus of this paper is theoretical analysis of a model, I think it is important to motivate it so that we understand this is a real problem and not just another mathematical model.\n\n2. The proof sketch section does not provide any intuition about the separation between GD and heteroSGD. Specifically, I'm looking for an explanation as to why $\\mathbf{Q}_t$ stays small when environment specific data is provided (as in heteroSGD) and it does not decay when the environment data is pooled (as in GD). (If the answer is already in the paper I might have missed it - please point me towards it).\n\n3. Simulations are run on synthetic examples. Are there any real datasets that you can demonstrate your methods on?\n\n4. The current theorem statements show that under conditions of heterogeneity, heteroSGD can learn invariant solutions. Can one show that heterogeneity is necessary for learning invariant solutions?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewers for the valuable feedback and insightful comments. We have carefully considered your comments and questions and have addressed them as below:\n\n> The problem of multi-environment matrix sensing seems under-motivated. Instances of problems and reasonability. \n\nA: We fully understand your concerns regarding our formulation. We will add more real scenarios to exemplify our abstract problem. Let us first explain two aspects.\n\n- About the formulation of multi-environment. In practice, the data are often collected from different environments. The spurious signal refers to the **endogenous spurious variables**, which inherit the dataset bias and so have non-zero associations that are unstable and thereby should be eliminated, such as the **intervened children of the response variable $Y$** in Structural Causal Model (SCM) [1]. While many existing works propose different methods to achieve invariance learning, the implicit regularization effect for achieving invariance learning is not well studied. This is also the starting point of our work. Additionally, besides eliminating the endogenous spurious variables, our model learns the sparsity, which also helps eliminate **exogenous spurious variables** which do not contribute to predicting the response variable $Y$.\n\n- About the matrix sensing problem. The matrix sensing problem is widely used in implicit bias contexts as a testbed for understanding the loss landscape and training dynamics of over-parameterized deep learning models since it behaves like neural networks with non-convexity and non-linearity. However, it can be solved efficiently under suitable conditions. We thus hope our work can provide insights into the implicit invariance learning abilities of deep learning models. The incoherence condition is common in matrix sensing or more broadly in high-dimensional feature selection problems.\n\n[1] Glymour, M., Pearl, J., & Jewell, N. P. (2016). Causal inference in statistics: A primer. John Wiley & Sons.\n\n> About the intuition about ``HeteroSGD``.\n\nA: Thanks for this question. For ``HeteroSGD``, as illustrated informally in **line 248**, the oscillation creates a contraction effect, which helps prevent the model from fitting the spurious signal. In contrast, for ``PooledGD``, there is no oscillation and the model is consistently driven towards the averaged signal and cannot distinguish between the invariance signal and the spurious signal. We will provide more intuitions.\n\n> About the simulations?\n\nA: Thanks for raising the question. There are some works that empirically achieve invariance learning from multi-environments (e.g. [2]). Our work attempts to theoretically reveal how models can learn invariance from the standard training procedure, and our simulations intends to verify our theoretical results, as the first work from this perspective. We will consider generalizing our theoretical findings to design more empirical methods for invariance learning in the future.\n\n[2] Simple data balancing achieves competitive worst-group-accuracy\n\n> Can one show that heterogeneity is necessary for learning invariant solutions?\n\nA: Yes, heterogeneity is essential for distinguishing between spurious and invariant components of the signal. Conceptually, the signal is said to be invariant as it does not change across the environments. If there is no heterogeneity, it means all the environments are the same or there is only one environment, so all the signal would be invariant. From the technical viewpoint, if our proposed heterogeneity conditions are severely violated, one can construct counterexamples where standard optimization algorithms fail to learn invariant solutions.\n\n> About trying searching for different learning rates for GD vs SGD, and the interaction between batch size and learning rate.\n\nA: Thanks for the insightful comment. For learning rate, our results point out the range of learning rate for ``HeteroSGD``. For ``PooledGD``, its failure is because the signal is averaged when calculating gradients. Therefore, tuning the learning rate will not prevent the model from learning the spurious solution.\n\nIn our current results of SGD, we do not treat the batch size as a tunable parameter. The batch size is set to satisfy the RIP condition, which is a very common condition in the field of matrix sensing. Emprically ``PooledSGD`` is not stable when batch size is very small. Here $(d,r_1,r_2,m)=(30,5,5,20)$:\n\n|  Learning rate $\\eta$ of ``PooledSGD``   | 0.001 | 0.005 | 0.01  | 0.05  | 0.1    |\n| ----- | ----- | ----- | ----- | ----- | ------ |\n| $\\min_{t\\le T}\\left\\|\\mathbf{U}_t\\mathbf{U}_t^\\top - \\mathbf{A}^* \\right\\|_F$ | 1.682 | 1.820 | 2.125 | 2.185 | Trajectory exceeds boundary. |\nWe will add representative figures in final version.\n\n> Does the separation holds when ``PooledGD`` is replaced by ``PooledSGD``?\n\nA: Yes. We will add the following theorem in the revised version of our paper. \n\nTheorem 3.1 (Negative Result for ``PooledSGD``). Under the assumptions of Theorem 3, if we perform SGD ending with $T=\\Theta(\\log d)$,  where each sample of a batch is randomly chosen from an environment w.r.t. $D$, and the linear measurements are symmetric Gaussian, and the batch size is $d\\cdot\\mathrm{poly}(r_1,r_2,M_1, \\log(d))$, then with probability over 0.99, \n$$\n\\left\\| \\mathbf{U}_T\\mathbf{U}_T^\\top - \\mathbf{U}^*{\\mathbf{U}^*}^\\top-\\mathbf{V}^*{\\mathbf{V}^*}^\\top \\right\\|_F = o(1),\n$$\nduring which for all $t=0,1,\\ldots, T$:\n$$\n\\left\\| \\mathbf{U}_t\\mathbf{U}_t^\\top - \\mathbf{A}^*\\right\\|_F\\gtrsim \\sqrt\n{r_1\\wedge r_2}.\n$$\nThe intuition is that the heterogeneity across environment is reduced. Hence the gradients will be very close to the gradients in ``PooledGD``. \n\n----\n\nWe again thank you for dedicating your time and effort to reviewing our manuscript. Your valuable questions and insightful comments have significantly improved our work. We hope our responses have addressed all your concerns.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9U02A0Jfti", "reviewer": "Reviewer_JQ35", "summary": "The paper studies the implicit bias of (Hetero) SGD towards learning invariant representations and discarding non-invariant features in a matrix sensing setup. The applicability of the setup to the 2-layer NN with quadratic activations is demonstrated. Finally, the authors demonstrate that pooledGD fails to recover the corresponding invariant representation. Experiments on synthetic data is used to demonstrate the validity of the claims.", "questions": "[1] Does the result in theorem 3 extend to pooledSGD as well?   \n\n[2] Typo: Missing Reference - line 200", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The paper is clearly written, and the motivation is clear. \n- The theoretical analysis seems sound.", "weaknesses": "- Claims for PooledSGD: The introduction claims that the authors show that PooledSGD fails in recovering the invariant representation. However, the analysis (theorem 3) and the experiments are only presented for PooledGD. It is quite well known that GD, when compared to SGD, results in worse generalization (see [1] and references therein), even in the homogeneous environment setting. Thus, given the context, the negative result for PooledGD is not very surprising. Additionally, the assumption in theorem 3 about that the expectation of the covariance matrix, that the training environments span all directions uniformly, seems rather restrictive. \n - Theorem 2 feasibility: The application of theorem 2 in eq. 10 requires assuming d log^2(d) << m (batch size), which seems unrealistic in standard machine learning scenarios. The authors can perhaps make the point more convincingly by demonstrating the applicability of the results of theorem 2 with more realistic values (d, m, C etc.).\n\nReferences:\n- [1] On the Generalization Benefit of Noise in Stochastic Gradient Descent - Smith et al, 2020", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewers for the valuable feedback and insightful comments. We have carefully considered your comments and questions and have addressed them as below:\n\n> About the findings on ``HeteroSGD`` verses ``PooledGD``? Does the result in theorem 3 extend to pooledSGD as well?\n\nA: Thanks for raising this. Our key point is that the separation comes from ``Hetero``, not from ``SGD`` itself. In fact, we can prove that ``PooledSGD`` will fail to learn the invariant signal as well as ``PooledGD``. The theorem is as follows, which will appear in the revised version of our paper:\n\nTheorem 3.1 (Negative Result for ``PooledSGD``). Under the assumptions of Theorem 3, if we perform SGD ending with $T=\\Theta(\\log d)$,  where each sample of a batch is randomly chosen from an environment w.r.t. $D$, and the linear measurements are symmetric Gaussian, and the batch size is $d\\cdot\\mathrm{poly}(r_1,r_2,M_1, \\log(d))$, then with probability over 0.99, \n$$\n\\left\\| \\mathbf{U}_T\\mathbf{U}_T^\\top - \\mathbf{U}^*{\\mathbf{U}^*}^\\top-\\mathbf{V}^*{\\mathbf{V}^*}^\\top \\right\\|_F = o(1),\n$$\nduring which for all $t=0,1,\\ldots, T$:\n$$\n\\left\\| \\mathbf{U}_t\\mathbf{U}_t^\\top - \\mathbf{A}^*\\right\\|_F\\gtrsim \\sqrt\n{r_1\\wedge r_2}.\n$$\nThe intuition is that the heterogeneity across environment is reduced. Hence the gradients will be very close to the gradients in ``pooledGD``. We remark that the counterpart for small batch size case even for single-environment case, to the best of our knowledge.\n\n> About the feasibility of Theorem 2 and assumptions.\n\nA: We fully understand your concerns regarding the assumptions. In our results, we set the batch size to satisfy the **RIP condition**, so that we can omit the randomness within the environment and focus on the randomness of the sampled environments. This condition is very common in the field of matrix sensing. However, it is interesting to study how small batch sizes affect the invariance learning for SGD. This problem requires much more involved analysis to deal with the randomness of data when the RIP is broken down. However, in practice, one does not often need such a large size of data to reach the conditions because the analysis is in the worst case. We will discuss the main difficulty in the revised version and leave the analysis as a future work.\n\nAs for the assumptions, in section D we briefly show how to generalize our results to the $\\kappa(\\mathbf A^*)>1$ case. In the main text we adopt this for clear presentation, since this involved technique mainly works for the invariance part $ \\mathbf{R}_t$, while we are more interested in the analysis of the **spurious part** $\\mathbf{Q}_t$.\n\n----\n\nWe again thank you for dedicating your time and effort to reviewing our manuscript. Your valuable questions and insightful comments have significantly improved our work. We hope our responses have addressed all your concerns.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "US7Yjc9aTx", "reviewer": "Reviewer_a5nG", "summary": "The authors show that in a matrix sensing context, and under a data distribution that includes invariant and environment-dependent components, SGD with successive batches from different environments lead to invariant features being provably learned. The authors show that SGD with mixed batches provably does not learn the invariant solution.", "questions": "- The authors results constitute a strong negative finding w.r.t. standard SGD-based training without environment annotation and presumably mixed minibatches. The authors should more clearly highlight this and explain the reasons for the success of standard SGD-based training in practical settings (_despite_ their findings).\n- Do the authors' results imply that in any dataset, even without access to environment labels, training with SGD with batch size=1 allows access to the favorable results they present in Thm 1? This is due to a single sample inevitably belonging to a single environment.\n- L47: \"Learning invariant predictions produces reliable, fair, robust predictions against strong structural mechanism perturbation.\" Please be more clear and specific.\n- L53: Typo: \"may not necessarily in practice\"\n- L58: Please describe matrix sensing problem either in the main paper, or refer to a relevant appendix section\n- L61-64: The introductory exposition is more confusing then helpful since the notation is not sufficiently introduced.\n- Similar to above, Figure 1 is a lavish use of authors' space without a strong unique contribution\n- L140: Please fix typo (combine sentences)\n- L152: Total dimension of spurious signals or total dimension of core + spurious signals?\n- L199: That the spurious features have large dispersion across environments is an important assumption for the present paper. Please provide positive and negative hypothetical examples for this in realistic settings (in what cases it is reasonable to expect such large heterogeneity? In what cases it is not?)\n- L200: Missing ref.\n- L280: The terms HeteroSGD and PooledSGD are proposed early in the paper but later neglected, please consistently refer to the algorithms as such for readability", "limitations": "The limitations implied by the choice of matrix sensing setting is not sufficiently discussed.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper addresses an important issue, the problem of learning invariant solutions, and targets a particular, well defined sub-problem within this issue.\n- The assumptions that allow the authors' exposition are clearly stated and well-justified, and their results justify the claims made.\n- The authors' argumentation is easy to follow and their results are situated well within the relevant literature.", "weaknesses": "- The authors limit their analysis to the setting of matrix sensing (and a 2-layer NN that is constructed to conform to it). This is a reasonable choice to enable their analyses, however the justifications and implications of this choice should be discussed more thoroughly. \n- Although their overall argumentation is easy to follow, certain gaps and/or inconsistencies in their exposition present difficulties (see below for details).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewers for the valuable feedback and insightful comments. We have carefully considered your comments and questions and have addressed them as below:\n\n> About discussing the adoption of matrix sensing problem.\n\nA: We fully understand your concerns regarding our model. We adopt the matrix sensing problem because it is widely used in implicit regularization contexts, and its non-convexity and over-parameterized nature reflects the complexity of deep learning models. From this perspective, this model is sufficient for us to convey our insights. We are willing to generalizing our findings to more models in future work.\n\n> Do the authors' results imply that in any dataset, even without access to environment labels, training with SGD with batch size=1 allows access to the favorable results they present in Thm 1? This is due to a single sample inevitably belonging to a single environment.\n\nA: This is a very good point. Our results **do not imply** that the ``HeteroSGD`` succeeds when batch size=1, as our results relies on the satisfaction of the RIP condition. It is interesting to study how small batch sizes affect the invariance learning for SGD. This problem requires much more involved analysis to deal with the randomness of data when the RIP is broken down. We will discuss the main difficulty in the revised version and leave the analysis as a future work.\n\n\n> About the presentation issues.\n\nA: Thanks for pointing out these issues. We have carefully resolved the presentation issues in the revised version.\n\n----\n\nWe again thank you for dedicating your time and effort to reviewing our manuscript. Your valuable questions and insightful comments have significantly improved our work. We hope our responses have addressed all your concerns.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9kgVXNCDFa", "reviewer": "Reviewer_hg76", "summary": "This paper studies \"implicit invariance learning\" within a simplified but meaningful setting---multi-environment low-rank matrix sensing problem. Authors show the implicit bias of SGD over heterogeneous data drives the model learning towards an invariant solution. The key insight is, through simply employing the large step size large-batch SGD sequentially in each environment without any  explicit regularization, the oscillation caused by heterogeneity can provably prevent model learning spurious signals, while the model learned using pooled GD over all data would simultaneously learn both the invariant and spurious signals.", "questions": "To clarify, I am not very familiar with implicit regularization of SGD and I do not check carefully the proof. My confidence score should be low.\n\n---\n\nFrom my sensible review (assuming correctness of proof), this submission is a clear accept, so my review would be short.\n\n- regarding: \"In this paper, we will show that surprisingly, if each batch is sampled from data in one environment rather than data in all the environments, the heterogeneity in the environments together with the implicit regularization effects in the SGD algorithm can drive ...\"\n\n  This finding is interesting and potentially quite helpful to invariant learning. From my own experience with many DG problems, I like to collect training data from each environment in a large batch, and this simple method serves a good baseline and usually has some improvement over ERM. This observation was also summarized in a paper by other researchers (\"Simple data balancing achieves competitive worst-group-accuracy\", CLEAR 2022), but is different from here. Of course,  the above observation is only empirical and does not necessarily hold in general. Nevertheless, I would like to see some discussions with this observation regarding more general representation-learning based DG problems. I also hope this information can help authors futher consider similar theoretic studies in the more general DG setting, and I look forward to new findings in expanded version or new works.\n\n- I believe the setting of applying SGD to varying environment should have been considered in the context of federated learning. Can you also provide additional discussion in your related work part?\n\n- minor suggestion: \n\n  - I don't think it necessary to put a summary in the abstract.; \n  - line 22:  ”??\" appear after \"conditions:\".", "limitations": "see above.", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- A novel perspective studying the invariant learning problem over multi environments.\n\n- Rigorous theoretic results with detailed proof. The theoretic findings are potentially very interesting and helpful.\n\n- Presentation is very good, and I am enjoy reading through the submission", "weaknesses": "The considered setting is somewhat simplified, but that is fine for the first work from this new perspective.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewers for the valuable feedback and insightful comments. We have carefully considered your comments and questions and have addressed them as below:\n\n> This finding is interesting and potentially quite helpful to invariant learning. From my own experience with many DG problems, I like to collect training data from each environment in a large batch, and this simple method serves a good baseline and usually has some improvement over ERM. This observation was also summarized in a paper by other researchers (\"Simple data balancing achieves competitive worst-group-accuracy\", CLEAR 2022), but is different from here. Of course, the above observation is only empirical and does not necessarily hold in general. Nevertheless, I would like to see some discussions with this observation regarding more general representation-learning based DG problems. I also hope this information can help authors further consider similar theoretic studies in the more general DG setting, and I look forward to new findings in expanded version or new works.\n\nA: Thanks for sharing this point. These empirical experiences strongly motivate us to consider generalizing our findings to other distribution-agnostic settings or models. The mentioned paper utilizes group reweighting and subsampling. It is indeed an effective approach in invariant learning contexts.  We will add some additional discussion in the revised version.\n\n> I believe the setting of applying SGD to varying environment should have been considered in the context of federated learning. Can you also provide additional discussion in your related work part?\n\nA: Thanks for pointing out this. Indeed federated learning naturally possesses the multi-environment structure. We would like to add some discussion about the relation to the works in federated learning. Federated learning ([1,2]) is a machine learning paradigm where data is stored separately and locally on multiple clients and not exchanged, and clients collaboratively train a model. Extensive work has focused on designing effective decentralized algorithms (e.g. [1,3]) while preserving privacy (e.g. [4,5]). The importance of fairness in federated learning has also garnered attention ([6,7]). One important issue in federated learning is to handle the heterogeneity across the data and hardware. Our work shows that by training with certain stochastic gradient descent methods, the system can automatically remove the bias from the individual environment and thus learn the invariant features. Our work provides insights into discovering the implicit regularization effects of standard decentralized algorithms. More discussion and related work will be added in the revised version.\n\n[1] Communication-efficient learning of deep networks from decentralized data.\n\n[2] Advances and Open Problems in Federated Learning\n\n[3] Scaffold: Stochastic controlled averaging for federated learning. \n\n[4] Our data, ourselves: Privacy via distributed noise generation.\n\n[5] On the upload versus download cost for secure and private matrix multiplication\n\n[6] Ditto: Fair and robust federated learning through personalization. \n\n[7] Personalized Federated Learning towards Communication Efficiency, Robustness and Fairness\n\n----\n\nWe again thank you for dedicating your time and effort to reviewing our manuscript. We appreciate that you highly acknowledge our work, and your suggestions will greatly help us improve our work. We will also carefully revise the paper according to your minor suggestions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This submission investigates the implicit bias of SGD in multi-environment matrix sensing. All reviewers rated this submission positively. They appreciated that the submission is well-written and contains solid theoretical analysis. The theoretical findings are interesting and would enhance our understanding of invariance learning. The reviewers made several suggestions to improve the submission (e.g., additional experiments when the assumptions are violated). But, based on the discussion, these suggestions do not seem to warrant an additional round of review to reach an acceptance recommendation.
Decision: Accept (poster)