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

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.

• Defining $\kappa, \sigma_{\min}$ in notation.

We 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.

• Dependence of the locality the requirement on $\kappa, \sigma_{\min}$.

The 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

Thus, 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)$.

• High-probability bound for SGD and role of Proposition 3.5.

We 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.

• Role of ridge parameter $\lambda_{\mathbf G}$ in linear representation experiments.

We 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). We observe that the regularization does not significantly alter the qualitative behaviors of the curves for a wide range of $\lambda_{\mathbf G}$.

[1] https://arxiv.org/abs/2102.10217

[2] https://arxiv.org/abs/2105.08306

[3] https://arxiv.org/abs/2308.04428

[4] https://arxiv.org/abs/2102.07078

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.

• Indirect support of certain claims.

We 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.

1. KF preconditioning is necessary vs element-wise.

Current 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.

2. Outperforming dense 2nd-order variant.

As 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].

• Correctness of rowspan.

The 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.

[1] https://arxiv.org/abs/1503.05671

[2] https://arxiv.org/abs/2201.12250

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.

• Anisotropy vs ill-conditioning.

We 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.

• Real-world experiments.

We 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.

• Relevance of Shampoo.

We 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). 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.

• Regularized $\mathbf P_{\mathbf G}$.

In 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.

• Negative result for 2nd-order methods.

We 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.

• Algorithm block and $\overline{\mathbf G}$.

We thank the reviewer for pointing these out. They have been fixed in the revision.

• One giant step vs multi-step.

The 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]).

• Is Adam a good feature learner?

For 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.

[1] arxiv.org/abs/2308.04428

[2] arxiv.org/abs/1503.05671

[3] arxiv.org/abs/2311.00636

[4] arxiv.org/abs/2309.06497

[5] arxiv.org/abs/2205.01445

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.

• More intuition for insufficiency of Adam and Batchnorm.

Regarding 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}$.

Regarding 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).

• Related work showing that Adam outperforms SGD when there is class imbalance.

We thank the reviewer for pointing out this paper. We will add it to the related works section.

• Adam in transfer learning settings.

We 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). 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).

The 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.