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

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:

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.

A: 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.

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?

A: 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.

[1] Communication-efficient learning of deep networks from decentralized data.

[2] Advances and Open Problems in Federated Learning

[3] Scaffold: Stochastic controlled averaging for federated learning.

[4] Our data, ourselves: Privacy via distributed noise generation.

[5] On the upload versus download cost for secure and private matrix multiplication

[6] Ditto: Fair and robust federated learning through personalization.

[7] Personalized Federated Learning towards Communication Efficiency, Robustness and Fairness

---

We 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.

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:

About discussing the adoption of matrix sensing problem.

A: 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.

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.

A: 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.

About the presentation issues.

A: Thanks for pointing out these issues. We have carefully resolved the presentation issues in the revised version.

---

We 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.

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:

The problem of multi-environment matrix sensing seems under-motivated. Instances of problems and reasonability.

A: 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.

• 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$.

• 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.

[1] Glymour, M., Pearl, J., & Jewell, N. P. (2016). Causal inference in statistics: A primer. John Wiley & Sons.

About the intuition about HeteroSGD.

A: 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.

About the simulations?

A: 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.

[2] Simple data balancing achieves competitive worst-group-accuracy

Can one show that heterogeneity is necessary for learning invariant solutions?

A: 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.

About trying searching for different learning rates for GD vs SGD, and the interaction between batch size and learning rate.

A: 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.

In 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)$:

Does the separation holds when PooledGD is replaced by PooledSGD?

A: Yes. We will add the following theorem in the revised version of our paper.

Theorem 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,

$$\left\| \mathbf{U}_T\mathbf{U}_T^\top - \mathbf{U}^*{\mathbf{U}^*}^\top-\mathbf{V}^*{\mathbf{V}^*}^\top \right\|_F = o(1),$$

during which for all $t=0,1,\ldots, T$:

$$\left\| \mathbf{U}_t\mathbf{U}_t^\top - \mathbf{A}^*\right\|_F\gtrsim \sqrt {r_1\wedge r_2}.$$

The intuition is that the heterogeneity across environment is reduced. Hence the gradients will be very close to the gradients in PooledGD.

---

We 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.

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:

About the findings on HeteroSGD verses PooledGD? Does the result in theorem 3 extend to pooledSGD as well?

A: 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:

Theorem 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,

$$\left\| \mathbf{U}_T\mathbf{U}_T^\top - \mathbf{U}^*{\mathbf{U}^*}^\top-\mathbf{V}^*{\mathbf{V}^*}^\top \right\|_F = o(1),$$

during which for all $t=0,1,\ldots, T$:

$$\left\| \mathbf{U}_t\mathbf{U}_t^\top - \mathbf{A}^*\right\|_F\gtrsim \sqrt {r_1\wedge r_2}.$$

The 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.

About the feasibility of Theorem 2 and assumptions.

A: 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.

As 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$.

---

We 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.

Clarifications Needed:

• I thank the authors for their response and clarifications. Regarding the first point, could you please provide a quick clarification for the following statement:

We remark that the counterpart for small batch size case even for single-environment case, to the best of our knowledge.

• And just to confirm, for theorem 3.1 on PooledSGD, I presume the authors still require the large batch size for the RIP condition to hold/gradients to behave similar to PooledGD?

Overall, I am unsure about the extent of the applicability of the matrix sensing paradigm to standard heterogeneous learning scenarios. However, I do believe the paper offers novel theoretical insights. Thus, I opt to increase my score by 1 point.

As a minor suggestion, even if the theoretical analysis is too involved for this work, I would request the authors to include numerical simulations for cases when some of the assumptions (eg. batch size) are violated, in order to help the wider audience better understand any potential limitations of this perspective.