The Feature Speed Formula: a flexible approach to scale hyper-parameters of deep neural networks

Authors: lénaïc  chizat
praneeth  netrapalli
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 21:06:23

Abstract: Deep learning succeeds by doing hierarchical feature learning, yet tuning hyper-parameters (HP) such as initialization scales, learning rates etc., only give indirect control over this behavior. In this paper, we introduce a key notion to predict and control feature learning: the angle $\theta_\ell$ between the feature updates and the backward pass (at layer index $\ell$). We show that the magnitude of feature updates after one GD step, at any training time, can be expressed via a simple and general *feature speed formula* in terms of this angle $\theta_\ell$, the loss decay, and the magnitude of the backward pass. This angle $\theta_\ell$ is controlled by the conditioning of the layer-to-layer Jacobians and at random initialization, it is determined by the spectrum of a certain kernel, which coincides with the Neural Tangent Kernel when $\ell=\text{depth}$. Given $\theta_\ell$, the feature speed formula provides us with rules to adjust HPs (scales and learning rates) so as to satisfy certain dynamical properties, such as feature learning and loss decay. We investigate the implications of our approach for ReLU MLPs and ResNets in the large width-then-depth limit. Relying on prior work, we show that in ReLU MLPs with iid initialization, the angle degenerates with depth as $\cos(\theta_\ell)=\Theta(1/\sqrt{\ell})$. In contrast, ResNets with branch scale  $O(1/\sqrt{\text{depth}})$ maintain a non-degenerate angle $\cos(\theta_\ell)=\Theta(1)$. We use these insights to recover key properties of known HP scalings (such as $\mu$P), and also introduce a new HP scaling for large depth ReLU MLPs with favorable theoretical properties.

Main paper: 
 introduction: 
 The ability of deep Neural Networks (NNs) to learn hierarchical representations of their inputs has been argued to be behind their strong performance in data-intensive machine learning tasks [LeCun et al., 2015 ]. Yet, the process via which gradient-based training leads to feature learning remains mysterious and defies our intuition; some architectures can even reach zero loss without feature learning at all [Jacot et al., 2018] . This limited understanding makes it difficult to design NNs architectures and hyper-parameters (HP) scalings, and begs the development of tools to quantify feature learning.In this paper, we demonstrate that the backward-feature angle (BFA) θ between the feature updates and the backward pass (at layer index ) is a central object in this quest. Indeed, we show that the magnitude of feature updates after one GD step, at any training time, can be expressed via a simple feature speed formula in terms of this angle, the loss decay and the magnitude of the backward pass. Given the knowledge of θ , this leads to a general approach to quantify key dynamical properties of the training dynamics of NNs -such as the speed of feature learning and loss decay -and to characterize the HP scalings satisfying these properties.Contributions Our contributions are the following:• We prove the feature speed formula (Thm 2.1) which quantifies the feature updates in terms of the BFA θ , the loss decay and the magnitude of the backward pass at layer . This formula, valid in any architecture and with an elementary proof, helps exploring the space of HP scalings, and understanding when feature learning arises.• In Section 3, we develop tools to quantify the BFA. In particular, we show that, in the batch-size 1 case, θ can be estimated in terms of the spectrum of the backward to feature kernel (BFK) K and is related to the conditioning of layer-to-layer Jacobians (Thm. 3.2).We study the case of MLPs and ResNets at random initialization, and obtain that for a depth L, cos(θ L ) = Θ(L -1 /2 ) for ReLU MLPs (Prop. 5.1, exploiting a result in Jelassi et al. [2023] ) and that cos(θ L ) = Θ(1) for linear ResNets with branch scale O(L -1 /2 ) (Prop. 5.2).• In Section 4, we consider several properties of NN training dynamics that can be conveniently studied with our tools, including feature learning and loss decay. Enforcing these properties leads to explicit contraints on the magnitude of the forward, backward pass and learning rates in general architectures (Prop. 4.1).• In Section 5, we show how various HP scalings for large width-then-depth MLPs and ResNets can be characterized by enforcing these properties. In particular we recover depth µP [Bordelon et al., 2023 , Yang et al., 2023b] for ResNets (Table 2 ) and, for ReLU MLP, we introduce a scaling with output scale √ depth width (Table 1 ) that does not suffer from vanishing loss decay, in contrast to the one studied in [Jelassi et al., 2023] .• Finally, in Section 6 we develop a more "axiomatic" approach: starting from a minimal list of desiderata, which include a notion of gradient stability, we show that we recover, in a certain extent, the convenient properties considered in Section 4. This section focuses on homogeneous architectures for which we show, along the way, a backward speed formula (Prop. 6.1) and an invariance under block-wise rescaling (Prop. 6.2).The theory of NNs has recently benefited from important insights from asymptotic analyses in the large width and/or depth limits. Our work is in the continuity of those.Analyses of wide and deep NNs at random initialization led to identifying critical initialization scalings that enable signal propagation [Poole et al., 2016 , Hanin and Rolnick, 2018 , Hanin, 2018] . They also identified dynamical isometry [Pennington et al., 2018] , namely the concentration of the singular spectrum of the layer-to-layer Jacobians around 1, as an important indicator of training performance. Our analysis gives a concrete justification of the link between dynamical isometry and successful training, as we show that it is related to the alignment between the backward pass and feature updates. These questions have also been studied in ResNets, see e.g. [Hayou et al., 2021 , Marion et al., 2022 , Li et al., 2021] for signal propagation and [Tarnowski et al., 2019, Ling and Qiu, 2019] for dynamical isometry.In 2018, two viewpoints for the dynamics of wide NNs were simultaneously introduced: a feature learning limit for two-layer MLPs [Mei et al., 2018 , Chizat and Bach, 2018 , Rotskoff and Vanden-Eijnden, 2018 ] and a limit without feature learning for general NNs [Jacot et al., 2018 , Du et al., 2018 , Allen-Zhu et al., 2019] . These works highlighted the crucial role of HP scalings -learning rates and initialization -in the behavior of large NNs [Chizat et al., 2019] .In order to classify HPs scalings, [Yang and Hu, 2021] formulated the maximal update µ-criterion (it is part of the properties we study in Section 4). This criterion led to a full classification of HP scalings in the infinite hidden width limit (at fixed depth), and singled-out the so-called µ-parameterization (µP) as ideal for this criterion. We note that, provided alignment holds, our analysis allows in particular to characterize µP in an elementary way. See also [Yang et al., 2023a] for another simple derivation of µP using matrix spectral norms (but that does not give tight control on the magnitude of feature learning and does not a priori apply to the large depth asymptotics). Several works have since shown the practical value of these analyses in predicting the behavior of NNs [Vyas et al., 2023] and improving HP tuning [Yang et al., 2021] .When restricted to the output layer of a NN, our notion of alignment/BFA coincides with that studied in Baratin et al. [2021] , Atanasov et al. [2021] , Lou et al. [2022] , Wang et al. [2022] and the BFK we consider coincides with the NTK [Jacot et al., 2018] . We extend these concepts to study and quantify feature learning at any layer (not just at the output layer). Here we focus on the batch-size 1 setting, but the large batch-size setting of these works is a natural next direction for our analysis.Finally, several recent works have studied feature learning in infinite width and depth NNs, starting with [Jelassi et al., 2023] for MLPs, and [Bordelon et al., 2023 , Yang et al., 2023b] for ResNets. The two latter identified the 1/ √ depth branch scaling as providing desirable properties, in particular that of HP transfer [Yang et al., 2021] . These works take the infinite width limit as a first step in their analysis, before studying the resulting objects, resulting in a technical analysis. In our approach, we first take the step-size to 0 (as in [Jelassi et al., 2023] ) and study in detail the structure of the back-propagation equations, before taking the large width-then-depth limit, as a last step.Notations For integers a, b ∈ Z, we write [a : b] = {a, . . . , b}. For any vector x ∈ R m we denote by x rms := m -1 /2 x 2 its root mean-square (RMS) norm. We use this as a proxy for the typical entry size of a vector, which is justified as long as that vector is dense.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification:Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
the feature speed formula: 
 Consider a depth-L NN architecture defined by the recursion, for∈ [1 : L], f 0 ∈ R m0 , f = T (f -1 , w ) ∈ R m , L = loss(f L ) ∈ R(1) where w ∈ R p are trainable parameters and we assume that the maps T : R m -1 × R p → R m admit elementary (log-exp) selections 1 [Bolte and Pauwels, 2020] . By flattening the tensors, one can encode most practical NN architectures in Eq. ( 1). For instance, m 0 is typically the product of batch-size (or context length) with input dimension. We denote by b = ∂L ∂f ∈ R m the vectors of the backward pass. A gradient descent (GD) step with layerwise learning-rate (LR) η • δt > 0 for ∈ [1 : L] consists in adding to each w the updateδw = -η • δt • ∇ L = -η • δt • ∂L ∂w .We are interested on the evolution of the NN over a single GD step with infinitesimally small step-size δt 1. For any quantity x associated to the NN, we denote ẋ its instantaneous velocity ẋ := lim δt↓0 δx δt when it exists. In particular, we have ẇ = -η ∇ L. The following identity is the seed of our approach. It expresses at any training time the speed of features in terms of other interpretable quantities, including the backward to feature angle (BFA) θ v . Theorem 2.1 (Feature speed formula).Let v ∈ [1 : L]. If ≤v η ∇ L 2 2 = 0 then ḟv = 0.Otherwise, the (non-oriented) angle θ v between ḟv and -b v is well defined in [0, π /2[ and it holdsEQUATION)Proof. By the chain rule, we have ḟv =≤v ∂fv ∂w ẇ = - ≤v η ∂fv ∂w ∂L ∂w . It follows -b v ḟv = ≤v η ∂L ∂f v ∂f v ∂w ∂L ∂w = ≤v η ∂L ∂w ∂L ∂w = ≤v η ∇ L 2 2 . (3) Clearly, if ≤v η ∇ L 2 2 = 0 then ẇ = 0 for ≤ v and thus ḟv = 0. Otherwise θ v is well defined and it holds b v 2 ḟv 2 cos(θ v ) = -b v ḟv = ≤v η ∇ L 22 and the claim follows. (In terms of the BFK defined below, Eq. ( 3) is equivalent to b v K v b v = ≤v η ∇ L 2 2 , for v ∈ [1 : L].)To better appreciate the content of Thm. 2.1, let us re-express it in terms of root mean-square (RMS) norms. Let L≤v := ≤v η ∇ 2 2 be the contribution to the loss decrease of all the parameters before f v in the forward pass, and note that L≤L = L. Then, the identity (2) rewrites asEQUATIONHere S v can be interpreted as the sensitivity (and, in the terminology of [Chizat et al., 2019] , 1/S v as the laziness) of the feature v: it is the proportionality factor between loss decay and feature speed. This formula is valid at any training time and involves three key quantities: the scale of the backward pass b v rms , the size of the feature m v , and the BFA θ v . Let us now build tools to quantify the BFA.3 Quantifying the backward-feature angles (BFA)Information about the BFA θ v can be gained from the Backward to Feature Kernel (BFK).Definition 3.1 (Backward to Feature Kernel). For v ∈ [1 : L], the BFK is the psd matrix defined asEQUATION)By construction, it holds ḟv = -K v b v . In other words, the BFK takes a backward pass vector as input and returns the (negative of the) feature velocity. For v = L, K v coincides with the Neural Tangent Kernel [Jacot et al., 2018] . We now show how the sprectrum of K v relates to BFA.Theorem 3.2. Let λ 1 ≥ • • • ≥ λ mv ≥ 0 be the sorted eigenvalues of K v and let M p := 1 mv mv i=1 λ p i be its spectral moments. It holdsλm v λ1 ≤ cos(θ v ) ≤ 1. Moreover, if b v is Gaussian and independent from K v , then as m v → ∞, cos(θ v ) pr. → M 1 √ M 2 .as soon as EQUATION)The first claim follows from Eq. ( 6) and the worst-case boundsK v b v 2 ≤ λ 1 b v 2 and K 1 /2 v b v 2 ≥ λ mv b v 2 .The second claim is related to the trace estimation method via random matrix-vector products [Martinsson and Tropp, 2020, Chap. 4 ]. We assume without loss of generality that E[ b v 2 2 ] = 1 and by Lem. 3.3, we can writeZ = K 1 /2 v b v 2 2 = a(1 + b) where a = E[Z|K v ] = M 1 and E[b 2 ] → 0 as m v → ∞. An analogous decomposition holds for K v (b v ) 2 2with a = M 2 and the second claim follows.Lemma 3.3. Let K ∈ R m×m be a (potentially random) psd matrix and a ∼ N(0,1 m I m ) be independent. Then E[ Ka 2 2 | K] = M 2 (K) and Var[ Ka 2 2 | K] = 2 m M 4 (K) where M p (K) := 1 m m i=1λ p i and λ 1 , . . . , λ m ≥ 0 are the eigenvalues of K.The second claim expresses the BFA in terms of the spread of the spectrum of the BFK, in an asymptotically exact way. Its assumptions hold at random initialization in the large width limit of typical NNs, provided f v is directly followed by a weight-matrix multiplication in the forward pass, so that b v is the output of a random matrix/vector multiplication. Asymptotic independence can be guaranteed in quite general contexts, see Yang [2020] . For MLP or ResNets with batch-size one, we show in Section 5 that cos(θ v ) is tightly related to the conditioning of layer-to-layer Jacobians, studied in the dynamical isometry literature [Pennington et al., 2017] .4 Ensuring feature learning in scaled NNsQuestion: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: Assumptions, that limit the scope of the theoretical results, are clearly stated.The limited practical relevance of some findings is mentioned. The conclusion discusses some limitations.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
theory assumptions and proofs: 
 Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: A link with the code to reproduce the experiments is provided.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
properties for scaled nns: 
 Consider a sequence of NNs and parameters as in (1) with some diverging architectural parameters such as depth or width. We refer to such a sequence as a scaled NN. In search of the optimal scaling of NNs, it is crucial to understand how HP scalings influence the properties of the training dynamics.In this section, we discuss the following properties:(SP) Signal propagation. It holdsf v rms = Θ(1) for v ∈ [1 : L -1].(FL) Feature learning. It holds ḟL-1 rms = Θ(1).(LD) Loss decay. It holds-L = Θ(1). (BC) Balanced contributions. It holds η ∇ L 2 2 = Θ(η ∇ L 2 2 ) for any , ∈ [1 : L].We discuss these specific properties because they are amenable to our tools and enforcing them requires (L -1) + 1 + 1 + (L -1) = 2L degrees of freedom, which exactly matches the number of free HPs if one counts one scale HP (such as the variance of the weights) and one LR per block. One may wonder if property (BC) is truly desirable: this is the topic of Section 6, where we adopt a more axiomatic approach and recover, for homogeneous architectures, (a more general version of) property (BC) as a consequence of enforcing gradient stability. Also, while enforcing these properties is reasonable when increasing depth and width, they should be rethought for other asymptotics.Property (SP) specifies L -1 scale HPs , but leaves the scale of f L free. The reason for not includingf L in (SP) is that f L rms = o(1)does not lead to vanishing gradient in general, so this behavior should not be excluded a priori. How should one then fix the scale of the output? The next proposition shows that for property (FL) to hold, the quantity that should be suitably normalized is the norm of the backward pass. EQUATIONandEQUATION)Proof. Property (LD) requires 1) and (BC) requires the terms in the sum to be balanced, this leads to Eq. ( 8). Now by Thm. 2.1, property (FL) requiresL =1 η ∇ L 2 2 = -L = Θ(EQUATIONwhich leads to (7). Conversely, it is clear that Eq. ( 8) and ( 7) imply (FL), (LD) and (BC).The criterion of Prop. (4.1), complemented with the property (SP), suggest a method to automatically adjust the scales and learning rates in any architecture. In general, properties (SP), (FL), (BC) and (LD) can be enforced as follows:• (SP): Forward layer normalization. Enforcing property (SP) can be done along with the computation of the forward pass, this is the usual layer normalization. • (FL): Backward layer normalization. Provided θ L-1 is known or measured, Eq. ( 7) can be enforced via a backward analog to layer normalization: one inserts a scaling factor in the forward pass between f L-1 and f L , adjusted so that Eq. ( 7) holds. • (BC) & (LD): Scale invariant learning rates. Directly tune the LRs via Eq. ( 8).We refer to the resulting scaling as FSC as it normalizes the Forward pass, the Sensitivities and the Contributions. Let us make some observations regarding the scale invariant LRs:• Link with Polyak step-size. In convex optimization, to minimize a convex and Lipschitz continuous function f : R d → R such that min x∈R d f (x) = 0, the Polyak-step-size [Polyak, 1987, Hazan and Kakade, 2019] for the GD algorithmx t+1 = x t -η t ∇f (x t ) is given by η t = f (xt) ∇f (xt) 2 2. With this step-size, GD achieves the optimal convergence rate for first order methods over the class of convex and Lipschitz functions. Eq. ( 8) require a layerwise version of this step-size schedule.• Interplay with adaptive methods (Adagrad [Duchi et al., 2011] , ADAM [Kingma and Ba, 2015] ). Adaptive gradient method typically divide the gradient by a quantity which grows linearly rather than quadratically with the norm of the gradient, such as δW = -η • ∇ L ∇ L 2 . For such an algorithm, properties (BC) and (LD), suggests the LR η = Θ(1 L• ∇ L 2 ), in place of Eq. ( 8).• Scale invariance and -2 homogeneity. These LRs arise naturally when one wants to make the gradient descent invariant to how scale is enforced (via initialization scale or via scaling factors). We show in App. C that any choice of LR that leads to this invariance must be a positively homogeneous function of the layer-wise gradient of degree -2, as in Eq. ( 8). We also show in Prop. 6.2 that these LRs make homogeneous architectures invariant to the choice of layer-wise scalings σ 1 , . . . , σ L , given a fixed global scaleL =1 σ .5 Scaling width and depth of MLPs and ResNets 
bfa for single input mlps and resnets at initialization: 
 Multilayer Perceptron Consider a ReLU MLP architecture with a single input x = g 0 ∈ R d and a forward pass given, for ∈ [1 :EQUATIONwhere φ(u) = max{0, u} is the ReLU nonlinearity and acts entrywise on vectors. The architecture HPs are the input width m 0 = d, the widths of the hidden layersm 1 = • • • = m L-1 = m (assumed equal), the output width m L = k. The trainable parameters are ∀ ∈ [1 : L], W ∈ R m ×m -1 . SuchNNs are of the form (1) and are thus covered by Thm. 2.1. Let us study their properties at random initialization under the following assumptions:(H1) the weights W are independent N(0, σ2 ) random variables for ∈ [1 : L].(H2) either k = Θ(1) or the loss is linear.In this setting, the following statements gather consequences of results from the literature on random NNs and of Thm. 3.2 to derive the scale of the forward and backward passes and the BFA. We require (H2) as a technical assumption to avoid dealing with cases where b L strongly depends on the forward pass, where different scalings may arise 2 .In what follows we write A = Θ(B) when there exists c, C > 0 independent of d, m, k, L, x 2 and b L 2 such that the probability that A/B ∈ [c, C] goes to 1 in the specified asymptotic. The new result in the following proposition is the BFA estimate, which relies crucially on a delicate computation due to [Jelassi et al., 2023] . Proposition 5.1 (Large width and depth MLP). Assume (H1-2) and for ∈ [1 :L -1], let σ = 2/m -1 . As m → ∞, it holds f v rms = Θ( x rms ), b v 2 = Θ( √ m σ L b L 2 ). (11) Moreover, if (BC) holds then cos(θ v ) = Θ(v -1/2 ).ResNets Consider now a ResNet with a branch scale parameter β ∈ [0, 1], as in Li et al. [2021] : with a single input x = f 0 ∈ R d , the forward pass is given, for ∈ [2 : L -1], by where φ(u) = u in our theoretical results. The architecture HPs are the input width m 0 = d, the widths of the hidden layersf 1 = W 1 x, f = 1 -β 2 f -1 + βW φ(f -1 ), f L = W L f L-1 , L = loss(f L ) (12)m 1 = • • • = m L-1 = m, the output width m L = k. The trainable parameters are ∀ ∈ [1 : L], W ∈ R m ×m -1 . When β = 1, we recover a MLP.Here we limit ourselves to the case of linear activation where we can directly apply a result from [Marion and Chizat, 2024] to estimate the BFA. We believe that the same result and proof technique extend to the ReLU activation and other variants of ResNets; these extensions are left to future work. Proposition 5.2 (Large width and depth linear ResNet). Assume (H1-2), let φ(x) = x, β = O(1/ √ L) and for ∈ [1 : L -1], let σ = Θ(1/ √ m -1 ).As m → ∞ it holds:EQUATION)Moreover, if (BC) holds then cos(θ v ) = Θ(1).We consider3 one GD step in the model ( 12) with ReLU nonlinearity, without training W 1 (input dimension d = 10, output dimension k = 1, master LR δt = 0.001). Fig. 1 , represent BFA, computed via θ v ≈ arccos(-b v δf v ) where δf v is the change of feature f v after one GD step. The results are consistent with Prop. 5.1 and 5.2. Interestingly, the last plot suggests that there exists a function ϕ : R + → ]0, π/2[ such that for a branch scale β = c/ √ L, the BFA converges to ϕ(c) (it can be conjectured numerically that cos(ϕ(c)) ≈ c -1/2 for c 1).Let us now discuss specific choices of HP scalings for single-input MLP architectures as in Eq. ( 10) (or Eq. ( 12) with β = 1) and at initialization. We consider 6 HPs: the scale of initialization σ 1 and LR η 1 of the input layer, the scaleσ hid := σ 2 = • • • = σ L-1 and LRs η hid := η 2 = • • • = η L-1of the hidden layers, and the scale σ L and LR η L of the output layer. The HP scalings mentioned in the next theorem are the following (see Table 1 ):• NTK: the standard scaling with LRs adjusted to satisfy (LD) and (BC) [Jacot et al., 2018] ;• MF+µP: the scaling proposed in [Jelassi et al., 2023] constructed by imposing the so-called "mean-field" output scale σ L ∝ 1/m and then enforcing (FL) by adjusting the learning rates; • FSC: the HP scaling singled-out by Prop. 5.3, obtained by adjusting the Forward scales, Sensitivities, and Contributions.The properties of HP scalings depend on x 2 and b L 2 . We consider two typical settings:• (Dense) Where x 2 = √ d and b L 2 = 1 √ k .This is representative of a dense whitened input and a RMS loss loss(f L ) = f L -y 2 2 /k for some dense signal y ∈ R k with y rms = Θ(1) as, e.g., in image generation applications.Table 1 : HP scalings for MLPs under the dense setting (for the sparse setting, replace k and d by 1). For L fixed, both MF+µP and FSC coincide with µP. Values in red are exact, the others are up to a multiplicative factor in Θ(1).Hidden Output FSC init. std.σ 1 / √ d 2 /m √ kL /m LR η m /L 2 d 1 /L 2 k /Lm MF+µP init. std. σ 1 / √ d 2 /m √ k /m LR η m /L 3/2 d 1 /L 3/2 k /L 3/2 m NTK init. std. σ 1 / √ d 2 /m 1 / √ m LR η 1 /Ld 1 /Lm k /LmTable 2 : FSC scalings identified in Prop. 5.4 for ResNets. All HPs are specified up to a multiplicative factor in Θ(1). When β = Θ(L -1 /2 ) and k = d = Θ(1), these scalings coincide with the so-called "depth µP" introduced in [Bordelon et al., 2023] and also studied in Yang et al. [2023b] .Input Hidden Output init. std. σ 1 / √ d 1 / √ m √ k /m LR η m /Ld 1 /β 2 L k /Lm• (Sparse) Where x 2 = 1 and b L 2 = 1. This is representative of a one-hot encoding input and the multiclass logistic loss (aka cross-entropy where b L 2 = Θ(log(k))). This setting is typical of natural language processing tasks.The scalings are reported in Table 1 . We have also introduced scalings in terms of output width k for NTK and MF + µP to ensure a non-degenerate behavior as k 1, although these are generally not written in the literature. Proposition 5.3 (MLP scalings). Under the assumptions of Prop. 5.1, the following hold at random initialization:(i) The scaling MF+µP satisfies (SP), (BC) , (FL) but not (LD);(ii) The scaling NTK satisfies (SP), (BC), (LD) but not (FL);(iii) Properties (SP), (BC), (LD), (FL) hold if and only if the scaling is FSC.This theorem identifies the new HP scaling FSC for deep ReLU MLP where the scale of the output layer depends on the depth. We compare empirically the sensitivities (Eq. ( 4)) of the various scalings in Fig. 2 , and the results are consistent with theory. Finally, let us mention that even though FSC fixes some degeneracies of deep MLPs, other problems arise when considering multiple inputs, such as degeneracy of the conjugate kernel and NTK [Hayou et al.] , which make ReLU MLPs a fundamentally flawed model at large depth. Hence, the analysis of large depth scalings for MLPs is mostly of theoretical interest.We now discuss HP scalings for single-input ResNets (Eq. ( 12)) withβ = O(1/ √ L).Proposition 5.4 (ResNets scalings). Take β = O(1/ √ L), consider the same 6 degrees of freedom as in the previous section and assume that the conclusions of Prop. 5.2 holds. Then properties (SP), (BC), (LD) and (FL) hold at initialization if and only if the scalings are those given in Table 2 .  
minimal desiderata and stability under homogeneity: 
 It is not clear a priori why the property (BC) studied in Section 4should be enforced. In this section, we consider homogeneous architectures, such as ReLU MLPs, and show that a slightly more general version of (BC) is a related to a notion of gradient stability.For a general architecture of the form Eq. ( 1), consider the following stability property (S), which is necessary if one wants to have comparable behavior between the first GD step and the next. It is related to the usual notion of smoothness in optimization:(S) Stability. It holds d dt ∇ L 2 / ∇ L 2 = O(1) for ∈ [1 : L].We will study this property in a ReLU MLP with a single input as in Eq. ( 10) (the extension to linear ResNets is simple as only the BFAs change). In this case we have ∇ L = b g -1 and thus∇ L F = b 2 g -1 2 . It follows d dt ∇ L F ∇ L F ≤ ḃ 2 g -1 2 + b 2 ġ -1 2 b 2 g -1 2 = ḃ 2 b 2 + ġ -1 2 g -1 2 .We can thus ensure the gradient stability by ensuring, for all ∈ [1 : L],(FS) Forward stability.EQUATION(BS) Backward stability. It holdsḃ 2 b 2 = O(1) for ∈ [1 : L].We focus on these simpler desiderata (FS) and (BS) instead of (S) for the rest of the discussion.To estimate ḃv , we rely on a "backward" version of the feature speed formula that holds in 1homogeneous NNs. Proposition 6.1 (Backward speed formula). Consider a general architecture of the form (1), take v ∈ [1 : L] and assume that the mapf v → f L is positively 1-homogeneous 4 . If -f L ∇ 2 loss[f L ] ḟL + >v η ∇ L 2 2 = 0 then ḃv = 0.Otherwise, the (non-oriented) angle θv between f v and ḃv is well defined in [0, π /2[ and it holdsḃv 2 = -f L ∇ 2 loss[f L ] ḟL + >v η ∇ L 2 2 f v 2 cos( θv ) .In the context of ReLU MLPs with a linear loss, we have by differentiating the back-propagation recursion and noticing that all terms involving φ are zero almost surely5 that:ḃv = >v η ∂g -1 ∂f v g -1 b b = >v η b 2 2 ∂g -1 ∂f v ∂g -1 ∂f v f v = Kv f v (14)where the last expression defines Kv . Reasoning as in Thm. (3.2), since f v is Gaussian at initialization and noticing that Kv has a structure similar to that of K L-v , we have that cos( θv ) = Θ( √ L -v), see the details in Lem. A.1. We can thus estimate ḃv just as well as ḟv .Homogeneous architectures such as ReLU MLP satisfy scale invariance properties that are important to take into account in our discussion. The following result presents a general invariance under blockwise rescaling, provided one uses scale-invariant LRs. This is related to known invariance results under global rescaling for scale invariant losses [Van Laarhoven, 2017 , Li et al., 2022 , Wan et al., 2020] . Proposition 6.2 (Invariance under block-wise rescaling). Consider a function f L (w 1 , . . . , w L ) (the NN, in our context) which is separately positively 1-homogeneous in each of its blocks of parameters w ∈ R p . Let θ 0 = (w 1 (0), . . . , w L (0)) and let θ0 = σ θ 0 := (σ 1 w 1 (0), . . . , σ L w L (0)) for some scale factors σ ∈ R L + . Let θ(t) and θ(t) be the iterates of GD on L : θ → loss(f L (θ)) with LR satisfying η (t) ∇ L(θ(t)) 2 2 = η (t) ∇ L( θ(t)) 2 2 and starting from θ 0 and θ0 respectively. If L =1 σ = 1 then θ(t) = σ θ(t) for all t ≥ 1.In view of Prop. 6.2, for homogeneous architectures, one can ignore (SP) since any GD dynamics is equivalent to a dynamics where (SP) holds at initialization. However, if the scale of the forward pass is not normalized, (FL) needs now to be adapted to a scale-free version, as follows:(RFL) Relative feature learning. It holds ḟL-1 2 / f L-1 2 = Θ(1).We are finally in position to gather all these insights and characterize all admissible scalings for ReLU MLPs, i.e. scalings that satisfy the minimal desiderata (RFL), (LD), (FS) and (BS) at initialization. Theorem 6.3 (Minimal desiderata for MLPs). Consider a ReLU MLP with 6 degrees of freedom: three initialization scales σ 1 , σ hid , σ L and three LRs η 1 , η hid , η L . Assume b L 2 = g 0 rms = 1 and a linear loss for simplicity. Then the minimal desiderata (RFL), (LD), (FS) and (BS) hold at initialization in the limit m → ∞ then L → ∞ if and only if( √ dσ 1 ) • ( m/2σ hid ) L-2 • σ L = Θ( √ L/m), C 1 + C hid = Θ(1) and C L = O(1),whereC 1 = η 1 ∇ 1 L 2 2 , C hid = L-1 =2 η hid ∇ L 2 2 and C L = η L ∇ L L 2 2 .In particular, the scaling FSC (Table 1 ) satisfies these desiderata.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: Experimental details are given either in the captions of the plots or in the main text.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
conclusion: 
 Starting from the feature speed formula, our approach allows to conveniently recover and characterize in an elementary fashion certain properties of existing HP scalings and to discover new ones, with essentially all the technical difficulty contained in the estimation of the BFA. The limitations of our approach are related to the blind spots of Thm. (2.1): it can only quantify feature speed for (S)GD (and does not apply to variants in its current form) and at "cut nodes" in the NN architecture, where all the signal goes through (in particular, it does not apply inside the blocks of a ResNet).In future works, besides removing these limitations, it would be interesting to have a better understanding of the BFA, both from a quantitative and a qualitative viewpoint.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: our plots include several random runs (which is more precise than error bars).Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [NA] Justification: the experiments are small scale and run in less than 10 minutes on a personal laptop.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification:Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [NA] Justification: Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. 
 
 Appendix: Proof of Lem. 3.3. Writing K = V DV with D = diag(λ 1 , . . . , λ m ) and V ∈ R m×m orthonormal, we have Ka 2 2 = a V D 2 V a. Conditioned on K, the vector u = V a is isotropic Gaussian soEu 2 i = 1 m for i ∈ [1 : m].Hence, on the one handE[ Ka 2 2 | K] = E m i=1 λ 2 i u 2 i | (λ i ) m i=1 = m i=1 λ 2 i E[u 2 i ] = 1 m m i=1 λ 2 i .On the other hand, using the fact that the variance of a chi-square random variable is 2,Var[ Ka 2 2 | K] = E m i=1 λ 2 i (u 2 i -1 /m) 2 | (λ i ) m i=1 = m i=1 λ 4 i E[(u 2 i -1 /m) 2 ] + m i =j λ 2 i λ 2 j E[(u 2 i -1 /m)(u 2 j -1 /m)] = 2 m 2 m i=1 λ 4 i .Proof of Prop.5.1. When β = 1, Eq. ( 11) is classical from the signal propagation literature [Poole et al., 2016 , Hanin and Rolnick, 2018 , Hanin, 2018 ] (the fluctuations around the limit have also been studied in Hanin and Nica [2020] ). Note that these results are proved with k = Θ(1), but Hanin [2018] allows to conclude as well when k diverges at least if the initial gradient b L = ∂L ∂f L ) ∈ R k is independent of the randomness of the weights, which is what (H2) guarantees. We note that analogous results have been derived for a variety of activation functions, and we focus on ReLU only for conciseness.Let us now discuss the BFAs, assuming for simplicity that η 1 = 0 as the contribution of w 1 to the BFK is asymptotically negligible assuming (BC). We consider the BFA at g v and denote z v := (∂L/∂g v ) . The main result of [Jelassi et al., 2023] can be restated as follows: with k = d = Θ(1), the choice σ L = 1 m and learning-rates η = Θ(L -3 /2 ), it holds ġL-1 rms = Θ(1). In view of (11), it holds in their setting for ∈ [2 :L -1] ∇ L 2 2 = b g -1 2 2 = g -1 2 2 b 2 2 = Θ(m -1 • m L-1 • σ 2 L ) = Θ(m -1 /m L-1 ).Using η = Θ(L -3 /2 ), it followsL-1 =2 η ∇ L 2 2 = Θ(L -1/2). By Thm. 2.1 and using m L-1 z L-1 rms = Θ(1), we getġL-1 rms = ≤L η ∇ L 2 2 cos(θ L-1 ) • m L-1 • z L-1 rms = Θ(1) ⇒ cos(θ L-1 ) = Θ(L -1/2 ).This shows the result for the BFA at g L-1 and the result holds as well for the BFA at f L-1 up to hidden constants.Interestingly, in view of Thm. (3.2), we can interpret the result of [Jelassi et al., 2023] as a computation on the spectral moments of a certain random matrix, as stated in the following lemma.Lemma A.1 (Spectrum of BFK and FBK in ReLU MLPs). For a ReLU MLP at random initialization satisfying (SP) and (BC), consider the BFK (at g v instead of f v ):K v = ≤v η ∂g v ∂w ∂g v ∂wand θ v the (non-oriented) angle between ġv and z := (∂L/∂g v ) . Then, in the notations of Thm. (3.2), it holds as hidden width diverges cos(θv ) = M 1 (K v )/ M 2 (K v ) = Θ(v -1/2 ).Consider also, for a linear loss, the kernel Kv such that ḃv = Kv f v (see Eq. ( 14)) and θv the (non-oriented) angle between f v and -ḃv . Then it holds, as hidden width diverges, cos( θv) = Θ M 1 ( Kv )/ M 2 ( Kv ) = Θ((L -v) -1/2 ).Proof. We have already seen in the proof of Prop.5.1 that cos(θ v ) = Θ(v -1/2 ). It thus remains to see that the assumptions of Thm. (3.2) are satisfied: the independence of z follows from [Hanin and Nica, 2020, Prop. 2 ] and the Gaussianity of z is direct since z = W +1 b +1 where W +1 is Gaussian and independent from b +1 . Also, we have the more explicit expression [Hanin and Nica, 2020, Prop. 2 ], these matrices can be taken as matrices with Bernoulli random variables on the diagonals, independent from everything else). Since under (SP) and (BC) we have that η g -1K v = v =1 η g -1 2 2 ∂g v ∂f ∂g v ∂f where ∂gv ∂f = D v W v . . . D +1 W +1 D and D i = diag(φ(f i )) (by2 2 is constant for ∈ [1 : L-1], it follows K v ∝ v =1 (D v W v . . . D +1 W +1 D )(D v W v . . . D +1 W +1 D ) .For the second claim, we have (see Section 6) Kv = L =v+1 η b 2 2 ∂g -1 ∂fv ∂g -1 ∂fv . Under (SP) and (BC), we have η b 2 2 is constant for ∈ [2 : L], hence it follows Kv ∝ L =v+1 (D -1 W -1 . . . W v+1 D v ) (D -1 W -1 . . . W v+1 D v ).By comparing the expressions for K v and Kv , we see that Kv has the same distribution of nonzero eigenvalues as K L-v (potentially up to a global multiplicative factor) and the conclusion follows.Proof of Prop. 5.2. The estimate for f rms is classical, see e.g. Li et al. [2021] . For the backward pass estimate, we rely on [Marion and Chizat, 2024, Lem. 3 ] (see also Zhang et al. [2022] for related results with the ReLU activation function), which implies that σ min ( → v) = Θ(1) and For the BFA, we will apply the first bound of Thm. 3.2, namely cos(θ v ) ≥ λ min (K v )/λ max (K v ) where λ min (K v ) and λ max (K v ) are the smallest, respectively largest, eigenvalues of K v . In the forward pass ( 12), let us write g = φ(f ) and h = W g -1 . By direct computations, it holds (here w is the vectorization of W ):σ max ( → v) = Θ(K v = v =1 η ∂f v ∂w ∂f v ∂w = v =2 η g -1 2 2 ∂f v ∂h ∂f v ∂h = v =2 η β 2 g -1 2 2 ∂f v ∂f ∂f v ∂f .Using the inequalitiesλ min (K) ≥ β 2 v =2 η g -1 2 2 σ min ∂f v ∂f 2 , λ max (K) ≤ β 2 v =2 η g -1 2 2 σ max ∂f v ∂f 2 we deduce cos(θ v ) ≥ λ min (K v )/λ max (K v ) = Θ(1).Proof of Prop. 5.3. In this proof, we say that a claim is found "by direct computation" when it can be directly deduced from the conclusion of Prop. 5.1. In particular, for the computation of scale invariant LRs, we use the fact that∇ L 2 = b g -1 F = b 2 • g -1 2. Also, by Prop. 5.1, under (SP) and (BC) it holds cos(θ L-1 ) = Θ(L -1 /2 ).(i) One has that MF+µP satisfies (SP), (BC) by direct computation, and (FL) by Prop. 4.1. We have seen in the proof of Prop. 5.1 that -L = Θ(L -1/2 ), so (LD) does not hold.(ii) For NTK, (SP), (LD) and (BC) can be checked by direct computation. For (FL), we have b L-1 rms = Θ(1/ √ m) so :ḟL-1 rms = Θ 1 cos(θ v ) • m • b L-1 rms = Θ 1 cos(θ L-1 ) • √ m . But in the considered asymptotics √ m • cos(θ L-1 ) = Θ( m/L) → ∞ so ḟL-1 rms = o(1).(iii) Properties (SP) specifies σ 1 and σ 2 = • = σ L-1 , and Prop. 4.1 gives, with (FL), b L-1 rms = Θ(1 cos(θ L-1 )m ) = Θ( √ L/m) which imposes σ L = √ kL/m.Then the LR are given by (8).Proof of Prop. 5.4. Properties (SP) specifies σ 1 and σ hid = σ 2 = • = σ L-1 , and Prop. 4.1 gives, with (FL), b L-1 rms = Θ(1 cos(θ L-1 )m ) = Θ(1/m) which requires σ L = √ k/m.Then the LR are characterized by (8).Proof of Prop. 6.1. By the chain rule and Euler's identity for positively 1-homogeneous functions, it holdsb v f v = ∂L ∂f v f v = ∂L ∂f L ∂f L ∂f v f v = ∂L ∂f L f L .Now, by differentiating in time both sides we getḃ v f v + b v ḟv = f L ∇ 2 loss[f L ] ḟL + ∂L ∂f L ḟL = f L ∇ 2 loss[f L ] ḟL + L. We have L = - L =1 η ∇ L 2 2 and moreover, from the proof of Thm. 2.1, -b v ḟv = ≤v η ∇ L 2 2 . So we deduce -ḃ v f v = -f L ∇ 2 loss[f L ] ḟL + >v η ∇ L 2 2 .We conclude by writing -ḃ v f v = ḃv 2 f v 2 cos( θv ) and rearranging.Proof of Prop. 6.2. By assumption at time t = 0, it holds θ(0) = σ θ(0) so let us prove the result by recursion. Assume that the claim is true at iteration t.Since σ = 1, it holds f L (θ(t)) = f L ( θ(t)).Moreover, since ∂f L ∂w is 0-homogeneous in w and separately 1-homogeneous in (w i ) i = . It follows∇ L( θ(t)) = ∂f L ∂w [ θ(t)] ∇loss(f L ( θ(t))) = ( i = σ i ) ∂f L ∂w [θ(t)] ∇loss(f L (θ(t))) = 1 σ ∇ L(θ(t)).In particular, we deduce that the LRs are related by η(t)η(t) = ∇ L(θ(t)) 2 2 ∇ L( θ(t)) 2 2 = σ 2 . It follows, for any ∈ [1 : L], w (t + 1) = w (t) -η(t)∇ L( θ(t)) = σ w (t) -σ 2 η(t) 1 σ ∇ L(θ(t)) = σ w (t + 1).This proves θ(t + 1) = σ θ(t + 1) and the claim follows by recursion.Let us mention an interesting degree of freedom for FSC in Table 1 : up to adjusting the initial LR, it is possible to initialize the output layer with 0 while still satisfying FSC at the next step. If one initializes the output layer W L with 0 then all gradients are 0 at time 0 except that for W L which leads to the update (non-infinitesimal in this paragraph):δW L (0) = -η L (0) • b L (0)g L-1 (0).The second forward pass is the same as the first one, with the only difference thatEQUATION1) (linear loss) for simplicity, this leads to a second backward pass:z L-1 (1) := ∂L ∂g L-1 (1) = (-η L (0)b L (0)g L-1 (0) ) b L (0) = -η L (0) • b L (0) 2 2 g L-1 (0).For the second GD step to satisfy (FL), we just need to ensurem z L-1 (1) rms = Θ( √ L) ⇔ η L (0) = Θ √ L m b L (0) 2 2 .This is the LR to be used at time 0, for the second step to satisfy (SP), (FL), (LD) and (BC).Proof of Thm. 6.3. Prop. 6.2 shows that in fact only the product σ 1 • σ L-2 h • σ L is a relevant degree of freedom of scale. We can thus fix (arbitrarily) 1) is not excluded. This is where (BS) comes into play. It requires in particularσ 1 = 1/ √ d and σ hid = 2/ √ m so that (SP) is satisfied; we then have b v 2 f v 2 = Θ(mσ L ) for v ∈ [1 : L -1]. Desideratum (RFL) requires ḟL-1 2 f L-1 2 = C 1 + C h cos(θ L-1 ) b L-1 2 f L-1 2 = Θ(1) ⇔ σ L (C 1 + C h ) √ L m . using that cos(θ L-1 ) = Θ(1/ √ L). Moreover, (LD) requires C 1 + C h + C L = Θ(1). At this stage, the output scale σ L is not yet entirely determined since C 1 + C h = o(ḃ1 2 b 1 2 = C hid + C L cos( θ1 ) b 1 2 f 1 2 = O(1) ⇔ (C hid + C L ) √ L m = O(σ L )using that cos( θ1 ) = Θ(1/ √ L) by Lem. A.1. Combining both conditions for σ L imply, on the one hand, thatC h + C L = O(C 1 + C hid ) hence C 1 + C hid = Θ(1) and σ L = Θ( √ L m), which are equivalent to the constraints written in the theorem. Conversely, it is not difficult to see that these constraints lead to satisfying (RFL), (LD), (FS) and (BS). One such LR is precisely that suggested by Prop. 4.1: η ∝ ∇ f (x) -2 2 .Proof. For ∈ [1 : L], it holdsα y = α y -α η (∇g(y))∇ g(y) = x -α 2 η (α • ∇f (x))∇ f (x). Then α • y = x for all α ∈ (R * + ) L is equivalent to α 2 η (α • ∇f (x)) = η (∇f (x)), ∀α ∈ (R * + )L which is the claimed homogeneity property.Justification: There is no societal impact of the work performed.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).This is a technical assumption that covers virtually all functions of interest in deep learning. In particular, the maps T admit selection derivatives that are compatible with the chain rule[Bolte and Pauwels,2020, Prop. 4] and that coincide almost everywhere with standard derivatives[Bolte and Pauwels, 2020, Prop. 3]. To simplify our exposition, we always implicitly assume that we are at a differentiability point of the maps T .Say, if loss(f ) = 1 2 f 2 2 , we have bL-1 rms = W L WLfL-1 rms = Θ(σL max{1, k/m} bL 2) (by Lem.3.3 and properties of the Marcenko-Pastur law), while under (H2) we have bL-1 rms = Θ(σL bL 2).Link to the Julia code to reproduce the experiments: https://github.com/lchizat/2023-BAFUFor Euler's identity to hold, we also assume that its selection[Bolte and Pauwels, 2020] is 0-homogeneous.A downside of this computation is that, because of its local nature, it ignores the contributions of the Jacobian's discontinuities to ḃL, while they do have a "macroscopic" effect with a non-vanishing step-size. For instance, taking first the infinite width and then the small step-size limit, would give a different expression. 


Decision: The paper proposes the a quantity named BFA to predict and control feature learning in deep neural networks. Based on this, the authors derive a feature speed formula which allows expressing the magnitude of feature updates after one GD step. This paper receives diverging scores: two reviewers appreciate the strength of this paper in terms of soundness and sufficient technical contributions, while the other two reviewers highlight the weakness of the paper, mainly in not providing sufficient experimental verifications. After careful discussions amoung the reviewers, and between the AC and SAC, I am current leaning towards accepting this paper. The authors should include additional experimental results in the final version of the paper.

Reviews: review: paper_summary: paper_summary: This paper studies the feature learning speed of the layers of Neural Networks (NNs). Specifically, it proposes to measure it through the quantity *Backward-Feature Angle* (BFA), denoted by $\theta_l$ for a layer $l$. This quantity is directly related to the layer-wise decomposition of the Neural Tangent Kernel (NTK). In practice, the BFA is measured experimentally and several properties (feature learning, signal propagation, etc.) can be related to the BFA.
strength_weakness: strength_weakness: # Originality

The BFA is closely related to the layer-wise decomposition of the NTK, which is already widely used in the NN optimization literature [1, 2, 3, 4]. Overall, the BFA does not contain any information that is not already available with previous objects.

# Significance

The benefits and the properties of the BFA are still unclear.

For instance, Section 5 proposes a new scaling of the hyperparameters, that is not clearly related to the BFA. Besides, the experimental validation of this new scaling is not provided.

# Quality

The contribution of this paper is unclear. The usefulness of the BFA, either theoretical or experimental, is still unclear, and the proposed hyperparameter scaling is not tested experimentally.

# EDIT: References

[1] Gradient descent provably optimizes over-parameterized neural networks (2018), Du et al.

[2] Gradient descent finds global minima of deep neural networks (2019), Du et al.

[3] A convergence theory for deep learning via over-parameterization (2019), Allen-Zhu et al.

[4] Stochastic Gradient Descent Optimizes Over-parameterized Deep ReLU Networks (2020), Zou et al.
questions: questions: How does the "ours" hyperparameter scaling compare to the others (usual, muP or NTK)?
limitations: limitations: Lack of experimental validation.
paper_summary: paper_summary: The authors propose a technical strategy for deriving neural net parameterizations that relies on controlling the angle between the activation gradient and the feature update. The authors derive various theoretical results about this quantity, including a formula for computing it, and some analyses in the context of MLPs and ResNets. The authors claim to use this principle to derive new parameterizations, but crucially they never test them in a real learning problem.
strength_weakness: strength_weakness: I need to preface this review by saying that this feedback is intended to be constructive and to help you improve the paper. My current impression is that the paper is not ready for publication. I strongly encourage you to keep working in this direction, and I hope this feedback will be useful for that.

With that said, the main issues I see with the paper are:

### **No real experimental evaluation**

My understanding is that the main practical outcome of your work and theoretical analysis is a new parameterization for training neural networks. I feel that it is really important for you to test this parameterization to check that it is useful, or at least to see what its properties are in an actual training situation. It's so easy to come by free cloud compute (e.g. Google Colab) that I can't really see a reason for not doing this.

I don't feel that the experiments in Figures 1 and 2 are enough to convince me of the utility of your framework. Also I'm not sure how to reproduce these experiments. For example, what dataset did you use? What is the loss function?

As a side note, I'm also a bit doubtful that you can even train MLPs effectively beyond depth 20 or so. I read the Jelassi et al paper (https://arxiv.org/abs/2305.07810) and noticed they don't test their parameterization either. I may be wrong here, but I don't think you can hope for some engineer or experimentalist to pick up the paper and implement things for you. I think you have to be proactive here.

### **Doesn't go that far beyond existing ideas**
A lot of the paper focuses on dealing with analyzing or re-deriving existing parameterizations---e.g. muP or the 1/sqrt(L) depth scaling in ResNets. But this is not so interesting because it has already been done and there are already ways to analyze these things. What does your analysis offer that prior analyses do not? I also want to point out that concurrent works to this paper go beyond 1/sqrt(L) depth scaling rules. For example arxiv.org/abs/2405.15712 and arxiv.org/abs/2405.14813. And these papers actually experimentally test these deviations. Clearly these are concurrent works, but I just mention it to demonstrate that there is more out there.

### **Paper only seems to focus on batch size one**

In my opinion, doing things only at batch size one is a bit toy, and it would be better to directly analyze larger batch sizes.
questions: questions: Please see the weaknesses section
limitations: limitations: Overall, I think it's hard to assess the limitations without having more thorough experimental evaluation. I would encourage you to work out how to streamline the mathematical exposition, and then to start testing these ideas.

I realize this feedback might be construed as being fairly negative, but I hope that it can help to improve the work.
paper_summary: paper_summary: The paper presents a novel perspective on infinite width and depth feature learning networks. It introduces the backward-to-feature kernel (BFK) as a central quantity determining the evolution of the intermediate layer features. The paper shows that the movement of the hidden layer features can be exactly related to an angle $\theta_\ell$ between the backward pass and the feature velocity, and uses insights on the scaling of the cosine of this angle with width to recover essentially all known infinite width and depth feature learning limits, as well as a novel large depth MLP limit.
strength_weakness: strength_weakness: There are no technical weaknesses that I have found, and I have gone through the derivations in detail. My only comment is expository:
In equation 1, the definition of the forward pass $T_{\ell}(f_{\ell-1}, w_\ell)$ as well as its discussion in terms of selection derivatives is quite technical and may confuse readers from outside sub-communities in machine learning. I recommend stating more clearly that this includes a simple forward pass such as $W_{\ell} \cdot f_{\ell-1}$ and perhaps adding a footnote to make this first paragraph a bit more readable.
questions: questions: As a simple clarification, I want to confirm that the ResNet scalings found precisely reproduce those predicted by Bordelon et al and Yang et al. Are there any additional ResNet scalings that have not been studied in prior work that this paper finds? 

In the second paragraph of the conclusion section "it can only quantify feature speed for (S)DG (and does not apply to variants, a priori) and at “cut nodes” in the NN architecture, where all the signal goes through (in particular, it does not apply inside the blocks of a
290 ResNet)" 

I assume you mean "(S)GD". Can you please elaborate a bit more on what you mean by cut nodes? Is this like a residual block with many layers? It will be interesting if you can derive a similar feature formula for more general differentiable circuits with branching.
limitations: limitations: Large width and depth limits may serve to determine the scaling laws for the next frontier of language and vision models, which may have major societal impact. However, the theoretical nature of the paper limits any major negative impacts.
paper_summary: paper_summary: The paper introduces the BFA - a novel quantity to predict and control feature learning in DNNs, as well as the feature speed formula which allows expressing the magnitude of feature updates after one GD step. The paper recovers key properties of known HP scalings, and also extends these results by introducing a new HP scaling for large depth ReLU MLPs.
strength_weakness: strength_weakness: Despite the strengths mentioned above, I did not give a higher score for the following reasons: 

1. Novelty for HP scaling: 
As far as I can see, the main takeaway regarding HP scaling is the extension of known results, such as muP, to the limit of large width-then-depth. While this is indeed new, this is a somewhat limited contribution. 

2. Applicability of results: 
While some of the results are rather general (like Thm 2.1), some other parts of the results seem to apply only under rather limited conditions, e.g. only a single input. 

3. Experimental findings: 
I found issues with some of the experimental findings: I did not find a mention of what is assumed about the data: is it synthetic, random, from some known benchmark etc. Also, by inspecting Fig 2b I was not convinced that the output sensitivity is bounded away from zero. 

4. I feel that the paper could be made less technical and more readable by delegating some of the proofs to the Appendix and using the space for some visualizations. 




typos: 
- Fig1 caption 1st line: witdh -> width
questions: questions: 1. in line 194 - BC condition - is this a strict equality? or is there some tolerance? 

2. In the Introduction you use the term "hierarchical features" - can you give a definition for that? 

3. in the BFK definition (eq. 5) - is this for multiple inputs? the NTK is defined for $x, x'$.. Is $m_v$ here the product of batch-size with input dimension to the layer?
limitations: limitations: The authors adequately addressed the limitations of their results.

Comments: comment: Thanks for these clarifications.

- I see now it's in the arXiv draft, but this feels quite an important omission from the submitted version
- I see what you mean that Theorem 2.1 is an equality so holds quite generally. A good way to convince people this is a useful reformulation is to demonstrate how it can be used to do a new scaling calculation. As you mention below this may not be so straightforward.

Pragmatically, I can't really engage with you further on this now (I still have two other papers in my batch to deal with). I've decided to maintain my score, and have the following suggestions:

- If you want to commit to the pure theory route, which I'm not sure is the best thing, but I would suggest making another pass through the paper to clarify things as much as possible, include the omission about bad depth limits in MLPs. Another worthwhile contribution would be clarifying the connection to the spectral Yang et al paper. I know you pointed out these are concurrent works, but clarifying the connection is a strict benefit to the community, and you likely understand how to do this better than anyone else
- My real suggestion is to do an actual practical training run of a neural network, and measure all your quantities
- In terms of getting the paper past NeurIPS reviewers, this may be slightly jaded advice, but picking an architecture that people don't already know how to scale (e.g. Mamba, graph neural net, etc, etc) could be a good way to do this

Good luck with your work
- I apologize for this: in fact, we forgot to include this comment about the degeneracy of MLPs in the submitted version; but it is in the arXiv version of the work, where we state: "let us mention that even though FSC fixes some degeneracies of deep MLPs, other problems arise when considering multiple inputs, such as degeneracy of the conjugate kernel and NTK [Hayou et al., 2019], which make ReLU MLPs a fundamentally flawed model at large depth." The case of MLPs is still an interesting case from a theoretical viewpoint because dynamical isometry does not hold, and it helps disentangling the pure effect of depth vs the effect of spectral degeneracy.
- Thm. 2.1 and 3.2 are *equalities* and Thm 2.1 is non-asymptotic besides the small learning-rate  (and Thm 3.2 requires the large width asymptotics, but in a robust way: it does not require width to be larger than any other quantity). So yes, we stand to the claim that anybody who attempts to quantify feature learning must be, implicitly or implicitly, trying to quantify the various terms in this formula.
- We sincerely appreciate all the time and energy that you have spent in this discussion. We agree that deep learning is fundamentally an empirical science and authors have the responsibility to validate their practical ideas. However, we disagree that it should be done for every single paper. There should also be room to develop intuitions into mature ideas, and to improve the understanding of known facts via alternative derivations, etc. Here, our paper's goal is solely to clarify the mechanism behind the scale of feature learning.
- Thanks, I understand our formulation "single batch size" is in fact confusing.
- I am aware that transformers are in a very different regime. This fact is one of the original motivations for our work ; to have a framework that does not relies on taking infinite width limit first (which is a strong limitation).
- The claim is perhaps huge, but it is indeed an equality, with a proof in the main text. The difficulty to make predictions in *new* settings are that (1) there is some theoretical difficulty in working out the various terms of the formula in general contexts (although it is not so difficult to get rules of thumb via, say, Marcenko-Pastur), and above all (2) it is not obvious what the desiderata should be, let's say when context length is very long. This is why we decided to make a paper where we focus on known results.
Thank you to the authors for writing this. I was planning to reply to this global response with something about the philosophy of science and the role of theory (which are both topics I've thought about quite hard over the years) but in the interest of time, let me get to the point.

**"as mentioned in the paper, deep ReLU MLPs cannot be trained"** I have gone over the paper three more times, and can nowhere find it mentioned that deep ReLU MLPs cannot be trained. On what line number do you say this? To me the paper gives the strong impression that you are proposing a new scaling for successively training deep ReLU MLPs. For example, you explicitly say that **"we introduce a new HP scaling (Table 1) that achieves both loss decay and feature learning"**

In the discussions below, the authors state that they have found the fundamental quantity **"needed to quantify feature learning, in any architecture, and any asymptotic."** yet they do not validate this claim empirically in any setting and only produce scaling rules that are either already known or that the authors now admit will not work.

### **Suggestion to the authors for improving their work**

This feedback is intended to be constructive. Of course, the authors may disagree with it and choose to ignore it, but here goes **I think you should take at least some responsibility for validating your ideas.** As I mentioned in my original rebuttal cloud compute is literally free at the moment (on Colab). All you need is an internet connection. You could run a simple test of your framework in under a day. You could then use this to sharpen your ideas, iterate on your ideas, and improve your science.

I'm sorry this may not be the type of feedback you're hoping for. But in my view I think it's really important that somebody says this.
Thanks for your response.

[Minor note: "at a single batch size" would usually be understood to mean "at a fixed batch size" and not "at a batch size of one". In other words, it would mean "at one value of the batch size" rather than "at a batch size of value one".]

[Also a second note, that you may find interesting, but in actual transformer training (e.g. Llama-2-13b) we may have a batch size of 4 million tokens and a width of 5000, so actual transformer training is way outside the realm of applicability of muP]

**"we feel that the strengths of our approach are not properly understood"** I take this seriously and I appreciate you've put a lot of work into your paper. I'll look over your paper again to check I haven't missed something.

**"This quantity is the fundamental spectral information that is needed to quantify feature learning, in any architecture, and any asymptotic."** This is a huge claim. If you could demonstrate this in practice, everybody would want to read your paper. But as it stands you're not demonstrating it in **any** new situation.

I won't reply further here as I want to move to writing a reply to your global response.
Thanks! Let us briefly comment on these points:

- The fixed batch size assumption (not the *single* batch size one, which I was referring to) is indeed critical for the current Tensor Program framework (an algorithmic formulation of DMFT). They can handle random matrix (RM)/ random vector products but not RM/RM products, which output RM).

- Regarding finite vs infinite width: our "feature speed formula" holds at any width as well. The "large size" asymptotics can then to be considered to obtain concentration of random quantities and sharp scalings (but, as we mention, it is not the only way one can use the formula).

- You seem to be appreciative of the "spectral criterion" approach. So perhaps we should a be a bit more explicit: the theoretical results in https://arxiv.org/abs/2310.17813 (which was contemporary to our preprint) are strictly weaker than ours. By assuming bounded RMS to RMS spectral norms of the weight matrices (the spectral criterion), we can directly deduce upper bounds on the feature speed via our formula, in the fixed depth setting. However, this approach does not provides lower bounds (ie sharp scalings), nor gives depth scalings (and in general it is not clear when it can or cannot be trusted for general asymptotics). We are very appreciative of this line of work, but we feel that the strengths of our approach are not properly understood. On a side note, our formula also relies on some "spectral" information (via the first two spectral moments of the kernel which determines the angle $\theta$). This quantity is the fundamental spectral information that is needed to quantify feature learning, in any architecture, and any asymptotic.
Thanks for engaging too. Okay I mainly agree with that characterization. The one remaining nitpick is that *"Works that prove both upper and lower bounds (Yang et al., Jelassi et al. ) consider a single batch size, generally for convenience"*. My understanding is that the fixed batch size assumption and infinite width assumption is critical to Tensor Programs analysis, since it essentially implies a low rank property of the gradients, giving you an a priori way to estimate gradient spectral norms. I can't speak for Jelassi et al, as I haven't been through that work carefully.

And sorry, my cryptic comment was just supposed to say that the bounding approach can easily deal with finite width.
Thank you for engaging with my response.
- Indeed, I should have been more precise, I was talking about all works that provide proof of the scale of feature learning (by this, I mean both upper and lower bounds, of matching scale). The line of work of arxiv.org/abs/2405.15712 (using dynamical mean-field theory) is concerning with proving the existence of dynamical limits. The existence of well defined limits  can be interpreted as the existence of upper bound, but does not, however, imply lower bounds. The line of work of arxiv.org/abs/2405.14813 (on the "spectral criterion") is also, as far as I am aware, proposing a method that only gives upper bounds. Works that prove both upper and lower bounds (Yang et al., Jelassi et al. ) consider a single batch size, generally for convenience. But as I mentioned in my review, this is not out of reach with our approach.
- I am not sure I understood the sentence : "my same comment applies again with the words "batch size" replaced by "width" (since width is studied in our work)
I'll reply in full to your top level comment---thanks for writing that. Just two nits with this rebuttal:

**"this is an assumption made by all related works on feature learning"** not true, for instance one of the papers I posted in my rebuttal (and it's non concurrent antecedents) does not make this assumption. It satisfies itself with proving bounds on feature learning that hold for any batch size.

**"and this is not accessible via other existing approaches, which start with the infinite width limit"** my same comment applies again with the words "batch size" replaced by "width"
thank you. 
I will keep my score. 
Good luck!
We apologize for the oversight. These illustrations are computed with a single sample as an input (batch-size $1$) which is a random Gaussian vector in $\mathbb{R}^d$, so this is synthetic data. Since those models are rotationally invariant -- because the distribution of the input layer's weight is Gaussian -- we would get the same result for any input vector, be it synthetic, or from a dataset. We plan to release the (simple) code with the final version of the paper.
thank you for your response. 
I still did not find a "mention of what is assumed about the data: is it synthetic, random, from some known benchmark etc"
It would be helpful to get your response to that.