On The Concurrence of Layer-wise Preconditioning Methods and Provable Feature Learning
Authors: Thomas TCK Zhang
Behrad Moniri
Ansh Nagwekar
Faraz Rahman
Anton Xue
Hamed Hassani
Nikolai Matni
Venue: ICML 2025 poster (2025)
Keywords: feature learning
representation learning
preconditioning
single-index models
two-layer networks
non-convex optimization
matrix sensing
quasi-newton methods
kronecker-factored approximate curvature
shampoo
Primary area: theory->optimization
TL;DR: Kronecker-Factored preconditioning is gaining popularity as an Adam/SGD alternative. We provide concrete evidence of their usefulness by analyzing how they uniquely enhance feature learning.

Abstract: 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 $\mathbf{x} \sim \mathsf{N}(\mathbf{0}, \mathbf{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.

On The Concurrence of Layer-wise Preconditioning Methods and Provable Feature Learning
Thomas T. Zhang" " Behrad Moniri 1
Ansh Nagwekar' Faraz Rahman
Anton Xue
Hamed Hassani'
Nikolai Matni
Abstract
Adam-like optimizers theoretically, especially from a statis- tical (e.g. generalization) perspective. In fact, there exist many theoretical settings where Adam and similar methods underperform in convergence generalization relative to, e.g., well-tuned SGD (see e.g. Wilson et al. (2017); Keskar & Socher (2017); Reddi et al. (2018); Gupta et al. (2021); Xie et (2022); Dereich al. (2024)), further complicating a principled understanding the role of Adam-like optimizers in deep learning. Given these challenges, is there an alternative algorithmic paradigm that is comparable the Adam fam- ily in practice that is also well-motivated from statistical learning perspective? Encouragingly, in a recent large-scale deep learning optimization competition, AlgoPerf (MLCom- mons, 2024), Adam and its variants were outperformed in various "hold-out error per unit-compute" metrics by method known Shampoo (Gupta al., 2018), member of layer-wise "Kronecker-Factorcd family of precondi- tioners, formally described in Section 2, contrasted with "diagonal" preconditioning methods like Adam.
Layer-wise preconditioning methods are fam- ily of memory-efficient optimization algorithms that introduce preconditioners per axis of each layer's weight tensors. These methods have seen recent resurgence, demonstrating impressive performance relative entry-wise ("diagonal") preconditioning methods such as Adam(W) 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 statisti- cal 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 ef- ficiently learn useful features to enable general- ization. In these problems, we show SGD is a suboptimal feature learner when extending be- yond ideal isotropic inputs x N(O, I) and well- conditioned settings typically assumed in prior work. We demonstrate theoretically and numeri- cally 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 wise preconditioning.
Notable members of the Kronecker-Factoree precondition- ing family include Shampoo and KFAC (Martens Grosse, 2015). as well as their many variants and descendants. These algorithms motivated from aoproximationntthoreee perspective, aiming to approximate some ideal curvature matrix (e.g. the Hessian or Fisher Information) in way that mitigates the computational and memory challenges associated with second-orde -order algorithms such as Newton's Method (NM) or Natural Gradient Descent (NGD). However, towards establishing the benefit of these preconditioners, an approximation viewpoint is bottlenecked by our limited understanding of how the idealized second-orde -oner methods perform on neural-network learning tasks, even disregarding the computational considerations. It in fact remains unclear whether these second-order methods are inherently supe- rior to the approximations designed to emulate them. For example, recent work has shown that, surprisingly, KFAC generally outperforms its ideal counterpart NGD in conver- gence rate and generalization error typical deep learning tasks (Benzing, 2022). Thus, key question remains:
1. Introduction
Well-designed optimization algorithms have been enabler to the staggering growth and success of machine learning. For the broader ML community, the Adam (Kingma & Ba, 2015) optimizer is likely the go-to scalable and performant choice for most tasks. However, despite its popularity in practice, it has been notoriously challenging understand
'quual contribution 'University of Pennsylvania. Correspon- dence to: T. Zhang, B. Moniri <(tt22, bemoniri) @wwas..upennn...... seess.upen..du>>.
'To be contrasted with "optimization" perspective, e.g. guar- antees of convergence to a critical point of training objective. "See Dahl et al. (2023, Section 4.2) for details.
Proceedings of the nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Layer-wise Preconditioning Methods and Provable Feature Learning
How do we explain the performance of Kronecker- preconditioned optimizers?
high-dimensional proportional limit. We show that for anisotropic covariates x N(O, Ex), SGD fails to learn useful features, whereas it is known that it learns suitable features in the isotropic setting. Furthermore, we show that KFAC is a natural fix to SGD, greatly enhancing the learned features in anisotropic settings (Section 3.2).
In a seemingly distant area, the learning theory community has been interested in studying the solutions learned by ab- stractions of "typical" deep learning set-ups, where the over- all goal is to theoretically demonstrate how neural networks learn features from data to perform better than classical "fixed features" methods, e.g. kernel machines. Much of this line of work focuses on analyzing the performance of SGD on simplified models of deep learning (see e.g. Collins et al. (2021); Damian et al. (2022); Ba et al. (2022); Barak al. (2022); Abbe et al. (2023); Dandi et al. (2024a); Berthier et al. (2024); Nichani et al. (2024b); Collins et al. (2024)). Almost invariably, certain imnocuous-looking assumptions are made, such as isotropic covariates x ~ N(0, 1). Under these conditions, SGD has been shown to exhibit desirable generalization properties. However, some works deviate from these assumptions in specific settings (Amari et al., 2021; Zhang et al., 2024b), and suggest that SGD exhibit severely suboptimal generalization. Thus, toward extending our understanding of feature learning, it seems beneficial to consider broader family of optimizers. This raises the following question:
Lastly, we carefully numerically verify our theoretical predictions. Notably, we confirm the findings in Benzing (2022) that full second-order methods heavily underper- form KFAC in convergence rate and stability. We also show standard tools like Adam-like preconditioning and batch-norm (Ioffe & Szegedy, 2015) do not fix the issues we identify, even for our simple models, and may even hurt generalization in the latter's case.
In addition to the works discussed earlier, we provide exten- sive related work and background in Appendix A. Notation. We denote vector quantities by bold lower-case, and matrix quantities by bold upper-case. We use @ to denote element-wise (Hadamard) product, @ for Kronecker product, and vec(.) the balumn-major vectorization operator. Positive -eeni-ideeinite matrices are denoted by Q > (z) O, and the corresponding partial order P v Q
Q P = O. We use Il-llap- |I-|/r denote the operator (spectral) and Frobenius norms. We use (x) to denote the expectation of fxx), and P[A((x) to denote the probability of event A(x). Given a batch {x. we denote the empirical expectation [[]]
What is a practical family of optimization algorithms that overcomes the deficiencies of SGD for standard feature learning tasks?
Di=1 Given indexed set of vectors, we use the upper case to denote the (row- wise) stacked matrix, e.g. X 2 [x]
We answer the above two questions by focusing on two prototypical problems used to theoretically study feature learning: linear representation learning and single-index learning. In both problems, we show that SGD is clearly sub- optimal outside ideal settings, such as when the ubiquitous isotropic data N(O, I) assumption is violated. By inspect- show
x,] E Ruxdx We reserve E (E) for (sample) covariance matrices, e.g. _ We use X 2 to omit universal numerical constants, and standard asymptotic notation o(-), ((),) (()) Lastly, we use the index shorthand [n] = (l.....,1 and subscript + to denote the "next iterate", e.g. G = Nex
Ex ExxxT], Σx -[xxx]]
ing the root cause behind these suboptimalities, that Cronecker-Factordd preconditioners arise naturally first-principles solution to these issues. We provide novel motivations for this class of algorithms, while establishing new and improved learning- theoretic guarantees. We hope that this serves as strong evidence of an untapped synergy between deep learning optimization and feature learning theory.
noo-approximaaioonttooooo
2. Kronecker-Factored Approximation
One of the longest-standing research efforts in optimization literature is dedicated to understanding the role of (local) curvature toward accelerating convergence rates of optimiza- tion methods. An example is Newton's 's method, where the curvature matrix (Hessian) serves as preconditioner of the gradient, enabling one-shot convergence in quadratic opti- mization, in which gradient descent enjoys at best a linear convergence rate dictated by the conditioning of the problem. However, for high- -igh-dimensional variables, computing and storing the full curvature matrix is often infeasible. Thus en- ter Quasi-Newton and (preconditioned) Conjugate Gradient methods, where the goal is to reap the benefits of curvature under computational or structural specifications, such as diagoadiagoonal, low-rank, sparsity, etc.) constraints (e.g. BFGS family (Goldfarb, 1970; Liu & Nocedal, 1989; No- cedal & Wright, 1999)), or accessing the curvature matrix
Contributions.
We study the linear representation learning problem under general anisotropic x ~ N(O, Σx) covariates and show that the convergence of SGD can be drastically slow, even under mild anisotropy. Also, the convergence rate suffers an undesirable dependence on the "conditioning" of the instance even for ideal step-sizes. We arrive at a variant of KFAC as the natural solution to these deficiencies of SGD, giving rise to the first condition condition-number-feee number-free convergence rate for the problem (Section 3.1).
Next, we consider the problem of learning a single- index model using two-layer neural network in the
2

---

Layer-wise Preconditioning Methods and Provable Feature Learning
only through matrix-vector products (see e.g. Pearlmutter (1994); Schraudolph (2002); Martens (2010)).
3. Feature Learning via Kronecker-Pactored Preconditioning
Nevertheless, the use of these methods for neural net- work optimization introduces new considerations. Con- sider an I.-layer fully-connected neural network (omitting (W (x)....,, where W. E Rddd E L] and 0 E LL is the concatenation of 0, _ voo( W2), E [Z]. Firstly, establishing convergence of SGD (Arora et al., 2019), NGD (Zhang et al., 2019), or Gauss-Newton (Cai et al., 2019) (or their corresponding gradient flows) global minima of the training objective is non-trivial, optimization over 0 is non-convex. Moreover, these results do not directly characterize the structure of the resulting features learned by the algorithms. Secondly, on the practical front, full preconditioners on O require mem- ory O(L2d*), which grows prohibitively with depth and width. Block-diagonal approximations (where one curva- ture block M2 E Rd xd² corresponds to a layer 0₂) still require (TTF)). Thus, entry-wise preconditioning as in Adam, with footprint ((Ld)) = dim(@), ; is usually consid- ered the only scalable class of preconditioners.
We present two prototypical models of feature learning, lin- ear representation learning and single-index learning, and demonstrate how typical guarantees for the features learned by SGD break down outside of idealized settings. We then show how to rectify these issues by deriving a modified algo- rithm from first principles, and demonstrate that both cases in fact coincide with a particular Kronecker-factoree pre- conditioning method. We now set-up the model architecture and algorithm primitive considered in both problems. We consider two-layer feedforward neural network predictors:
biases) (x) - W, 1.-1...0
fr, = Fo FGx),
(2)
where F E [Rdvxd, G E Rddxxx denote the weight ma- trices and (-)) a predetermined activation function. For scalar outputs dy 1, we use fr, f o (Gx). For our purposes, we omit the bias vectors from both layers. We further denote the intermediate covariate pre- and post- activation h - Gx, 2 a(x)). We consider stan- dard mean-squaared-rror (MSE) regression objective and its (batch) empirical counterpart:
However, a distinct notion of Croonecker-Factored pre- conditioning emerged approximately concurrently with Adam, with representative examples such as KF AC and Shampoo. As its name suggests, since full biocka approximations are too expensive, a Kronecker-Factored approximation is made instead, where Vv, (Q Px)- ((0), Pr; Q2 O. Using properties of the Kronecker product (see Lemma E.1), this has the conve- nient interpretation of pre- and post-multiplyin the weights in their matrix form:
L(F, G) _ 2 (x [lly - fx)a(()]13]
(3)
Ž(F, G) & Elly fr. (x)((x)
Given a batch of inputs {x:}== we define the left and right preconditioners of the two layers (recall equation (I)):
Qe Σx = =E|xx 1,
QF En
E[zz J,
(4)
( Oh
T
OFF.G Oh
P. =
(or Ia,), Pr Lay.
(Q, 8 (p) Vo. = P Cw, (1)
We introduce the flexibility of Pa La, for when Pa does not play a significant role; notably, this recovers cer- tain Kronecker-Factored preconditioners that avoid extra backwards passes (see Appendix D). We consider a stylized alternating descent primitive, where we iteratively perform
As such, the memory requirement of Kronecker-Factored Se
layer-wise preconditioning is O [dd), matching that of entry-wise preconditioning. The notion of curvature dif- fers from case to case, e.g., for KFAC, this is the Fisher Information matrix corresponding to the distribution pa- rameterized by fx()), whereas for Shampoo this is the full-matrix Adagrad preconditioner, in turn closely related to the Gauss-Newton matrix.3 We provide some sample derivations and background in Appendix D. However, as aforementioned, an approximation viewpoint falls short of explaining the practical performance of Kronecker-Factored methods, as they typically converge faster than their cor- responding second-order method (Benzing, 2022) on deep learning tasks. This motivates understanding the unique benefits of layer-wise preconditioning methods from first principles, which brings us to the following section.
G = G - no PoP VoL(F, G) (Qa + Aa F+
1/ T.PP VV G4) G.) Q
where na. 7/2 > 0 are layer-wise learning rates, and A E R is a regularization parameter. In line with most prior work, we consider an alternating scheme for analytical conve- nience. We also assume that F. are computed on independent batches of data, equivalent to sample-splitting strategies in prior analysis (Collins et al. (2021); Zhang et al. (2024b); Ba et al. (2022); Moniri et al. (2024) etc).
The preconditioners (4) and update (5) bear a striking re- semblance to KFAC (cf. Appendix D). In fact, the precon- ditioners align exactly with KFAC if we view L(F, G) as a
*This is itself a positive-definite approximation of the Hessian.
3

---

Layer-wise Preconditioning Methods and Provable Feature Learning
negative log-likelihood of conditionally Gaussian model with fixed variance: fxx) N(fr, (x),,.a(),,)) o²T). This is in some sense a coincidence (and a testament to the pre- science of KFAC's design): rather than deriving the above preconditioners via approximating the Fisher Information matrix, we will show shortly how they arise as a natural adjustment to SGD in our featured problems. We note that Bonnecker-Factreed preconditioning methods often involve further moving parts such as damping exponents P- ad- ditional ridge parameters on various preconditioners, and momentum. Though of great importance in practice, they are beyond the scope of this paper, 4 and we only feature parameters that play a role in our analysis.
equivalent to determining its row-space rowsp G.). Re- covering F,,, G. is an ill-posed problem, as for any invert- ible L E R*** the matrices F, F,L, L GG, remain optimal. Therefore, we measure recovery of rowsp(G.) via a sub- space distance.
Definition 3.2 (Subspace Distance (Stewart & Sun, 1990)), Let G, G. E Rkxxx be matrices whose rows are orthonor- mal. Let P. E pdxxdx be the projection matrix onto rowsp (G.) Define the distance between the subspaces spanned by the rows of G and G. by
dist(G, GG G.) - |GPP, llop
convenient observation that unifies various stylized al- gorithms on two-layer networks is that the F-update in (5) can be interpreted as exponential moving average (EMA) over least-squares estimators conditioned on z.
The subspace distance quantitatively captures the align- ment between two subspaces, ranging between 0 (occur- ring iff rowsp(G.) = rowsp(G)) and 1 (occurring iff rowsp (G.) 1 rowsp(G)), We further make the follow- ing non-degeneracy assumptions.
Lemma 3.1. Given G, define the least-squares estimator:
Assumption 3.3. We assume G. is row-orthonormal, and F. is full-rank, rank(F.) = k S dy. This is without loss of generality: if & > dy, then recovering a k-dimensional row- space from y, is underdetermined. If rank(F.) .) = K' < k, then it suffices to consider G. E IR xdx
1 E[lly - Fo 1 = Y - ZZ as an EMA of F iss i.e., F+ = (1 1/1 F + The Fis.
Fis A arginn F
Given The E (0, 11. then the F-update in (5) can be re-written
The linear representation learning problem has often been studied in the context of multi-task learning (Du et al., 2021; Tripuraneni et al., 2020; Collins et al., 2021; Thekumpara- mpil et al., 2021; Zhang et al., 2024b).
In particular, many prior works (e.g. Collins et al. (2021); Nayer & Vaswani (2022); Thekumparampil et al. (2021): Zhang et al. (2024b)) consider an alternating "minimization- descent" approach, where out of analytical convenience F is updated by performing least-squares regression holding the hidden layer G fixed. In light of Lemma 3.1, this corre- sponds to the case where The = 1.
Remark 3.4 (Multi-task Learning). Multi-task learning considers data generated as y! F(!) G + e!") for distinct tasks t = I....,T, with the same goal of recovering the shared representation G.. Our al- gorithm and guarantees naturally extend here, see Ap- pendix B.3 for full details. In particular, by embedding F, = FFOTT lent to the "task-diversity" conditions in the above works: rank(F.) = rank C("
3.1. Linear Representation Learning
F(UT] Assumption 3.3 is equiva-
Assume we have data generated by the following process
x; l.d. N(O, Ex); y; = F,G.x, + +e;
(6)
k.
We maintain the "single-task" setting in this section for concise bookkeeping while preserving the essential fea- tures of the representation learning problem. Various al- gorithms have been proposed toward provably recovering the representation G.. A prominent example is al- ternating minimization-SGD scheme (Collins et al., 2021; Vaswani, 2024). In the cited works, a local convergence result is established for isotropic data Σx = Lax- In Zhang et al. (2024b), it is shown that using SGD can dras- tically slow convergence even under mild anisotropy; their proposed algorithmic adjustment equates to applying the rreondrrecon Qa = Σx- However, their local convergence result suffers a dependence the condition
where Ex is the input covariance, F, Rdrx G. E Rexdy are (unknown) rank-k matrices, and E, N(0, Σx) is additive label noise independent of all other randomness. We consider Gaussian data throughout the paper for conciseness; all results in this section can be ex- tended to subgaussian x, E via standard tools, affecting only the universal constants (see Appendix E). Let us define o 2 Amax (E.). Accordingly, our predictor model is two-layer linear net (2) with F E dy G E R
i.d.
The goal of linear representation learning is to learn the low-dimensional feature space that G. maps to, which is
WWe refer the interested reader to Ishikawa & Karakida (2023) for discussion of these settings in the -upditel-update parameteri- zation" framework (Yang& & Hu, 2021b).
'By local convergence here we mean dist (G, G.) is suffi- ciently (but non-vanishingly) small.

---

Layer-wise Preconditioning Methods and Provable Feature Learning
number of F., slowing the linear convergence rate for ill- conditioned F.. Let now specify the algorithm template used in this section, that also encompasses the above work:
Proposition 3.5. Let Σx Ldx nu xc. Choose any dx k; d, 2 k 2 2. Let the learner be given knowl- edge of F., G. and dist(Go, G.). However, assume the learner must fix aa > O before observing Go. Then, there exists F- E Rdyxk, G., Go E TRkkdx, such that Gy = SGD(GG) T)) satisfies:
G via (5), G = Ortho(G+), F+ via (5).
(8)
Notably, we row-orthonormalize the representation after each update. Besides case of analysis, we have observed this numerically mitigates the elements of G from blow-up when running variants of SGD. The alternating min-SGD al- gorithms in Collins et al. (2021); Vaswani (2024) are equiv- alent to iterating (8) setting P. Qe = I, 1/r = 1 in (4), whereas Zhang et al. (2024b) use Pa = I, Qa = Σx. ? = 1. Let us now write out the full-batch gradient update.
dist(GT, G.) 2 11-4 Amin (F. F.)
(F.F.) Amax (F.
T
4
dist(Go, G.).
The proof can be found in Appendix B.I. Since we set Σx = I, the lower bound also holds for the algorithm in Zhang al. (2024b). We remark departing from a worst- case analysis random initialization or varying step-sizes, is nuanced topic even for the simple case of convex quadratics; see e.g. Bach (2024); Altschuler & Parrilo (2024). In light of Proposition 3.5 and (9), a sensible alteration might be to pre- and post- multiply V VaL(F, G) by (F and Σx! These observations bring us to the proposed recipe in (5).
generic performance lower bound, e.g.
fresh batch of data ((x, and current weights G), we have the representation gradient and corresponding SGD step:
Full-Batch SGD. Given
(FGX T x VoÎ(F, G) = F
x
G = - l/c V C(F(F G).
Stylized KFAC. By analyzing the shorteomings of the SGD update, we arrive at the proposed representation up- date:
When x is isotropic Σx Idxe the key observation is that by multiplying both sides of (9) by Ptt, recalling Y F. P.G.X +ET we have
G - FF
G.P, = G GG (FG- - _ p. -(G
n
We can verify from (4) and (5) that Pa F and Qa Σx Thus, we have recovered stylized variant of KFAC as previewed. Our main result in this section is a local convergence guarantee.
Ta T x 2 (I& - Me F)
n
where the approximate equality hinges covariance con- centration Ex Idx and G.P 0. Therefore, in the isotropic setting, for sufficiently large n 2 dy, and appro- priately chosen 7/ ~ (F F.J. then (omitting many de- tails) we have the one-step contraction (Collins et al., 2021; Vaswani, 2024): defining A O (oevdu/n)
Theorem 3.6. Consider running (8) with da 0. na [0, 1]. and 1)F = 1. Define 72 & Fain (2x) As long as dist(G, 1 0.01x~1(Ex)) (F.) and 2 max (1, (1) (dg 1og(115)) we have:
dist(G G.) S (1 0.9m ddsst G.).G.)) nn
dist(G+, G.) (1- Amin FF. G.) (10) Amaa F.)
with probability 2 1 1-6 8, where A 2 O TV dxttog(1/88)
Therefore, in low-noise/lurge-btee settings, this demon- strates SGD the representation G converges geometri- cally to G. (in subspace distance). However, there are clear suboptimalities to SGD. Firstly, the above analysis critically relies on Σx= = Lax such that G, E.P. x 0. As afore- mentioned, this is demonstrated to be crucial in Zhang et al. (2024b) for dist G.) to converge using SGD. Secondly, the convergence of SGD is bottlenecked by the conditioning of F.. In fact, we show the dependence on F, in the con- traction rate bound (10) cannot be improved in general, even under the most benign assumptions. Following Collins et al. (2021): Vaswani (2024), we define Gr = SGD( Go; Tl, T) as the output of alternating min-SCD, i.e. iterating (8) set- ting P. = Q. = I, 7 = 1 in (4), for T steps with fixed step-size a/ starting from Go-
Crucially, our contraction factor is condition-number-eeeee subverting the lower bound in Proposition 3.5 for suffi- ciently ill-conditioned F.. Therefore, setting Vaa near 1 ensures a universal constant contraction rate. Curiously, our proposed stylized KFAC (8) aligns with an alternat- ing "min-min" scheme (Jain al., 2013; Thekumparampil et al., 2021), where F, G are alternately updated via solv- ing the convex quadratic least-squares problem, by setting 1/c 1. However, our experiments (see Figure 5) demonstrate ee = is generally suboptimal, highlighting the flexibility of viewing KFAC as a descent method.
1. TRANSFER LEARNING
The upshot of representation learning is the ability to trans- fer (e.g. fine-tune) to a distinct, but related, task by only
5

---

Layer-wise Preconditioning Methods and Provable Feature Learning
retraining F (Du et al., 2021; Kumar et al., 2022). Assume now have target data generated by:
(2021a); Moniri et al. (2024); Moniri & Hassani (2024a),
etc.) that ensures the covariates x; alone do not carry any information about the target direction.
(t) lid. N(O, )( yi F..) G.X!"
F.DGG. + ei (11)
Assumption 3.8. The vector B. is drawn from G. (( 'xx) independent of other sources of randomness.
where E; Ltd. N(O, Σe), F E Rdyxk Notably, G. is
shared with the "training" distribution (6). Given G (e.g. by running (8) on training task), we consider fitting the last layer F given a batch of (1( data from the target task (11).
In this section, we study the problem of fitting a two-layer feedforward neural network ft.G for prediction of unseen data points drawn independently from (12) at test time, When G is kept at random initialization and f is trained using ridge regression, the model coincides with a random features model (Rahimi & Recht, 2007: Montanari et al., 2019; Hu & Lu, 2023) and has repeatedly used as toy model to study and explain various aspects of practical neural networks (see Lin & Dobriban (2021); Adlam & Pennington (2020); Tripuraneni et al. (2020); Hassani & Javanmard (2024); Bombari et al. (2023); Lee et al. (2023a); Bombari & Mondelli (2024a:b), etc.).
Lemma 3.7. Let F(2) = argming [("[)ly( Fz 12).
(()) A Gx(() be the optimal F on the batch of n(() tar- get data (11) given G. Defining v= dist(G, G.). given
((r) 2k k + log(1/0), we have with probability 2 1 8:
(F (t) (), [1x(t)) - xx
< || Amax (D) + o2(d,k + log(1/6))
n!==)
When the covariates are isotropic Σx = Lax. it is shown that single step of full- SGD update on G can drastically improve the performance of the model over random features as result of feature learning by aligning the top right- singular-vector of the updated representation layer G with the direction G. (Damian et al., 2022; Ba et al.. 2022; Moniri et al., 2024; Cui et al., 2024; Dandi et al., 2024a;b;c). In this section, we assume that the covariates are anisotropic and show that in this case, the one-step full batch SGD is suboptimal and can learn an ill-correlated direction even when the sample size is large. We then demonstrate that the KFAC update with the preconditioners from (4) is in fact the natural fix to the full batch SGD,
As hoped, the MSE of the fine-tuned predictor decom- poses into a bias term scaling with the quality of G. and noise term scaling with dim(F)) (F)/(()). We comment the required data is x k rather than dx resulting from do- ing regression from scratch (Wainwright, 2019). Addi- tionally, the noise term scales with (F) = dy rather than dydx of the full predictor space. The transfer learn- ing set-up (11) also reveals why data normalization (e.g. whitening, batch-norm (loffe & Szegedy, 2015)) can be counterproductive. To illustrate this, consider perfectly whitening the training covariates v= = Σx 1/2x. x By this change of variables, the ground-truth predictor changes y 2 F.G.x = F.G.E.² v. This is unproblematic so far-in fact, since the covariates v are isotropic, SGD now may converge. However, instead of rowsp(G,) the representation now converges to rowsp G.E, Ex/3). De- ploying on the target task, since Ex # Σx we have rowsp(G) & (Gwsp(G, x²² + rowsp( G. ( In other words, in return for stabilizing optimization, normal- izing the data destroys the shared structure of the predictor model! We illustrate this effect in Figure 3.
Full-Batch SGD. Following the prior work, at initializa- tion, we set f = d 11/2f with fo ~ N(O, d, Tay), and G ==Go Go with i.i.d. N(0, 1) entries. We update G with one step of full batch SGD with step size Te = n i.c.
Go nVdh VoL(fo, Go).
SGn
In the following theorem, we provide an approximation of the updated first layer G which is generalization of (Ba al., 2022, Proposition 2.1) for Σx + Lax-
3.2. Single Index Learning
Theorem 3.9. Assume that the activation function a is (1)--iipshhitz and that Assumption 3.8 holds. In the limit where n, d2y du lend to infinity proportionally, the matrix GoeD. with probability 1 - o(1), satisfies
Assume that we observe n i.i.d. samples generated accord- ing to the following single-index model:
X, j.idd. N(O, Ex): 3/4 = (3.x,(3Tx,)+ +
(12)
||Go + on anfo.B3.
GGooo ll.p - 0,
where Ex E dx is the input covariance, o. : R + R is the teacher activation function, B. dRdx is the (un- known) target direction, and E; is an additive noise E ~ N(O, o²) independent of all other sources of randomness. We also make the following common assumption on ,,. (Dicker (2016); Dobriban& Wager (2018); Tripuraneni et al.
in which a = E, [=)] with = ~ N(O, d Tr(xx)), and the vector Bseu is given by BoeD = n X y.
i.d.d.
This thcorem shows that one step of full batch SCD update approximately adds a one component a: To to the
6

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Subspace Distance
1.00
6x:00
Train Loss
4x10
0.75
1PV CVC
Thansfer Loss
4x:00
0.50
2x30
0.25
10"
100
500
1000
1500
2000
500
1000
1500
2000
500
1000
1500
2000
Number of iterations
Number of iterrtios
Number of Iterations K
Figure 1. From left to right: the training loss, subspace distance, and transfer loss induced by various algorithms on a linear representation learning task. We note that various algorithms converge in training loss, but negligibly in subspace distance, and thus transfer loss.
initialized weights Go. Thus, the pre-activation features for a given input x E RRxx after the update are given by
Stylized KFAC. This time, we update G using the styl- ized KFAC update from (5) with the regularized Pa- We use the same initialization as full-batch SGD. The updated representation layer in this case is given by
hasc = GGGD se Gox + a (g fo E 12ch
where the first and second term correspond to the random feature, and the learned feature respectively. To better un- derstand the learned feature component, note that defin- ing C,,1 = E ==-N(0.47 TT(E,)) (od, (=)], the target function a, (B. x) can be decomposed as
Guzac - Gu nvd Vo Go) (Qc + Aatax)
The preconditioning factor (Qa + Aa with QG Σx is precisely the factor required so that the direction learned by the one-step update match the ridge regression estimator with ridge parameter À as shown in the following immediate corollary of Theorem 3.9.
x.(,,x))c.,1333 = x to
(3. (5
Corollary 3.11. Under the same set of assumptions as Theorem 3.9, the matrix GREAC, satisfies
satisfying Ex [B., T. (3, x)] = 0. Therefore, when C,,1 * 0, the target function has a linear part. Full batch SGD is estimating the direction of B. using this lin- ear part with the estimator ASED
Go + aon fog - GREAC
X y/n. However, the natural choice for this task is in fact ridge regression Bx (Ex Aaxxaa)-11 X -X y/n, and BoGo is missing the prefactor (Exx + -axx) In the isotropic case Σx Lax. we expect Σx R: Iax when n >> dx. Thus, in this case the estimator Bac is roughly equivalent to the ridge estimator and can recover the direction B.. However, in the anisotopic case, Asco is biased even when n >> dx. To make these intu- itions rigorous, we characterize in the following proposition the correlation between the learned direction 3scn and the true direction B..
with probability 1- o1), where a is defined in Theorem 3.9, and BaFAc (Qa AaTax) - y/n
Bsa.
Because BxPac is equivalent to ridge regression, expect it to align well with 3. even for anisotropic Σx, given a proper choice of Aa. The following lemma formally characterizes the correlation between BREAC and B., for any Xa E TR.
Lemma 3.12. Under the ussumptions of Theorem 3.9, the correlation between B, and Brenc satisfies
B BkFaC B.
Cx,1 V
Lemma 3.10. Under the assumptions of Theorem 3.9, the correlation between B, and Baco satisfies
B 1|211/2111
C U/2 + dx
+2)
with probability 1 - o(1), where 2. =E(o(()))))) ct E (=)) with = ~ N(O, Tr(Ex)), and 11, V2, V are defined in (34) and depend on Ex, dx/n,, and Xa In particular, as No - and dx/n n - 0, we have
B. 9900
(xxx
dx
0
|scD 1 1211,3111
cl+02 Tr [x) +
E.a
Tr Tr(E2)
dx
with probability 1- o(1), in which C..1 = Ex(o' and 2² [z)] with N(O, d Tr(Ex)).
B 3./ ll 11. + 1.
This lemma shows that when n >> dx. and A - 0, the one- step stylized KFAC update- unlike the one-step full- pero-peffectly recovers the target direction G., fixing the issue with full batch SGD with anisotropic covariances.
This lemma shows that the correlation increasing in the strength of the linear component C, while keeping the sig- nal strength c, fixed. Also, based on this lemma, when n >>
dx. the correlation is given by d₂ Tr(Ex)/ / Tr which is equal to one if and only if Σx o² Iax for some E R. This means these are the only covariance matrices for which applying one step of full batch SGD update learns the correct direction of B..
Remark 3.13. It well-known that, given features that align with 3., applying least-squares on Z o (Gurae X), which from Lemma 1 is equivalent to the KFAC f-update with
M) = 1, leverages the feature to obtain a solution with good generalization. See Appendix C.4 for more details.

---

Layer-wise Preconditioning Methods and Provable Feature Learning
0.82
0.950
Correlation
0.80
0.79
0.900
KFAC Theary
KFAC Theory
0.77
5GD Theory
0.850
SCD Theory
SGD Simulation
SGD Simulation
0.76
KFAC Simulation
KFAC Simulation
6.75
0.800
log(Ac)
Degree of Anisotropy (e)
Figure 2. The correlation of the direction leamed by SGD and KFAC with the the true direction by numerical simulations averaged over 30 trials, and theoretical predictions. (Left) For different values of λa the theoretical predictions match the simulations very well. (Right) The alignment of the feature learned by SGD deteriorates as anisotropy is increased (larger s), whereas the KFAC update remains accurate.
4. Numerical Validation
1.00
100
Subspace Distance
4.1. Linear Representation Learning
0.75
6x100
Train Loss
We numerically study the behavior of different algorithms for transfer learning setting (11), where the model is to be trained on data generated by (Ftrain G.), and the transfer task has data generated by (Ftest, G,), i.e. the embedding G. is shared, but the task heads Ftrain and Fisst are dif- ferent. The training and test covariates have anisotropic covariance matrices Ex,train and Exttest respectively. Our data generation process for the training task and the transfer task are follows:
APGO ISbiT
0.50
4x100
0.25
2x10
0.00
500
1000
1500
2000
Number of Iterations k
Figure 3. Subspace distance and the training loss of KFAC and AMGD (with and without batch-norm). Notably, batch-norm en- ables AMGD's 's train loss to converge. but not its subspace distance.
y = F + tej
xi lid. DV/ Unif({=1)dk),
4.2. Single-Index Learning
o N(O, o Iav), S (test, train),
Lid. ~
(13)
Consider the single-index function learning setting of Sec- tion 3.2 with (=) + 2=22 - -1), and ar
where triiii 0.1 and Te,test 1. We use dg = 100, dy = 15, k 8, and batch size 1024. We present additional experiments and details in Appendix F, including discussions on the learning rates, and how Fs,, G Σx, are precisely generated.
Different Levels of Anisotropy. In this experiment, we set d = 200, = 6000, and dp = 1000 and set X + 0. For parameter R, we define Σx Σ with
diag( [ E,...
E, E). (14)
Head-to-head Evaluations. We track the training loss, subspace distance, and transfer loss of different algorithms during the update (Figure 1). Alongside SGD, KFAC, Adam, and NGD, we also consider Alternating Min-SGD (AMGD) (Collins et al., 2021; Vaswani, 2024), and De-bias & Feature- Whiten (DEW) (Zhang et al., 2024b) (corresponding to (5) with P.- I,y), two algorithms studied in linear represen- tation learning. The transfer loss is the loss incurred by fitting least-squares Fiest on the current G iterate (see Lemma 3.7). Although various algorithms converge s train, KFAC outperforms all others in terms of sub- space distance and transfer loss, as suggested by the theory.
dx/2
dx/2
For different values of e, we simulate the KFAC and SGD updates numerically and compute their correlation with the true direction. We also theoretically predict the correlation using Lemma 3.10 and 3.12; see Figure 2 (Right). The SGD update fails to recover the true direction in highly anisotropic settings (large E), whereas the one-step KFAC update remains accurate.
Theory vs. Simulations. We set d, 900, 5000, dn 1000, and Σx (00-)) For different λς we simu- late the correlation between the directions learned by KFAC and SCD with the true direction and compare it with predic- tions of Lemma 3.10 and 3.12; see Figure 2 (Left). We see that the theoretical results match very well with numerical simulations, even for moderately large n, dy, and dn. The direction learned by KFAC has a larger correlation with the true direction compared to that learned by SGD, as predicted.
Effect of Batch Normalization. We track the subspace distance and the training loss of AMGD (with and without batch-norm) and KFAC, see Figure 3. As theoretically pre- dicted in Section 3.1.1, since batch-norm approximately whitens x|rain, MMGD++batc-noom converges in training loss. However, as predicted, it does not recover the correct repre- sentation, whereas KFAC does.
8

---

Layer-wise Preconditioning Methods and Provable Feature Learning
5. Discussion
In Advances in Neural Information Processing Systems, 2020.
We study two models of feature learning in which we iden- tify key issues of SGD-based feature learning approaches when departing from ideal settings. We then present Kronecker-Factored precondiiiooingg-rrevveienn of KFAC-to to provably overcome these issues and derive im- proved guarantees. Our experiments these simple models also confirm the suboptimality of full second-order methods, as well as the marginal benefit of Adam preconditioning and data normalization. We believe that analyzing properties of statistical learning problems can lead to fruitful insights into optimization and normalization schemes.
Altschuler, J. and Parrilo, P. Acceleration by stepsize hedg- ing: Multi-step descent and the silver stepsize schedule. Journal of the ACM, 2024.
variants
Amari, S.-i., Ba, J., Grosse, R. B., Li, X., Nitanda, A., Suzuki, T., Wu, D., and Xu, J. When does preconditioning help or hurt generalization? In International Conference on Learning Representations, 2021.
Amid, E., Anil, R., and Warmuth, M. Locoprop: Enhancing backprop via local loss optimization. In International Conference on Artificial Intelligence and Statistics, 2022.
Acknowledgments
Anil, R., Gupta, V., Koren, T., Regan, K., and Singer, Y. Scalable second order optimization for deep learning. arXiv preprint (eXii:2002.09015 2020.
Thomas Zhang and Behrad Moniri gratefully acknowledge gifts from AWS Al to Penn Engineering's ' ASSET Center for Trustworthy AI. The work of Behrad Moniri and Hamed Hassani supported by The Institute for Learning-enale Optimization Scale (TILOS), under award number NSF- CCF-2112665, and the NSF CAREER award CIF- 1943644 Thomas Zhang and Nikolai Matni are supported in part by NSF Award BLES-2331880, NSF CAREER award ECCS- 2045834, NSF HECS- 2231349, and AFOSR Award FA9550- 24-1-0102.
Arnaboldi, L.,, Dandi, Y., Krzakala, F., Pesce, L., and Stephan, L. Repetita iuvant: Data repetition allows SGD to learn high-dimensional multi-inder funetions. arXiv preprint arXiv:2405.15450 2024.
Arora, S., Cohen, N., Hu, W., and Luo, Y. Implicit regu- larization in deep matrix factorization. In Advances in Neural Information Processing Systems, 2019.
Impact Statement
Ba, J., Grosse, R. B., and Martens, J. Distributed second- order optimization using Kronecker- factorrd approxima- tions. In International Conference on Learning Represen- tations, 2017.
The focus of this paper is on theoretical aspects of opti- mization and feature learning. We expect the results to be illuminating for the optimization and learning theory com- munity. We do not anticipate any negative societal impact.
Ba, J., Erdogdu, M. A., Suzuki, T., Wang. Z., Wu, D., and Yang. G. High-dimensional asymptoties of feature learn- ing: How one gradient step improves the representation. In Advances in Neural Information Processing Systems, 2022.
References
Abbasi-Yaddes Yadkori, Y. and Szepesvári, C. Regret bounds for the adaptive control of linear quadratic systems. In Con- ference on Learning Theory, 2011.
Ba, J., Erdogdu, M. A., Suzuki, T., Wang. Z.,, and Wu, D. Learning in the presence of low-dimensional structure: spiked random matrix perspective. In Advances in Neural Information Processing Systems, 2024.
Abbe, E., Adsera, E. B., and Misiakiewicz, T. The merged- staircase property: a necessary and nearly sufficient con- dition for SGD learning of sparse functions two-layer neural networks. In Conference on Learning Theory, 2022.
Bach, F. URL
Scaling laws of optimization, 2024. https://Erancisbact: com/
Abbe, E., Adsera, F. B., and Misiakiewicz, T. SGD learning on neural networks: leap complexity and saddle-to- toaleettossee saddee dynamics. In Conference on Learning Theory, 2023.
Bai, Y. and Lee, J. D. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. In International Conference on Learning Representations, 2020.
Absil, P.-A,, Mahony, R., and Sepulchre, R. Optimization algorithms on matrix manifolds. Princeton University Press, 2008
Barak, B., Edelman, B., Goel, S., Kakade, S., Malach, E., and Zhang, C. Hidden progress in deep learning: SGD learns parities near the computational limit. In Advances Neural Information Processing Systems, 2022.
Adlam, B. and Pennington, J. Understanding double de- scent requires fine-grained bias variance decomposition.

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Ben Arous, G., Gheissari, R., and Jagannath, A. Online stochastic gradient descent non-convex losses from high-dimensional inference. Journal of Machine Learn- ing Research, 22(106):1-51, 2021.
Collins, L., Hassani, H., Soltanolkotabi, M., Mokhtari, A., and Shakkottai, S. Provable multi-task representation learning by two-layer ReLU neural networks. In Interna- tional Conference on Machine Learning, 2024.
Benzing, F. Gradient descent on neurons and its link to approximate second-order optimization. In International Conference on Machine Learning, 2022.
Cui, H., Pesce, L., Dandi, Y., Krzakala, F., Lu, Y., Zde- borova, L., and Loureiro, B. Asymptotics of feature learning in two-layer networks after one gradient-step. In International Conference on Machine Learning. 2024.
Bernstein, J. and Newhouse, I.. Modular duality in deep learning. arXiv preprint arXiv:2410.2/265,, 2024a.
Dahl, G. E., Schneider, F,, Nado, Z., Agarwal, N., Sastry, C. S., Hennig, P., Medapati, S., Eschenhagen, R., Kasim- beg, P., Suo, D., al. Benchmarking neural network training algorithms. arXiv preprint arXiv:2306.07179 2023.
Bernstein, J. and Newhouse, L. Old optimizer, new norm: An anthology. arXiv preprint arXiv:2409.20322. 2024b.
Berthier, R., Montanari, A., and Zhou, K. Learning time- scales in two-layers neural networks. Foundations of Computational Mathematics, pp. 1-84, 2024.
Damian, A., Lee, J., and Soltanolkotabi, M. Neural net- works can learn representations with gradient descent. In Conference on Learning Theory, 2022.
Bollapragada, R., Nocedal, J., Mudigere, D., Shi, H.-J., and Tang, P. T. P. A progressive batching I-bfgs method for machine learning. In International Conference on Machine Learning, 2018.
Dandi, Y., Krzakala, F., Loureiro, B., Pesce, L., and Stephan, L. How two-layer neural networks learn, one (giant) step at a time. Journal of Machine Learning Research, 25 (349):1-65, 2024a.
Bombari, S. and Mondelli, M. How spurious features are memorized: Precise analysis for random and NTK fea- tures. In International Conference on Machine Learning, 2024a.
Dandi, Y., Pesce, L., Cui, H., Krzakala, F., Lu, Y. M., and Loureiro, B. A random matrix theory perspective on the spectrum of learned features and asymptotic gener- alization capabilities. arXiv preprint 2rXiv:2410.18938 2024b.
Bombari, S. and Mondelli, M. Privacy for free in the over- parameterized regime. arXiv preprint arXiv:2410. 14787, 2024b.
Dandi, Y.. Troiani, E., Arnaboldi, L.,, Pesce, L., Zdeborova, L., and Krzakala, F. The benefits of reusing batches for gradient descent in two-layer networks: Breaking the curse of information and leap exponents. In International Conference on Machine Learning, 2024c.
Bombari, S., Kiyani, S., and Mondelli, M. Beyond the universal law of robustness: Sharper laws for random features and neural tangent kernels. In International Conference Machine Learning, 2023.
Dangel, F., Harmeling, S., and Hennig, P. Modular block- diagonal curvature approximations for feedforward ar- chitectures. In International Conference on Artificial Intelligence and Statistics, 2020.
Botev, A., Ritter, H., and Barber, D. Practical Gauss-Netwon optimisation for deep learning. In International Confer- ence on Machine Learning, pp. 557-565, 2017.
Dereich, S., Graeber, R., and Jentzen, A. Non-convergence of Adam and other adaptive stochastic gradient descent optimization methods for non-vanishing learning rates.
Byrd, R. H., Hansen, S. L., Nocedal, J., and Singer, Y. A stochastic quasi-Networ method for large-scale optimiza- tion. SIAM Journal Optimization, 26(2): 2016.
arXiv preprint arXiv:2407.08100 2024.
Dicker, L. H. Ridge regression and asymptotic minimax estimation over spheres of growing dimension. Bernoulli, pp. 1-37, 2016.
Cai, T., Gao, R., Hou, J., Chen, S., Wang. D., He, D., Zhang, Z., and Wang. L.. Gram-Gauss- Gauss- Netwon method: Learning overparameterized neural networks for regression prob- lems. arXiv preprint
Dobriban, E. and Wager, S. High-dimensional asymptotics of prediction: Ridge regression and classification. The Annals of Statistics, 46(1):247-279, 2018.
2019.
Collins, L., Hassani, H., Mokhtari, A., and Shakkottai, S. Exploiting shared representations for personalized feder- ated learning. In International Conference on Machine Learning, 2021.
Du, S. S., HU, W., Kakade, S. M., Lee, J. D., and Lei, Q. Few-shot learning via learning the representation, provably. In International Conference on Learning Rep- resentations, 2021.
10

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Duchi, J., Hazan, E., and Singer, Y. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(7), 2011.
Hassani, H. and Javanmard, A. The of over- parametrization in adversarial training: Precise analysis of robust generalization for random features regression. The Annals of Statistics, 52(2):44-440 2024.
Frerix, T., Möllenhoff, T., Moeller, M., and Cremers, D. Proximal backpropagation. In International Conference on Learning Representations. 2018.
Horn, R. A. and Johnson, C. R. Matrix analysis. Cambridge university press, 2012.
Fu, H., Wang. Z., Nichani, E., and Lee, J. D. Learning hier- archical polynomials of multiple nonlinear features with three-layer networks. arXiv preprint arXiv:2411.17201, 2024.
Hsu, D., Kakade, S. M., and Zhang, T. Random design analysis of ridge regression. In Conference on Learning Theory, 2012.
Hu, H. and Lu, Y. M. Universality laws for high-dimensional learning with random features. IEEE Trunsactions on Information Theory, 69(3), 2023.
Ghorbani, B., Mei, S., Misiakiewicz, T., and Montanari, A. Linearized two-layers neural networks in high dimension. The Annals of Statistics, 49(2): 2021a.
loffe, S. and Szegedy, C. Batch normalization: accelerating deep network training by reducing internal covariate shift. Proceedings of the 32nd International Conference on International Conference on Machine Learning Volume 37, pp. 448-456, JMLR.org, 2015.
Ghorbani, B., Mci, S., Misiakiewicz, T., and Montanari, A. When do neural networks outperform kernel methods? Journal of Statistical Mechanics: Theory and Experiment, 2021(12), 2021b.
Goldfarb, D. A family of variable-metric methods derived by variational means. Mathematics of computation, 24 (109):23-26, 1970.
Ishikawa, and Karakida, R. On the parameterization of second-order optimization effective towards the infinite width, arXiv preprint arXiv:2312. 2023.
Goldfarb, D., Ren, Y., and Bahamou, A. Practical quasi- Netwon methods for training deep neural networks. In Ad- vances Neural Information Processing Systems, 2020.
Jacot, A., Gabriel, F., and Hongler, C. Neural tangent ker- nel: Convergence and generalization neural networks. In Advances in Neural Information Processing Systems, 2018.
Goldt, S., Loureiro, B., Reeves, G., Krzakala, F., Mézard, M., and Zdeborová, L. The Gaussian equivalence of gen- erative models for learning with shallow neural networks. In Mathematical and Scientific Machine Learning, pp. 426-471, 2022.
Jain, P.,, Netrapalli, P., and Sanghavi, Low-rank matrix completion using alternating minimization. In ACM Sym- posium Theory of Computing, pp. 665-674. 2013.
Jordan, K., Jin, Y., Boza, V., Jiacheng, Y., Cesista, F., New- house, L., and Bernstein, J. Muon: An optimizer for hidden layers in neural networks, 2024. URL https:
Guionnet, A., Ko, J., Krzakala, F., Mergny, P., and Zde- borová, L. Spectral phase transitions in non-linear wigner spiked models. arXiv preprint arXiv:2310. 14055, 2023.
/ //kellerjordart git io http/posss/////e http///sssssms.....////
Gupta, A., Ramanath, R., Shi, J., and Keerthi, S. S. Adam vs. SGD: Closing the generalization gap on image clas- sification. In OPT2021: 13th Annual Workshop on Opti- mization for Machine Learning, 2021.
Keskar, N. S. and Socher, R. Improving generalization performance by switching from Adam to SGD. arXiv preprint arXiv:1712.07622 2017.
Gupta, V., Koren, T., and Singer, Y. Shampoo: Precondi- tioned stochastic tensor optimization. In International Conference on Machine Learning, 2018.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Hanin, B. and Nica, M. Finite depth and width corrections to the neural tangent kernel. In International Conference on Learning Representations, 2020.
Kumar, A., Raghunathan, A., Jones, R. M., Ma, T., and Liang, P. Fine-tuning can distort pretrained features and underperform cittof-distribution. In International Con- ference on Learning Representations,, 2022.
Hanson, D. L. and Wright, F. T. A bound on tail probabili- ties for quadratic forms independent random variables. The Annals of Mathematical Statistics, 4279):177-11083, 1971.
Large, T., Liu, Y., Huh, M., Bahng, H., Isola, P., and Bern- stein, Scalable optimization in the modular norm. arXiv preprint arXiv:2405. /48/3, 2024.

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Laurent, B. and Massart, Adaptive estimation of quadratic functional by model selection. The Annals of Statistics, pp. 1302-1338, 2000.
MLCommons. Announcing the results of the inaugural
AlgoPerf: Training algorithmms benchmark competition,
2024. m m.--aaoopper--becccaa---oooooopo wtw.eeccccaa.cccmmmtmiiooooo.oo
org/2024/088
Lee, D., Moniri, B., Huang, X., Dobriban, E., and Hassani, H. Demystifying disagreemeent-on-the--lle in high dimen- sions. In International Conference on Machine Learning. 2023a.
Moniri, B. and Hassani, H. Asymptotics of linear re- gression with linearly dependent data. arXiv preprint arXiv:2412.03702, 2024a.
Lee, J. D., Oko, K., Suzuki, T., and Wu, D. Neural net- work learns low-dimensional polynomials with SGD near the thforeition-theoreticc limit. arXiv preprint arXiv:2406.01581, 2024.
Moniri, and Hassani, H. Signal-plus-noise decomposition of nonlinear spiked random matrix models. arXiv preprint arXiv: arXiv:2405. 20..2245528277 18274, 2024b.
Moniri, B., Lee, D., Hassani, H., and Dobriban, E. theory of non-linear feature learning with one gradient step in two-layer neural networks. In International Conference on Machine Learning. 2024.
Lee, Y., Chen, A. S., Tajwar, F., Kumar, A., Yao, H., Liang, P., and Finn, C. Surgical fine-tuning improves adaptation to distribution shifts. In International Conference on Learning Representations, 2023b.
Montanari, A., Ruan, F., Sohn, Y., and Yan, J. The gen- eralization error of max-margin linear classifiers: High- dimensional asymptotics in the overparametrized regime. arXiv preprint arXiv: 1.1154444,44 2019.
X.-L. Preconditioned stochastic gradient descent. IEEE Iransactions neural networks and learning systems, 29 (5):144-1146, 2017.
Li, X.-L. Preconditioner on matrix lie group for sgd. arXiv preprint 100.::1800.102.22 2018.
Morwani, D., Shapira, I,, Vyas, N., Malach, E., Kakade, S.. and Janson, L. A new perspective shampoo's preconditioner. arXiv preprint arXiv:2406. 17748, 2024.
Lin, I.. and Dobriban, What causes the test error? going beyond bias-variance via ANOVA. Journal of Machine Learning Research, 22:155-1, 2021.
Mousavi-Hosseini, A., Wu, D., Suzuki, T., and Erdogdu, M. A. Gradient-basee feature learning under structured data. In Advances in Neural Information Processing Systems, 2023.
Lin, W., Dangel, F., Eschenhagen, R., Bae, J., Turner, R. E., and Makhzani, A. Can we remove the square-root in adaptive gradient methods? second-order perspective. In Forty-first International Conference on Machine Learn- ing, 2024.
Nakhleh, J., Shenouda, J., and Nowak, R. D. The effects of multi-task learning ReLU neural network functions. arXiv preprint arXiv:2410.. 40.. 2666, 2024.
Liu, D. C. and Nocedal, J. On the limited memory bigs method for large scale optimization. Mathematical pro- gramming, 45(1):503-555 1989.
Nayer, and Vaswani, N. Fast and sample-efficient feder- ated low rank matrix recovery from column-wise linear and quadratic projections. IEEE Transactions Infor- mation Theory, 69(2):1177-1202 2022.
Martens, Deep learning via Hessian free optimization. In International Conference on Machine Learning, 2010.
Nichani, E., Damian, A., and Lee, J. Provable guaran- tees for nonlinear feature learning in three-layer neural networks. In Advances in Neural Information Processing Systems, 2024a.
Martens, J. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1-76, 2020.
Nichani, E., Damian, A., and Lee, J. D. How transformers learn causal structure with gradient descent. In Interna- tional Conference on Machine Learning, 2024b.
Martens, J. and Grosse, R. Optimizing neural networks with Kronecker-factored approximate curvature. Interna- tional Conference on Machine Learning, 2015.
Nocedal, J. and Wright, S. Numerical optimization. Springer, 1999.
Maurer, A., Pontil, M., and Romera-Parede B. The benefit of multitask representation learning. Journal of Machine Learning Research, 17(81):1-32, 2016.
Pearlmutter, B. A. Fast exact multiplication by the hessian. Neural computation, 6(1):147-160, 1994.
Mci, S. and Montanari, A. The generalization error of random features regression: Precise asymptotics and the double descent curve. Communications on Pure and Applied Mathematics, 75(4)6667-766, 2022.
Pooladzandi, O. and Li, X.-L. Curvature-informat sgd via general purpose lie-group preconditioners. arXiv preprint arX5v:2402.04553, 2024.
12

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Rahimi, A. and Recht, B. Random features for large-scale kernel machines. In Advances in Neural Information Processing Systems, 2007.
Tripuraneni, N., Adlam, B., and Pennington, J. Overpa- rameterization improves robustness covariate shift in high dimensions. In Advances in Neural Information Processing Systems, 2021a.
Reddi, S. J., Kale, S., and Kumar, S. On the convergence of Adam and beyond. In International Conference Learning Representations, 2018.
Tripuraneni, N., Jin, C., and Jordan, M. Provable meta- learning of linear representations. In International Con- ference on Machine Learning, 2021b.
Rudelson, M. and Vershynin, R. Manson-wright inequality and sub-gaussian concentration. Electronic Communica- tions in Probability, 18:1-9, 2013.
Troiani, E., Dandi, Y., Defilippis. L., Zdeborová, L., Loureiro, B., and Krzakala, Fundamental limits of weak learnability in nigh-dimensional multi-index mod- els. arXiv preprint arXiv:2405./5444 15480, 2024.
Schmidt, R. M., Schneider, F,, and Hennig, P. Descend- ing through crowded valley- beneemannnc deep learn- ing optimizers. In International Conference on Machine Learning, 2021.
Vaswani, N. Efficient federated low rank matrix recovery via alternating GD and minimization: A simple proof. IEEE Transactions on Information Theory, 2024.
Schraudolph, N. N. Fast curvature matrix-vecte products for second-order gradient descent, Neural computation, 14(7):123-1138, 2002.
Vershynin, R. Introduction to the non-asymptotic analysis of random matrices. In Eldar, Y. C. and Kutyniok, G. (eds.), Compressed Sensing: Theory and Applications, pp. 210-268. Cambridge University Press, 2012.
Shi, H.-J. M., Lee, T.-H., Iwasaki, S., Gallego-Posade J., Li, Z., Rangadurai. K., Mudigere, D., and Rabbat, M. A distributed nata-parallel pytorch implementation of the distributed Shampoo optimizer for training neural net- works scale. arXiv preprint arXiv:2309.0644 2023.
Vershynin, R. nigh-dimensional probability: An introduc- tion with applications in data science, volume 47. Cam- bridge University Press, 2018
Vyas, N., Morwani, D., Zhao, R., Shapira, I,, Brandfon- brener, D., Janson, L., and Kakade, S. M. Soap: Improv- ing and stabilizing Shampoo using Adam. In 2024: Optimization for Machine Learning. 2024.
Shi, Z., Wei, J., and Liang, A theoretical analysis on feature learning in neural networks: Emergence from imputs and advantage over fixed features. In International Conference on Learning Representations, 2022.
Wainwright, M. J. High-dimensional statistics: A non- asymptotic viewpoint, volume 48. Cambridge university press, 2019.
Silverstein, J. W. and Choi, S.-.. Analysis of the limiting spectral distribution of large dimensional random matri- ces. Journal of Multivariate Analysis, 54(2):295-309, 1995.
Wang. Z., Nichani, E., and Lee, J. D. Learning hierarchi- cal polynomials with three-layer neural networks. In International Conference on Learning Representations, 2024.
Stewart, G. W. and Sun, g.gg Matrix perturbation theory. Academic press, 1990.
Thekumparampil, K. K., Jain, P., Netrapalli, P., and Oh, S. Sample efficient linear meta-learning by alternating minimization. arXiv preprint arXiv:2105.0833 2021.
Wilson, A. C., Roelofs, R., Stern, M., Srebro, N., and Recht, B. The marginal value of adaptive gradient methods in machine learning. In Advances in Neural Information Processing Systems, 2017.
Tieleman, T. and Hinton, G. Lecture 6.5 6.5-rmsprop: Divide the gradient by running average of its recent magnitude. Coursera: Neural Networks for Machine Learning, 4(2): 26, 2012.
Woodbury, M. Inverting modified matrices. Department of Statistics, Princeton University, 1950.
Xie, Z., Wang. X., Zhang, H., Sato, I., and Sugiyama, M. Adaptive inertia:
Trefethen, L. N. and Bau, D. Numerical linear algebra, volume 181. SIAM, 2022.
Disentangling the effects of adaptive learning rate and momentum. In International Conference on Machine Learning, 2022.
Tripuraneni, N., Jordan, M., and Jin, C. On the theory of transfer learning: The importance of task diversity. In Advances in Neural Information Processing Systems, 2020.
Yang, G. and Hu, E. J. Feature learning in infinite- width neural networks. In International Conference on Machine Learning, 2021a.
13

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Yang, G. and Hu, J. Tensor programs iv: Feature learn- ing in wwith neural networks. In International Conference Machine Learning, 2021b.
Zhang, G., Martens, J., and Grosse, R. Fast convergence of natural gradient descent for over-parameterized neural networks. In Advances in Neural Information Processing Systems, 2019.
Zhang. T. T., Lee, B. D., Ziemann, I., Pappas, G. J., and Matni, N. Guarantees for nonlinear representation learn- ing: Non-identical covariates, dependent data, fewer plcs. In International Conference Machine Learning, 2024a.
Zhang, T. T., Toso, L. F., Anderson, J., and Matni, N. Sample-efficient linear representation learning from non- IID non-isotropic data. In International Conference Learning Representations, 2024b.
Ziemann, I., Tsiamis, A., Lee, B., Jedra, Y.. Matni, N., and Pappas, G. J. A tutorial on the non-asymptotic theory of system identification. In IEEE Conference on Decision and Control, 2023.
14

---

Layer-wise Preconditioning Methods and Provable Feature Learning
A. Extended Background and Related Work
Preconditioners for Neural Network Optimization. A significant research effort in neural network optimization has been dedicated to understanding the role of preconditioning in convergence speed and generalization. Perhaps the most widespread paradigm falls under the category of entry-wise ("diagonal") preconditioness,, whose notable members include Adam (Kingma & Ba, 2015), (diagonal) AdaGrad (Duchi al., 201 I), RMSprop (Tieleman & Hinton, 2012), and their innumerable relatives and descendants (see e.g. Schmidt et al. (2021); Dahl et al. (2023) for surveys). However, diagonal preconditioners inherently do not fully capture parem-rrreeeter dependencies, which are better captured by stronger curvature estimates, e.g. Gauss- Newton approximations (Botev et al., 2017; Martens, 2020), L-BFGS (Byrd et al., 2016; Bollapragada et al., 2018; Goldfarb et al., 2020). Toward making non- non-diagonal preconditioners scalable to neural networks, many works (including the above) have made use of layer-wise Kronecker- Kronecker-Fuctored approximations, where each layer's curvature block factored into a Kronecker product Q @ P. Perhaps the two most well-known examples are Kronecker-Factore Approximate Curvature (KFAC) (Martens & Grosse, 2015) and Shampoo (Gupta et al., 2018; Anil et al., 2020), where approximations made to the Fisher Information and ideas, such as by improving practical efficiency (Ba et al., 2017; Shi et al., 2023; Jordan et al., 2024; Vyas et al., 2024) and exploring generalized constructions (Dangel et al., 2020; Amid et al., 2022; Benzing, 2022; Li, 2017; 2018; Pooladzandi & Li, 2024). An interesting alternate view subsumes certain preconditioners via steepest descent with respect to ayer-wise ("modular") norms (Large et al., 2024; Bernstein & Newhouse, 2024a;b). We draw a connection therein by deriving the steepest descent norm that Cronecker-Factoree preconditioners correspond to; see Appendix D.3.
curvature, respectively. Many works have since expanded these
Multi-task Representation Learning (MTRL). Toward a broader notion of generalization, the goal of MTRL is to characterize the benefits of learning shared representation across distinct tasks. Various works focus the generalization properties given access to empirical risk minimizer (ERM) (Maurer et al., 2016; Du et al., 2021; Tripuraneni et al., 2020; Zhang et al., 2024a), with the latter work resolving the setting where distinct tasks may have different covariate distributions. Closely related formulations have been studied in the context of distribution shift (Kumar et al., 2022; Lee et al., 2023b). While these works consider general non-linear representations, access to an ERM obviates the (non-convex) optimization component. As such, multiple works have studied algorithms for linear representation learning (Tripuraneni et al., 2021b; Collins et al., 2021: Thekumparampil et al., 2021; Nayer & Vaswani, 2022) and specific non-linear variants (Collins et al., 2024; Nakhleh al., 2024). In contrast the ERM works, which are mostly agnostic the covariate distribution, all the listed algorithmic works assume isotropic covariates N(O, I). Zhang et al. (2024b) show that isotropy is in fact key enabler, and propose an adjustment to handle general covariances. In this paper, we show that many prior linear representation learning algorithms belong to the same family of (preconditioned) optimizers. We then propose an algorithm coinciding with KFAC that achieves the first conition-nuumer-fiee convergence rate.
Nonlinear Feature Learning. In the early phase of training, neural networks are shown to be essentially equivalent to the kernel methods, and can be described by the neural tangent kernel (NTK). See Jacot et al. (2018); Mei & Montanari (2022); Hu & Lu (2023). However, kernel methods are inherently limited and have sample complexity superlinear in the input dimension d for learning nonlinear functions (Ghorbani et al., 202 a0)). The main reason for this limitation is that kernel methods use set of fixed features that are not task specific. There has been a lot of interest in studying the benefits of feature learning from a theoretical perspective (Bai & Lee (2020); Hanin & Nica (2020); Yang & Hu (2021a): Shi et al. (2022); Abbe et al. (2022), etc.). In a setting with isotropic covariates N(O, I), it is shown that even a one-step of SGD update on the first layer of a two-layer neural networks can learn good enough features to provide a significant sample complexity improvement over kernel methods assuming that the target function has some low-dimensional structure (Damian et al., 2022: Ba et al., 2022; Moniri et al., 2024; Cui et al., 2024: Dandi et al., 2024a;bxe; Arnaboldi et al., 2024; Lee al., 2024) and this has became very popular model for studying feature learning. These results were later extended to three-layer neural networks in which the first layer is kept at random initialization and the second layer is updated using one step of SGD (Wang et al., 2024; Nichani et al., 2024a; Fu et al., 2024). Recently, Ba et al. (2024); Mousavi-Hossein et al. (2023) considered an anisotropic case where the covariance contains a planted signal about the target function and showed that a single step of SGD can leverage this to better learn the target function. However, the general case of anisotropic covariate distributions remains largely unexplored. In this paper, we study feature learning with two-layer neural networks with general anisotropic covariates in fanll--inex, models and that one-step of SGD update has inherent limitations in this setting, and the natural fix will coincide with applying the KFAC laier-wise preconditioner.
15

---

Layer-wise Preconditioning Methods and Provable Feature Learning
B. Proofs and Additional Details for Section 3.1
B.1. Convergence Rate Lower Bound SGD
Our goal is to establish the following lower bound construction.
k2 2. Let the learner be given knowledge of F., G. and dist Go. G.). However, assume the learner must fix 7/ > before observing Go- Then, there exists F. E Ryyyk G., Go E Rexd, such that Gr = SGD Go: T)) satisfies:
Proposition 3.5. Let Ex = Lax. n = co. Choose any dx > k, dy k
(1-4
Amin F. F.) mmax FFFF.)
T
dist GT, G+)
dist G.).
Proof of Proposition 3.5. We prove the lower bound by construction. First, we write out the one-step SGD update given step size Ia.
G G 13 VoC(F.G)
1 F [FGX X
G
Y
F (FG F.G.).
(YT = =FG.X +ET = 0c, Ex I)
G Ortho(G+)..
We recall F is given by the F-update in (5) with !) 1, which by Lemma 3.1 is equivalent to setting F to the least-squares solution conditional on G:
F Y - (GE,G))
(z Gx, n oo)
F.G.G' (G
(G row-orthonormal)
Therefore, plugging in into the SGD update yields:
G = G 11a F (FG-F,G.)
G Ta GG.
Before proceeding, let us present the construction of F.; G.. We focus on the case dy 3, 2, will be clear the
construction is trivially embedded to arbitrary d2 > k: 2 2. We observe that F, only appears in the SGD update in the form F. E k thus dy 2 k can be set arbitrarily long F F. satisfies specifications. Set F., G. such that
1 A 0
0 X
1 0 0 AE (0,1/2). G.
F.F.
0
0
Accordingly, the initial representation Go (which the learner is initially given) will have form
0
V1
=3
80 0
G =
V1
Eo
or
We prove all results with the first form of Go. as all results will hold for the second with the only change swapping X, 1 A. is clear that we may extend to arbitrary dx > 2 by setting:
(1
a 0 0 A
TA-1 0
F F.F.
G. = [ Odx-
Go =
0
1
Eo
0
Returning to the dy 3 case, we first prove the following invariance result.
[
0
Lemma B.1. Given Go
then for any t 2 0, G
for some +f 1. Furthermore,
V1
Eo
C1
C2
we have dist(Gr, G.) 102
16

---

Layer-wise Preconditioning Methods and Provable Feature Learning
1 0
Proof of Lemma B. This follows by induction. The base case follows by definition of Go. Now given G. for some C12 cz. we observe that
0
c1
1 0 F.F.
1
G.G.
0 X
0
Notably, we may write
G G. 1o G.G. G.G. (F,GG, F.G.)
[Ig ?a G.G.F.FF..GGG
G. + Ma G.G.G.GG
I - no 1
1-1
c.c.
G. +o
no
G.
ad
(15)
0 ai + (1++)
al1-7oGfA)
G+41 Ortho(G.+1).
Therefore, G shares the same support as G and by the orthonormalization step, the squared entries of the second row of G:+1 equal 1, completing the induction step.
To prove the second claim, we see that P.- = 002
1 ] and since G, is by assumption row-orthonormal, we have
dist Gr,, G.) = ||. = [ o c2] ] | =
completing the proof.
With these facts in hand, we prove the following stability limit of the step-size, and the consequences for the contraction rate.
Lemma B.2. no 2 I-x thenfor any given dist (Go; G.) we may find Go such that lim sup, dist (G., G.) 2 1 2
Proof of Lemma B.2. By assumption
< 1/2 and thus S 1 A. Evaluating Lemma B.1 instead on Go
0 Eo writing out (15) yields symmetrically:
1
dist G.) = |02|
G [c - 7Tc3(1 )))
0
c2(1 noef(1 - A))
0
1
G Orthh
We first observe that regardless of 7a. the norm of the first row G+41 is always greater than 1 pre-orthonormalization. Let us define w Va (1 A). Then, the squared-norm of the first row satisfies:
(c+
(1(+))))) w + (c wei)) =1+
Therefore, the norm is strictly bounded away from when > 0 and either C1 * 0 by the constraint of + cτ 1. importantly, this implies that regardless of the step-size taken, the resulting first-row norm of G must exceed 1 prior to orthonormalization Given this property, we observe that for w 2 1/c2, we have:
weel
wef -1
+ we
1 we
When this ratio is greater than 1, we are guaranteed that the first-row coefficients oi ch of G post-orthonommaization satisfy <//c/ > c2/01, and recall from Lemma B.I dist(Gr41, G.) c₂ and thus dist (ast(G,1, G.) > dist (G G.). Rearranging the above ratio, this is equivalent to the condition
na (1 X) > 2
no A) 1/c
Setting oî 3/4, o 1/4, this implies for 7/6 2 -- the moment dist G. G.) S 02 1/2, then are guaranteed dist(Gr+1, G.) > dist Gr G.), and thus lim sup, dist dist(G,, G.) 2 2, irregardless of dist dist(Go, G.).
17

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Now, finish the construction of the lower bound, Lemma B.2 establishes that llc
T is necessary for convergence
(though not sufficient!). This implies that when we plug back in Go =
we have G
V % 60
0 G++1 0 c1 + 7c (222
02
1
0
G141 OrthhoGG++1 =
0
di
d2
We have trivially 1 - n 2 1- naA. Therefore, for À I 1/5 such that - S 1/4, we have 1 - 11 As shown in the proof of Lemma B.2, the norm of the second row pre-orthonoormalizatinn is strictly greater than 1, and thus:
dist(Gr+1: G.) = c2 2 c2 - - 2 (1 - maA)dist(G G.) > (1 1-4 12
dist(G G.).
--X
Applying this recursively to G. yields the desired lower bound.
B.2. Proof of Theorem 3.6
Recall that running iteration of stylized KFAC (5) with Xr, Ao 0r yields:
G G 7e P VaL(F, G) (Qc + datax)
G 1a (F F) (FG - F,G. Ex 1 X)
(16)
G Na (F F) (FG F.G.) + (F FF F
X(X
where the matrix F is given by
F = F prey TPPP VrC(F prev, G) (Qr + ApIa,)
Y Z(Z Z)-1
= =F.G.X z(z Zj + E Z(Z Z) -1
recalling that z A Gx. Focusing on the representation update, we have
G_P. = (1 - 7a GP, + Ma (F' F) X(X X)-1
Therefore, to prove one-step contraction of rowsp G+) toward rowsp( G.), we require two main components:
Bounding the noise term Va FF E X X)-1
Bounding the orthonormalization factor; the subspace distance measures distance between two orthonormalized bases (a.k.a. elements of the Stiefel manifold (Absil et al., 2008)), while a step of SGD or KFAC does not inherently conform to the Stiefel manifold, and thus the "off-manifold" shift must be considered when computing dist(G., G.). This amounts to bounding the "R" "R"-factor of the QR-decomposition (Trefethen & Bau, 2022) of G
Thanks to the deft-reeconditioning by F)-1, the contraction factor essentially determined by (1 - Na); however, the second point about the "off ooff-manifold" shift is what prevents us from setting 7k
Bounding the noise term
We start by observing T T X = exx] xxxx)x,,) (and thus & is least-
squares error-like term, and thus can be bounded by standard self-normalized martingale arguments. In particular, defining
F A F) FF we may decompose
X < X XXx Amin ( (xTx)-1/2
where the first factor is the aforementioned eeff-normalized martingale (see e.g. Abbasi- Yadkori & Szepesvári (2011); Ziemann et al. (2023)). and the second can be bounded by standard covariance lower-tail bounds. Toward bounding the first factor, we invoke a high-pun high-probbbility self elf-noorallized bound:
18

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Lemma B.3 (ef. Ziemann al. (2023, Theorem 4.1)). Let (v,, wa)121 be IRd x Rev-valued process and (Fitiz1 be filtration such that (vi)121 is adapted (Fi)e2, {w:)221 is adapted (Fi) and (wi)121 is a o² -ubgauussian martingale difference sequence.. Fix (non-randon) pasitive-definiee matrix Q. For l 2 1, define Σg - Eis v.v. Then, given any fixed N4, with probability at least 6:
det (Q (Q+Ein)
)
n
-1/2
Σ w.v. (Q+
S 2 log
+ 1. + 80 log(1/8).
(17)
det(QQ)
Instantiating this for Gaussian W; Lid. Ew). V: Lid. N(O, Ev), may set Q Ey to yield:
Lemma B.4. Consider the quantities defined Lemma B.3 and assume wi i.dd Ew). V4 iidd N(O, Ey). defining
2 A ( 22 Amas (Ev). Then, as long as n Z o²
18.27 (d, + log(1/6)), with probability at least 8:
Σ wiV. T (.))
-1/2
S 8d, log
c 1-c)
oh + 26d + 160 log(1/0)
lo
Amil mx) 2 (1 c) Amin Ev).
Proof of Lemma B.4. We observe that if En, 2 Q, then
2Σ, - Q + Σn = (Σm) (5m)²² 32 Q <2 E.
This implies
{ = 2}
Σ
n
w.v, T.
En
-1/2
S 21 Σ ビ 0)
Σ w.v. T (Q+ Σn
n
-1/2 2
(18)
Let us consider the event:
(1-c)E, 3 Σ, 3 +1++0) E)
which by Lemma E.2 occurs with probability at least 5 as long as n 2 18.27 (dg + log(1/5)). This immediately
establishes the latter desired inequality. Setting Q (1 c)nE,v and conditioning on the above event, observe that by definition Σx ł Q, and
det (Q + Em. (Q)
log
log det (Id. + Enn (Q) +En(())
< log det
(1+ 1-c Idh
(1
1+c c
(+c)
S dy log
1-c)
Plugging this into Lemma B.3, applied to the RHS of (18), we get our desired result.
Therefore, instantiating w - Fe, v + x, (x;, 6:)121 1 is a R x & ----uue process. Furthermore, since we assumed out of convenience that F, G. are computed on independent batches of data, we have that ~ N(O, FE,F", ). In order to complete the noise term bound, it suffices to provide uniform bound on F 1/ m (F) in terms of F.,
"See Appendix E.3 for discussion formalism. It suffices for our purposes to consider w Lidd E.))
19

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Lemma B.S. Assume the following conditions hold:
n Z max h: + log(1/0), o² dy log(1 (1/)) }
m (F,) Amin Ex)
dist(G, G.) (xn))))(F.)) (x()))<<<(xx))
2
then with probability least 1 6, we have F 1/ m (F) min (F.)
Proof of Lemma B.5. Recall we may write F
Z(Z Z) Z) F = F=F.G.X Z(Z Z)
= =F + F.,,, G G Z)Z z) 3 E 2(Z Z)-
(19)
By Weyl's inequality for singular values (Horn & Johnson, 2012), we have
min (F) 2 miin (FG.G max (Iax G
Z(Z
Since G. T an orthogonal matrix, the first term equal m (F.). On the other hand, applying triangle inequality on the second term, for 2 k + log(1/6) we have:
F. G T G) Z(Z Z) 'lop
S Fx Fx || G. - G TG) x 2 Z) *l
S F llop dist (G, (a) |Eillop Amin /G
< F. Flop dist(G, G.)x(Ex,
where used covariance concentration for the second inequality Lemma E.2 and the trivial bound Amin (AEAT) > Amin (2) for the last inequality. In turn, we may bound:
|| Z(Z 2)
<
3 Z(Z
mii ZZ TZZ (Z) -1/2
3 Z(Z Z) 1/2
1/2
Amin 1xx)1/2
(Lemma E.2)
lop
(dy + k + log(1/8)
100
Amin Ex)n
Therefore, setting dist(G, < xx),, and 2 0% dy +x+ we have a (F) 2 Omin (F.), which leads to our desired bound on |||
Fnin (F.)² (2x)
With a bound on |F|| bounding the noise term is straightforward application of Lemma B.4. Proposition B.6 (KFAC noise term bound). Let the conditions in Lemma B.5 hold. In addition, assume Z dy + log(1/8).
Then, with probability at least - 8:
dx + log(1/8) FE X(X XX x) '1lo - S det mmin (F.)² Amin (2x)r )n
Proof of Proposition B.6. Condition the event of Lemma B.5, Then, assuming 2 d2 + log(1/8), may apply covariance concentration (Lemma E.2) on Ex and Lemma B.4 bound the noise term by:
| X(X T x) -
< FE X(X x) 1/2
Amin X (X)
F op
[ X(X X
11/2 Amin ) -1/2
(Lemma E.2)
dy + k + log (1/8) Omin 1F)22x. Amin Σxx )
$ _=
(Lemma B.4)
which completes the proof.
20

---

Layer-wise Preconditioning Methods and Provable Feature Learning
This completes the bound on the noise term. We proceed the orthonormalization factor.
Bounding the orthonormalization factor
Toward bounding the orthonormalization factor from (8). Defining G+ as the updated representation pre-orthonormalization, write G+ RG+, where G+ the orthonormalized representation and R R**k is the corresponding orthonormal- ization factor. Therefore, defining the shorthand Sym(A)
A AT, we have
RR RG T
G lc (F
F)
F (FG - F, G.) (F E X (1.))
(from (16))
L Sym
F(FG - F.G....
Sym
FE X(X -G
er
n Sym (F(FG F.G.) (EE (FE'X(XTx))))
where the strictly inequality comes from discarding the positive-defimg "diagonal" terms of the expansion. Therefore, by Weyl's inequality for symmetric matrices (Horn & Johnson, 2012), have:
min (RR ))
2ne (Ir'illop + [| To [|Fillop 11T2llop).
Toward bounding
[T' Tillop the conditions of Lemma B.5 hold. Then,
[ llop
|[FG - F.G.) llop
[I FF= FGG -
F (F..G. G G) - Z(Z ">)
(from (19))
5
k log(1/0)
$ 4 min (F) |F dist(G,
+ Onin (F)0
Amin (Ex)
5 (FF)ddist G.)x(Ex) / Tmin (F,)2 Amin Ex)ri
k + log(1/8)
2
(Lemma B.5)
Similarly, letting the conditions of Proposition B.6 hold, have
[ 12lppp = FE X(X X -
dx + k log( (18) Omin (F.)² Amin Σx )
S /
(Proposition B.6)
72-
We observe that 1/2 will always dominate the second term the bound on [[T llop: and therefore:
Aminn (RR)) )) 2 - 2na (71 2n/2 + na(1 + 72)7)).
Therefore, we have the following bound on the orthonormalization factor. Proposition B.7, Let the following conditions hold:
a² dx + log n Z max { ddx de + log(1/0),
72 min (F.) Amii (xx)
dist G.) S 5m
xxFF))
Then, with probability at least
8, we have the following bound on the orthonormalizztion facto
Gmin (R) 2 v1-21a(m 27/ + Ma (71 + 72)))))
21

---

Layer-wise Preconditioning Methods and Provable Feature Learning
The constants 71 72 will be instantiated control the deflation of the contraction factor 1 1/a = 1 o/ec due the orthonormalization factor.
Completing the bound
We almost ready to complete the proof. By instantiating the noise bound Proposition B.6 and the orthonormalization factor bound Proposition B.7, have:
|++ op
R
((1 11a) + 1la (F "FF(FFF)=FF 8T X(X XX
7c R
|GP, 1| !lop + the FF F 3 x(X T x) 1
min
1 1 1/6
GGP op
dx + k + log(1
V1 (m + 2 + na (n + 72)
G (F.)2 Amiin (xx n
To understand the effective deflation of the convergence rate, we prove the following numerical helper lemma. Lemma B.8. Given d E (0, 1) and E (0, 1/2), i 2 c, then the following holds:
1 - d
< 1 - (1 -1-8)) E)d.
vI cd
Additionally, as long e < l-VI-d then
(1 - e)d < v1 - d.
Proof of Lemma B.8: squaring both sides of the desired inequality and re-arranging some terms, we arrive at
(1 - d1²
- =)d)
1
1 d
1 - d 1+
1
d
e)d
-edd
To certify the above inequality, it suffices to lower-bound the RHS. Since c, d E 1), the last factor is at least 1, such that we have
1 (1 -1-i-(11 11--11 1-d (1-1-11-1) (1+
1 1- d
1 d
8)d
1- 1- s)d
1- (1 - 2)d
(1 (1-E)dd
(1 8)d
Therefore, c S & is sufficient for certifying the desired inequality. The latter claim follows by squaring and rearranging terms to yield the quadratic inequality:
(1- - 2(1 5) + +1 0, (--vT= 1+yl-
Setting := 1 - s, the solution interval is À E E (0, 1), leaving the lower limit as the condition on proposed in the lemma.
The upper limit is redundant as it exceeds I and
d
Plugging in no d E (0, and 2(11 + 272 na(1 72)22) (72 272 (71 + 72h2 c, try candidate values ?1 1/40, 7/2 1/100 and set to get:
1 1- 7l
< (1 - 0.9ma).
VI 2ma - 272 + na(11 + 72)
Plugging in our candidate values 71, 72 into the burn-in conditions of Proposition B.7 finishes the proof of Theorem 3.6.
22

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Theorem 3.6. Consider running (8) with A = 0. 7/ E [0, 1I. and 1/e = 1. Define o² dist(G, G.) & 0.01x (Ex)x=) (F.) and n 2 ma )²) (dx + log(1/5)), we have: GG+ G.) = (1 0.9ma ).9n)) dist
o²
As long as
ouin Auin (2x)
G, G.) + 7,
with probability 2 1 - 5, where A 0 T
dx dx+ 1og(110))
B.3. Multi-Tast and Transfer Learning
We first discuss how the ideas in our "single-task" setting directly translate to multi-task learning. For example, taking our proposed algorithm template in (8), an immediate idea is, given the current task heads and shared representation ({F()), G), to form task-specific preconditioners formed locally on each task's batch data:
P
E(t)
[z
P
F(t)
F Qi Exx(xxx
and perform a local update on F(t) G before a central agent averages the resulting updated G:
F p(t)
V
(((()))
(( Q
G = - naP
V
G) , E
T
G T t=1
(9
However, this presumes F(t) are invertible, i.e. the task-specific dimension dy k. As opposed to the single-task setting, where as stated in Remark 3.4 we are really viewing F as the concatentation of of F make recovering the representation well- well-posed problem, in multi-task settings d, may often be small, e.g. d, (Tripuraneni et al., 2021b; Du et al.,
2021; Collins et al., 2021; Thekumparampil et al., 2021). Therefore, (pseudo)-inverting away (F(() p( may be highly suboptimal. However, we observe that writing out the representation gradient (9), as long as invert away Q first, then we have:
(2) -1
_ (t)
G
7e
V
T Σ
G.P, T
t=l
T T
I
1l0 T
Σ
F(t)
F(t)
GP, + (task-averaged) noise term.
Since by assumption 4 ΣT F(O) ((t) is full-rank (otherwise recovering the rank k representation is impossible), then suggestively, may instead invert away the lask-averaged preconditioner Pa Σ p(t) F(()) the task-averaged G descent direction before taking representation step G summarize, we propose the following two-stage preconditioning:
F(!) p(t) - 7r V.L G O
(20)
( Vo 2 G)Q (0-
EE
(21)
G, G M/ P-1
M D(
(22)
such that G.P. = (1 (1-1)a) GP. + (task-averaged) noise term.
(23)
The exact same tools used the proof of Theorem 3.6 apply here, with the requirement a few additional standard tools to study the "task-averaged" noise term(s). As an example, we refer to Zhang et al. (2024b) for some candidates. However,
23

---

Layer-wise Preconditioning Methods and Provable Feature Learning
note the qualitative behavior unchanged. As such, since are using data points per each of T tasks to update the gradient, the scaling of the noise term goes from Ox/=/ in our bounds O( dd/nT).
We remark that in the multi-task setting, where each task may have differing covariances and task-heads F(.) the equivalence of our stylized KFAC variant and the alternating min-min algorithm proposed in Jain et al. (2013); Thekumparampi et al. (2021) breaks down. In partieular, the alternating min-min algorithm no longer in general admits iterates that can be expressed as a product of matrices as in (5) or (20), and rather can only be stated in vectorized space vec(G). This means that whereas (20) be solved as T parallel small matrix multiplication problems, the alternating min-min algorithm nominally requires operating in the vectorized-space dyk.
Transfer Learning
We first prove the proposed fine-tuning generalization bound.
Lemma 3.7. Let 1s argmine [(y(( F2
[1] z(0) & Gx(() be the optimal on the batch of n()) target data (11) given G. Defining = dist G. G.). given (e) 2 k + log(1 /8), we have with probability 2 1 - 6:
c (
(2)
G - E ||s - -
Gx 13
|| Amax D 22(crk: + log(1 (1/6))
nl))
Proof of Lemma 3.7. We observe that we may write:
(
G = E -[1x(()) - G
(t
= E {[ICALE ls F G.)x
= =|| G F G. ز
(()) 1/2
Now writing out the definition of (=)) ls
defining (t)) Gx(L), we have
F
argmin
E [1y
(
(Z())
F..GG.GTT + F G.P. -
where P/ = Iax G is the projection matrix onto the rowspace of G. using the fact that G is row-orthonormal (8).
Therefore, plugging in the last line into error expression, we have
C (f()) G p F. f G.) (1)
F. M(:)
+2 +2| (Z(())))))))))))))))))))))))))))
((a + 6)2 S 20² 262)
Focusing the first term, we have:
F
F.
-
G
= = G.
x
(U)
7 Z (() lax)
By a covariance concentration argument Lemma E.2, since (() (t) and Z())(()) rank-k matrices, long as (=) Z k + log(1/8), we have with probability at least - d:
(t) (t)) R (t) &T (Z (1) G
24

---

Layer-wise Preconditioning Methods and Provable Feature Learning
and thus
G.P.
(Z
Lax
x 1|
GD G
(1)/// - Lix ||5
G.
| c.
IIx op
| F dist(G: G, G.) Amax 1//)
where in the last line we applied the definition dist G. G.) = G. G. = G and the fact that the matrix
op
p à Σx 10 G (
can be verified to be a projection matrix p P. PT P, such that p - -- I= pl is also an orthogonal projection and || Ilop = 1. Now, we analyze the noise term:
33 (Z (() GG
-G 1/2 (Z
(GIAGT) 1/2 11/
GI
where we observed (t)) (t²) n G and applied covariance concentration. Now, defining the (compact) SVD of
G/ V. we find
F Vtt
11 (Z(
V.
|| (())()) (2((((())))))))))))))))))))))( (Z(T)) (2))))))))
Vn((t
24
Vn(t) Te d log(1/5),
for (1)) 2 h: + +lgg(1/8) The last line comes from the Frobenius norm variants of Lemma B.3 and Lemma B.4 (see Ziemann et al. (2023, Theorem 4.1) or Zhang et al. (2024b, Lemma A.3) for details). Putting the two bounds together yields the desired result.
C. Proofs and Additional Details for Section 3.2
C.1. Proof of Theorem 3.9
Theorem 3.9. Assume that the activation function a is D(1)-Lipschitz and that Assumption 3.8 holds. In the limit where n, dx, dn tend to infinity proportionally, the matrix GeGD. with probability 1 o(1), satisfies
| + an fo sCD - GGoo - O,
in which a = [z]] with z ~ N(O, - and the vector Asco is given by BsCO = n -
Proof. To prove this theorem, we first note that
Vdn fg (Gxx;) ) (V), fo o' (Gox,) x
VoL (fu, Go) = Σ
yi
n 1=1
Adopting the matrix notation X = 1x.l... |x,]
E gnxdx and y = [y....., E R", we can write
VaL Go) - [1,, 1/200 tofo o )ox()) x.
(24)
25

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Let N(O, d( - and define a =E [=)(=),, and To R - R o [(2) az. This function satisfies
Es [o [z)] = 0. With this, we decompose the gradient into three components as VoL (fo, Go) = -T T2 + T3 with
y y
X
T
T -a 1/2 to
T2 -n Loy ool X,
T3 = n (fot) (Go
of x
We will analyze each of these components separately.
Term 1: For this term, using the facts that fu f2 (11) and |X y/nl12 O(1), we have
IT Ilop =0 =0(1n (d)
Term 2: To analyze this term, note that
fo y @ = diag(fo) of (Gox") diag(y),
which gives
fl y of 1lpp S |1000=0 Illpp
Using basic concentration arguments, we have ||foll = Ö(1, 1/2) and || = (1), with probability o(1). By construction of a (-), the matrix o'1 (GoXT) has mean zero entries, thus using (Vershynin, 2012, Theorem 5.44), we have ||o` (G Olop Ö(ap + 1/2) with probability 1 o(1) Thus, the norm of T2 can be upper bounded as
T2llop δ ( dh
dn
dx n
)
1
+
Ô (4,
7
Term 3: Similar to the second term, note that
(fo fu (GoX )) @ (GoxT) = diag(fo) o'Gox diag (fo oGox
Thus, the norm of the third term can be upper bounded as
1 | nd (fof) (GoXT) o (Go
IT op
Sn -n-
X Goll o (G, Ib I (Go Ilop
lpp
analyze the right hand side, note that assuming that (ip)-Lipsehitzz, the entries of a' (Gx are bounded by the Lipschitz constant, and have ||o Olop = (( dn). Also, using simple orderwise analysis we have [lf] (GxXX 1l O(1,, which gives
dx ) = -(dh 1+
=(1
1 dh
Ta Tallop = =
To wrap up, note that d/ / dlop = O()), whereas / lop and d//22 TTallp = o(1). Thus, with probability - o(1) we have
Gsm = Go + 17 1 VaC Go)
Go) Go + am -1=x'² +A
with ||p = o(1), finishing the proof.
26

---

Layer-wise Preconditioning Methods and Provable Feature Learning
C.2. Proof of Lemma 3.10 Lemma 3.10. Under the assumptions of Theorem 3.9, the correlation between Ba and feen satisfies
3. BSGD
Ca,1
Tr(Ex)
dx
+0
G 12
22 Tr [xx
+
G.
Tr(=2)
n
dx
where we have used the fact that e mean zero. Thus, using the weak law of large numbers, with probability (1), in which cx,1 (z)) and ct [x)] with N(O, Tr(Ex)).
(((E))
Proof. Recall that BaGu 1. X y and y (x(3..) where E ( Ex]" Therefore, with probability d(1)
have
3, 1 (x..) y ox..)),.(X33.) + (1)) 11
71
Bscn 3. + E, [20, d Tr(Ex) (o( c, TTE,.
(25)
in probability, where N(O, 100, Tr(Ex)) Similarly, ||G 11 can be written as
B =
XX
y =
xx o XX (X33.) + o1))
We will analyze each of the two remaining term separately. For the first term, recall that is independent of X. Using the Hanson- Hanson-Wright inequality (Theorem E.6) we have
Tx(XX /n) + o(1) -x=n- - xx E = =
Tr(Ex) o(1).
For the second term, note that X,3, is vector with i..d. elements x, B., each of them distributed according to N(O. 3 Ex Let be random variable distributed as ~ 3, T, Ex We decompose the function a, into a linear and a nonlinear part as
o.(2) C2,1 0.,11
(26)
This decomposition satisfies
E, [a = (2)]= [x] =
E, [zo,, (=) = E, 20. c., E. [z²] E, [z). cx,1 1, = 0,
where the last equality due Stein's lemma (Lemma E.7). This shows that the random variables and o. (2) are uncorrelated. With this, have
n 22, (XB.) XX o. (XB.)) = = (c,, + -. (X3.)) XX Xxx (c,, + +0..
= 2 23, (xT)) + 2 2 2(X3,) XXT a (X33.) + 2
XXT (X3.). (27)
For the first term in this sum, by assumption 3.8 and the Hanson-Wright inequality (Theorem E.6), write
- (xT)² d Tr(n
Tr(n (x + o(1) = ",, Tr(E 2 -d TT²(EE))))))
where in the last we plugged in the second Wishart moment. For the second term in (27), although by construction X,3, and _..1 (X,3.) are uncorrelated, the vector X,3, and the matrix XX T are dependent, which complicates the analysis. To
resolve this issue, we define X = X - X which satisfies X,3, 1 X and write
XX
XxT + x X33,3.
T
x TT
27

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Thus, the second term in (27) be written
n (X,,))
TX XX 0x,1 (X,3. n 7 (X,9.) XX + X33,3. X x33,8, X 0x,1 (XB3.)
n =
XX _. ox (X + t (X
X + X X (x,3))
Ti
72
The term 7i can be shown to be o(1) by using the Hanson-Wright inequality (Theorem E.6) and noting that XXT is independent of XB.. and also the fact that XB, and T.,1 (X,3.) are orthogonal, by construction. Similarly, 72 can also be shown to be (1) using similar argument. For this, we also use that fact that by construction we have X,3, 1. X,3, which alongside E [X3,]= = O, proves that n (XX3.) X3, o(1) and n (X.3)" 0-,1 (X33) = o(1). Hence,
2c,,1n PXB3)) XX o. (X.3.) 0.
For the third term in (27), we can similar argument and replace X with X to show that
20 (X3.) XXT ox, 0x,33 [=)0 n Tr(Ex).
Putting everything together, we have
B300 = d Tr =
Tr
+ toē n Tr(Ex) + Es
+ 2 dd,' Ix2(2x)) ( + +11)
+ n +n (o + Ezlo,, +2 di Ir(Ex)) + o(1).
Note that given the decomposition (26), have
Eg [x] = [o (z)] + dr Ir(Ex)
given the orthogonality of the linear and nonlinear terms. Hence,
B 112 = d Tr(E2)) + Tr(Ex) (o? + +2) +o(1),
which alongside (25) proves the lemma.
C.3. Proof of Lemma 3.12 Lemma 3.12. Under the assumptions of Theorem 3.9, the correlation between 3r and BRrac satisfies
3 T 3.
C.1 V
+0
|| Catrrcll2110 12122,112
Vet, U1 + dx (c2 1 + o2) 11
with probability 1 o(1), where c2 E2(o((=)), 2 Eslo?, [=)) with N(O, Tr(xx and 11, Dz V are defined in (34) and depend Ex, dx/n, and Xa. In particular, as Xo and dx/n - have
Toaac B.(((),rrxclllllBB....
Proof. Recall that Qo = and let = (Qo + do Aa -1,
The inner product of B, and BRrxc is given by
3. Barae = n -1 3. RX o..... + o(1),
where we have used the fact that is mean zero and is independent of all other randomness in the problem. Defining we can use the push-through identity to rewrite the inner product as
R (XX7 /nn + daInn) -1
B. BREAC n (XX3) (X,(. o()).
28

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Note that X,3, is vector with i.i.d. elements x, B., each of them distributed according to N, 3, Ex3.). Using the same decomposition for o, the used in the proof of Lemma 3.10 in (26), we have
1 -x33.) ( 71
B_ BReaC
(c.,1X3,, + T..o (X.3.) o(1)
= c
Tr (x [Xx AdmIn) X) + in-' (0-'((33)) (x,,1(X3..) + o(1),
(28)
where for the first term we have used Assumption 3.8 and the Hanson-Wright inequality (Theorem E.6). To analyze the second term, note that although by construction X,3 and (X,3), are uncorrelated, the vectors X,3, and R. are dependent, which complicates the analysis. To resolve this issue, we use the same trick used in the proof of Lemma 3.10 and set X = X - -33,3, which satisfies X,3, 1 X, and use it to write
--==n-XX + dcIn = -1 1-11x++ (x+.3. (X+
d
= -n1-1XX + AoJ1.) +1 (X3.)((3)3) (XB.) ( (X..)
Defining R = (XXT / +
E ERaxn V -IIn-1/2XX0
n 1/1 E ER and [n-1/2X,3,
1 D
0
have R R VDVT. Using Woodbury matrix identity (Theorem E.8), R is given by
R - R
(29)
and plugging this expression into the second term in (28) gives
n (x3.) Ro.,, (XB.) = n - RR-RVDD + R) (x3.) +
=n R Ro,, -n R
Ro., (XB,) o(1)
The first term can be shown to be o(1) probability by using fact that R independent of X,3, and orthogonality of X,3, and T..L (X,33) To analyze the second term, first note that the elements of the matrix (D- -1 R22x can all be shown to be ((1) by simple norm argument. Moreover,
(X3,) RV(D- + V RV) -V ' Ro.,(XX3.)
(Q =
(X,,) + S12 (X3.) + 9221 (x,.) + 922 (x3.))))) R 1x.11(8,3)
where 2i are the elements of the matrix Ω. We analyze each term in this sum separately and show that all of them are o(1).
First Term. Using simple norm argument, n1 -XX33 ((X3.)) = ((). Also, by construction of o. have ((X,)) Ra. (X,.)) 0.
Thus, the whole term is o(1).
Second Term. Similar to the first term, we have (X.3.) (X,3)) O()) Also, n (X,3,) T (XB.)) in probability, using the weak law of large numbers by noting that 0., (X33)) is independent of n (X,3.) R by construction and that has mean zero. Hence, the whole term o(1)
Third Term. By construction, the vector XB3, independent of R(X,,)) and has mean zero, which gives n (XB..) ( (X.3.) + 0. Also, using a simple norm argument, we have n -XX3,)) (X3.) = O(1) which proves that the third term is also o()).
29

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Fourth Term. This term can be shown to be o(1) with argument very similar the argument for the third term.
Putting these all together and using (28), have
B. BReac = Id Tr ((X) X/n) + o11)
(30)
Next we move to the analysis of the squared fz-norm of the vector BREAC. By decomposing the function into linear and an orthogonal nonlinear component similar to the one used for the analysis of the inner product term above, we write
1| n n XR y = n XR E c 33 X XR X,8
0 XR o ( + 2c,,, ²3, XR X
+
We will analyze each of these terms separately.
First Term. Recalling that e ~ N(O, In) independent of all randomness the problem, using the Hanson- Hanson-Wright Wright inequality (Theorem E.6) we have
E = (<(((XX/))))R)) X/n)R²) + o(1)
Second Term. Using Assumption 3.8, and by the Hanson-Writ Wright inequality we have
n 3 B. XR x,3, = ,
d Tr [(X X/n) R + o(()
Third Term. Note that R. and X,3, are dependent. Note that XR = XX R = nR >. Thus, an almost
identical argument to the argument used above for the analysis of B. BEFAC using X X -X,33, gives
20 (XG.) (X..)xXXXXXX XR o. (X33) E. (=)] n Tr ((x X R²) o
Fourth Term. This term can readily be shown to be o(1) in the analysis of BT_ Braaci i.e.
2c,, n 23 XR2xT A (x3.) = =())
Putting everything together, we find
|| 112 = d Tr (XTX/1) R + (o? + ((XXx//)))))
(31)
Now, given (30) and (31), defining c. = E. (z)]. we have
3 BRPaC
1d Tr (X X/ X)n)R)
(32)
1Kraall 13,11
Ver, Tr I [ R + (o2 + +t, n Ir ((XTX/n) R²)
Thus, noting that if dx/n + and de - we have R Ex we find
BRFAC
lim
lim
1,
dx
BREEC
[.
proving the second part of the lemma. For the first part, we define m(2): IR R the limiting Stieltjes transform of the the empirical eigenvalue distribution of n n X;
m(z) lim - Tr (x X/n Iax))
(33)
where the limit is taken under the assumption that dy/n + & > O. For a general covariance matrix Ex, m(x) does not have closed form except for very special cases; however, it can be efficiently computed. See Section E.8 for more details, The derivative of the function m is given by
(x) = lim dz Tr (x T X/n - Lax) -9.
200
30

---

Layer-wise Preconditioning Methods and Provable Feature Learning
We write all the traces appearing in (32) in terms of the function and its derivative:
T X//) R) = d Tr ((X X/n + AaIax AGTax) R) = dx Tr (Idx AGR) = 1- - do ml m(->a),
dx d Tr (X"X)) R²) = dx Tr ((X X /n - Aa \c R?² - 1 - -Xa Xam'(-Xo) - -2 m2-am
n- Tr (xx X n) R²) n T (R (R- \a) - mm + 2 da mmm'm-Nall (->a).
With these, the correlation is given by
3 B. XPAC
Cx [1 Ao (nd-)
B KFAC 3+122
V2,11 X2 m2mm--oo - 2A ml m( (->a)] ] + 2(2,,> + mz) + Aam'(--ee
which defining
1 Ay mm-da))))
U2 = 2a m'm (-aa) - 27a ml mm-)aa)
V3 m(-da) Aam'(->aa
(34)
concludes the proof.
L
C.4. From Feature Learning to Generalization
In Section 3.2, we showed that after one step of SGD and KFAC, the first layer weights will become approximately equal to
G2 = Go + on T a E (SGD,KFAC)
(35)
Given Lemma 3010 and Lemma 3.12, we argued that compared to SGD, the weights obtained by the KFAC algorithm more aligned to the true direction B.. Given nontrivial alignment between the weights and the target direction, the second layer f can be trained using least squares (or based on Lemma 3.1, equivalently using one step of the KFAC update on f with V/ 1) with ((d) samples to achieve good generalization performance (See e.g., Ba et al. (2022, Theorem 11) and Dandi et al. (2024c, Section 3.4)). The existence of nontrivial alignment of the learned weights and the true direction in a single index model is often called weak recovery and has been subject to extensive investigation (see c.g., Ben Arous et al. (2021); Dandi et al. (2024c); Troiani al. (2024); Amaboldi al. (2024), etc.).
To see this, consider the feature matrix Z, Raxds as Za a (X where the activation function is applied element- wise. Based on equation (35), this matrix can be written
Zn R a XGG + an (x3.)6))
This is an example of a random matrix in which a nonlinear function is applied element-wise to a random component plus a rank-one signal component which has been studied in the literature (Guionnet et al., 2023; Moniri et al., 2024; Moniri & Hassani, 2024b). particular, by Taylor expanding the activation function, the feature matrix Zg can be written as
Zo (XG Σ at1"" (o ((x3.) rot E
where denotes element-wise power and Eg is the reminder term. Let 1) = no for some E [0, 0.5). Given a, the integer is chosen to be large enough so that the operator norm of & is o(d and the reminder term is negligible compared to (XG ). By simple concentration argument, the matrix (k) (XG can be replaced with its mean (a (XG])) ,111 to get
Σ a a
Zo o(G))
(x,6)) pakT.
k!
31

---

Layer-wise Preconditioning Methods and Provable Feature Learning
The first term a ( is the feature matrix of a random feature model and based on the Gaussian Equivalence Theorem (GET) (see e.g., Goldt al. (2022); Hu & Lu (2023); Dandi et al. (2024a); Moniri et al. (2024)), we can replace it with aXG, + N where N is properly scaled independent Gaussian noise. The vectors (X,3,) )* are nonlinear functions of the covariates with different degrees. The least squares estimator f. is then fit on the features Zn in way that L(f,, Go) is minimized; i.c.
linearize it; i.e.,
a
(fo* fa)
y o (XG
Σ
k! (X,3.)
(36)
Based on the GET, the random feature component can only learn linear functions with sample complexity of learning n = (dnn) = (d2)) When is large enough and 3g is aligned to B... with the finite dimensional correction to the random features model, the model can also represent nonlinear functions (æ: Ba)* x of degree k < é by matching the coefficients a²1/" (fg* fa)/k! with the Taylor coefficients of the teacher function (x ( B.).
Although we have provided complete proof sketch for providing generalization guarantees given weight alignment, complete analysis require tedious computations and is beyond the scope of this work as we mainly focus on feature learning properties different optimization algorithms.
D. Additional Information Regarding Kronecker-Factored Preconditioners
Here, provide additional background information regarding key Kronecker- preconditioning methods, including their derivation and relations to various methods in the literature. We recall the running example of a fully- connected net omitting biases, introducing layer-wise dimensionality and a final non-linear layer (e.g. softmax) for completeness:
fo = (WW
We Rdxde- do dx.
(37)
As before, we define 0 as the concatenation of 0, = vec (ec(W)), & [Z]. We define an expected loss induced by the neural network ((0) = (()) and its batch counterpart C(0). Here, we define the family of Cronecker-Factored preconditioned optimizers as those that update weights in the following fashion:
We = W - P Vw, (E
where Py Rdexxe, Qe E 1dde- [L] are square matrices. For simplicity, we ignore moving parts such as momentum, damping exponents, adaptive learning rate schedulesiregularizztion etc. We now demonstrate the basic principles and derivation of certain notable members of these preconditioning methods on the feedforward network (37).
D.1. Kronecker-Factored Approximate Curvature KFAC
As described in the main paper, KFAC (Martens & Grosse, 2015) is at its core an approximation to natural gradient descent. Given that we are approximating NGD, a crucial presumption on fx() and L(0) that the network output fo (x) parameterizes a conditional distribution p(y) x; 0), and L(O) x [ log p(y) x; 0)] is the corresponding negative log- likelihood. As such, KFAC is technically only applicable to settings where such an interpretation exists. However, this notably subsumes cases (()) E(x.y) [e(fo y)], where E(-) is strictly convex function in fe (x), ), this admits interpretation as fo(x) parameterizing exponential family distribution. In particular, the square-loss regression case E(),y) || y corresponds to conditionally- GGaussian predictive distribution with fixed variance (xx) N(fo(x), T²), and if d(.) is softmax layer and E(ý.y) = CrossEnt(ý,y) the multi-class classification case corresponds to a conditionally- multinomial predictive distribution.
Defining h Weze 1, 24 = a (h), 2n x, the Fisher Information of the predictive distribution p(y x; 0) at O can be
32

---

Layer-wise Preconditioning Methods and Provable Feature Learning
expressed in block form:
Öp(y x 0)
Opp(||x;)) 0)
T
FI(0) 2 Ex
(recall 0 is vec-ed parameters)
00
00
Ex
(p(r)xx0) 80,
Op(y|x:0) x:0)
Ex
@p(y)x:0) 00:
Op(y xx:))
00,
80.
Op(ylx:0) 80
(Oply|x:)) xx:0)
@p(y)x:0)
(pplyix:0) 80L
Ex
Ex
30,
80
Looking at the (ih block, we have
Op(y|x:0) 0)
(Op(yix:0)
(Op(y|x:0) 0)
(Op(y|x:0) 0)
Ex
Ex
vec
vec
00,
00,
OW;
aw,
(z gi) gi) 3 B Ex
(ge
Oply the
Ex [ x (g.g,) )
(Lemma E. item 2)
where the second line comes from writing the backpropagation formula. KFAC makes two key approximations:
1. The matrix FI(0)-+ -1 is approximated by block-diagonal, and hence FI(0). We note the original formulation of KFAC in Martens & Grosse (2015) also supports tridiagonal inverse approximation.
2. The vectors z and ge are independent for all [L]. such that
Ex [Zz- x & (g+ ge
E(2, 27,1 @ Elgr 82
Now replacing the true expectation with the empirical estimate, and defining P = Elg g, Q Elzz 2T-1 completes the Kronecker-Factorts approximation to the Fisher Information. It is clear to see from the derivation that, as we previewed the introduction and expressed emphatically in Martens & Grosse (2015), this approximation is never expected be tight.
Some related preconditioners
Having introduced KFAC, we introduce some related preconditioners. Notably, has been noted that computing gr requires backwards gradient computation, whereas zę only requires forward pass. particular, various works have recovered the right preconditioner Qe of KFAC via various notions of "local" (layer-wise) losses. Notably, these alternative views allow KF KFAC-like preconditioning extend beyond the negative-log-likee likelihood interpretation.
LocoProp, square-loss case (Amid et al., 2022):
Update rule: WW+ argmin 1 [ W1 hell
W
Well
W
2m
W. W.-mee VV- ((0) (Iu... + n E(Zz
As noted in Amid et al. (2022), this update is also closely related ProxProp (Frerix al., 2018).
FOOF (Benzing, 2022):
g: al(So(x),)) Th.
Update rule: AW, argmin
È [ 7
- + 2
AWI
AW
= v, 2(0) [Eiz-12Zz]] Aldee) ))-
W4+ = A
33

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Interestingly, note that these right-preconditione-onlly variants subsume the DFW algorithm for two-layer linear represen- tation learning proposed in Zhang et al. (2024b); thus may see the guarantee therein support of the above algorithms from feature learning perspective, albeit weaker than Theorem 3.6.
D.2. Shampoo
Fronecker-Factored Factored approximation of the full AdaGrad preconditioner, which recall is the running sum of the outer-product of loss gradients. Turning off the AdaGrad accumulator and instead considering the empirical batch estimate E & (x), Vo (x)(())))) "J. the curvature matrix being estimated also be viewed as the Gauss-Newton matrix E(x.y) Vo (x)(x)))))) y) Ve (x)) y) As documented in various works (see e.g. Martens (2020)), the (generalized) Gauss-Newton matrix in many cases is related or equal to the Fisher Information, establishing link between the target curvatures of KFAC and Shampoo.
Shampoo is designed to be Kronecker-
However, the Shampoo preconditioners differ from KFAC's. Let define the Ze, he before, and g: (6(fa(x).)) the Shampoo preconditioners are given by
Then,
oh,
P E [g
(
giz
(8:2
Q: =E [z =E[2z-- (zg--ie (ze- )T]
Notably, Shampoo takes the fourth root in the preconditioners, as its target is the AdaGrad preconditioner which is (modulo scaling) the square-root of the empirical Gauss- matrix- analogouu to the square-root of the second moment in Adam. Whether the target curvature should be the square-root or not of the Gauss-Newton matrix is the topic of recent discussion (Morwani et al., 2024; Lin et al., 2024).
D.3. Kronecker-Factored Preconditioners and the Modular Norm
The "modular norm" (Large et al., 2024; Bernstein & Newhouse, 2024a;b) is a recently introduced notion that provides a general recipe for producing different optimization algorithms that act layer-wise. By specifying different norms customized for different kinds of layers (e.g. feed-forward, residual, convolutional etc.), one in principle has the flexibility to customize an optimizer handle the different kinds of curvature induced by different parameter spaces. Given choice of norm on the weight tensor We, the descent direction is returned by steepest descent with respect to that norm. To introduce steepest descent, require few definitions (cf. Bernstein & Newhouse (2024a)):
Definition D.1 (Dual norms, steepest direction). Given a norm
||-|| defined over a finite-dimensional real vector space
The dual norm ||-||; defined by
|v]|+
max (u, v).
With g E V and a "sharpness" parameter 7) > the steepest direction(s) are given by the following variational representation:
argmin
(g. d) +
2n ||| = n 7/gg|1
argmax (g. u)
||u|==
Here we focus on unite-dimension normed spaces, but note that these concepts extend mutatis mutandis to general Banach spaces. The aforementioned works derive various standard optimizers by choosing different norms, including induced matrix norms [Well-+8 maxx
[Waxlle applied to given layer's weight space, for example (Bernstein & Newhouse, 2024b):
|x]|..
SGD: induced by Frobenius (Euclidean) norm ||-|| = ||-1| Note the Frobenius norm is not an induced matrix norm.
Sign-descent ('ideal" Adam with EMA moments turned off): induced by ||-I| |1
Shampoo ("ideal" variant with moment accumulator turned off): induced by ||||| = |1-1le,++6 = =Il-llop-
Therefore, in light of this characterization, natural question to ask is what norm induces a given Kronecker- Factored preconditioner (which includes Shampoo). We provide simple derivation that determines the norm.
34

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Proposition D.2 Kronecker-Factoree matrix norm), Recall the fully-connected network (37). Given preconditioners {(Pr; Qe) where Ps E Rededd, Qe E R [T. are invertible square matrices. Then, the layer-wise Kronecker-Factored update:
W4 W. Vw. QE", [
equivalent to layer-wise steepest descent with norm
|M&]
P ||:
argmin M
M
- Pi C(,) Q |M|
27)
Proof of Proposition D.2. It straightforward to verify |M]
- |P T || for invertible P, Q satisfies the axioms of a
norm. It remains to verify the steepest descent direction:
(V), 1 ||M 21]
argmin
Vw.c(()ll
argmax ((w, PT'Vw,C(0) Q
M
M
We start by writing:
Vw, C(0))+ & max ||M
(Vw,CiO),
max |P/ MQ, 1=
TrMM Vw,((0)
max Tr P (ww,(0)
(P DD)
[D
= [P (()) Q l|F.
(trace cyclic property, ||-I| self-dual)
Similarly, it straightforward to verify that the maximizing matrix is:
P V ((0) Q
argmax (Vw.C(0),
|P, V ((0) Q ||
||||=I
such that plugging it into the steepest descent expression yields:
argmin
V
((,), M)
1
|M|
= P. V Vw,(()) 0 1
P. V ((,) Q
M
2n
P. V w. C(B) Q ||
=n P Vw, Q
as required.
We remark that for complex-valude matrices, the above holds without modification for the Hermitian transpose A Notably, the layer-wise norm corresponding to Kronecker-Factored preconditioning not an induced matrix norm, though modified optimizers can certainly be derived via induced-norm variants, such
Maahalooobsa induced norm:
|M A max ((x) P(Mx)
(P, Q.Q 0)
Vx 1
1 /2x]
1/2
max x||=1
-|12//MQQ
35

---

Layer-wise Preconditioning Methods and Provable Feature Learning
E. Auxiliary Results
E.1. Properties of Kronecker Product
Recall the definition of the Kronecker Product: given A Raxn B E Rexy
Anp B
A B
A A&B B
E R mmpnnn
Ami
Amn B
Complementarily, the vectorization operator vec(A) is defined by stacking the columns of A top of each other (i.e. volumn-major order)
vec AA [An
Am1
Ain
A Amn
R
We now introduce some fundamental facts about the Kronecker Product.
Lemma E.1 (Kronecker-Prodieee Product Properties). The following properties hold:
1. (A @ B) A < B-1 Holds for Moore-Penrose pseudoinverse as well.
2. For size-compliant A, B, C, we have (A & B) D) (AC) @ B D).
3. vec(AXB) = (B*
A)ve
E.2. Covariance Concentration
We often use the following Gaussian covariance concentration result.
Lemma E.2 (Gaussian covariance concentration). Let x; i.d. N(O, Ex) for = I,....n, where x; Rdd Defining the
empirical covariance matrix Ex
E i=1 x,x;" as long as n 2 18:27 (d + (lodlog(1111) 1og(1/6)), we have with probability at least
1 1-6, 8,
(1- (1+0
Proof of Lemma E.2. The result follows essentially from combining by-now standard concentration inequality for Gaussian quadratic forms and a covering number argument. To be precise, we observe that
Σx
Σx op
I Σx (1 - c) 5x
Therefore, it suffices to establish a concentration bound on 15x Exll and invert for Ex To do we recall standard
covering argument (see e.g. Vershynin (2018, Chapter 4)) yields: given an E-covering of gd-1 , -N(S--1 (l, the operator norm of a symmetric matrix Σ bounded by
E
< 1 - 28 ueN
max u Eu,
where the corresponding covering number bounded by:
1 1++2 11+))
W(S 11-112,)) 8)1 <
As such it suffices to provide a concentration bound on Eu for each E V and then union-bound. Toward establishing this, we first state the Gaussian quadratic form concentration bound due to Hsu et al. (2012), which is in turn an instantiation of chi-squared concentration bound from Laurent & Massart (2000).
Proposition F.3 (Prop. in (Hsu et al., 2012)). Let E Ramxd be fixed matrix. Let g N(O, Ia) be a mean-zero, isotropic Gaussian random vector. For any E (0, 1), we have
2|||Ag| TF(ATA) + 2 Tr((AT (ooooo/)((++//////// A log(1/8)] <
36

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Now, given u E gd-1, setting A = u IT 1//² such that u E// g = u x, instantiating Proposition E.3 yields:
(g((1/)) + 2u Exu
log(1/6)
P u Ex E
Exu+ 2u /x/V
S 8.
Put another way, this says with probability at least 8:
(108(110)
log(1/6)
u E Eju
2u Exu
n
n
Taking a union bound over u E N,, we get with probability at least 8:
(V
Log((NV/88)
log( (VV(6))))
max u T (Ex
Eju
max 2u T Exu uEN
ueN
"
dlog (1 + ?) + log(1/6)
dloog (1 + +2) + log (1/6)
< 2 Ex llop
n
2 |Ex1lop v d+ log(1/5) (1/5)
4 log
1+
as long as 2 dllog (1+ 2) + loo(l1/ (18). Chaining together inequalities, this yields with probability least ô under the
same condition nt:
Σx
Σx op
2
log 1
2 8 V
d log(1/5)
lp
Minimizing the RHS for 0.0605 yields the result.
E.3. Extensions to subgaussianity
As previewed, many results be extended from the Gaussian setting subgaussian random vectors.
Definition E.4. A (scalar) random variable X subgaussian with variance proxy o² the following holds on its moment- generating function:
(X202 2
E(exp(AX)] s exp
mean-zero random vector x E Rdd E(x] O, is subgaussian with variance proxy o² if every linear projection is 2² subbaaussia random variable:
Elexp(Av x)] S exp
for all v e Rd
2
With this in hand, we may introduce the subgaussian variant of covariance concentration and the Hanson-Wright inequality,
E.4. Subgaussian Covariance Concentration
We state the subgaussian variant of Lemma F.2, whose proof is structurally the same, replacing the x² random variables with generic subexponential (c.f. Vershynin (2018, Chapter 2)) random variable, and using generic Bernstein's 's inequality rather than the specific x² concentration inequality. The result is qualitatively identical, sacrificing tight/explicit universal numerical constants. The result is relatively standard, and can be found in e.g., Vershynin (2018, Chapter 5) or Du et al. (2021, Lemma A.6).
37

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Lemma E.5 (Subgaussian covariance concentration). Let x, be i.i.d. zero-meun o² o--subgaussian random vectors for i ,...., n, where X, TRd and E[xx"] Ex Defining the empirical covariance matrix Ex +E x:X. there exists a universal constant C₁ > 0 such that with probability at least 20:
(V=+
log(1/0) + log(1/8)
Σx Ex
Co |Ex|lop
n
Therefore, as long as n 2 O2 (d++ogg( +ogg(1000) we have with probability at least 6.
=) Ex (1++
E.5. Hanson- Hrriinn-Wright Inequality
We often use the following theorem to prove the concentration inequality for quadratic forms. A modern proof of this theorem can be found Rudelson & Vershynin (2013).
Theorem E.6 Hansonn-righht Wright Incquality (Hanson & Wright, 1971)). Let x
(XI,...,XA)) E Rdd be a random vector with independent sub- gaussian components X; with EX; = Let D be an n x n matrix. Then, for every 2 we have
t2
Dx > + S 2 exp
-emin
D |
where e is a constant that depends only on the subgaussian constants of X
E.6. Stein's 's Lemma
We use the following simple lemma which is application of integration by parts for Gaussian integrals.
Lemma E.7 (Stein's 's Lemma). Ler X be random variables drawn from N(H, and g R + RR be a differentiable function. We have E [[x()((X p = o² (X))
E.7. Woodbury Matrix Identity
In the proofs, we use the following elementary identity which states that the inverse of a rank-k correction of a matrix is equal to rank-k correction to the inverse of the original matrix.
Theorem E.8 (Woodbury Matrix Identity (Woodbury, 1950)), Let A E REXn The following matrix identity holds:
C R bx
UE R and V E Rkxn
(A + UCV) ------------CCC--- -
assuming that the inverse matrices in the expression exist.
E.8. Stieltjes Transform of Empirical Eigenvalue Distribution
For distribution y over R. its Stieltjes transform is defined
m2(2) = / du(x)
x
Let Ha be the (discrete) empirical eigenvalue distribution of Σx E Rddxd and let Fa be the (discrete) empirical eigenvalue distribution of the sample covariance matrix Σ E dxd Consider the proportional limit where d, n + co with d n + 6 > 0. Suppose that the eigenvalue distribution Ha converges to limit population spectral distribution HΣ,; i.e. H = HE, in distribution. Given the definition of m(x) from equation (33), we have m(²) mp(z). The following theorem characterizes mr terms of HE
Theorem E.9 (Silverstein Equation (Silverstein & Choi, 1995)). Let vp(x) ((m 1/2) 1/2. The function vF the solution of the following fixed-poin equation:
1
t 1+(ur(() | dlt (t).
①
(2)2)
38

---

Layer-wise Preconditioning Methods and Provable Feature Learning
Thus, using this theorem, given Hz can numerically compute up (and hence, my) using fixed-point iteration. For example, for xe from equation (14), we have F' = 1/2 + 1/2814e.
F. Additional Numerical Results and Details
F.I. Details of the Experiment Setups
In the experiments, we generate Fo d with i.idd N(O, 1) entries. Then, for each task we randomly draw matrix B, E Xdy and set F; exp (0.005(B, - B) Fu, where exp(-) is the matrix exponential. The shared representation matrix G, E xx is generated by sampling uniformly from the space of row-orthonormal matrices in &
We consider two settings for the covariance matrices Ex,si the low-anisotropic, and the high-anisotropic settings. In the now-anisotropic setting, we define E = 5 Lax +N where N E dxxx has i..d. N(0, 1) entries, and set Σx,= = 0.5 (E+ET). For the high-anisotropic setting, we first sample uniformly a rotation matrix O E d xdx and set Σx,, = ODO where D diag(1 (iggacog 5, d.)). In the experiments for the main paper, we always consider the high-anisotropic setting.
In the following experiments, in addition the data generation process in equation (13) used in the experiment in the main paper, we also consider a Gaussian data setup where samples for task s generated according to
l.d. y +.G.Xi+ x N(0, Ex.)), of ~ N(O, s E (test, train).
(38)
F.2. Additional Experiments
F.2.1. EFFECT OF BATCH NORMALIZATION
We used the same experiment setting described in Section 4 to generate the plots in Figure 9. Explicitly, we use data dimension dx 100, task dimension dy 15, and representation dimension k: We use the same learning rate 10-2 for each optimizer except for NGD, in which used 10-4 The batch size is 1024. In Figure 3 we considered the Uniform data (13) with high anisotropy. Here, consider the other three setting: Uniform data (13) with low anisotropy, Gaussian data (38) with low anisotropy, and Gaussian data (38) with high anisotropy.
1.00
1.00
3.00
Sulopace Blstance
Subuspace Distance
Sub-mace Distance
0.75
10
Train Loss
0.75
10
0.75
Train Loos
0.50
0.50
Train
0.25
0.25
0.25
0.00
500
10000
1500
2000
500
10000
2500
2000
500
1000
2500
2000
Number of Ieeratios
Number of Iterations
Number of Iterations k
Figure 4. The effect of batch normalization (on AMGD) vs. KFAC in our experiment settings (Left) Uniform with low anisotropy. (Middle) Gaussian with low anisotropy. (Right) Gaussian with high anisotropy.
As discussed in the main paper Section 3.1.1, we expect AMGD with batch-norm to converge in training loss but to perform poorly with respect to the subspace distance from the optimal in settings in the case with high anistropy (Right). However, in the experiment settings with low anisotropy (Left and Center), we expect reasonable performance from this algorithm because rowsp G, Ex) is close the target rowsp(G.) (G.).
F.2.2. LEARNING RATE SWEEP
We further test the performance of each learning algorithm at different learning rates from 10-6, 155 ..... 10 10-0.5 0.5 10 with results shown in Figure 5, where we plot the subspace distance at 1000 iterations for different algorithms. If the algorithm encounters numerical instability, then we report the subspace distance as the maximal value of 1.0. We observe that KFAC and DFW coverage to a solution with small subspace distance to the true representation for a wide range of step sizes, whereas the set of suitable learning rates for other algorithms is much narrower. Furthermore, we observe the poor performance of various algorithms in Figure and Figure is not due to specific choice of learning rate.
39

---

Layer-wise Preconditioning Methods and Provable Feature Learning
SS Dist at 1000 Iters
0.75
0.50
500
Azam
0.25
10-6
10
10
10-3
10-2
10
10"
Learning Rate (log scale)
Figure 5.
The subspace distance of representations learned by different algorithms after 1000 iterations and the true representation as function of learning rate.
F,2.3. HEAD-TO HEAD EXPERIMENTS
We again consider the same experimental setting used for Figure In particular, we use data dimension d, = 100, task dimension d, = 15, and representation dimension = 8. We use the same learning rate 10-2 for each optimizer except for NGD optimizer, in which we used 10-4. The batch size 1024. In Figure we considered the Uniform data (13) with high anisotropy. Here, we consider the other three setting: Uniform data (13) with low anisotropy, Gaussian data (38) with low anisotropy, and Gaussian data (38) with high anisotropy. We plot the training loss, subspace distance the ground truth shared representation, and the transfer loss obtained by different algorithms. See Figure 9, We observe that in all three settings, various algorithms converge in training loss. In the case with high anisotropy (second row), methods other than KF AC do not converge to the optimal representation in subspace distance and transfer loss. However, in the low anisotropy settings (first and third rows), the performance of other algorithms also improve, but are notably still suboptimal relative KF AC, confirming the theoretical results showing that anisotropy root cause behind the sub-optimality of prior algorithms and analysis.
40

---

Layer-wise Preconditioning Methods and Provable Feature Learning
10a
Sutspace Distance
Alse AMG
0.75
Transfer Loss
Train Loss
16a
0.00
10a
1000
1500
2000
Number of iterations
Number iterations
Figure Gaussian with low anisotropy
1.00
Subspace Distance
6x 10
Ttain 1036
4x
0.75
Trarsfer Loss
0.50
20a
500
1000
1500
2000
0
500
1000
15:00
2000
500
1000
1500
2000
Number iterations k
Number iterations K
Number of iterations
Figure 7. Gaussian with high anisotropy
161
Subsuce Distance
10a
Transfer Loss
Train Loss
0.50
109
0.00
10a
500
1000
1300
2000
1000
1500
500
1000
1500
2000
Number of ittrations
Number steratons
Number of iterations k
Figure Bernoulli with low anisotropy
Figure From left to right: the training loss, subspace distance, and transfer loss induced by various algorithms linear representation learning task.


Reviews:
[{"review_id": "Xg7fkeQAuZ", "reviewer": "Reviewer_RnJ3", "summary": "This paper shows that layer-wise preconditioning is statistically necessary for efficient feature learning using two common models: linear representation learning and single-index learning. They prove that SGD struggles in non-isotropic inputs, and demonstrate theoretically and experimentally that this suboptimality is fundamental. They show that layer-wise preconditioning naturally addresses this issue. In addition, they show experimentally that Adam and batch normalization help only slightly, and that layer-wise preconditioning is uniquely beneficial even compared to the actual second-order methods that it is approximating.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable comments. We are very glad to see that you find our paper to be a solid contribution  and a valuable step toward understanding the limitations of SGD and the necessity of layer-wise preconditioning. Below please find our response to the comments and questions.\n\n- **More intuition for insufficiency of Adam and Batchnorm.**\n\nRegarding batchnorm, the discussion following Lemma 3.7 can be mathematically formalized. We will reformat this to be a separate lemma that demonstrates how batchnorm/whitening causes convergence to the wrong feature space $\\mathrm{rowsp}(\\tilde{\\mathbf G}) \\triangleq \\mathrm{rowsp}(\\mathbf G_\\star \\Sigma_{\\mathbf x}^{1/2})$, and further provide a simple estimate of $\\mathrm{dist}(\\tilde{\\mathbf G}, \\mathbf G_\\star)$ in terms of the anistropy of $\\Sigma_{\\mathbf x}$.\n\nRegarding the insufficiency of Adam, one of the motivating factors to this paper is that Adam has proven difficult to analyze in *feature learning* contexts, i.e. fine-grained models of neural-network learning, thus motivating the search and analysis of optimizers that are naturally amenable to the compositional / layered structure of neural networks (see e.g. Eq (1)). In our opinion, the reason Adam is not optimal even in our idealized neural network learning settings is that diagonal preconditioning is not optimal for dealing with sources of ill-conditioning arising from the compositional structure of the parameter-space, e.g. $f_{\\mathbf F, \\mathbf G}(x) = \\mathbf F \\mathbf G x$. Since Adam is a diagonal preconditioner, it necessarily suffers from these issues. Furthermore, Adam takes a \"-1/2\" power of its curvature estimate, rather than the full inverse, which necessarily slows its convergence for cases where curvature estimates are reliable (though certainly has other practical merits). Initial probes into the relative effectiveness of diagonal preconditioning have been proposed in concurrent studies, see e.g. https://nikhilvyas.github.io/SOAP_Muon.pdf, https://arxiv.org/abs/2411.12135 (cf. Section 4.3).\n\n- **Related work showing that Adam outperforms SGD when there is class imbalance.**\n\nWe thank the reviewer for pointing out this paper. We will add it to the related works section.\n\n- **Adam in transfer learning settings.**\n\nWe emphasize that our transfer learning set-up is a proxy for more robust notions of generalization, which allows us to diagnose the quality of solutions returned by various algorithms in a more informative way than comparing training loss convergence. Not accounting for curvature (e.g. SGD) or using weaker forms of it (Adam) are therefore exposed by either slow convergence rate overall or by poor \"generalization\" (exposed by poor convergence in subspace distance); see additional plots [(anon. link)](https://drive.google.com/file/d/1yitFmd-Kvqr_bBmrb_L7OhCKQd9RUMyC/view?usp=sharing). Notably, we believe the latter notion is a promising way to diagnose whether curvature-aware/adaptive descent methods capture the geometry of neural-network optimization correctly. For example, despite Adam being an adaptive gradient method, it exhibits poor subspace distance recovery due to still being biased by the (spurious) curvature introduced by the anisotropy of the covariates (compare e.g. to low anisotropy settings in Figure 8). \n \nThe essence of why Adam (or diagonal preconditioners) is a suboptimal feature-learner is the following: many adaptive methods are able to reduce training loss quickly by fitting the \"low-frequency\" features that explain most of the data. Therefore, many candidate solutions can attain low training loss. However, the \"high-frequency\" directions/features that explain less of the *training* objective correspond to *sharper* curvature; weak preconditioners have trouble smoothing these directions, and thus make very slow progress therein, reflected by the subspace distance plot. Mild distribution shifts (e.g. transfer learning) can lead to significant shift in which features are relevant, and thus of the many candidate solutions that get low training loss, it is important to also accurately detect all the possible relevant features like KFAC does in this setting.", "claims_and_evidence": "The main claims about layer-wise preconditioners versus SGD are clear and well-supported by theoretical analysis, with well-explained intuition and implications. Overall, it seems to be a solid contribution.  \n\nPerhaps one limitation would be that the experiments are somewhat limited, and while the paper shows that Adam (diagonally preconditioned methods) and batch norm are insufficient, it lacks deeper intuition or analysis on why that’s the case.", "methods_and_evaluation_criteria": "The theoretical examples and analysis and interesting and I believe appropriate for the question. Experimentally, the subspace distance seems to correlate well with whether transfer learning is successful.", "theoretical_claims": "I didn't check the proofs line by line, but they seem reasonable. The authors explained the intuitions and implications of each lemma/theorem well.", "experimental_designs_or_analyses": "I didn't rerun the experiments to check correctness, they seem reasonable. My complaint is that the experiments seem to be limited to small-scale examples and transfer learning settings. Though I understand empirical results are not the main focus of this paper, it would be very interesting to see if these observations hold in bigger, more practical cases.", "relation_to_broader_scientific_literature": "It’s observed in existing literature and practice that K-FAC-like methods (Shampoo for example) can outperform SGD in training and generalization. While this paper’s theoretical analysis is limited to Gaussian data and small-scale tasks (e.g., two-layer networks for linear representation learning and single-index models), it’s still a valuable step toward understanding the limitations of (S)GD and the necessity of layer-wise preconditioning. The paper also briefly touches on why Adam's advantage is limited in such settings, though more in-depth study of that could be helpful.", "essential_references_not_discussed": "I find the related work pretty well cited in general.\n\nI have seen in the following work that Adam can outperform SGD when there is a class imbalance in the data, which could be related to the problem condition. Would be interesting to see the discussion on whether this is related (but feel free to use your judgment).\n\n[F. Kunstner et.al, Neurips 2024. Why Adam Outperforms Gradient Descent on Language Models: A Heavy-Tailed Class Imbalance Problem]", "other_strengths_and_weaknesses": "Other strength: The paper is educational and well-written, providing a clear discussion of optimization history and why the problem is well-motivated. \n\nOther weakness: The analysis is limited to specific settings, though this is fairly common in theoretical work. Additionally, the experiments are somewhat limited in scope.", "other_comments_or_suggestions": "No", "questions_for_authors": "The experiments that contain Adam seem to be limited to the transfer learning settings. Is there any intuition on why layer-wise preconditioning is better than diagonal (Adam) in this setting? Is there a difference in general, outside of the transfer learning setting?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XQLeLgyzxk", "reviewer": "Reviewer_NYms", "summary": "This paper demonstrates KFAC is better at feature learning than vanilla SGD with two model examples. For linear representation learning, they show the convergence rate of SGD will suffer from the condition number while KFAC gets rid of it. For single-index learning, they show that one-step update of SGD can only learn the correct direction with isotropic data while the one-step update of KFAC can learn the direction under isotropy.  They conduct experiments to show the effectiveness of their theory.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable comments. We are glad that you find our results rigorous, and that our results make a good connection between feature learning and KF algorithms.\n\n- **Anisotropy vs ill-conditioning.**\n\nWe clarify that we make two distinct claims in Section 3.1. We identify *two* sources of suboptimality in SGD for linear rep. learning: the first sourced from the anisotropy of the inputs, and the second coming from the ill-conditioning of the output layer $\\mathbf F$. The bias from anisotropy is critical, and can prevent convergence *altogether* (as seen in AMGD/SGD lines in Figure 1). As pointed out in our literature review, this observation has escaped notice, apart from a recent work [1], whose algorithm is a special case of our recipe in Eqs (4) and (5); see DFW in Fig. 1. The second issue is what the lower bound in Prop. 3.5 concerns. We assume isotropy therein (and multiple favorable assumptions) to make the problem *as benign as possible* for SGD, which we showed does not converge without isotropy. The point of Prop. 3.5 is that *even under ideal conditions* the convergence rate of SGD and related methods is suboptimal compared to KFAC as shown in Theorem 3.6. We hope this clarifies our separate contributions.\n\n- **Real-world experiments.**\n\nWe believe it is well-documented that KFAC performs well in practice (see e.g. [2, Sec. 13] and [3]). In fact, KFAC was designed as a *practical* approximation of the Natural Gradient method. Our main goal is to resolve the limitations of SGD in prominent feature learning theory setups. As a byproduct, we derive KFAC as the natural solution, providing an alternate *feature-learning* justification for KFAC. Our experiments are solely to verify and complement our theoretical findings.  \n\n- **Relevance of Shampoo.**\n\nWe are aware that methods such as Shampoo and Muon have received a lot of attention lately. However, we should note that up until last year, KFAC has been largely the more well-known method, seeing many extensions and applications. We also mention that making Shampoo outcompete the Adam family (e.g. the winning submission in AlgoPerf) takes additional heuristics like learning-rate grafting [4, Sec. 2.4]. However, **we have implemented basic versions of Shampoo**, and run it in the set-up for Figure 1 [(anon link)](https://drive.google.com/file/d/1MS4PjZXTXVidzvzjmoM1P1X7mKz7_Ih6/view?usp=sharing). SHMP(G) corresponds to updating $\\mathbf G$ with Shampoo, and $\\mathbf F$ with least-squares (like AMGD and KFAC), and SHMP is Shampoo on both layers. We observe that Shampoo performs roughly on par with NGD, but not as well as KFAC. This is expected, since we derived KFAC to provably get *constant-factor* linear convergence. We believe that exploring the feature learning properties of Shampoo is an interesting future direction.\n\n- **Regularized $\\mathbf P_{\\mathbf G}$.** \n\nIn the single-index setting, it is more natural to consider regularized $\\mathbf P_{\\mathbf G}$ because analyses are traditionally in the proportional limit $d_x \\asymp n$, where a non-zero regularizer is known to be optimal. In fact, the result we provide (Lemma 3.10 and 3.12) are more general and hold even for $\\lambda_{\\mathbf G} \\to 0$. In Figure 3 (Left), we still see that for $\\lambda_{\\mathbf G} \\to 0$, our theory matches the simulations, and that the alignment of the direction learned by KFAC to the true direction is still significantly larger than that of SGD.\n\n- **Negative result for 2nd-order methods.**\n\nWe clarify that the motivation for our paper is that many algorithms are designed to approximate 2nd-order methods, but neural network analyses for 2nd-order methods are still lacking. Therefore, our goal is not to debunk 2nd-order methods, but to suggest a direct motivation for the approximation methods themselves. We refer the reviewer to **Outperforming dense 2nd-order variant** of our response to Reviewer AMdp for more discussions.\n\n- **Algorithm block and $\\overline{\\mathbf G}$.** \n\nWe thank the reviewer for pointing these out. They have been fixed in the revision.\n\n- **One giant step vs multi-step.** \n\nThe one-step update for learning a single index model is a popular framework in deep learning theory and is the state of the art model to theoretically analyze feature learning properties of two-layer networks (cf. line 283 in the submission). In this setting, features are learned by taking a single (giant) step on the first layer as opposed to multiple small steps. A single step can often be shown to be equivalent to taking multiple steps with smaller step sizes (Section B.1.3 of [5]).\n\n- **Is Adam a good feature learner?** \n\nFor insights on why Adam has fundamental limitations in our feature learning setting, we refer to **Adam in transfer learning settings.** in our response to Reviewer RnJ3.\n\n[1] arxiv.org/abs/2308.04428\n\n[2] arxiv.org/abs/1503.05671\n\n[3] arxiv.org/abs/2311.00636\n\n[4] arxiv.org/abs/2309.06497\n\n[5] arxiv.org/abs/2205.01445", "claims_and_evidence": "The author claims they show SGD can be drastically slow on linear representation learning under anisotropy. But I only saw they cite previous papers on this claim and proposition 3.5 assumes isotropy. It can not viewed as one contribution by this paper without further explanation.", "methods_and_evaluation_criteria": "The two representation learning models are commonly used in this area. There are two synthetic experiments for each models. They are already reasonable but I am hoping for some real world dataset.", "theoretical_claims": "I roughly checked the proof of proposition 3.5 and didn’t find any obvious mistakes.", "experimental_designs_or_analyses": "I checked the experiment setting and results and didn’t find any serious issue.", "relation_to_broader_scientific_literature": "It can help understand the difference between optimization algorithms and why some algorithms can succeed on some tasks.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strength:\n1. This paper makes a good connection between feature learning and kronecker-factored algorithms. \n2. The theoretical claims and proof look well-written and relatively rigorous. \n\nWeakness:\n1. The theoretical results are only for one specific layer-wise preconditioning method. Since Shampoo is a more popular kronecker-factored optimization algorithm than KFAC, it should be more interesting to get the results for Shampoo. \n2. Different $P_G$ are used for KFAC in the two models. It would be better if lemma 3.12 is proved with unregularized $P_G$. \n3. The negative result is only shown for SGD. As mentioned in the introduction, it is unknown why kronecker-factored algorithm can be better than the idealized second-order methods. If we want to claim the reason might be KFAC is better at feature learning, then the negative results of second-order methods are expected.", "other_comments_or_suggestions": "1. It’s better to explicit show each algorithm with some algorithm block and keep a consistent notation for $G_+$. It is defined differently in (5) and (8) and looks very confusing every time it is mentioned later. For example, I think it should be $\\bar{G}_+$ in the LHS of (9). \n2. Typo in line 280: $F_\\star^{(t)}$ should be $F_{ls}^{(t)}$?\n3. Typo in line 401: The distribution of $x_i^s$ seems unfinished.", "questions_for_authors": "1. I am not convinced why we only care about the one-step update in the single index learning. Even if the one-step update of SGD is bad, we almost never update the parameter only for once so SGD might be able to learn the correct direction after several steps. \n2. Do you try Shampoo for the first experiment? If the theoretical results can not be easily extended to Shampoo, I am curious whether Shampoo can learn features well in practice. That can help understand whether the connection to feature learning is applicable to a broader class of kronecker-factored algorithms rather than only KFAC. \n3. I am also curious whether there can be any negative theoretical results for Adam. From figure 1, Adam can learn features much better than any other algorithms and is only slightly worse than KFAC. If there is no obvious negative theoretical result, is there any insight why Adam can be good at learning features?\n4. As mentioned in weakness 2, I wonder there is any obstacle in proving lemma 3.12 with unregularized $P_G$.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9HuVSvRfKo", "reviewer": "Reviewer_AMdp", "summary": "The authors aim to study the effectiveness of layer-wise preconditioning methods with per-axis preconditioning. They show that the Kronecker-Factored structure naturally occurs for the problem of linear representation learning and single index learning. This provides a potential explanation for why these methods can outperform their exact second-order counterpart.\n\n## update after rebuttal\n\nIn the rebuttal, the authors state that they avoid making general claims about Kronecker-Factored preconditioning being better than element-wise/diagonal preconditioning for feature learning, or that factored curvature estimates always lead to faster convergence than the dense versions.\n\nIf that is the case, then I believe some parts of the writing can be improved to avoid giving out this impression. For example, in the introduction section, they first state that KFAC generally outperforms its ideal counterpart NGD and then state that they want to give explanation to the performance of Kronecker-Factored methods. This could give the impression that they want to claim evidence for KFAC better than not only SGD but also NGD.\n\nIn either case, this means the paper only try to show that Kronecker-Factored methods are better than SGD, which is solid work, but also much more expected so I keep more original score.\n\nI also verified the authors claim about rowspan is correct. I was originally thinking about left-multiplication instead of right multiplication.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable comments. We are glad that you find that these results provide explanations for the empirical success of Kronecker-Factored preconditioning methods, and that it can potentially guide algorithm design. Below please find our response to the comments and questions.\n\n- **Indirect support of certain claims.**\n\nWe would like to clarify that the main goal of our paper is to address key issues of relying on SGD in feature learning theory, and in doing so establish a feature-learning motivation for layer-wise preconditioned optimizers. We avoid making general claims about Kronecker-Factored preconditioning being better than element-wise/diagonal preconditioning for feature learning, or that factored curvature estimates always lead to faster convergence than the dense versions. Rather, specific observations of these phenomena in prior literature are what *motivate* our study of Kronecker-Factored preconditioning. We provide some additional information and context for each.\n\n1. **KF preconditioning is necessary vs element-wise.**\n\nCurrent literature notably lacks work showing provable feature learning capabilities of diagonal preconditioning methods like Adam, despite a huge amount of work studying SGD in this context. Anecdotally, this is because the entry-wise (non-linear) operations that Adam performs are somewhat unamenable to analyze given the linear-operator structure of weights in neural networks. Therefore, by straightforwardly deriving feature learning properties of a representative Kronecker-Factored method, we see our work as evidence that they are a more natural NN-oriented optimizer class. On the other hand, we note diagonal preconditioning is *not* a strict subset of Kronecker-Factored preconditioning: for a given entrywise product $\\mathbf M \\odot \\mathbf G$, generally there does not exist $\\mathbf P, \\mathbf Q$ such that $\\mathbf M \\odot \\mathbf G = \\mathbf P \\mathbf G \\mathbf Q$, which is why we do not claim that diagonal preconditioning is generally worse than KF preconditioning for feature learning, even if KF preconditioning may be more mathematically natural.\n\n2. **Outperforming dense 2nd-order variant.**\n\nAs stated in our introduction, many deep learning optimizers, including Adagrad (Adam), KFAC, Shampoo etc. were initially derived as computationally-efficient approximations to full 2nd-order methods like (Gauss-)Newton method and Natural Gradient Descent. Therefore, a natural conclusion is that *if one had the resources*, the full 2nd-order method should yield better convergence. However, the behavior of full 2nd-order methods on neural networks is notoriously poorly understood from theory, with basic questions such as \"are negative curvature directions (negative Hessian eigs) in Newton's method good?\" and \"is Gauss-Newton (guaranteed psd) preferable to Newton?\" having no conclusive answers [1]. Therefore, our contribution can be seen as *directly* deriving the benefits of the approximating method *without* reasoning about the full 2nd-order method. The fact that the approximant numerically *outperforms* the full method is only further evidence that this is a fruitful path of analysis, though we note this observation is not original to us [2].\n\n- **Correctness of rowspan.**\n\nThe math in line 302 should be correct. To see how right-multiplication by a matrix can change the rowspace, consider the simple example $\\mathbf G = [1 \\quad 0] \\in \\mathbb R^{1 \\times 2}$, whose rowspace is  $\\{[c \\quad 0], c \\in \\mathbb R\\}$. Multiplying by a (psd) matrix $\\Sigma = \\begin{bmatrix}2 & 1 \\\\\\\\ 1 & 2 \\end{bmatrix}$ yields $\\mathbf G \\Sigma = [2\\quad 1]$, which is not contained in $\\mathrm{rowsp}(\\mathbf G)$. To tie this back to batchnorm, our observation is that whitening the training-distribution covariates allows SGD to converge (since the covariates are now isotropic) but changes the rowspace, breaking the shared structure between the source and transfer distributions.\n\n[1] https://arxiv.org/abs/1503.05671\n\n[2] https://arxiv.org/abs/2201.12250", "claims_and_evidence": "It feels like some of the claims are only supported indirectly without being articulated clearly. Specifically, two of the more interesting claims,\n\n1. Kronecker-Factored structured preconditioning is necessary and thus element-wise preconditioning like Adam is not enough,\n\n2. There are theoretical justification for why Kronecker-Factored structured preconditioning can outperform their exact second-order counterpart,\n\nare discussed very little in the paper. Most of the main paper is about how SGD can fail, which is more expected and less interesting to the readers.", "methods_and_evaluation_criteria": "The evaluation setting is quite simple so that it aligns with the analyzed theoretical setting. Since the purpose of the experiments is just to verify the theory, this seems reasonable.", "theoretical_claims": "I do not find any issues with their theoretical claims.", "experimental_designs_or_analyses": "I do not find any issues with their analyses.", "relation_to_broader_scientific_literature": "The paper provides potential explanations to the empirical success of Kronecker-Factored preconditioning methods. This might lead to the development of new optimization methods.", "essential_references_not_discussed": "I do not find any essential references that are missed by the authors.", "other_strengths_and_weaknesses": "The writing seems to be a weakness of the paper. Some of the claims seem not to be clearly articulated.", "other_comments_or_suggestions": "I have no other suggestions.", "questions_for_authors": "1. I am not sure if the equality about rowspan on line 302 is correct. It seems to me that right multiplication with a matrix will not change the rowspan. (The authors might be considering some noise term to the solution of G, which cannot be clearly seen from the equation)", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WCdtKuUG6f", "reviewer": "Reviewer_aaxq", "summary": "The author show theoretically how, in the settings of linear representation and single index learning with data with non-trivial covariance structures using a two-layer network, SGD suffers from fundamental statistical biases and limitations, and how the use of a simple Kronecker-factored preconditioning scheme allows to overcome those issues. More precisely, in the case of linear representation learning, they establish how the preconditioned dynamics are endowed with condition number-free guarantees, in contrast to SGD, which is suboptimal. For single-index learning, they study the spike developed by the first layer weights following a single gradient step. They show that a SGD step leads to a biased spike compared to the target weights, while a preconditioned step does not. They corroborate their findings in numerical experiments, and include a comparison with Adam and batch-normalization schemes, which are suboptimal in the considered settings.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback. We are glad that you found our results regarding the fundamental limitations of SGD and the natural emergence of preconditioning schemes as solutions to address these challenges interesting. Please find our detailed response to the comments and questions below.\n\n- **Defining $\\kappa, \\sigma_{\\min}$ in notation.**\n\nWe thank the reviewer for pointing this out. The parameters $\\kappa$ and $\\sigma_{\\min}$ correspond to condition number and the smallest singular value respectively. We have added these to the revision.\n\n- **Dependence of the locality the requirement on $\\kappa, \\sigma_{\\min}$.**\n\nThe presence of $\\kappa, \\sigma_{\\min}$ in the initialization requirement is rather subtle. Firstly, we note that locality requirements are distinct from the rate; an analogous bound for SGD in the isotropic setting will have both the initialization requirement *and* a rate that contains $\\kappa(\\mathbf F)$ [1, Thm 3.1-3.2], not to mention that SGD is non-convergent under anisotropy $\\kappa(\\Sigma_{\\mathbf x}) \\gg 1$ (cf. Figure 1). Secondly, similar requirements are ubiquitous in linear representation learning guarantees; however, it is generally believed that initialization is not required for this problem [2, Sec. 6], [3, Remark 3.2]. On the other hand, linear representation learning can be viewed as a low-rank matrix-sensing problem [1], and as far as we know there does not exist a general *global* convergence guarantee for that problem, let alone for the anisotropic sensing matrices our setting induces. We lastly remark that the initialization requirement does not appear in a stylized, noiseless setting $\\sigma_{\\varepsilon} = 0$, where we fix $\\mathbf F_0$ at initialization and solely perform KFAC updates on $\\mathbf G$. In particular, for each iteration $t$ we have\n$$\n\\mathbf G_t = \\mathbf R_0 \\mathbf G_\\star + \\left(1-\\eta_{\\mathbf G}\\right)^t \\left(\\mathbf G_0 - \\mathbf R_0 \\mathbf G_\\star\\right), \\\\;\\\\; \\mathbf R_0 = (\\mathbf F_0^\\top \\mathbf F_0)^{-1}\\mathbf F_0 \\mathbf F_\\star.\n$$\nThus, for $\\eta_{\\mathbf G}$ close to $1$, $\\mathbf G_t$ converges linearly to $\\mathbf R_0 \\mathbf G_\\star$, and for full-rank $\\mathbf F_0$, $\\mathrm{rowsp}(\\mathbf R_0 \\mathbf G_\\star) = \\mathrm{rowsp}(\\mathbf G_\\star)$.\n\n\n- **High-probability bound for SGD and role of Proposition 3.5.**\n\nWe note that we identify two suboptimalities of SGD. Firstly, given anisotropic covariates, SGD can be non-convergent. However, even given isotropic covariates, SGD is unavoidably affected by the conditioning of the output layer $\\mathbf F$. Convergence upper bounds analogous to Theorem 3.6 for SGD can be found in [1, 4]. The role of Proposition 3.5 is therefore to show that *even under extremely ideal conditions* such as isotropic inputs, SGD necessarily suffers when $\\mathbf F_\\star$ is ill-conditioned. This bound is indeed worst-case in that the adversary chooses $\\mathbf G_0$, but should be interpreted as \"SGD generally suffers from ill-conditioning\", whereas KFAC's convergence rate is always condition-number-free. As noted after Proposition 3.5, an instance-specific SGD convergence rate *could* in principle be estimated and will almost surely involve a function of the spectrum of $\\mathbf F_\\star$, but is besides the point of our paper.\n\n- **Role of ridge parameter $\\lambda_{\\mathbf G}$ in linear representation experiments.**\n\nWe have set $\\lambda_{\\mathbf G} \\to 0$ in Figures 1, 5, and 8. We conduct additional experiments for the same experimental setup, but with varying degrees $\\lambda_{\\mathbf G}$ [(anon. link)](https://drive.google.com/file/d/1n6OmlzOpo0hRFp2YOwkHDun2bo63HjK3/view?usp=sharing). We observe that the regularization does not significantly alter the qualitative behaviors of the curves for a wide range of $\\lambda_{\\mathbf G}$.\n\n[1] https://arxiv.org/abs/2102.10217\n\n[2] https://arxiv.org/abs/2105.08306\n\n[3] https://arxiv.org/abs/2308.04428\n\n[4] https://arxiv.org/abs/2102.07078", "claims_and_evidence": "The theoretical claims of Theorem 3.9 are quantitatively verified in Fig. 2, where they are compared to numerical experiments, with convincing agreement. The theoretical insights on the suboptimality of SGD and benefits of KFAC are illustrated in numerical experiments in Figs. 1,3.", "methods_and_evaluation_criteria": "The paper is primarily theoretical in nature, and all experiments are conducted within the two stylized settings (linear representation task and single-index learning) considered in the theoretical parts.", "theoretical_claims": "From my reading, the claims look reasonable and sound, but I did not check the proofs. \n\nI have minor concerns and questions regarding Proposition 3.5 and Theorem 3.6. In Theorem 3.6, the notation $\\kappa$ and $\\sigma_{min}$ are, unless I am mistaken, not introduced. I assume they refer to the condition number and the $\\lambda_{min}$. In this case, would the range of application of Theorem 3.6 not also suffer from the conditioning number, as an increasingly close initialization to the ground truth is required? If this is the case, the statement under Theorem 3.6, although still valid, would gain to be nuanced by discussing this fact. \nAs a minor question, is it possible to reach a similar high-probability local convergence for SGD? Even if impossible or unclear, I believe it would aid to grasp better the comparison with KFAC, as for now Proposition 3.5 is a worst-case result. On the other hand, I do believe the experiments of Figs. 1,5,9 compellingly show the superiority of KFAC. This is meant as a minor suggestion.", "experimental_designs_or_analyses": "I did not check the experimental designs in detail, however, the exposition of the setup provided in the main text is sufficiently detailed and I do not see any issue.\n\n(Minor question) What is the value of $\\lambda_G$ in Figs. 1,5,8? Do the qualitative behaviors of the curve depend on this parameter ?", "relation_to_broader_scientific_literature": "The alternating SGD scheme studied is related to algorithms appearing in (Collins et al., 2023) or (Zhang et al., 2024), but seems more general, although I do not have a large familiarity with this line of works.\n\nThe single-step analysis of Theorem 3.9 is a generalization of Proposition 2 and Theorem 3 of (Ba et al., 2022) to anisotropic covariates, which is to the best of my awareness a novel result.", "essential_references_not_discussed": "I did not identify any essential reference which the authors fail to discuss. However, I have limited familiarity with the literature on alternating or preconditioned descent methods, so it is possible I overlooked some works.", "other_strengths_and_weaknesses": "I am overall in favor of acceptance. The results contained in the paper are interesting, as they highlight the fundamental limitations of vanilla SGD and how preconditioning schemes naturally emerge as pathways to mitigate these issues. Due to my limited expertise, I am unable to assess confidently the novelty or significance of the results presented in section 3.1. I have left a few questions and comments in the above sections.", "other_comments_or_suggestions": "I do not have particular comments or suggestions.", "questions_for_authors": "I do not have any particular question, beside those I listed in \"Theoretical Claims\" and \"Experimental Design\"", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Dear Authors,

Thank you for submitting your paper to ICML and for contributing a theoretically grounded and well-motivated analysis of layer-wise preconditioning in the context of feature learning. Your work explores the statistical and algorithmic advantages of Kronecker-factored preconditioning methods, demonstrating that they uniquely resolve fundamental limitations of SGD.

The reviewers generally agreed that this is a well-written and impactful contribution to our understanding of optimization in deep learning. Several noted that your theoretical results were rigorous and complemented by clearly designed experiments on certain stylized tasks. While some reviewers expressed interest in further empirical validation on larger or real-world datasets, the conceptual clarity and theoretical insights were recognized as a core strength. The rebuttal further addressed reviewer questions about generalization to Shampoo, the role of initialization, and the scope of your claims regarding second-order methods, preconditioners etc.

Given the strength of the theoretical contributions, the clarity of exposition, and the positive consensus among reviewers, I am happy to recommend acceptance to ICML 2025. I encourage you to use the camera-ready version to further refine the exposition, clarify your scope of claims (particularly relative to full second-order methods), and highlight directions for extending the work to broader optimizer families or empirical settings.

Congratulations on your contribution.

Best regards,
AC
Decision: Accept (poster)