Research map / Optimization / Implicit bias and generalization
The Implicit Bias of Heterogeneity towards Invariance: A Study of Multi-Environment Matrix Sensing
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows theoretically that large-step-size, large-batch SGD run sequentially over heterogeneous environments (HeteroSGD) implicitly learns the invariant low-rank signal in multi-environment matrix sensing, whereas pooled SGD learns invariant plus spurious components.
Problem
Whether standard training algorithms (SGD) on heterogeneous multi-environment data exhibit an implicit bias toward invariant solutions, when association and invariance differ, in over-parameterized low-rank matrix sensing.
Approach
Environment-sequential large-step-size SGD makes the spurious component oscillate multiplicatively across environments, so it is contracted and cannot grow, while the invariant component grows steadily; pooled SGD instead converges to the biased solution containing the spurious signal.
Key results
- HeteroSGD recovers the invariant solution with ||U_T U_T^T - A*||_F = O(log(d) sqrt(d/m)) under conditions on heterogeneity and step size, while the spurious signal is eliminated.
- Pooled SGD converges to U*U*^T + V*V*^T, with error at least order sqrt(r1 ∧ r2) from the invariant solution (negative result).
- Simulations show a phase transition in heterogeneity M around 5 and that spurious signal is eliminated when step size eta >= 0.05.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use sequential batches from individual environments instead of pooled batches | Over-parameterized matrix sensing with SGD · Multi-environment matrix sensing | Frobenius error to invariant signal: HeteroSGD: o_{\mathbb{P}}(1); PooledSGD: \Omega_{\mathbb{P}}(1) | Helped |
| Increase environment heterogeneity (ablation) | HeteroSGD · Matrix sensing simulation | Signal recovery proportion | Helped |
| Increase the step size (ablation) | HeteroSGD · Matrix sensing simulation | Signal recovery proportion | Helped |
| Use a small step size (ablation) | HeteroSGD · Matrix sensing simulation | Spurious signal recovery | Hurt |
Surprising findings
- Heterogeneity combined with large step size, rather than explicit regularization, removes the spurious signal; the paper notes over-parameterization may not be necessary.
Limitations
- Analysis restricted to matrix sensing with Gaussian/RIP measurements and a two-layer quadratic-activation example.
- Assumes equal singular values of the invariant signal in the main theorem; general case deferred to appendix.
- Requires strong conditions (RIP, diagonal dominance, small overlap of invariant and spurious subspaces).
Experimental setup
Metrics: signal recovery proportion Frobenius norm error to invariant solution
Baselines: PooledSGD
Builds on: implicit regularization of gradient descent for matrix sensing invariant risk minimization benign oscillation of SGD with large learning rate
Closest papers by meaning
- Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruction (NeurIPS 2021 · similarity 0.87)
- Towards Self-Supervised Covariance Estimation in Deep Heteroscedastic Regression (ICLR 2025 · similarity 0.86)
- Learning with Exact Invariances in Polynomial Time (ICML 2025 · similarity 0.86)
- TIC-TAC: A Framework For Improved Covariance Estimation In Deep Heteroscedastic Regression (ICML 2024 · similarity 0.85)
- Learning Conditional Invariances through Non-Commutativity (ICLR 2024 · similarity 0.85)
- Implicit Bias of Spectral Descent and Muon on Multiclass Separable Data (NeurIPS 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.