Research map / Optimization / Implicit bias and generalization
On the Parameterization of Second-Order Optimization Effective towards the Infinite Width
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.
Derives a maximal-update (muP)-style parameterization for second-order optimizers K-FAC and Shampoo via a one-step update analysis in the infinite-width limit, covering learning-rate and damping scaling and showing HP transfer across widths.
Problem
Hyperparameters (initialization scale, learning rate, damping) for second-order optimizers do not transfer across network widths under standard parameterization, so training larger models is unstable and requires re-tuning.
Approach
Extends the maximal update parameterization (muP) from first-order and entry-wise adaptive methods to K-FAC and Shampoo by analyzing the first one-step update with push-through identities, yielding width-scaling rules for learning rates and damping, plus a rescaled damping that is width-consistent for K-FAC.
Key results
- Under muP with K-FAC and rescaled damping, the optimal learning rate and damping remain fixed across widths, while SP shows shifting optima.
- Table of test accuracy shows muP gains over SP grow with width, e.g., ResNet18 on CIFAR100 with K-FAC at width 16 gains +4.00 points over SP.
- Zero-initialized last layer with K-FAC under muP-style analysis converges toward an NNGP solution; b_L=64 lowers K-FAC accuracy (e.g., 84.29 vs 82.63 at batch 64 for b_L=1 vs 64) while SGD is unaffected.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use μP with K-FAC (b_L=1.0) instead of SP (b_L=0.5) | K-FAC, 3-layer CNN · FashionMNIST test, batch size 1024 | accuracy: 80.25 -> 83.21 | Helped |
| Use μP with K-FAC (b_L=1.0) instead of b_L=0.5 | K-FAC, 3-layer CNN · FashionMNIST test, batch size 16 | accuracy: 84.17 -> 84.16 | No clear effect |
| Set b_L=64.0 instead of μP's b_L=1.0 (ablation) | K-FAC, 3-layer CNN · FashionMNIST test, batch size 1024 | accuracy: 83.21 -> 75.37 | Hurt |
| Set b_L=64.0 instead of b_L=1.0 (ablation) | SGD, 3-layer CNN · FashionMNIST test, batch size 1024 | accuracy: 73.40 -> 74.63 | Helped |
| Use μP instead of SP | K-FAC, VGG19 · CIFAR100 test, width 16 | accuracy change vs SP: +4.00 | Helped |
| Use μP instead of SP | Shampoo, VGG19 · CIFAR100 test, width 1 | accuracy change vs SP: -0.28 | Hurt |
| Use μP instead of SP | K-FAC, ResNet18 · CIFAR100 test, width 16 | accuracy change vs SP: +4.00 | Helped |
| Use μP instead of SP | Shampoo, ResNet18 · CIFAR100 test, width 8 | accuracy change vs SP: +1.04 | Helped |
| Use μP instead of SP | K-FAC, ResNet50 · ImageNet test, width 8 | accuracy change vs SP: +0.13 | Helped |
| Replace K-FAC damping heuristic with rescaled damping consistent with μP | K-FAC · CNN across widths | optimal damping transfer across widths | Helped |
| Use μP rather than SP for learning-rate scaling | K-FAC · models across widths | learning-rate transfer | Helped |
| Use the standard largest-eigenvalue damping heuristic | Shampoo · infinite-width parameterization | valid damping scale | No clear effect |
Surprising findings
- Zero initialization of the last layer can cause K-FAC to get stuck near an NNGP solution, an implicit bias specific to K-FAC, while SGD does not show this degradation.
- K-FAC does not require learning-rate width scaling under muP (c_l=0), unlike SGD.
Limitations
- Derivation is restricted to fully-connected networks in the infinite-width limit and a one-step update; depth scaling is not addressed.
- Transformers are not established; finite-width effects may matter.
- Analysis focuses on the valid regime where preconditioner eigenvalues match damping scale; Tensor Program theory does not yet cover second-order methods.
Experimental setup
Datasets: CIFAR10 CIFAR100 FashionMNIST WikiText2 ImageNet
Benchmarks: CIFAR10 CIFAR100 FashionMNIST WikiText2 word analogy ImageNet
Metrics: test accuracy training accuracy test loss
Baselines: standard parameterization SGD with muP K-FAC with heuristic damping
Builds on: maximal update parameterization K-FAC Shampoo Tensor Programs
Closest papers by meaning
- Exact, Tractable Gauss-Newton Optimization in Deep Reversible Architectures Reveal Poor Generalization (NeurIPS 2024 · similarity 0.90)
- Studying K-FAC Heuristics by Viewing Adam through a Second-Order Lens (ICML 2024 · similarity 0.89)
- $\boldsymbol{\mu}\mathbf{P^2}$: Effective Sharpness Aware Minimization Requires Layerwise Perturbation Scaling (NeurIPS 2024 · similarity 0.88)
- Deep Linear Network Training Dynamics from Random Initialization: Data, Width, Depth, and Hyperparameter Transfer (ICML 2025 · similarity 0.88)
- The Feature Speed Formula: a flexible approach to scale hyper-parameters of deep neural networks (NeurIPS 2024 · similarity 0.88)
- Scaling Exponents Across Parameterizations and Optimizers (ICML 2024 · similarity 0.87)
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.