The Feature Speed Formula: a flexible approach to scale hyper-parameters of deep neural networks
Authors: Lénaïc Chizat
Praneeth Netrapalli
Venue: NeurIPS 2024 poster (2024)
Keywords: Feature Learning
Deep Neural Networks
One SGD step
Hyperparameter scaling
dynamical isometry
Primary area: optimization_for_deep_networks
TL;DR: We propose an elementary approach to quantifying feature learning in wide and deep neural networks and to derive hyperparameter scalings.

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
Lénaïc Chizat
Praneeth Netrapalli Google DeepMind pnetrapall1@google.co com
Institute of Mathematics, EPFL Lausanne, Switzerland lenaic. www.ccciiztttiiiee
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, introduce a key notion to predict and control feature learning: the angle 0g between the feature updates and the backward pass (at layer index E). We show that the magnitude of feature updates after one GD step, at any training time, can be expressed via simple and general feature speed formula in terms of this angle 0g, the loss decay, and the magnitude of the backward pass. This angle 0g is controlled by the conditioning of the layer-to- layer Jacobians and at random initialization, it is determined by the spectrum of certain kernel, which coincides with the Neural Tangent Kernel when = depth. Given 8g, the feature speed formula provides 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 cos (e) ((1/VE). In contrast, ResNets with branch scale (de/vt)phh) maintain nen-degene angle cos(O,) (()). We use these insights to recover key properties of known HP scalings (such nP), and also introduce new HP scaling for large depth ReLU MLPs with favorable theoretical properties.
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 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 all (Jacot al., 2018). This limited understanding makes difficult design NNs architectures and flyer-parameeter (HP) scalings, and begs the development of tools to quantify feature learning.
In this paper, we demonstrate that the backward-feann angle (BFA) 0g between the feature updates and the backward pass (at layer index f) is central object in this quest. Indeed, show that the magnitude of feature updates after one GD step, at any training time, can be expressed via a simple feature speed formula terms of this angle, the loss decay and the magnitude of the backward pass. Given the knowledge of 0,, this leads to a general approach to quantify key dynamical properties of the training dynamics of NNs such the speed of feature learning and loss decay and to characterize the HP scalings satisfying these properties.
Contributions
Our contributions the following:
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

We prove the feature speed formula (Thm 2.1) which quantifies the feature updates in terms of the BFA 8,, the loss decay and the magnitude of the backward pass at layer f. 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 we develop tools to quantity the BFA. In particular, we show that, in the batch-size I case, O2 can be estimated in terms of the spectrum of the backward to feature kernel (BFK) K2 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 depth L, co,) = OL for ReLU MLPs (Prop. 5.1 exploiting result in Jelassi et al. (2023) and that cos(Or) (1) for linear ResNets with branch scale V22) (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 S 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 uP [Bordelon et al.) [023] Yang et al. 2023b) for ResNets (Table and, for ReLU MLP, we
introduce scaling with output scale d'depth
width (Table that does not suffer from vanishing loss decay, in contrast to the one studied in [Jelassi et al. 2023).
Finally, in Section we develop more "axiomatic" approcac: starting from 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. 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 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-layee Jacobians around 1, performance. Our analysis gives a concrete justification of the link between dynamical isometry and successful training, as show that it is related to the aligmment between the backward pass and feature updates. These questions have also been studied in ResNets, see e.g. (Hayou et al.] 2202
important indicator of training
Marion et al. 2022 Li et al. [021] for signal propagation and (Tamowski al. 2019 Ling and Qiu 2019] for dynamical isometry.
In 2018, two viewpoints for the dynamics of wide NNs were simultancously introduced: feature learning limit for two-layer MLPs [Mei et al.] 2018 Chizat and Bach, 2018 Rotskoff and Vanden-
Eijnden 2018) and limit without feature learning for general NNs (Jacot et al. [2018] Du et al.
2018 (188|Alles AAlele------n -hu 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 202T] formulated the maximal update p-criterion (it is part of the properties we study in Section(4). 4). This criterion led to full classification of HP scalings in the infinite hidden width limit (at fixed depth), and singled-out the so-called j-parameeterization
(uP) 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 μPP 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. 2202T), Atanasov et al. 202 Lou et al. 20022. J. 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 setting,
but the large batch-size setting of these works is 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 /eeth branch scaling as providing desirable properties, in particular that of HP transfer [Yang et al. 202T]. These works take the infinite width limit as first step in their analysis, before studying the resulting objects, resulting in 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-propagann equations, before taking the large width-then-depet limit, last step.
Notations For integers a, E 2, write (a: bj = {a...., b). For any vector I we denote by ||=||rms m |/|==/ 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 depth-L. NN architecture defined by the recursion, for & E [1 L.].
fo mo
Se = Te ws) E R
loss( E
(I)
where wy E 12 are trainable parameters and we assume that the maps T3 1 R + admit elementary (log-exp) selection [Bolte and Pauwels, 2020j. By flattening the tensors, one encode most practical NN architectures in Eq. (1). For instance, mo is typically the product of batch-size (or context length) with input dimension. We denote by by = (afe) E R the vectors of the backward pass. A gradient descent (GD) step with layerwise learning-rate (LR) ye ōl > 0 for L E [1 L] consists in adding to each wy the update
mm
00
Swe - 5t Veh = -Tpe 5t
We are interested on the evolution of the NN over single GD step with infinitesimally small step-size öt < 1. For any quantity r associated to the NN, we denote its instantaneous velocity = limm bg 32
when it exists. particular, we have the - v,C.
The following identity is the seed of our approach. It expresses at any training time the speed of features terms of other interpretable quantities, including the backward feature angle (BFA) 0g.
Theorem 2.1 (Feature speed formula). Let v E [1: If Erso V 0 then fo 0. Otherwise, the (non-oriented) angle Og between f. and bb is well defined in [0, 7/2[ and it holds
llful12= Eis, Well V
cos(O.). -
Proof. By the chain rule, we have f. = =Ess
as. wi = -Ees
as.
( (.
) It follows
T
Bwee
me
thre
Buee
-o fo
Σ
Ml of Ow: Owp
ac 0f. dL
mi
dL
)
dC
Σ ESC
V
(3)
Owi
Clearly, if Eis Me vellV = 0 then wi = 0 for l < and thus fo = O. Otherwise 0, is well defined
and it holds ||b..laal| 12 cos (.,) - b. = E Ees nel V and the claim follows. (In terms of the BFK defined below, Eq. 3 is equivalentt to b. K.b. Eeso Me V for v E []: 0
To better appreciate the content of Thm. 2.1 let us re-express it in terms of root mean-square (RMS)
norms. Let Lego = Eiso Me Vell be the contribution to the loss decrease of all the parameters
before fo the forward pass, and note that İgr. c. Then, the identity rewrites as
|.
S..
(4)
Ago
cos(O,)) m, b I|rms
'This technical assumption that covers virtually all functions interest in deep learning. In particular, the maps Te admit selection derivatives that are compatible with the chain rule [Bolte and Pauwels| (2200 Prop. 4] and that coincide almost werywhere with standard derivatives [Bolte and Pauwels] [020 Prop. 3J. To simplify exposition, always implicitly assume that at differentiability point of the maps Tt.
3

---

Here Sy can be interpreted as the sensitivity (and. in the terminology of [Chizat et al. 2019], 1/S, as the laginess) of the feature e: 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 [|bellmms the size of the feature my, and the BFA 0g. Let now build tools to quantify the BFA.
Quantifying the backward-feature angles (BFA)
Information about the BFA 0, can be gained from the Backward to Feature Kernel (BFK). Definition 3.1 (Backward to Feature Kernel). For E [1: I.J. the BFK is the psd matrix defined as
O5u Owe
af
K2 Σ Me 85
E R mm..
(5)
Ows
By construction, it holds f. --=KKubb In other words, the BFK takes a backward pass vector as input and returns the (negative of the) feature velocity. For u = L, K, coincides with the Neural Tangent Kernel (Jacot al. 2018). We now show how the sprectrum of K2 relates to BFA.
Theorem 3.2. Let \1 2 Am. 2 0 be the sorted eigenvalues of Ky and let Mp = Th TH X be its spectral moments. holds from K., then as m, - 00,
Am S cos(O,)) S X1
Moreover, ifb, is Gaussian and independent
Pr M cos(O,) MM2
as soon as VMaafan Mi and Ma are uniformly bounded (i.e. are upper bounded by some C' with probability going to 1 as my - 20).
Proof of Thm.
3.2
By the chain rule, it holds
fo =
ne Σ SS
0fe O. Owg!
Σ ES
Öfe [85. T 0C
(
ne Owr Owc! Of.
Owe
hence f. = --KKb Kbbb. Denoting K the unique psd square-root of K,, it follows
-bT
|K 60
cos(O,) ||Bollalbbo11 I1, 112
(6)
The first claim follows from Eq. (6 and the worst-case bounds |,,bb,11 S Aillbolle and
b 2 /mm 112. lby.lll The second claim related to the trace estimation method via ran- dom matrix-vector products (Martinsson and Tropp. 2020, Chap. 4J. We assume without loss of generality that E([lb,,1l
write Z ||/]/// b, =a( b) where a E Z K, = Mj and E > 0 as m. > oc. An analogous decomposition holds for K (b.) [
1 and by Lem. 3.3
with M2 and the second claim follows.
Lemma 3.3. Let K. Rmxmm be a (potentially random) psd matrix and a ~ N(O, 1 Im) be independent. Then E[]] K]= M2(K) and Varl]] K] 2 M&(KK where Mp(K)= = 1 Em X and >i...., Am 2 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 random initialization in the large width limit of typical NNs, provided fo is directly followed by a weight-matrix matrix multiplication in the forward pass, so that b, guaranteed in quite general contexts, see Yang [2020) For MIP or ResNets with batch-size one, we show in Section 5 that cos(O,)) ) tightly related to the conditioning of layer-to-layer Jacobians, studied in the dynamical isometry literature [Pennington et al. 2017).
output of random matrix/vector multiplication. Asymptotic independence can be

---

Ensuring feature learning in scaled NNs
4.1 Properties for scaled NNs
Consider a sequence of NNs and parameters as in ① with some diverging architectural parameters such depth or width. We refer to such sequence as scaled NN. 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 [ = (1) for v E [1:L-1] -1..
(FL) Feature learning. It holds 1lfit-llemm ((1)
(LD) Loss decay. holds - (1).
(BC) Balanced contributions. It holds me V @(ne (V for any e, E E [1:
We discuss these specific propertics because they are amenable to our tools and enforcing them requires (L 1) + 1++++ (L (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 where we adopt more axiomatic approach and recover, for homogeneous architectures, (a more general version of) property (BC) 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 7. 1 scale HPs, but leaves the scale of fL free. The reason for not including ft in (SP) is that l.ffllens = o(1) does not lead to vanishing gradient in general, so this behavior should not be excluded 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 suitably normalized the norm of the backward pass.
Proposition 4.1. A scaled NN I satisfies (FL), (LD), and (BC) if and only if
164-1|1mm = 0 1---1 cos (1-1))
(7)
and
V [1: ne = -
(8)
Proof. Property (LD) requires Σ
Me V V
-
(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
ΣE Mh V
Irms cos cos(0L-1)
mL-1 | (1)
which leads Conversely, is clear that Eq. (8) and ( imply (FL), (LD) and (BC).
4.2 Towards automatic HP scaling
The criterion of Prop. 4.1., complemented with the property (SP), suggest 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 0z-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 fi-1 and ft. adjusted so that Eq. 77) holds.
(BC) & (LD): Scale invariant learning rates. Directly tune the LRs via Eq.
We refer the resulting scaling as FSC as it normalizes the Forward pass, the Sensitivities and the Contributions, Let make some observations regarding the scale invariant LRs:
5

---

Link with Polyak step-size. In convex optimization, to minimize a convex and Lipschitz continu- ous function f Rd + R such that min f(z) the stlyak-step-size PPolyak 1987
Hazan
and Kakade, 2019] for the GD algorithm xx+1 = 2t - 1/t (x)e) is given by 1 [(
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 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 quantity which grows linearly rather than quadratically with the norm of the gradient, such OW; m
VzC For such
IV rm.
an algorithm, properties (BC) and (LD), suggests the LR 7i = O | e mmz place Eq. B
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. that any choice of LR that leads to this invariance must be 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 the choice of layer-wise scalings 01,..., given a fixed global scale I oz.
5
Scaling width and depth of MLPs and ResNets
5.1 BFA for single input MLPs and ResNets at initialization
Multilayer Perceptron Consider ReLU MLP architecture with single input go E IRd and forward pass given, for l E [1: - 1]. by
fe = Wege-1,
gt = x(fe.
fo = f WL9-11,
= loss (fL)
(10)
where d()) = max (0, 2) the ReLU nonlinearity and acts entrywise on vectors. The architecture HPs are the input width man = d. the widths of the hidden layers my
M_-11 m (assumed
equal), the output width my, = k. The trainable parameters are Vl E [1: L], W, E pm Such
NNs are of the form ① and thus covered by Thm. 2.1 Let study their properties random initialization under the following assumptions:
(H1) the weights We are independent N(0, ((0,endden((((,,, (H2) either 1: O(1) or the loss is linear.
random variables for E [1:/A].
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) technical assumption to avoid dealing with cases where by strongly depends on the forward pass, where different scalings may arise")
In what follows we write A = ( when there exists C' > O independent of d, m, k. L, ||0/|2 and lobz.l2 such that the probability that A/B B [c, C] goes to I in the specified asymptotic. The new result in the following proposition is the BFA estimate, which relies crucially on a delicate computation due to JJelassi al. 2023
Proposition 5.1 (Large width and depth MLP). Assume (I11-2) and for é E [1:L 11. lei o 22/me-1- As m oo, it holds
llfellems = m([
|bb.l22 1 l
(11)
Moreover, (BC) holds then cos(O.) = (1-1/2))
ResNets Consider
ResNet with branch scale parameter 3 E [0, 11, as in
Li al. 2021
with single input x = fo E RRd, the forward pass is given, for E [2: T, 1]. by
f1 = Wiz, ft = / f + 3W @Weo(fe-1). fL = (:=WifL-n L = (f) (12)
2Say, if loss( f) = 11 we have I1bz_ = [[WI Weft-1llmms O(o max /ko /k/m/ m/ bzll2)
(by Lem 3.3 and properties of the Marcenko- -osttuu law). while under (H2) we have | || = bbo

---

175
100
(a) BFA vs. layer (7. 200)
(b) Output BFA depth
(c) Output BFA vs. scale
Figure I: Backward-Feature Angle (BFA) 0, observed at initialization in MLPs (3 1) and ResNets (width m = 200), for few random realizations. (a) for all architectures, BFA 0,, varies the first few layers and then stabilizes. (b) BFA at output layer 0 is asymptotically independent of depth, with a non-trivial angle only when 3 x 1/vI (same color scheme as (a)). (c) for a branch scale s/vI, the factor directly determines asymptotic output BFA 0 (averaged over 5 draws).
where d(u) in theoretical results. The architecture HPs are the input width mo d, the widths the hidden layers m1 ... = mL-1 m. the output width m.z = k. The trainable parameters are Vl [1: LJ, We mmmexme-s When 8 1, recover MLP.
Here we limit ourselves to the case of linear activation where we can directly apply result from Mar ion and Chizat 2024] 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 (HI-2), les 6(x) x, B O11/VT))
for E [1 L 1), let oe ((1/ me- As m + co it holds:
[[felrmm @([[=|1mms))
llbell2 = ((vorrllbll)))
(13)
Moreover, if (BC) holds then cos(O.) = +(1)
Numerical experiments We consider one GD step in the model with ReLU nonlinearity, without training W1 (input dimension
10, output dimension k master LR 8t 0.001). Fig. represent BFA, computed via 0, x arccos(-6o ff)) where f. is the change of feature fs after one GD step. The results consistent with Prop. 5. and 5.2 Interestingly, the last plot suggests that there exists function & R4 - j0, 2/ such that for branch scale c/ /L, L, the BFA converges to (p() (it can be conjectured numerically that cos((p(c)) c-1/2 1/2 for > 1).
5.2 Characterizing HP scalings for MLPs
Let us now discuss specific choices of HP scalings for single-input MLP architectures in Eq. (TO (or Eq. (12) with 3 = 1) and at initialization. We consider 6 HPs: the scale of initialization 01 and
LR ?/ of the input layer, the scale ani = 02 the hidden layers, and the scale σ, and LR of the output layer. The HP scalings mentioned in the
OL-1 and LRs 7hid == 72
7/-11 of
next theorem are the following (see Table T:
NTK: the standard scaling with LRs adjusted to satisfy (LD) and (BC) Jaacot et al. 2018
MF++PP the scaling proposed [Jelassi et al. 2023 constructed by imposing the so-called "mean-field" output scale o, 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 ||l|| and lbbll22 We consider two typical settings:
(Dense) Where ||~||| = vd and | = vे" This is representative dense whitened
input and a RMS loss loss(fz) llf yllž/k: for some dense signal y R with [|y]] (1) e.g., in image generation applications.
Link to the Julia code reproduce the experiments: https://github... httpps///ggtabbb.mo////ha//////t BAFU

---

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+ and FSC coincide with JPP. Values in red are exact, the others are up to multiplicative factor in (1).
Input
Hidden
Output
NTK MF+pP FSC
init. std.
1/va
V 2/m
vkt.mm
LR me
m/L'a
1/12
k/Lm
init. std. oe
1/va
vR/m
LR m)
m//3/2a
1/23/2
k/19/3m
init. std.
1/va
V2mm
1/m
LR n/
1/La
1/Lm
k/Lm
Table 2: FSC scalings identified in Prop. 5.4 for ResNets. All HPs are specified up multiplicative
factor in (1). When 3 = (/-= and k d (I), these scalings coincide with the so-called
"depth yP" introduced in [oodeloon et al. 2023 and also studied in Yang et al. [2023b].
Input
Hidden
Output
init. std. de
1/va
1/vm
vK/m
LR me
m/Ld
1/82L L
k/Lm
(Sparse) Where ||=]||a and |bb.||2 This is representative of one-hot encoding input and the multiclass logistic loss (aka cross-entropy where ||bı.1|2 = (llog(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 pP to ensure a non-degenerate behavior as 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+pP satisfies (SP), (BC), (FL) but not (LD);
(ii) The scaling NTK satisfies (SP), (BC), (L.D) but not (FL);
(iii) Properties (SP), (BC), (LD), (FL) hold 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
various scalings consistent with theory. Finally, let mention that even though FSC
in Fig. [2] and the results fixes some degeneracies of deep MLPs, other problems arise when considering multiple inputs, such degeneracy of the conjugate kernel and NTK [Hayou et al.), which make ReLU MLPs fundamentally flawed model at large depth. Hence, the analysis of large depth scalings for ML.s 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
Proposition 5.4 (ResNets scalings). Take B = O(1/VL), consider the same 6 degrees of freedom as the previous section and assume that the conclusions of Prop. 5.2holds Then properties (SP), (BC), (LD) and (FL) hold at initialization if and only if the scalings are those given in Table 2

---

SAnGR outour feature
MP-+pp +p0
senstwiy of
10.0
7.5
5.0
Mr1
2.5
0.0
depth 1.
depth
ReLU MLP (B =
(b) ReLU ResNet
Figure 2: Sensitivities SL-1 (see Eq. (4 of the last layer of activations (OL-1 11 in the MLP and fL-1 in the ResNet) computed via the formula 1llf.
where á denotes the change after one
input sample on the unit sphere and k: 1). (a) ReLU MLP of width m = 400. From our theory we have for NTK, S = (1/v (close to 0 and constant with depth); for MF+uP S = ((vL) and for the FSC S = (1) (b) ReLU ResNet of width m 400: for both choices of branch scale, the sensitivities stable around nonzero value.
GD step (master learning rate 81 0.01,
10,
Minimal desiderata and stability under homogeneity
It is not clear priori why the property (BC) studied in Section 4hould be enforced. In this section, we consider homogeneous architectures, such as ReLU MLPs, and show that a slightly more general version of (BC) is related notion of gradient stability.
6.1 Stability and backward speed formula
For general architecture of the form Eq. [T. consider the following stability property (S), which 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. holds | Ve 12/1//// = O(1) for []:
We will study this property in ReLU MLP with single input in Eq. (10) (the extension to linear ResNets is simple as only the BFAs change). In this case we have V,L begé-1 and thus It follows
|| |F =
de
Vel
[|bella [|bel2219e--11] + (eellallee1lls
libel
190 1|2
199
|
1gc-1
We can thus ensure the gradient stability by ensuring, for all é [1:
(FS) Forward stability. It holds
l1de- O(1) for é E [1 T.]. and
li ill2
(BS) Backward stability. It holds
lb O(1) for é [1 L].
Tb
We focus these simpler desiderata (FS) and (BS) instead of (S) for the rest of the discussion.
To estimate by. we rely on "backward" version of the feature speed formula that holds in homogencous NNs.
Proposition 6.1 (Backward speed formula), Consider a general architecture of the form T take V [1: L] and assume that the map fo fL is positively 1-homogeneous If -f Floss + Eiso Mell V C
[ = 0 then b. = 0. Otherwise, the (non-oriented) angle 0, between fo and by is well defined in [0, #/2| and it holds
SI V (osss + Σi>o ne V |I cos(
|b 12
'For Euler's identity hold, also assume that its selection Bolte and Pauwels 2020 is 0 O--homoeenoous.

---

In the context of ReLU MLPs with linear loss, we have by differentiating the kack-propagation recursion and noticing that all terms involving d" are zero almost surely that:
by =
Σ> E>
dge-1 Of
De ge-1b = Σ Te E>
nellbell
( (20r.) 3f.
) O5v
KuJo (14)
where the last expression defines K.. Reasoning in Thm. (3.2). since fo is Gaussian initialization and noticing that K, has structure similar that of K we have that cos see the details Lem. A.1 We can thus estimate by just as well as fe.
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 Wan
results under global rescaling for scale invariant losses (Van Laarhoven 2017 Li al. 2022
et al. 2020
Proposition 6.2 (Invariance under block-wis, rescaling). Consider a function fr fu(w...... wL) (the NN, in our context) which is separately positively 1-homogeneous in each of its blocks of parameters We E RPc Let 0g = (wi (O). ..... Wr. (0)) and let on a @ 0n = (a1 W1 (0) ..... Wr (0)) for some scale factors E 14. Let O(t) and ((t) be the iterates of GD on & 0 + loss(fr with LR satisfying [
= (te(t) & and starting from Gg and Oo respectively. If
I1 og = then ((t) = a @ 0(t) for all 2
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 scale-free version, follows:
(RFL) Relative feature learning. It holds f 1 /tt//Ilfflllle = O(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 01, Thid, OL and three LRs m, Thid, ME- Assume |16y.12 llgo Imss 1 and a linear loss for simplicity. Then the minimal desiderata (RFL), (LD), (FS) and (BS) hold at initialization in the limit m 3o then L
and only if
(Wdo1) (m//oonaa) L-2 oL = O(VL/m))
C1 + Chid =(1 and Cy 0(1),
where C1 = 7/ 1 LIg. Cnid Σt=2 Thial Ve and Cr. n/. V In particular, the
scaling FSC (Table D 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. 1 it can only quantify feature speed for (S)GD (and does not apply to variants in its current form) and at "cut nodes" the NN architecture, where all the signal goes through (in particular, it does not apply inside the blocks of ResNet).
In future works, besides removing these limitations, would be interesting to have better under- standing of the BFA, both from a quantitative and a qualitative viewpoint.
'A downside of this computation is that, because of its local nature, it ignores the contributions of the Jacobian's 's discontinuities to by. while they do have "macroscopic" effect with a non- vanishini step-size. For instance, taking first the infinite width and then the small step-size limit, would give a different expression.
10

---

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 (ood::ICTS/dspoo22237))
References
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via overr--arametteizztionnn In International conference 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 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. mathematical model for automatic differentiation in machine learning. Advances Neural Information Processing Systems, 33: 10809 10819, 2020.
Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hamin, and Cengiz Pehlevan, Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit. arXiv preprint arXiv:2309, 11099/66666 /6620, 2023.
Lénaïcc Chizat and Francis Bach. On the global convergence of gradient descent for wver-parameterieed 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 over- parameterized neural networks. In International Conference 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 neural information processing systems, 31, 2018.
Boris Hanin and Mihai Nica. Products of many large random matrices and gradicnts in deep neural networks. Communccations in Mathematical Physics, 376( 1):2777322 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 machine learning, pages 2672 -6800.
Soufiane Hayou, Eugenio Clerico, Bobby He, George Deligiannidis. Amaud Doucet, and Judith Rousseau. Stable ResNet. In International Conference Artificial Intelligence and Statistics, pages 1324-1332. PMIR, 2021.
Elad Hazan and Sham Kakade, Revisiting the Polyak step size. arXiv preprint arXiv: 2019.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kernel: Convergence and generalization 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.0788 2023.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International ConfaattionaCCooeeeeneeee on Learning Representations (ICLR), San Diega, CA, USA, 2015.
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. Nature, 426 2015.
Mufan Li, Mihai Nica, and Dan Roy. The future is log-Gaussian: ResNets and their seinitt-ddeth--ane aad--indddwdidd width limit at initialization. Advances in Neural Information Processing Systems, (1-7782277644 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énaic Chizat, Deep linear networks for regression are implicitly regularized towards flat minima. arXiv preprint arXiv:2405. 2024.
Pierre Marion, Adeline Fermanian, Gérard Biau, and Jean-Philippe Vert. Scaling ResNets in the large-depth regime. arXiv preprint ar arXiv:2206,06922 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, 1 115633))E7665-7777 2018.
Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli. Resurrecting the sigmoid deep learning through dynamical isometry: theory and practice. Advances in Neural Information Procexsing 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
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 Warchol, Stanislaw Jastrzehski, 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. L.2 regularization versus batch and weight normalization. arXiv preprint 17X6vv1706.053350 2017.
Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cengiz Pehlevan. Feature-Icaming networks 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. 08409, 2020.
Zhichao Wang, Andrew Engel, Anand Sarwate, Ioana Dumitriu, and Tony Chiang. Spectral evolution and invariance linear-width neural networks. arXiv preprint arXiv:2211.06500 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 Bahuschkin, 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, 178477 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 NewrIPS 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 T. Machine Learning, 11 111(9):355-3339 2022.
12

---

A Proofs omitted from the main text
Proof of Lem.
3.3 Writing K = VDVT with D = diag(A!....., Am) and E R xm orthonormal, K a T VDDVTa a. Conditioned on K, the vector a = VTa isotropic Gaussian so for E [1 ml. Hence, on the one hand
we have
Euf
m
m
E| Ka 2 K|
Σ
(A:)
Σ
E
Σ
X
Y71
On the other hand, using the fact that the variance of chi-square random variable is 2,
22 - -1m)) (A))
Varil| Kall| K] = E
1=1
m
2 m²
m
Σ
[uu -\mm)
Σ X E - 1/m)(u]-1////
Proof of Prople 5.7 When 3 Eq. (11) is classical from the signal propagation literature (Poole et al., 6 Hanin and Rolnick] 2018 Hamin 2018] (the fluctuations around the limit have also been studied in Hanin and Nica (.020). ). Note that these results are proved with k +(1), but Hanin (2088) allows to conelude as well when & diverges at least if the initial gradient b, [(,)* E 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 7n = as the contribution of w1 to the BFK asymptotically negligible assuming (BC). We consider the BFA at g. and denote 2, (AC/0g.) The main result of [Jelassi et al. 2023) can be restated as follows: with
= +(1), the choice it holds 11 = O(1). In view (1I. holds in
or = 1 and learning-rates f) ( 1/2,. their setting for L [2: L 1]
VV = |bag 1| l1 1 be 0(m=-1 mi-- of. 0(me-1/mm.-))
Using m/
it follows
Σ ne V ((T.-1/2). By Thm. 2.1 and using
m =1-1111111 rms = (11), we get
111--11mma cOs(OL-1) -L-11 ||
9(1)
This shows the result for the BFA 9L-1 and the result holds as well for the BFA at fi-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 certain random matrix, stated in the following lemma.
Lemma A.1 (Spectrum of BFK and FBK in ReLU MLPs), For ReLU MLP at random initialization satisfying (SP) and (BC), consider the BFK (at Iu instead of fw):
Ogo
Ogn
K
Σ ME E<C
Ows
Ows
and 0, the (non-oriented) angle between åo and 25 = (TL/Og.)T. Then, in the notations of Thm. (3.)). holds as hidden width diverges cos(f.) (M, (M, 112)))).
Consider also, for a linear loss, the kernel K, such that b, Kofo (see Eq. 14 and 0, the (non-oriented) angle between fo and -b.. Then it holds, as hidden width diverges, cos(O.) = (M.((K.) M2(KK)) MM(Kk)) @( 11)-1/2.
13

---

Proof. We have already seen in the proof of Prop 5.1 that cos(6.) (vv-1/2. It thus remains to see that the assumptions of Thm. (3.2) are satisfied: the independence of žy follows from [[Hami and Nica. 2020, Prop. 2] and the Gaussianity of ze is direct since W bet1 where W4+1 is Gaussian and independent from bę+1. Also, we have the more explicit expression
dg. Of.-
dge Öff
K2 Σ
nellac-1l1 112
where 0g. öfc
D.W. De+1 We+1 De and D; (ff)) (by (Hanin and Nica 2020 Prop. 2), these matrices can be taken matrices with Bernoulli random variables the diagonals, independent from everything else). Since under (SP) and (BC) we have that V 1 is constant for / E [1 L-1), it follows
Kr E(D.W. We
De+1 Det
For the second claim, we have (see Section 6) K2 E melbbell (SP) and (BC), we have ne nel1be|8 12 is constant for f E [2: L). hence it follows
Oge 05
T(
(gge af.
Under
K, x Σ t=0+1
W..1 (De- Wo41
By comparing the expressions for K2 and Kg, we see that K, has the same distribution of nonzero cigenvalues as - (potentially up to global multiplicative factor) and the conclusion follows. O
Proof of Prop. 5.2 The estimate for [| 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 (E + v) O and
mmax( (E 0) (1). where Tmiu (E + v) and F (e + v) are the smallest, respectively largest singular value of afc The estimate on by (8fz. by. directly follows. afe"
(af.)
For the BFA, we will apply the first bound of Thm. 3.2 namely cos(O.) 2 Amin(Ky)) Amxx (wx)
where Amia (K,) and Amax Amax(Ke) are the smallest, respectively largest, eigenvalues Kg- In the forward pass (12), let us write ge p(ff) and he Wege-1- By direct computations, it holds (here wy is the vectorization We):
O5 T1 1=1
Of Owee
( The
0f.
05. @5O) T
Of Ofe
05 Ofe
K, Σ
Σ E=2
ne ge gelge-
Σ
The!
6=2
Using the inequalities
Of.
Of. 05,
Amin (K) > 32
Σ nel ge- 11200ii
Amax S 82
Σ
ne 1 max
aepe
8=
we deduce cos( ) 2 Amin (Kx)) Amax (K,) (().
Proof of Prop. 5.3 In this proof, we say that claim is found "by direct computation" when it be
directly deduced from the conclusion of Prop. In particular, for the computation of scale invariant
LRs, use the fact that ||V ||2 lbbgg |I [blll2 11 2|2- Also, by Prop. 5.1 under (SP) and (BC) it holds cos(O,,-1) =
(i) One has that MF+ satisfies (SP), (BC) by direct computation, and (FI.) by Prop. seen in the proof of Prop. 5.1 that -C (0(L-1/)) so (LD) does not hold.
4.
We have
(ii) For NTK, (SP), (LD) and (BC) can be checked by direct computation. For (FL), we have ||by-Illmms = (11/ /m) so:
1 lms -0
@
(cos((,)) m bb
(0z-1) v
1rmm
14

---

But in the considered asymptotics vm cos((L-1) ((/m/L) + ao so Ilft- =(1).
(iii) Properties (SP) specifies on and Ø2 01.1, and Prop. 4.1 gives, with (FL.), 110 mmrms
① cn 1)ml (VLL/m) which imposes 0y = VkL/m. Then the LR are given by
Properties (SP) specifies 01 and hhid = 02 = L- and Prop. 4.1 gives, with (FL), | rms = O cos(01-1)m = ((1/m) which requires aL = v Vk/m. Then the LR are
Proof of Prop. 5.4
characterized by ().
Proof of Prop. 6.1
By the chain rule and Euler's identity for positively 1-homogeneous functions,
holds
0C b. = ofe fo Of 0f. dfL
ac
a. 0fl fo fL.
Now, by differentiating in time both sides we get
0C j SL V &.
by + byg so S (o
0fl
We have 2 = - E m V and moreover, from the proof of Thm. 2.1 Eeso nell V... So we deduce
b. f =
-b = - vosss
EndVve
E>v
We conclude by writing -b, fo
|| 22 foll co.) and rearranging.
By assumption time it holds ((0) 000(0) so let us prove the result by recursion. Assume that the claim is at iteration t. Since π[ oz 1, it holds (e(t)) fz (0(t). 0fc. is O-homogencous in we and separately 1-homogeneous in (wi)ige. It follows
Proof of Prop. 6.2
Moreover, since Dw,
VeL((()) = OS (0(O)) Vloss((fr(0())))
Twe
III a.) @S (0(O)) (loss(((((0()))) iffe
v.C(O(1)- de
Owee
In particular, we deduce that the LRs are related by 7(E)
of. It follows, for any T(E) (V
[V
l E [I:
we(t + 1) (t( -(tVVrc((()))) = (t)
veL(O(t)) (nn
orwe (t + 1).
d2
This proves 0(t 1) = 2 0t + 1) and the claim follows by recursion.
B Initializing with zero output weights
Let us mention an interesting degree of freedom for FSC Table [ 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 Wz with then all gradients are at time 0 except that for Wz which leads to the update (non-infinitesimal in this paragraph):
SW, (0) = (0) by (0)
The second forward pass is the same the first one, with the only difference that
fr.(1) = - (0) (gr (0)
Assuming (0) by (1) (linear loss) for simplicity. this leads to second backward pass:
-1-1(1)= = (āg--1 (1)) T (m-m (0mm000) (0) = (0) |b2 (0)]2ggg))11
15

---

For the second GD step to satisfy (FL), we just need to ensure
VL < (0)
11) O(VL)
(0) θ
This is the LR to be used time O, for the second step satisfy (SP), (FL), (LD) and (BC).
Proof of Thm. 6.3 Prop. 6.2
shows that in fact only the product 01 o o is relevant degree of freedom of scale. We can thus fix (arbitrarily) 01 1/vd and Thid 2/ 3m/mmmm that (SP) satisfied; ((mrr) for E [1: L - 1]. Desideratum (RFI.) requires
we then have ||b.|12l|
11 fL-1112 cos(0L-1) (bL-
C1 C Ch
VL m
O(1)
or = (C1 + Ch.
using that cos(01-1)= O(IvVD). Moreover, (LD) requires C₁ + Ch + C2 (1). At this stage, the output scale or is not yet entirely determined since C1 + Cr o(1) is not excluded. This is where (BS) comes into play. It requires in particular
l|b la bb 2
Chid + CC = O(1)
VL
(Chid + Ct) m
0
cos c0s(@)))lb/laall/ill
using that cos(01) = LLIVVL) by Lem. A.1 Combining both conditions for σL imply, on the one hand, that C₁ + CL OC1 + Chia) hence C1 Chid ((1) and dg = O VI which equivalent to the constraints written the theorem. Conversely, 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 II 1. - R admitting a (selection) derivative and, for a fixed scale vector (R*)) ) consider the function g(v) = f'a y) where a a denotes "er,x,,.,,,711 Consider one step of GD the two functions, given for & E [1 L). by
de = 7/e vz)
Ve = = T/e Vegly)
with identical starting points, that is It = ag Ye for E []:
Proposition C.1. Consider adaptive learning rates, which are of the form me = 7e (V(()) Then x' a of for all a E (RR,)* if and only if me is (-2)-hommgenneou in V (z)) and O-homogeneous
Ve [(((() for l &.
One such LR precisely that suggested by Prop.
4.1 Ie x [+](())12²
Proof. For l [1 LJ. it holds
aryi Oefl CreT ( (Voly))
aznela (1)
Then a y = x all a E (R L is equivalent to
Va e E n vf(x)) (VV(())),
which is the claimed homogeneity property.
16

---

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 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 other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations 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 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 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 (c.g.. independence assumptions, noiseless settings, model cell-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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails 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 faimesss.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 complete (and correct) proof?
17

---

Answer:
[Yes]
Justification:
Some of the proofs are deferred to the appendix, but all 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 cross- referenced.
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 provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix 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 ex- perimental results of the paper to the extent that 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 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 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 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 large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily 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 new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome 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., 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
18

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
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.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.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.c.// cttpsscccpbbbcc//cceeceeeccccccc httpe//giideee/eeeeebbsssesssss /ac////odebiisssseeeeeee wwwssiiiccciiyiicccooooocociiiii 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 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 the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
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 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, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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, confi-
dence intervals, statistical significance tests, 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).
19

---

The method for calculating the error bars should be explained (closed form formula, call 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 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 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
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 into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https::/neurips..cm http:///wrrpp...cc//////ccccccca
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal 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 why the paper does not address societal impact.
20

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that 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 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 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 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe 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)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 best faith effort.
12. Licenses for existing assets
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?
Question:
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 dataset.
The authors should state which version of the asset used and, possible, include 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 use in package should be provided. For popular datasets, paperswithcone com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
21

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 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 the dataset/code/moodl as part 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. Crowdsoureing 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 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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included 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 equivalent approval/review based 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 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.
22


Reviews:
[{"review_id": "gSuvCHRD0D", "reviewer": "Reviewer_jmiK", "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.", "questions": "How does the \"ours\" hyperparameter scaling compare to the others (usual, muP or NTK)?", "limitations": "Lack of experimental validation.", "rating": 3, "confidence": 4, "soundness": 1, "presentation": 3, "contribution": 2, "strengths": "# 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.", "weaknesses": "# 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We have replied to this review in the main comment.\n\nCan you please specify what are the references [1,2,3,4] in your review? To the best of our knowledge our approach to derive HP scalings is new (the closest work being Jelassi et al), but we would appreciate precise pointers to related results. \n\nWe are aware that the NTK and related objects have been heavily used in NN theory. This is expected since it is a fundamental characteristic of the training dynamics. However, we do not think any work has used these objects the way we do.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zQ9Hz9oNCQ", "reviewer": "Reviewer_gz2f", "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.", "questions": "Please see the weaknesses section", "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.\n\nI realize this feedback might be construed as being fairly negative, but I hope that it can help to improve the work.", "rating": 3, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "- 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", "weaknesses": "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.\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We have replied to your first 2 criticisms in the main rebuttal. Concerning the limitation to batch-size 1: this is an assumption made by all related works on feature learning. However, note that the \"feature speed formula\" applies to any architecture -- and in particular to any batch size. With our approach, it is in fact possible to obtain the scalings in the *joint limit* of width and batch size -- which is more powerful than considering a large width limit with fixed batch size (and this is not accessible via other existing approaches, which start with the infinite width limit). \n\nFor instance, a direct application of the feature speed formula in a two-layer MLP (one hidden layer) shows that to ensure feature learning, the scale of the output layer $\\sigma_{L}$ should be scaled as \n$$\n\\sigma_L \\asymp \\frac{1}{\\cos(\\theta_1)\\cdot \\text{width}} \\asymp \\frac{\\sqrt{M_2}}{M_1}\\frac{1}{\\text{width}}.\n$$\nwhere $M_p$ are the spectral moments of the input covariance. (To derive the expression of $\\cos(\\theta)$, we use that in this case, the BKF is a linear operator that acts on a matrix $B\\in \\mathbb{R}^{\\text{width}\\times \\text{batch}}$ as $K(B) = BX^\\top X$, writing the linear case for clarity). This differs from the usual $\\frac{1}{\\text{width}}$ mean-field scaling (this new insight is not discussed in the manuscript to avoid overcrowding, but we mention it here to highlight the flexibility of our approach).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "L8DSAtaRCj", "reviewer": "Reviewer_PDm8", "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.", "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? \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.", "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.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The 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.", "weaknesses": "There 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for his/her detailed and encouraging review.\n- yes the scalings for ResNets are precisely those predicted by Bordelon et al, and Yang et al, this is mentioned in the manuscript but we'll make this more visible (note that the first version of our work appeared in November 2023, the same month as Yang et al).\n- yes, we mean \"SGD\". By \"cut nodes\" we mean \"cut nodes in the DAG computational graph where each node represents a tensor\". This means that a ``cut node'' is any intermediate computations in the forward pass that is not skipped by any other computation, such as a skip connection in a Resnet. For instance, in a ResNet, the formula applies just after each residual connection, but not just before. We will clarify that (with a mathematical definition). We also think that it would be interesting and important to extend this approach to any node in any computational graph!\n- we agree with the comment that we should be a bit more concrete in the beginning of the paper, when introducing the forward pass and what this can represent. We will add examples of how to instantiate the formulas.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CYk0INQOVj", "reviewer": "Reviewer_PezB", "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.", "questions": "1. 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?", "limitations": "The authors adequately addressed the limitations of their results.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. 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", "weaknesses": "Despite 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for his/her detailed and encouraging review! Here is our answer to the main weaknesses raised by the reviewer:\n- in this paper, we focus on introducing a theoretical methodology and do not introduce new (useful) HP scalings indeed. In current work we are tackling other asymptotics via this approach (see answer to R#2 for an example).\n- the assumption of a single input is also made by all related works, but we do not believe that it is too problematic. The scalings are the same for a finite batch size as long as the NTK does not degenerate (eg on a ReLU ResNet -- a setting that we are developing in a work in progress). However, the situation would change and become more interesting for joint limits such as batch-size and width jointly going to $\\infty$ (see also answer to R#2).\n- there was indeed a problem in the code of Fig.2b which will be fixed in the revision.\n- since our focus is on exposing a theoretical approach, we would like to keep the proofs in the main paper as much as possible. We will use the additional page to introduce illustrations and expand the discussions.\n\nAnswers to questions:\n- the equality is indeed a typo, this should be a $\\Theta(\\cdot)$. Thank you!\n- we have removed the discussion on \"hierarchical features\" which we believe is not central to our goals and was too vague. By this term, we meant that each layer learns a representation of the input that relies on the learnt representation of the previous layer, and so on.\n- Yes the BFK in Eq 5 is for multiple inputs (in fact for essentially any architecture and tensor dimensions). In a vanilla MLP and Resnet, $m_v$ would indeed be the product of batch-size with the input dimension to the layer. It is this ability to cover any architecture and tensor shapes that we think make our approach a good starting point to derive HP scalings.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their time and their comments and we appreciate their encouraging remarks. We also disagree with a few comments which, we believe, result from a misunderstanding of the content of our paper, the state of the theory on feature learning, and perhaps from a disagreement on the role of theory, in general.
- *lack of experimental validation*: this paper is theoretical and the usefulness of the hyper-parameter (HP) scalings that we discuss has been already demonstrated in prior works (which themselves followed from previous purely theoretical works on infinite width neural nets). The only new scaling that is introduced is the one for deep ReLU MLP but this just for the purpose to illustrate how to apply our theory in various contexts : as mentioned in the paper, deep ReLU MLPs cannot be trained (the NTK degenerates, a problem comes up when considering more than 1 sample) so there is no relevant "real data" experiment to do in this case.
- *lack of significance*: We sincerely believe that our approach is a significant advance in the theory of feature learning and HP scalings. The existing approaches to HP scalings are the following:

  (1) *write infinite-width limits dynamics first*, and see whether "features move" (eg Chizat et al, Yang et al, Bordelon et al, among many): this is the original approach to design HP scalings, but this leads to proofs which are technical and that help intuition in a limited way. For instance, in their breakthrough paper, Yang & Hu obtain a scaling for feature learning but only for specific architectures (namely finite depth tanh & gelu MLP) at the end of quite technical computations and use of a heavy random matrix machinery (see their appendix H.7.2/3, let us mention that this paper was obviously very influential to us). Note that the recent extensions of this approach to large depth limit are written only at a formal level, due to their high technicality. Besides, this approach is intrinsically limited to "infinite width first" limits and cannot deal with joint limits (e.g. with depth, batch-size or context length). All this facts make it crucial to dispose of alternative theoretical approaches to HP scalings.

  (2) *heuristics*: because of the difficulty of the above approach, other researchers typically rely on heuristics involving alignment/CLT/LLN etc. These heuristics can be useful as long as one lacks of rigorous and simple derivations (which we provide). For instance the "spectral condition" in (Yang, Simon, Bernstein) is a heuristic that recovers the good scalings in terms of "hidden width", but it fails to give the correct scalings for other asymptotics (such as depth or large batch-size). 

In contrast, our approach is simple, rigorous and general. Since the "feature speed formula" is a non-asymptotic equality, it can be used as a starting point to obtain the good scalings in any asymptotic, provided each term is worked out : we have done this for width and depth in the current submission, and are currently extending it to other asymptotics (see also answer to R#2).

Meta-review: 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.
Decision: Accept (poster)