The Implicit Bias of Heterogeneity towards Invariance: A Study of Multi-Environment Matrix Sensing

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.

Strengths

- A novel perspective studying the invariant learning problem over multi environments. - Rigorous theoretic results with detailed proof. The theoretic findings are potentially very interesting and helpful. - 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.

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. --- From my sensible review (assuming correctness of proof), this submission is a clear accept, so my review would be short. - 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 ..." 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. - 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? - minor suggestion: - I don't think it necessary to put a summary in the abstract.; - line 22: ”??" appear after "conditions:".

Rating

7

Confidence

2

Soundness

4

Presentation

4

Contribution

3

Limitations

see above.

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.

Strengths

- The paper addresses an important issue, the problem of learning invariant solutions, and targets a particular, well defined sub-problem within this issue. - The assumptions that allow the authors' exposition are clearly stated and well-justified, and their results justify the claims made. - 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. - Although their overall argumentation is easy to follow, certain gaps and/or inconsistencies in their exposition present difficulties (see below for details).

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). - 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. - L47: "Learning invariant predictions produces reliable, fair, robust predictions against strong structural mechanism perturbation." Please be more clear and specific. - L53: Typo: "may not necessarily in practice" - L58: Please describe matrix sensing problem either in the main paper, or refer to a relevant appendix section - L61-64: The introductory exposition is more confusing then helpful since the notation is not sufficiently introduced. - Similar to above, Figure 1 is a lavish use of authors' space without a strong unique contribution - L140: Please fix typo (combine sentences) - L152: Total dimension of spurious signals or total dimension of core + spurious signals? - 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?) - L200: Missing ref. - 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

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The limitations implied by the choice of matrix sensing setting is not sufficiently discussed.

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.

Strengths

- The paper is clearly written, and the motivation is clear. - 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. - 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.). References: - [1] On the Generalization Benefit of Noise in Stochastic Gradient Descent - Smith et al, 2020

Questions

[1] Does the result in theorem 3 extend to pooledSGD as well? [2] Typo: Missing Reference - line 200

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

Yes

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.

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. 2. 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). 3. Simulations are run on synthetic examples. Are there any real datasets that you can demonstrate your methods on? 4. 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?

Questions

Some questions in weaknesses. Others below: 1. 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? 2. 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.

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes