# On The Concurrence of Layer-wise Preconditioning Methods and Provable Feature Learning

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 3 Feb 2025.

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

- [Preconditioning Benefits of Spectral Orthogonalization in Muon](https://mlsanity.com/paper/2601.13474/) (2026)
- [Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner](https://mlsanity.com/paper/2506.03595/) (2025)
- [Preconditioned Norms: A Unified Framework for Steepest Descent, Quasi-Newton and Adaptive Methods](https://mlsanity.com/paper/2510.10777/) (2025)
- [On the Mechanisms of Weak-to-Strong Generalization: A Theoretical Perspective](https://mlsanity.com/paper/2505.18346/) (2025)
- [Better Hessians Matter: Studying the Impact of Curvature Approximations in Influence Functions](https://mlsanity.com/paper/2509.23437/) (2025)
- [From Information to Generative Exponent: Learning Rate Induces Phase Transitions in SGD](https://mlsanity.com/paper/2510.21020/) (2025)
- [Design Criteria for SGD Preconditioners: Local Conditioning, Noise Floors, and Basin Stability](https://mlsanity.com/paper/2511.19716/) (2025)
- [Leveraging Second-Order Curvature for Efficient Learned Image Compression: Theory and Empirical Evidence](https://mlsanity.com/paper/2601.20769/) (2026)
- [Reassessing Muon for Matrix Factorization](https://mlsanity.com/paper/2607.13246/) (2026)
- [How Does Preconditioning Guide Feature Learning in Deep Neural Networks?](https://mlsanity.com/paper/2509.25637/) (2025)
- [Feature Learning in Linear-Width Two-Layer Networks: Two vs. One Step of Gradient Descent](https://mlsanity.com/paper/2605.17767/) (2026)
- [Spectral Asymptotics of Neural Network Loss Landscapes: An Exact Decomposition of the Curvature Exponent](https://mlsanity.com/paper/2606.02596/) (2026)

## Abstract (arXiv)

Layer-wise preconditioning methods are a family of memory-efficient optimization algorithms that introduce preconditioners per axis of each layer's weight tensors. These methods have seen a recent resurgence, demonstrating impressive performance relative to entry-wise ("diagonal") preconditioning methods such as Adam(W) on a wide range of neural network optimization tasks. Complementary to their practical performance, we demonstrate that layer-wise preconditioning methods are provably necessary from a statistical perspective. To showcase this, we consider two prototypical models, linear representation learning and single-index learning, which are widely used to study how typical algorithms efficiently learn useful features to enable generalization. In these problems, we show SGD is a suboptimal feature learner when extending beyond ideal isotropic inputs x ~ mathsfN(0, I) and well-conditioned settings typically assumed in prior work. We demonstrate theoretically and numerically that this suboptimality is fundamental, and that layer-wise preconditioning emerges naturally as the solution. We further show that standard tools like Adam preconditioning and batch-norm only mildly mitigate these issues, supporting the unique benefits of layer-wise preconditioning.
