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

arXiv: https://arxiv.org/abs/2403.01420 · HTML: https://mlsanity.com/paper/2403.01420/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 53 references that are in the map.
- First posted on arXiv 3 Mar 2024.

## Builds on (most cited first, 30 of 37)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Ditto: Fair and Robust Federated Learning Through Personalization](https://mlsanity.com/paper/2012.04221/) (2020)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability](https://mlsanity.com/paper/2103.00065/) (2021)
- [Implicit Regularization in Matrix Factorization](https://mlsanity.com/paper/1705.09280/) (2017)
- [The Risks of Invariant Risk Minimization](https://mlsanity.com/paper/2010.05761/) (2020)
- [Implicit Bias of Gradient Descent on Linear Convolutional Networks](https://mlsanity.com/paper/1806.00468/) (2018)
- [SGD on Neural Networks Learns Functions of Increasing Complexity](https://mlsanity.com/paper/1905.11604/) (2019)
- [Gradient descent aligns the layers of deep linear networks](https://mlsanity.com/paper/1810.02032/) (2018)
- [Directional convergence and alignment in deep learning](https://mlsanity.com/paper/2006.06657/) (2020)
- [Simple data balancing achieves competitive worst-group-accuracy](https://mlsanity.com/paper/2110.14503/) (2021)
- [Implicit regularization for deep neural networks driven by an Ornstein-Uhlenbeck like process](https://mlsanity.com/paper/1904.09080/) (2019)
- [Label Noise SGD Provably Prefers Flat Global Minimizers](https://mlsanity.com/paper/2106.06530/) (2021)
- [Does Invariant Risk Minimization Capture Invariance?](https://mlsanity.com/paper/2101.01134/) (2021)
- [Convergence of Gradient Descent on Separable Data](https://mlsanity.com/paper/1803.01905/) (2018)
- [What Happens after SGD Reaches Zero Loss? --A Mathematical Framework](https://mlsanity.com/paper/2110.06914/) (2021)
- [Invariant Causal Prediction for Block MDPs](https://mlsanity.com/paper/2003.06016/) (2020)
- [Shape Matters: Understanding the Implicit Bias of the Noise Covariance](https://mlsanity.com/paper/2006.08680/) (2020)
- [Understanding the Generalization Benefit of Normalization Layers: Sharpness Reduction](https://mlsanity.com/paper/2206.07085/) (2022)
- [Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruction](https://mlsanity.com/paper/2106.15013/) (2021)
- [The Implicit Bias of Depth: How Incremental Learning Drives Generalization](https://mlsanity.com/paper/1909.12051/) (2019)
- [ZIN: When and How to Learn Invariance Without Environment Partition?](https://mlsanity.com/paper/2203.05818/) (2022)
- [Understanding Causality with Large Language Models: Feasibility and Opportunities](https://mlsanity.com/paper/2304.05524/) (2023)
- [Nonlinear Invariant Risk Minimization: A Causal Approach](https://mlsanity.com/paper/2102.12353/) (2021)
- [Understanding Incremental Learning of Gradient Descent: A Fine-grained Analysis of Matrix Sensing](https://mlsanity.com/paper/2301.11500/) (2023)
- [(S)GD over Diagonal Linear Networks: Implicit Regularisation, Large Stepsizes and Edge of Stability](https://mlsanity.com/paper/2302.08982/) (2023)

## Built on by (most cited first, 2 of 2)

- [Escaping Local Minima Provably in Non-convex Matrix Sensing: A Deterministic Framework via Simulated Lifting](https://mlsanity.com/paper/2602.05887/) (2026)
- [Towards Understanding Momentum Acceleration in River-Valley Loss Landscape](https://mlsanity.com/paper/2609.30957/) (2026)

## Abstract (arXiv)

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.