Research map / Architectures and efficiency / RNNs and normalization
On The Concurrence of Layer-wise Preconditioning Methods and Provable Feature Learning
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 SGD is a suboptimal feature learner under anisotropic inputs in linear representation learning and single-index models, and that a Kronecker-factored (K-FAC-like) layer-wise preconditioner provably fixes these issues.
Problem
Understanding why SGD fails to learn good features when inputs are anisotropic (non-isotropic covariance), and what optimizer provably fixes this in feature learning models.
Approach
Derives a stylized K-FAC update (left-preconditioning by F^T F and right-preconditioning by the input sample covariance) from first principles in linear representation learning and single-index learning, showing it removes the condition-number dependence and anisotropy-induced misalignment of SGD.
Key results
- Stylized K-FAC achieves a condition-number-free contraction factor (1 - 0.9 eta_G) in subspace distance, whereas SGD has a lower bound of (1 - 4 lmin/lmax)^T on its convergence.
- In single-index learning, the one-step K-FAC direction's correlation with the true direction tends to 1 as lambda_G -> 0 and dx/n -> 0, while one-step SGD correlation is (tr Sigma/dx)/sqrt(tr Sigma^2/dx), equal to 1 only for isotropic covariance.
- In linear representation experiments, K-FAC outperforms SGD, Adam, NGD, AMGD and DFW on subspace distance and transfer loss, and batch-norm does not fix AMGD's subspace distance.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace SGD with the stylized KFAC representation update | Linear representation learning · Linear representation learning | subspace distance: (1 - λ_min(F*ᵀF*)/λ_max(F*ᵀF*)) -> (1 - 0.9η_G) | Helped |
| Use unpreconditioned SGD instead of KFAC | Linear representation learning · Linear representation learning | subspace distance: (1 - 0.9η_G) -> (1 - λ_min(F*ᵀF*)/λ_max(F*ᵀF*)) | Hurt |
| Replace SGD with one-step KFAC | Two-layer network · Anisotropic single-index learning | direction correlation: tr(Σ_x)/√(d_x⁻¹ tr(Σ_x²)) -> 1 | Helped |
| Increase covariance anisotropy while using one-step SGD | Two-layer network · Single-index learning | direction correlation | Hurt |
| Increase covariance anisotropy while using one-step KFAC | Two-layer network · Single-index learning | direction correlation | No clear effect |
| Use KFAC instead of SGD, Adam, NGD, AMGD, or DFW | Linear representation learning · Anisotropic transfer-learning experiment | subspace distance | Helped |
| Add batch normalization to AMGD (ablation) | AMGD · Anisotropic linear representation learning | training loss | Helped |
| Add batch normalization to AMGD (ablation) | AMGD · Anisotropic linear representation learning | subspace distance | No clear effect |
| Whiten the training covariates (ablation) | SGD-based representation learning · Transfer learning with differing train and test covariances | transfer loss | Hurt |
| Set η_G to 1 (ablation) | Stylized KFAC · Linear representation learning | subspace distance | Hurt |
| Use Adam-like preconditioning instead of layer-wise preconditioning | Feature-learning methods · Linear representation learning | generalization | No clear effect |
Surprising findings
- Batch-norm/whitening stabilizes training loss convergence but converges to the wrong representation (rowspace of G* Sigma^{1/2}), hurting transfer.
- Full second-order methods (NGD) underperform K-FAC in convergence rate and stability in experiments.
Limitations
- Analysis restricted to two-layer linear representation learning and two-layer single-index models with Gaussian inputs
- Stylized K-FAC omits damping exponents, momentum, and other practical components
- Single-index results assume n, dx, dhid proportionally large limit and one-step updates
Experimental setup
Datasets: synthetic Gaussian linear representation data synthetic single-index data
Benchmarks: synthetic linear representation learning synthetic single-index learning
Metrics: subspace distance training loss transfer loss correlation with true direction
Baselines: SGD Adam natural gradient descent AMGD DFW batch-norm
Builds on: K-FAC Shampoo alternating minimization-SGD for linear representation learning one-step SGD feature learning (single-index)
Closest papers by meaning
- SGD with Adaptive Preconditioning: Unified Analysis and Momentum Acceleration (ICLR 2026 · similarity 0.89)
- Preconditioners for the Stochastic Training of Neural Fields (CVPR 2025 · similarity 0.87)
- Adam Reduces a Unique Form of Sharpness: Theoretical Insights Near the Minimizer Manifold (NeurIPS 2025 · similarity 0.87)
- Layerwise LQR for Geometry-Aware Optimization of Deep Networks (NeurIPS 2026 · similarity 0.86)
- Task-Specific Preconditioner for Cross-Domain Few-Shot Learning (AAAI 2025 · similarity 0.86)
- A Unified Stability Analysis of SAM vs SGD: Role of Data Coherence and Emergence of Simplicity Bias (NeurIPS 2025 · similarity 0.86)
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.