The Implicit Bias of Heterogeneity towards Invariance: A Study of Multi-Environment Matrix Sensing
This provides a theoretical insight into how data heterogeneity and modern algorithms can implicitly promote invariance, potentially benefiting robustness and fairness in ML applications, though it is incremental as it builds on existing invariance learning frameworks.
The paper tackles the problem of invariance learning in multi-environment matrix sensing, showing that standard SGD with large step sizes and sequential training across heterogeneous data can implicitly drive models to invariant solutions, preventing spurious signal learning, unlike pooled SGD which learns both.
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.