Title: The Feature Speed Formula: a flexible approach to scale hyper-parameters of deep neural networks
Venue: NeurIPS 2024
Accepted: yes
Keywords: Feature Learning

Deep Neural Networks

One SGD step

Hyperparameter scaling

dynamical isometry
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.

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

Abstract

Deep learning succeeds by doing hierarchical feature learning, yet tuning hyperparameters (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 θ glyph[lscript] between the feature updates and the backward pass (at layer index glyph[lscript] ). 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 θ glyph[lscript] , the loss decay, and the magnitude of the backward pass. This angle θ glyph[lscript] is controlled by the conditioning of the layer-tolayer Jacobians and at random initialization, it is determined by the spectrum of a certain kernel, which coincides with the Neural Tangent Kernel when glyph[lscript] = depth. Given θ glyph[lscript] , 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( θ glyph[lscript] ) = Θ(1 / √ glyph[lscript] ) . In contrast, ResNets with branch scale O (1 / √ depth ) maintain a non-degenerate angle cos( θ glyph[lscript] ) = Θ(1) . We use these insights to recover key properties of known HP scalings (such as µ P), and also introduce a new HP scaling for large depth ReLU MLPs with favorable theoretical properties.

1 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) θ glyph[lscript] between the feature updates and the backward pass (at layer index glyph[lscript] ) 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 θ glyph[lscript] , 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 θ glyph[lscript] , the loss decay and the magnitude of the backward pass at layer glyph[lscript] . 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 , θ glyph[lscript] can be estimated in terms of the spectrum of the backward to feature kernel (BFK) K glyph[lscript] 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).

Related work 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 VandenEijnden, 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.

2 The Feature Speed Formula

Consider a depthL NN architecture defined by the recursion, for glyph[lscript] ∈ [1 : L ] ,

where w glyph[lscript] ∈ R p glyph[lscript] are trainable parameters and we assume that the maps T glyph[lscript] : R m glyph[lscript] -1 × R p glyph[lscript] → R m glyph[lscript] 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 glyph[lscript] = ( ∂ L ∂f glyph[lscript] ) glyph[latticetop] ∈ R m glyph[lscript] the vectors of the backward pass. A gradient descent (GD) step with layerwise learning-rate (LR) η glyph[lscript] · δt > 0 for glyph[lscript] ∈ [1 : L ] consists in adding to each w glyph[lscript] the update

We are interested on the evolution of the NN over a single GD step with infinitesimally small step-size δt glyph[lessmuch] 1 . For any quantity x associated to the NN, we denote ˙ x its instantaneous velocity ˙ x := lim δt ↓ 0 δx δt when it exists. In particular, we have ˙ w glyph[lscript] = -η glyph[lscript] ∇ glyph[lscript] 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 ∑ glyph[lscript] ≤ v η glyph[lscript] ‖∇ glyph[lscript] L ‖ 2 2 = 0 then ˙ f v = 0 . Otherwise, the (non-oriented) angle θ v between ˙ f v and -b v is well defined in [0 , π / 2 [ and it holds

Proof. By the chain rule, we have ˙ f v = ∑ glyph[lscript] ≤ v ∂f v ∂w glyph[lscript] ˙ w glyph[lscript] = -∑ glyph[lscript] ≤ v η glyph[lscript] ∂f v ∂w glyph[lscript] ( ∂ L ∂w glyph[lscript] ) glyph[latticetop] . It follows

Clearly, if ∑ glyph[lscript] ≤ v η glyph[lscript] ‖∇ glyph[lscript] L ‖ 2 2 = 0 then ˙ w glyph[lscript] = 0 for glyph[lscript] ≤ v and thus ˙ f v = 0 . Otherwise θ v is well defined and it holds ‖ b v ‖ 2 ‖ ˙ f v ‖ 2 cos( θ v ) = -b glyph[latticetop] v ˙ f v = ∑ glyph[lscript] ≤ v η glyph[lscript] ‖∇ glyph[lscript] L ‖ 2 2 and the claim follows. (In terms of the BFK defined below, Eq. (3) is equivalent to b glyph[latticetop] v K v b v = ∑ glyph[lscript] ≤ v η glyph[lscript] ‖∇ glyph[lscript] 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 := ∑ glyph[lscript] ≤ v η glyph[lscript] ‖∇ glyph[lscript] ‖ 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 as

Here 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 as

By construction, it holds ˙ f 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 ≥ · · · ≥ λ m v ≥ 0 be the sorted eigenvalues of K v and let M p := 1 m v ∑ m v i =1 λ p i be its spectral moments. It holds λ mv λ 1 ≤ cos( θ v ) ≤ 1 . Moreover, if b v is Gaussian and independent from K v , then as m v →∞ ,

as soon as √ M 2 / M 1 and √ M 4 / M 2 are uniformly bounded (i.e. are upper bounded by some C > 0 with probability going to 1 as m v →∞ ).

Proof of Thm. 3.2. By the chain rule, it holds

hence ˙ f v = -K v b v . Denoting K 1 / 2 v the unique psd square-root of K v , it follows

The first claim follows from Eq. (6) and the worst-case bounds ‖ K v b v ‖ 2 ≤ λ 1 ‖ b v ‖ 2 and ‖ K 1 / 2 v b v ‖ 2 ≥ √ λ m v ‖ 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 write Z = ‖ 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 2 with 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 NNs

4.1 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 holds ‖ f v ‖ rms = Θ(1) for v ∈ [1 : L -1] .

(FL) Feature learning. It holds ‖ ˙ f L -1 ‖ rms = Θ(1) .

(LD) Loss decay. It holds -˙ L = Θ(1) .

(BC) Balanced contributions. It holds η glyph[lscript] ‖∇ glyph[lscript] L ‖ 2 2 = Θ( η glyph[lscript] ′ ‖∇ glyph[lscript] ′ L ‖ 2 2 ) for any glyph[lscript], glyph[lscript] ′ ∈ [1 : L ] .

We discuss these specific properties because they are amenable to our tools and enforcing them requires ( L -1) + 1 + 1 + ( L -1) = 2 L 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 including f 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.

Proposition 4.1. A scaled NN (1) satisfies (FL), (LD), and (BC) if and only if

Proof. Property (LD) requires ∑ L glyph[lscript] =1 η glyph[lscript] ‖∇ glyph[lscript] L ‖ 2 2 = -˙ L = Θ(1) and (BC) requires the terms in the sum to be balanced, this leads to Eq. (8). Now by Thm. 2.1, property (FL) requires

which leads to (7). Conversely, it is clear that Eq. (8) and (7) imply (FL), (LD) and (BC).

4.2 Towards automatic HP scaling

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 F orward pass, the S ensitivities and the C ontributions. Let us make some observations regarding the scale invariant LRs:

and

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 algorithm x t +1 = x t -η t ∇ f ( x t ) is given by η t = f ( x t ) ‖∇ f ( x t ) ‖ 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 glyph[lscript] = -˜ η glyph[lscript] · ∇ glyph[lscript] L ‖∇ glyph[lscript] L ‖ rms . For such an algorithm, properties (BC) and (LD), suggests the LR ˜ η glyph[lscript] = Θ( ‖∇ glyph[lscript] L ‖ rms L ·‖∇ glyph[lscript] L ‖ 2 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 scale ∏ L glyph[lscript] =1 σ glyph[lscript] .

5 Scaling width and depth of MLPs and ResNets

5.1 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 glyph[lscript] ∈ [1 : L -1] , by

where φ ( 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 layers m 1 = · · · = m L -1 = m (assumed equal), the output width m L = k . The trainable parameters are ∀ glyph[lscript] ∈ [1 : L ] , W glyph[lscript] ∈ R m glyph[lscript] × m glyph[lscript] -1 . Such NNs 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 glyph[lscript] are independent N (0 , σ 2 glyph[lscript] ) random variables for glyph[lscript] ∈ [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 glyph[lscript] ∈ [1 : L -1] , let σ glyph[lscript] = √ 2 /m glyph[lscript] -1 . As m →∞ , it holds

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 glyph[lscript] ∈ [2 : L -1] , by

(a) BFA vs. layer (

L

= 200

)

(b) Output BFA vs. depth

(c) Output BFA vs. scale

n/2

n/2

L =500

L=100

0E = 7

β=

L

Output BFA

rt/4

n/4

/

Q

25

50

75

Layer index v

100

125

150

175

200

25

50

75

Depth L

100

125

150

175

200

10-1

Normalized branch scale /L - β

100

101

where φ ( u ) = u in our theoretical results. The architecture HPs are the input width m 0 = d , the widths of the hidden layers m 1 = · · · = m L -1 = m , the output width m L = k . The trainable parameters are ∀ glyph[lscript] ∈ [1 : L ] , W glyph[lscript] ∈ R m glyph[lscript] × m glyph[lscript] -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 glyph[lscript] ∈ [1 : L -1] , let σ glyph[lscript] = Θ(1 / √ m glyph[lscript] -1 ) . As m →∞ it holds:

Moreover, if (BC) holds then cos( θ v ) = Θ(1) .

Numerical experiments We consider 3 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 glyph[latticetop] 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 glyph[greatermuch] 1 ).

5.2 Characterizing HP scalings for MLPs

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 -1 of 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 F orward scales, S ensitivities, and C ontributions.

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.

(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 glyph[greatermuch] 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.

5.3 Characterizing HP scalings for ResNets

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.

√

4

sensitivity of output feature

+

NTK

+

MF+μP

3

+

ours

2

T

十

0

5

10

15

T

20

25

T

30

35

40

depthL

(a) ReLU MLP ( β = 1 )

(b) ReLU ResNet

20.0

outputfeature

17.5

15.0

12.5

10.0

sensitivity

7.5

5.0

+

β=1/√L

2.5

+

β = 1/L

0.0

20

30

40

50

60

70

80

06

100

depthL

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

6.1 Stability and backward speed formula

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:

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 ∇ glyph[lscript] L = b glyph[lscript] g glyph[latticetop] glyph[lscript] -1 and thus ‖∇ glyph[lscript] L ‖ F = ‖ b glyph[lscript] ‖ 2 ‖ g glyph[lscript] -1 ‖ 2 . It follows

We can thus ensure the gradient stability by ensuring, for all glyph[lscript] ∈ [1 : L ] ,

(FS) Forward stability . It holds ‖ ˙ g glyph[lscript] -1 ‖ 2 ‖ g glyph[lscript] -1 ‖ 2 = O (1) for glyph[lscript] ∈ [1 : L ] , and

(BS) Backward stability . It holds ‖ ˙ b glyph[lscript] ‖ 2 ‖ b glyph[lscript] ‖ 2 = O (1) for glyph[lscript] ∈ [1 : L ] .

We focus on these simpler desiderata (FS) and (BS) instead of (S) for the rest of the discussion. To estimate ˙ b v , we rely on a 'backward' version of the feature speed formula that holds in 1 -homogeneous NNs.

Proposition 6.1 (Backward speed formula) . Consider a general architecture of the form (1) , take v ∈ [1 : L ] and assume that the map f v → f L is positively 1 -homogeneous 4 . If -f glyph[latticetop] L ∇ 2 loss[ f L ] ˙ f L + ∑ glyph[lscript]>v η glyph[lscript] ‖∇ glyph[lscript] L ‖ 2 2 = 0 then ˙ b v = 0 . Otherwise, the (non-oriented) angle ˜ θ v between f v and ˙ b v is well defined in [0 , π / 2 [ and it holds

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 surely 5 that:

where the last expression defines ˜ K v . Reasoning as in Thm. (3.2), since f v is Gaussian at initialization and noticing that ˜ K v 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 ˙ b v just as well as ˙ f v .

6.2 Scale invariance for homogeneous architectures

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 glyph[lscript] ∈ R p glyph[lscript] . Let θ 0 = ( w 1 (0) , . . . , w L (0)) and let ˜ θ 0 = σ glyph[circledot] θ 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 η glyph[lscript] ( t ) ‖∇ glyph[lscript] L ( θ ( t )) ‖ 2 2 = ˜ η glyph[lscript] ( t ) ‖∇ glyph[lscript] L ( ˜ θ ( t )) ‖ 2 2 and starting from θ 0 and ˜ θ 0 respectively. If ∏ L glyph[lscript] =1 σ glyph[lscript] = 1 then ˜ θ ( t ) = σ glyph[circledot] θ ( t ) for all t ≥ 1 .

6.3 Characterization of admissible scalings for ReLU MLPs

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:

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

where C 1 = η 1 ‖∇ 1 L ‖ 2 2 , C hid = ∑ L -1 glyph[lscript] =2 η hid ‖∇ glyph[lscript] L ‖ 2 2 and C L = η L ‖∇ L L ‖ 2 2 . In particular, the scaling FSC (Table 1) satisfies these desiderata.

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

Acknowledgments

This work was supported in part by the International Centre for Theoretical Sciences (ICTS) for participating in the meeting -Data Science: Probabilistic and Optimization Methods (code:ICTS/dspom2023/7)

References

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization. In International conference on machine learning , pages 242-252. PMLR, 2019.

Alexander Atanasov, Blake Bordelon, and Cengiz Pehlevan. Neural networks as kernel learners: The silent alignment effect. In International Conference on Learning Representations , 2021.

Aristide Baratin, Thomas George, César Laurent, R Devon Hjelm, Guillaume Lajoie, Pascal Vincent, and Simon Lacoste-Julien. Implicit regularization via neural feature alignment. In International Conference on Artificial Intelligence and Statistics , pages 2269-2277. PMLR, 2021.

Jérôme Bolte and Edouard Pauwels. A mathematical model for automatic differentiation in machine learning. Advances in Neural Information Processing Systems , 33:10809-10819, 2020.

Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit. arXiv preprint arXiv:2309.16620 , 2023.

Lénaïc Chizat and Francis Bach. On the global convergence of gradient descent for over-parameterized models using optimal transport. Advances in Neural Information Processing Systems , 31, 2018.

Lénaic Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. Advances in Neural Information Processing Systems , 32, 2019.

Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes overparameterized neural networks. In International Conference on Learning Representations , 2018.

John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research , 12(7), 2011.

Boris Hanin. Which neural net architectures give rise to exploding and vanishing gradients? Advances in neural information processing systems , 31, 2018.

Boris Hanin and Mihai Nica. Products of many large random matrices and gradients in deep neural networks. Communications in Mathematical Physics , 376(1):287-322, 2020.

Boris Hanin and David Rolnick. How to start training: The effect of initialization and architecture. Advances in Neural Information Processing Systems , 31, 2018.

Soufiane Hayou, Arnaud Doucet, and Judith Rousseau. On the impact of the activation function on deep neural networks training. In International conference on machine learning , pages 2672-2680.

Soufiane Hayou, Eugenio Clerico, Bobby He, George Deligiannidis, Arnaud Doucet, and Judith Rousseau. Stable ResNet. In International Conference on Artificial Intelligence and Statistics , pages 1324-1332. PMLR, 2021.

Elad Hazan and Sham Kakade. Revisiting the Polyak step size. arXiv preprint arXiv:1905.00313 , 2019.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kernel: Convergence and generalization in neural networks. Advances in Neural Information Processing Systems , 31, 2018.

Samy Jelassi, Boris Hanin, Ziwei Ji, Sashank J Reddi, Srinadh Bhojanapalli, and Sanjiv Kumar. Depth dependence of µ -P learning rates in ReLU MLPs. arXiv preprint arXiv:2305.07810 , 2023.

Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR) , San Diega, CA, USA, 2015.

Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. Nature , 521(7553):436-444, 2015.

Mufan Li, Mihai Nica, and Dan Roy. The future is log-Gaussian: ResNets and their infinite-depth-and-width limit at initialization. Advances in Neural Information Processing Systems , 34:7852-7864, 2021.

Zhiyuan Li, Srinadh Bhojanapalli, Manzil Zaheer, Sashank Reddi, and Sanjiv Kumar. Robust training of neural networks using scale invariant architectures. In International Conference on Machine Learning , pages 12656-12684. PMLR, 2022.

Zenan Ling and Robert C. Qiu. Spectrum concentration in deep residual learning: a free probability approach. IEEE Access , 7:105212-105223, 2019.

Yizhang Lou, Chris E Mingard, and Soufiane Hayou. Feature learning and signal propagation in deep neural networks. In International Conference on Machine Learning , pages 14248-14282. PMLR, 2022.

Pierre Marion and Lénaïc Chizat. Deep linear networks for regression are implicitly regularized towards flat minima. arXiv preprint arXiv:2405.13456 , 2024.

Pierre Marion, Adeline Fermanian, Gérard Biau, and Jean-Philippe Vert. Scaling ResNets in the large-depth regime. arXiv preprint arXiv:2206.06929 , 2022.

Per-Gunnar Martinsson and Joel A Tropp. Randomized numerical linear algebra: Foundations and algorithms. Acta Numerica , 29:403-572, 2020.

Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences , 115(33):E7665-E7671, 2018.

Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice. Advances in Neural Information Processing Systems , 30, 2017.

Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. The emergence of spectral universality in deep networks. In International Conference on Artificial Intelligence and Statistics , pages 1924-1932. PMLR, 2018.

Boris T. Polyak. Introduction to optimization. 1987.

Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. Advances in Neural Information Processing Systems , 29, 2016.

Grant M. Rotskoff and Eric Vanden-Eijnden. Neural networks as interacting particle systems: Asymptotic convexity of the loss landscape and universal scaling of the approximation error. stat , 1050:22, 2018.

Wojciech Tarnowski, Piotr Warchoł, Stanisław Jastrze ¸bski, Jacek Tabor, and Maciej Nowak. Dynamical isometry is achieved in residual networks in a universal way for any activation function. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 2221-2230. PMLR, 2019.

Twan Van Laarhoven. L2 regularization versus batch and weight normalization. arXiv preprint arXiv:1706.05350 , 2017.

Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cengiz Pehlevan. Feature-learning networks are consistent across widths at realistic scales. arXiv preprint arXiv:2305.18411 , 2023.

Ruosi Wan, Zhanxing Zhu, Xiangyu Zhang, and Jian Sun. Spherical motion dynamics: Learning dynamics of neural network with normalization, weight decay, and SGD. arXiv preprint arXiv:2006.08419 , 2020.

Zhichao Wang, Andrew Engel, Anand Sarwate, Ioana Dumitriu, and Tony Chiang. Spectral evolution and invariance in linear-width neural networks. arXiv preprint arXiv:2211.06506 , 2022.

Greg Yang. Tensor programs III: Neural matrix laws. arXiv preprint arXiv:2009.10685 , 2020.

Greg Yang and Edward J Hu. Tensor programs IV: Feature learning in infinite-width neural networks. In International Conference on Machine Learning , pages 11727-11737. PMLR, 2021.

Greg Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hyperparameter transfer. Advances in Neural Information Processing Systems , 34:17084-17097, 2021.

Greg Yang, James B Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint arXiv:2310.17813 , 2023a.

Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Feature learning in infinite-depth neural networks. In NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning , 2023b.

Huishuai Zhang, Da Yu, Mingyang Yi, Wei Chen, and Tie-Yan Liu. Stabilize deep ResNet with a sharp scaling factor τ . Machine Learning , 111(9):3359-3392, 2022.

A Proofs omitted from the main text

Proof of Lem. 3.3. Writing K = V DV glyph[latticetop] with D = diag ( λ 1 , . . . , λ m ) and V ∈ R m × m orthonormal, we have ‖ Ka ‖ 2 2 = a glyph[latticetop] V D 2 V glyph[latticetop] a . Conditioned on K , the vector u = V glyph[latticetop] a is isotropic Gaussian so E u 2 i = 1 m for i ∈ [1 : m ] . Hence, on the one hand

On the other hand, using the fact that the variance of a chi-square random variable is 2 , glyph[negationslash]

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 ) glyph[latticetop] ∈ 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 ) glyph[latticetop] . 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 η glyph[lscript] = Θ( L -3 / 2 ) , it holds ‖ ˙ g L -1 ‖ rms = Θ(1) . In view of (11), it holds in their setting for glyph[lscript] ∈ [2 : L -1]

Using η glyph[lscript] = Θ( L -3 / 2 ) , it follows ∑ L -1 glyph[lscript] =2 η glyph[lscript] ‖∇ glyph[lscript] L ‖ 2 2 = Θ( L -1 / 2 ) . By Thm. 2.1 and using m L -1 ‖ z L -1 ‖ rms = Θ(1) , we get

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 ):

and θ v the (non-oriented) angle between ˙ g v and z glyph[lscript] := ( ∂ L /∂g v ) glyph[latticetop] . 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 ˜ K v such that ˙ b v = ˜ K v f v (see Eq. (14) ) and ˜ θ v the (non-oriented) angle between f v and -˙ b v . Then it holds, as hidden width diverges, cos( ˜ θ v ) = Θ ( M 1 ( ˜ K v ) / √ M 2 ( ˜ K v ) ) = Θ(( 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 glyph[lscript] follows from [Hanin and Nica, 2020, Prop. 2] and the Gaussianity of z glyph[lscript] is direct since z glyph[lscript] = W glyph[latticetop] glyph[lscript] +1 b glyph[lscript] +1 where W glyph[lscript] +1 is Gaussian and independent from b glyph[lscript] +1 . Also, we have the more explicit expression

where ∂g v ∂f glyph[lscript] = D v W v . . . D glyph[lscript] +1 W glyph[lscript] +1 D glyph[lscript] and D i = diag( φ ( f i )) (by [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 η glyph[lscript] ‖ g glyph[lscript] -1 ‖ 2 2 is constant for glyph[lscript] ∈ [1 : L -1] , it follows

For the second claim, we have (see Section 6) ˜ K v = ∑ L glyph[lscript] = v +1 η glyph[lscript] ‖ b glyph[lscript] ‖ 2 2 ( ∂g glyph[lscript] -1 ∂f v ) glyph[latticetop] ( ∂g glyph[lscript] -1 ∂f v ) . Under (SP) and (BC), we have η glyph[lscript] ‖ b glyph[lscript] ‖ 2 2 is constant for glyph[lscript] ∈ [2 : L ] , hence it follows

By comparing the expressions for K v and ˜ K v , we see that ˜ K v 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 glyph[lscript] ‖ 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 ( glyph[lscript] → v ) = Θ(1) and σ max ( glyph[lscript] → v ) = Θ(1) , where σ min ( glyph[lscript] → v ) and σ max ( glyph[lscript] → v ) are the smallest, respectively largest singular value of ∂f v ∂f glyph[lscript] . The estimate on b glyph[lscript] = ( ∂f L ∂f v ) glyph[latticetop] b L directly follows.

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 glyph[lscript] = φ ( f glyph[lscript] ) and h glyph[lscript] = W glyph[lscript] g glyph[lscript] -1 . By direct computations, it holds (here w glyph[lscript] is the vectorization of W glyph[lscript] ):

Using the inequalities

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 ‖∇ glyph[lscript] L ‖ 2 = ‖ b glyph[lscript] g glyph[latticetop] glyph[lscript] -1 ‖ F = ‖ b glyph[lscript] ‖ 2 · ‖ g glyph[lscript] -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 :

But in the considered asymptotics √ m · cos( θ L -1 ) = Θ( √ m/L ) →∞ so ‖ ˙ f 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 holds

Now, by differentiating in time both sides we get

We have ˙ L = -∑ L glyph[lscript] =1 η glyph[lscript] ‖∇ glyph[lscript] L ‖ 2 2 and moreover, from the proof of Thm. 2.1, -b glyph[latticetop] v ˙ f v = ∑ glyph[lscript] ≤ v η glyph[lscript] ‖∇ glyph[lscript] L ‖ 2 2 . So we deduce

We conclude by writing -˙ b glyph[latticetop] f = ‖ ˙ b ‖ ‖ f ‖ cos( ˜ θ ) and rearranging.

Proof of Prop. 6.2. By assumption at time t = 0 , it holds ˜ θ (0) = σ glyph[circledot] θ (0) so let us prove the result by recursion. Assume that the claim is true at iteration t . Since ∏ σ glyph[lscript] = 1 , it holds f L ( θ ( t )) = f L ( ˜ θ ( t )) . Moreover, since ∂f L ∂w glyph[lscript] is 0 -homogeneous in w glyph[lscript] and separately 1 -homogeneous in ( w i ) i = glyph[lscript] . It follows glyph[negationslash]

glyph[negationslash]

In particular, we deduce that the LRs are related by ˜ η ( t ) η ( t ) = ‖∇ glyph[lscript] L ( θ ( t )) ‖ 2 2 ‖∇ glyph[lscript] L ( ˜ θ ( t )) ‖ 2 2 = σ 2 glyph[lscript] . It follows, for any glyph[lscript] ∈ [1 : L ] ,

This proves θ ( t +1) = σ glyph[circledot] θ ( t +1)

˜ and the claim follows by recursion.

B Initializing with zero output weights

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):

The second forward pass is the same as the first one, with the only difference that

Assuming b L (0) = b L (1) (linear loss) for simplicity, this leads to a second backward pass:

For the second GD step to satisfy (FL), we just need to ensure

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 = 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

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) is not excluded. This is where (BS) comes into play. It requires in particular

using that cos( ˜ θ 1 ) = Θ(1 / √ L ) by Lem. A.1. Combining both conditions for σ L imply, on the one hand, that C 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).

C Characterization of reparameterization invariant LR

Consider a function f : ∏ L glyph[lscript] =1 R p glyph[lscript] → R admitting a (selection) derivative and, for a fixed scale vector α ∈ ( R ∗ + ) L consider the function g ( y ) = f ( α · y ) where α · x denotes ( α 1 x 1 , . . . , α L x L ) . Consider one step of GD on the two functions, given for glyph[lscript] ∈ [1 : L ] , by

with identical starting points, that is x glyph[lscript] = α glyph[lscript] · y glyph[lscript] for glyph[lscript] ∈ [1 : L ] .

glyph[negationslash]

Proposition C.1. Consider adaptive learning rates, which are of the form η glyph[lscript] = η glyph[lscript] ( ∇ f ( x )) . Then x ′ = α · y ′ for all α ∈ ( R ∗ + ) L if and only if η glyph[lscript] is ( -2) -homogeneous in ∇ glyph[lscript] f ( x ) and 0 -homogeneous in ∇ glyph[lscript] ′ f ( x ) for glyph[lscript] = glyph[lscript] ′ .

One such LR is precisely that suggested by Prop. 4.1: η glyph[lscript] ∝ ‖∇ glyph[lscript] f ( x ) ‖ -2 2 .

Proof. For glyph[lscript] ∈ [1 : L ] , it holds

Then α · y ′ = x ′ for all α ∈ ( R ∗ + ) L is equivalent to

which is the claimed homogeneity property.

NeurIPS Paper Checklist

1. Claims

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.

2. Limitations

Question: 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.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Some of the proofs are deferred to the appendix, but all are provided.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Experiments (with synthetic data) corroborate closely the theoretical results.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

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.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

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.

7. Experiment Statistical Significance

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.

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

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

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

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

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

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. The BFA and BFK are interesting objects to study and the geometrical picture that arises (mentioned in the introduction) gives a nice intuition.\n\n2. The main results (Thms 2.1 and 3.2) are clearly stated and the proofs are straightforward. \n\n3. The contributions are clearly stated and the relation to previous work distinguishes these contributions. \n\n4. Earlier results are recovered here with a transparent derivation, but Ref [1] also provided quite an intuitive derivation, as you mentioned. \n\n\n\n\n\n[1] https://arxiv.org/abs/2310.17813\n\nWEAKNESSES:\nDespite the strengths mentioned above, I did not give a higher score for the following reasons: \n\n1. Novelty for HP scaling: \nAs 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. \n\n2. Applicability of results: \nWhile 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. \n\n3. Experimental findings: \nI 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. \n\n4. 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. \n\n\n\n\ntypos: \n- Fig1 caption 1st line: witdh -> width\n\nQUESTIONS:\n1. in line 194 - BC condition - is this a strict equality? or is there some tolerance? \n\n2. In the Introduction you use the term \"hierarchical features\" - can you give a definition for that? \n\n3. 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?\n\nLIMITATIONS:\nThe authors adequately addressed the limitations of their results.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe paper studies an extremely important topic in deep learning theory. Given the technically challenging nature of the study of large width and depth limits, the paper is superbly well-written and accessible. Prior papers by Yang et al and Bordelon et al have done important work in developing the study of large width and depth limits, but their derivations are either very dense or otherwise rely on non-rigorous methods to derive the scalings. This paper manages to both rigorously motivate feature learning at infinite width and depth while simultaneously making the paper short and accessible. This is a major strength and no easy feat. I commend the authors on it. \n\nBeyond this, there are several results of strong technical merit that will be of value for researchers studying infinite width and depth limits. The infinite depth MLP and scale invariant learning rate discussions are particularly interesting. The authors do a good job placing their work in context by presenting tables comparing their parameterization to others. \n\nUltimately, I believe that this paper is not only technically novel and sound, but is also a service to the community. I strongly recommend it for acceptance.\n\nWEAKNESSES:\nThere are no technical weaknesses that I have found, and I have gone through the derivations in detail. My only comment is expository:\nIn 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.\n\nQUESTIONS:\nAs 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? \n\nIn 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\n290 ResNet)\" \n\nI 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.\n\nLIMITATIONS:\nLarge 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.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- the authors propose an interesting notion and derive interesting analyses surrounding it\n- the parts of the math I checked seem rigorous and sound\n- the authors do a good job of connecting their work to related work\n- the ideas are quite creative\n\nWEAKNESSES:\nI 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.\n\nWith that said, the main issues I see with the paper are:\n\n### **No real experimental evaluation**\n\nMy 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.\n\nI 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?\n\nAs 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.\n\n### **Doesn't go that far beyond existing ideas**\nA 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.\n\n### **Paper only seems to focus on batch size one**\n\nIn my opinion, doing things only at batch size one is a bit toy, and it would be better to directly analyze larger batch sizes.\n\nQUESTIONS:\nPlease see the weaknesses section\n\nLIMITATIONS:\nOverall, 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.\n\nI realize this feedback might be construed as being fairly negative, but I hope that it can help to improve the work.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n# Originality\n\nThis paper tackles an important problem: the relation between the optimal hyperparameters of a neural network and its architecture.\n\n# Clarity\n\nThis paper is easy to read and the statements are clear.\n\nWEAKNESSES:\n# Originality\n\nThe 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.\n\n# Significance\n\nThe benefits and the properties of the BFA are still unclear.\n\nFor 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.\n\n# Quality\n\nThe 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.\n\n# EDIT: References\n\n[1] Gradient descent provably optimizes over-parameterized neural networks (2018), Du et al.\n\n[2] Gradient descent finds global minima of deep neural networks (2019), Du et al.\n\n[3] A convergence theory for deep learning via over-parameterization (2019), Allen-Zhu et al.\n\n[4] Stochastic Gradient Descent Optimizes Over-parameterized Deep ReLU Networks (2020), Zou et al.\n\nQUESTIONS:\nHow does the \"ours\" hyperparameter scaling compare to the others (usual, muP or NTK)?\n\nLIMITATIONS:\nLack of experimental validation.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2024-07-17"}