Super Consistency of Neural Network Landscapes and Learning Rate Transfer

Recently, there has been growing evidence that if the width and depth of a neural network are scaled toward the so-called rich feature learning limit (\mup and its depth extension), then some hyperparameters -- such as the learning rate -- exhibit transfer from small to very large models. From an optimization perspective, this phenomenon is puzzling, as it implies that the loss landscape is consistently similar across very different model sizes. In this work, we study the landscape through the lens of the loss Hessian, with a focus on its largest eigenvalue (i.e. the sharpness), and find that certain spectral properties under $μ$P are largely independent of the size of the network, and remain consistent as training progresses. We name this property Super Consistency of the landscape. On the other hand, we show that in the Neural Tangent Kernel (NTK) and other scaling regimes, the sharpness exhibits very different dynamics at different scales. But what causes these differences in the sharpness dynamics? Through a connection between the Hessian's and the NTK's spectrum, we argue that the cause lies in the presence (for $μ$P) or progressive absence (for the NTK scaling) of feature learning. We corroborate our claims with a substantial suite of experiments, covering a wide range of datasets and architectures: from ResNets and Vision Transformers trained on benchmark vision datasets to Transformers-based language models trained on WikiText.

Paper

References (97)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer ssWb8/10 · confidence 2/52024-07-10

Summary

The paper investigates the loss landscape of the model with scaling width or depth, through observing the largest eigenvalue of the Hessian marix and the NTK matrix. Authors show empirically that the loss Hessian evolve almost identically for different model sizes (which is named Super consistency), however, ntk accumulates finite-size effects over time. Authors also validate their empirical findings using theory in a two-layer NN with linear activations.

Strengths

- The paper discusses an interesting phenomenon that the loss landscape gradually becomes stable, this might explain the transfer of learning rate - The paper shows NTK (lazy learning) has different behavior when comparing with Hessian of the loss, this suggests NTK is insufficient to explain the behavior - The paper gives an explicit evolution law under 2-layer linear NN setting,

Weaknesses

- The authors do not conduct larger-scale environment since needing to compute Hessian

Questions

- I am confused about Line 209, this line mentions we can decompose H into G + R, while G = K^T K and the dynamics are mainly driven by G. NTK and G share the same nonzero eigenvalues, doesn't this mean H and NTK should have similar behavior? Or feature learning happens in H which makes the difference between NTK and Hessian?

Rating

8

Confidence

2

Soundness

4

Presentation

4

Contribution

4

Limitations

The authors have fully addressed the limitations in their paper

Reviewer 47kG7/10 · confidence 3/52024-07-10

Summary

The authors argue that the top eigenvalues of the loss Hessian stabilize throughout training under width and depth muP scaling. This phenomenon is called the Super Consistency of the loss landscape. The authors provide theoretical convergence guarantees and empirical experiments supporting their claims. The learning rate transfer under muP scaling also correlates with the super consistent landscape, i.e. optimization follows the trajectories of sharpness (top eigenvalue of the loss Hessian). Under NTK or other suboptimal scaling, the super consistency is also violated and thus learning rate transfer failed. The authors further show that the dynamics of the weights are fully specified by a set of governing equations and thus one can derive a theoretical edge of stability result under scaling. The phenomenon of progressive sharpness towards stability also happens along with the NTK evolution with finite-size effects, suggesting other factors contributing to the super consistency of the sharpness. **[raising score from 6 to 7 after rebuttal]**

Strengths

1. The observation about the super consistency of the loss landscape and its relation to learning rate transfer under muP scaling is quite an important contribution to the community. I’m not an expert in this field, but this result seems novel to me. 2. The paper is well-written and easy to follow.

Weaknesses

The empirical evaluations are a bit limited. - First, a lot of the claims are made under the setting of ConvNet on CIFAR-10. It’s understandable from a compute perspective, but it’s not clear if super consistency scales to even larger models. - Second, the GPT-2 results on WikiText seem to break the super consistency, though it could be possible that it’s due to the Transformer block itself as the authors explained. So it’s not clear if the super consistency results hold across different models and modalities.

Questions

1. How are you getting the eigenspectrum of the loss Hessian? It seems possible to get at least a few top eigenvalues using iterative Lanczos methods for models larger than 124 million parameters GPT-2 with A100-80GB. 2. I don’t find enough evidence that the super consistent sharpness causally helps the learning rate transfer under muP scaling. In line 178, “optimization happens along trajectories of super consistent sharpness λ_max”. Is there evidence that the optimization does happen first along this direction, i.e. some kind of spectral bias during learning?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

See weakness and questions

Reviewer 11Ly6/10 · confidence 4/52024-07-12

Summary

This paper proposes the concept of super consistency, which describes the stable properties of the loss landscape during training. By analyzing the maximum eigenvalue of the Hessian matrix, it is found that the sharpness under the μP and Depth-μP frameworks remains super-consistent and stable near the threshold. In contrast, NTK and other frameworks show significant differences.

Strengths

This paper proposes a new concept of "super consistency", which provides a new perspective for understanding the behavior of models of different scales. Through a large number of experiments, the learning rate migration phenomenon under the μP and Depth-μP frameworks is verified, and the consistency of these frameworks on different tasks and datasets is demonstrated, including ResNets, Vision Transformers, and GPT-2.

Weaknesses

This paper provides some theoretical analysis, mainly focusing on two-layer linear networks and failing to fully verify the theoretical applicability in nonlinear networks and complex structures.

Questions

Is it feasible to extend the results of this article to other algorithms, such as AdamW?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The theoretical setting i simple. The author can consider trying to expand the two-layer linear network to two-layer ReLU or deep linear network.

Reviewer f2zq7/10 · confidence 5/52024-07-12

Summary

The authors conduct a series of experiments in which they investigate for which attributes of the neural network the gap between its infinite-width (or infinite depth) value and the finite counterpart grows or shrinks during training. Among other attributes, they look at the training loss, largest loss Hessian eigenvalue, and largest NTK eigenvalue. They investigate how the results for different parameterisations (such as muP, neural tangent parameterisation, and some other unstable parameterisations with no well-defined infinite width behaviour).

Strengths

The asymptotic scale properties of neural network training are an important area of research. This paper investigates an interesting question: how do the properties of the loss landscape, such as the curvature along the training trajectory, evolve as one scales up the neural network size. Furthermore, the experimental results are interesting, and well-presented in the figures. I liked Figure 3 in particular, which most convincingly illustrates the claim that the gap between finite-width and infinite-width attribute value shrinks/grows throughout training. In fact, I wish for every hypothesis investigated, a plot like that in Figure 3 was shown. Lastly, it appears the authors did proper ablations, investigating multiple datasets and models, to verify their empirical conclusions.

Weaknesses

The paper falls short in presentation and the formalism. One of the largest issues is the definition of the term “super consistency”. Definition 3.1 has several issues: - The authors defines $S_N(t)$ as a function of the predictor $f_t(x)$. This is pretty vague, and not sure matches what authors are trying to do. If the predictor $f_t$ is interpreted to be a function $f_t:\mathcal{X}\to\mathcal{Y}$ from some input space to some outputs space implied by the neural network architecture and weights at time $t$, then $S_N(t)$ cannot capture something like the curvature of the loss with respect to the weights. I'm convinced this hurdle can be overcome by carefully defining all objects, and what they are (e.g. a ‘predictor’). - If $S_N(t)$ depends on the weights (such as when considering the spectral norm of the loss Hessian), then it's a random variable. Hence, all the notions of distance and convergence need to be defined for random variables for the expressions in Definition 3.1 to make sense. - It's not clear the limit $S_\infty(t)=\lim_{N\to\infty}S_N(t)$ exists for many properties being considered. In fact, this already precludes the authors from talking about super-consistency of parameterisations that do not have well-defined infinite-width limit training dynamics (e.g. SP). - line 127 “if at any finite $N\geq N_0$:” – did the authors mean to say ‘there exists some $N_0$ such that for all $N\geq N_0$’? Otherwise this sentence doesn't make sense. - Does the symbol $\sim$ here represent asymptotic equivalence as $t\to\infty$? This is something that should have been defined. - “$\sim$ denotes the finite-time behaviour during training” is not a formal definition. I have no idea what it's meant to say. - Also, the whole condition of “$|S_N(t)-S_\infty(t)|\sim g(t)$ where $g(t)$ is a non-increasing function of time...” seems like it could have been equivalently stated much more simply as $|S_N(t)-S_\infty(t)|=\mathcal{O}(1)$ as $t\to\infty$. - This definition seems different from how super consistency was described in the abstract introduction: “certain [...] properties [... are largely independent] of the width and depth[...]. - We name this property super consistency of the landscape”. The concept of super consistency, as defined in Definition 3.1, is simple enough to express in two sentences, that I don't see a reason it can't be explained in the introduction and/or abstract properly. Later in the paper, the authors proceed to use ‘super consistency’ in a sense completely disjoint from that in Definition 3.1. On lines 168-159 they say: “The optimal learning rate is preserved across widths/depths, indicating very fast convergence with respect to the scaling quantity (i.e. it's super consistency).” Definition 3.1. has seemingly nothing to do with the speed of convergence in the scaling quantity (width/depth). In fact, I think the paper would have been stronger had the authors cut back on formalising things like “super consistency”, and just presented the empirical results for what they are. The term “super consistency” doesn't strictly seem necessary to convey the take-aways of the paper, and, at the moment, makes the paper more convoluted. Of course, a thorough formalism might be preferred, but it its current form it detracts from the paper. Others: - I find the phrase “optimization happens along trajectories of super consistent top Hessian eigenvalues” quite confusing. It took me a while to realise the authors are saying that the ‘Hessian eigenvalues are super consistent along the optimisation trajectory’, which is the way I'd recommend they phrase it throughout. - There are several typos throughout the paper. Sometimes, seemingly parts of mathematical expressions are missing (e.g. line 161).

Questions

- Why are there two sets of lines for each width on Figure 1, one solid one dashed? - Is the Appendix G, which is meant to discuss “the effect of the lower order eigenvalues”, empty? - In Figure 2, the authors look at 4-th and 10-th largest eigenvalue. How come the authors decided to look at a fixed N-th largest eigenvalue? Given the scaling in width, wouldn't it be more interesting too look at say 4th and 10th percentile largest eigenvalue, given that the number of eigenvalues grows with depth/width?

Rating

7

Confidence

5

Soundness

2

Presentation

1

Contribution

3

Limitations

- Definition 3.1 concerns asymptotic properties in *both* width/depth and training time, which are then doubly difficult to establish with certainty with finite width, depth, and training time experiments.

Authorsrebuttal2024-08-07

Rebuttal by Authors

Finally, we thank the reviewer for their valuable and precise feedback, and we would like to gently push back on the score of 1 given to the paper presentation. The general consensus amongst the other reviewers is that the main body of work is well-written and well-presented and we would kindly ask the reviewer to reconsider their score.

Reviewer f2zq2024-08-09

Response 1

Thank you for engaging on the points regarding presentation, and for trying to work with the feedback in the review to improve it. I'll try and go through the author's rebuttals, and point out where I think issues still remain. > Super Consistency refers to when certain aspects of the loss landscape and of the predictor $S_N(t)$ (in this paper $S_N(t)$ refers to the NTK's and loss Hessian's eigenvalues or the loss itself) [...] I think that is an improvement. I'm familiar with the (Vyas et al. 2023) work. As far as I know, they don't have a formal or semi-formal definition of consistency, but my understanding is it just colloquially means that some quantity is close to the infinite-width counterpart, whenever that limit is well-defined (e.g. muP, NTP). The key interesting part of their paper is not that some quantities become consistent (that directly follows from what it means for a limit to exist), but that they do so at realistic widths and depths. I think the definition you ascribe to the term consistency: >  At large N, $S_N(t)$ does not deviate significantly from its limit. This is what is referred to as consistency in Vyas et al, 2023. is a non-sequitur. By definition, for **any** $S_N(t)$ that has a limit it is true that $S_N(t)$ does not deviate significantly from the limit at large $N$. I guess a way to phrase this that would make more sense would be to just say: > $S_N(t)$ has a limit. This is what is referred to as consistency in Vyas et al, 2023. But at that point you might as well just say that $S_N(t)$ has a limit, rather than that it is consistent. That being said, I think informally referring to some quantities as being “consistent” (meaning they are either “close” to one another, or are “close” to some limit), like is done (Vyas et al. 2023) is perfectly clear, and doesn't need a formal or informal definition. The reminder of the changes also sound great. I would maybe slightly reword certain parts: > When $g(t)$ is an increasing function (up to fluctuations), Super Consistency is violated. nit: I would change this to “when $g(t)$ increases over time (up to fluctuations)...” just because “increasing function” is a mathematically commonly used term that implies monotonicity. --- ### Clarifying the definition in the abstract: Even given the clarification, I still don't think the line in the abstract is particularly clear: > [...] **find that certain spectral properties under μP are largely independent of the width and depth of the network along the training trajectory. We name this property super consistency of the landscape.** I think a reader, after reading only the abstract, would have no idea that what you have in mind is what you later describe as super-consistency in Section 3. “properties being largely independent of size along the training trajectory” could just mean that they are consistent, or close to the limit at reasonable sizes. It's absolutely not clear that super-consistency encompasses whether the gap grows or shrinks as the training progresses. Here is a suggested alternative: > we find that certain spectral properties under μP are largely independent of the width and depth of the network along the training trajectory, **and they become more consistent as the training progresses**. We name this property super consistency of the landscape.

Reviewer f2zq2024-08-09

Looking at the N-th largest eigenvalue or N-th percentile eigenvalue

### Re: looking at the N-th largest eigenvalue or N-th percentile eigenvalue I realised after the rebuttal that looking at eigenvalues of the Hessian, although interesting, is kind of difficult to interpret, since the Hessian is composed of gradients for many parameters that grow at different rates with width (e.g. input layer vs middle layer weights). This is in contrast to hidden-layer weights in the infinite width limit. By RMT results, the per-layer eigenvalue (or singular value) spectrum of the weights should converge to a fixed spectrum in muP as width goes to infinity (and I think this holds throughout training). This means that top eigenvalue, but also something like the 10th percentile eigenvalue (or any other percentile) all have a well-defined limit. I think it would have been very interesting to see such results for changes in the weight eigenvalues and gradient eigenvalues on a per-layer basis, for different percentiles (i.e. comparing convergence for different parts of the limiting eigenvalue spectrum). This is a very late request, so I cannot ask the authors to implement it, and will not count this point in the final scoring. That being said, I'd personally be very much looking forward to seeing such experiments, and am disappointed I didn't notice the distinction sooner. I hope the authors consider running them for a camera-ready.

Reviewer f2zq2024-08-09

Response 3

I am looking forward to the authors' response. I am currently leaning towards increasing my score, but look forward to discussing the above points with the authors.

Authorsrebuttal2024-08-10

looking at the N-th largest eigenvalue or N-th percentile eigenvalue

This is a very interesting point and we thank the reviewer for bringing it up. Indeed it would be interesting to visualize the evolution of the whole spectrum (or some percentiles) at increasing $N$. > gradients for many parameters that grow at different rates with width In $\mu$P for gradient descent, the learning rate is set to $\eta_0 \gamma^2$, where $\gamma$ is $\mathcal{O}(\sqrt{N})$. This parametrization of the network is such that the magnitude of feature updates is $\mathcal{O}(1)$ for all layers. What does the reviewer mean that different weights move at different rates in $N$? Perhaps a simple example would help us understand. However, after the initial rebuttal phase, we thought about the convergence of the Hessian eigenspectrum in $N$ and its evolution during training. We set up the code to compute it and we are now in the process of estimating the spectral density and percentiles. We are partially confident that we will get these results before the end of the rebuttal and will keep the Reviewer updated on this matter. However, anonymous links are unfortunately not allowed and to our knowledge, we do not have a way to share the plots. We are also in the process of computing the per-layer Hessian. Unfortunately, we probably cannot deliver it by the end of the discussion phase, but will surely include it in the camera-ready version. > “gradient eigenvalues on a per-layer basis”. A clarification: what does the Reviewer mean by eigenvalues of a gradient here (which is a vector)? Finally, we would like to stress that the focus of our work is on the top $k$ eigenvalues (where $k$ is a fixed quantity equal to the number of classes) because in SGD the gradient lies in the top $k$ Hessian subspace, and on the sharpness because of its relevance in step size selection (and thus on learning rate transfer). We have also performed an experiment where we computed the curvature along the gradient direction (Figure 3 of the one-page pdf).

Authorsrebuttal2024-08-10

Answer to Response 1

We thank the reviewer for the additional valuable feedback. Indeed some of the extra points (e.g. convergence of the Hessian spectrum) have been in our minds after the first rebuttal and we have been thinking about that experiment. 1. **Clarifying the definition in the abstract**: By *certain spectral properties under μP are largely independent of the width and depth.* we mean that they are consistent (i.e. the finite width object is “close” to the infinite width one). And by *along the training trajectory* we mean that it remains independent of the size along the training trajectory (i.e. **super** consistent). However, we agree that the fact that we often observe the sharpness curves getting even closer to their large width limit is not stressed by this phrasing. Thus, we agree that “[...] become more consistent as the training progresses” better captures this intuitive meaning and have updated the abstract accordingly. 2. **On the meaning of “Consistency”** We partially disagree on the interpretation of “Consistency” in the work of Vyas et al (2023). The Reviewer (we apologize if we misunderstood) interprets consistency in a similar way as having a well-defined limit. In our interpretation, Consistency refers to the fact that at realistic widths (the word realistic is crucial here) the object of interest is (informally) practically converged to its limiting object. This is a very important result, as it implies that the infinite width model (proxied by a large width model in their work) is a good model for finite-width neural networks at realistic scales. In this sense, the existence of the limit alone does not imply consistency of the dynamics. In fact, under the NTK parametrization, consistency of the dynamics is not observed despite having a well-defined limit. Quoting from Vyas et al, 2023: “*We stress that this observed **consistency** is a property of networks in mean field/μP parameterization but is not present in other parameterizations which also give an infinite width limit like NTK parameterization*". We are slightly modifying our phrasing from "*At large $N$, $S_N(t)$ does not deviate significantly from its limit*" to "*At **realistically** large $N$, $S_N(t)$ does not deviate significantly from its limit*" to match the wording and meaning of consistency in Vyas et al (2023). Finally, we modified the phrasing of $g(t)$ increasing with time to the Reviewer’s suggestion of “when $g(t)$ increases over time (up to fluctuations)...”

Reviewer f2zq2024-08-10

Response

Thank you for the response, and being open to taking the feedback on-board. > We partially disagree on the interpretation of “Consistency” in the work of Vyas et al (2023). The Reviewer [...] interprets consistency in a similar way as having a well-defined limit. In our interpretation, Consistency refers to the fact that at realistic widths (the word realistic is crucial here) the object of interest is (informally) practically converged to its limiting object. Yes, I don't necessarily disagree, I think to a certain extent both meanings of consistency are used in Vyas et al. (2023). I think the “at realistic scales” is not always implied by the term “consistency”, otherwise the “At Realistic Scales” in “Feature-Learning Networks Are Consistent Across Widths At Realistic Scales” would be superfluous in the title. I think sometimes in the text they just drop the “at realistic scales” and then ‘constistency’ implicitly captures that aspect as well. That being said, this current remark is pretty pedantic from my side, so I consider this settled. I agree with the authors that *“the word realistic is crucial”*, and its addition to their description of consistency makes it a non-trivial statement and a useful definition. Together with some extra clarification in the abstract (and hopefully similarly in the introduction wherever applicable) the authors mentioned above, I am happy to consider these points resolved.

Reviewer f2zq2024-08-10

Response 2

> In $\mu$P for gradient descent, the learning rate is set to $\eta_0 \gamma^2$, where $\gamma=\mathcal{O}(N^{0.5})$. This parametrization of the network is such that the magnitude of feature updates is order 1 for all layers. Depends on the implementation; when not using any multipliers (Table 3 implementation variant in [1]) the learning rates scale as $\Theta(N)$, $\Theta(1)$, or $\Theta(1/N)$ for input, hidden and output layers respectively. > What does the reviewer mean that different weights move at different rates in ? Perhaps a simple example would help us understand. Generally, since the changes input to any layer (activations/features) in $\mu$P will be order $\Theta(1)$, and there are $\Theta(N)$, that means that the change to the weights $\Delta W$ (resulting from a gradient step with a muP learning rate) must be roughly $\Theta(1/\sqrt{N})$ to ensure the output pre-activations are also $\Theta(1)$. However, this is not the case for the input layer; there are only $\Theta(1)$ inputs (rather than $\Theta(N)$) for the input layer. Hence, the size of the change $\Delta W$ to the input weights should also be order $\Theta(1)$ to ensure the entries of the output of that layer are $\Theta(1)$. That means there should be a discrepancy in the size of the changes $\Delta W$ to the weight matrices in the hidden and last layer ($\Theta(1/\sqrt{N})$) and input matrices $\Theta(1)$. Similar reasoning applies to biases, where the change after a gradient update $\Delta b$ will generally be $\Theta(1)$ as well. Of course, the above applies to changes to weights and biases (i.e. gradients multiplied by per-layer learning rates), and not to the gradients (which haven't been multiplied by a learning rate yet). In general, these will also have different scales for different layers (input, hidden, last). Now, this becomes doubly weird when considering something like the gradient vector or Hessian matrix for the entire network. There, you will have elements that grow at different rates with width (depending on whether they are weights or biases for an input, hidden or last layer), and also the *proportion* the elements will be different for different layers. For example, the number of the elements in the gradient vector corresponding to weights of a hidden layer grows as $N^2$, but for the input and last layers the number of elements grows as $N^1$. Hence, for large enough widths almost all elements of a gradient vector or Hessian matrix will correspond to hidden layer weights. --- > A clarification: what does the Reviewer mean by eigenvalues of a gradient here (which is a vector)? I mean the gradient $\frac{\partial L}{\partial W^{(\ell)}}$ for a particular weight matrix $W^{(\ell)}$, reshaped to have the same shape as that weight matrix. If it's not square, I mean the singular values. To be more specific, I did actually have the weight update $\Delta W^{(\ell)}$ in mind (i.e. the gradient multiplied by the learning rate in the case of SGD), since I think in that case, the $\mathcal{\Theta}(1/\sqrt{N})$ scaling for hidden layer weights means that the spectrum should converge to something well-defined, and for others it might not (although I'm not 100\% sure about this claim, I haven't worked through the maths in detail). ### muP check As an aside, did the authors do a muP coordinate check, as recommended in Appendix D.1 and Figure 5 of [1]? Speaking from personal experience, it's quite easy to mess-up the implementation of muP as there are many small caveats to the general rules (e.g. how to initialise LayerNorm parameters), and sometimes one can get learning rate transfer even when there are minor mistakes in the implementation. The coordinate checks usually reveal those. I think adding these to an appendix would make the authors' empirical claims significantly more resistant to any potential readers' doubts. [1] Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer

Authorsrebuttal2024-08-10

Response

**On the definition of Super Consistency** We are glad that the issue of the definition (which was the main issue raised in the initial review) is resolved, and we sincerely thank the Reviewer for the feedback, which has ultimately helped our paper. We will make it in clear what we mean by Super Consistency whenever it applies. **Equivalent parameterizations of $\mu$P** > Depends on the implementation; when not using any multipliers (Table 3 implementation variant in [1]) the learning rates scale as … We haven’t included the experiments, but as a sanity check, we have tried to repeat some experiments with the code of Yang et al [1] (https://github.com/microsoft/mup), which adopts a different parametrization (which is equivalent in terms of the predictor’s dynamics). In this case, instead of computing the largest eigenvalue of $\gamma^2 N$, we have computed the largest eigenvalue of $DH$, where $D$ is a diagonal matrix containing the width-dependence of the learning rate for each parameter. In both $\mu$P parametrizations, we get super consistency of the sharpness, converging at the EoS value of $2/\eta_0$. Thus, the largest eigenvalue stays the same regardless of the (rescaled) Hessian. This is as expected, as the sharpness should be invariant to the equivalent $\mu$P reparametrizations. This also makes us confident about our implementation. Furthermore, in the experiments of the attached one-page pdf, we also use the code in https://github.com/microsoft/mup, and compute the largest eigenvalue of $DP^{-1}H$, where $P$ is the diagonal preconditioner computed by Adam (calculated as in the AEoS paper of Cohen et al (2022)). We have added a section in the appendix describing the Hessian computation for Adam + width scaling. > change to the weights $\Delta W$ is different for different layers and > Now, this becomes doubly weird when considering something like the gradient vector or Hessian matrix for the entire network. There, you will have elements that grow at different rates with width … We think the above discussion answers this question, i.e., we observe Super Consistency of the Hessian’s largest eigenvalue when the Hessian is rescaled appropriately to account for the width dependence of the learning rates and preconditioner (as a subtlety, notice that with the same meaning of $P$ as above, $P^{-1/2}HP^{-1/2}$ has the same eigenvalues as $P^{-1}H$ (Cohen et al (2022))). We will make more clear in the camera-ready version. > spectrum of $\Delta W^{\ell}$ It would be of course of crucial interest to derive the eigenvalue (or singular value) distribution of something like $\Delta W^{\ell}$ under $\mu$P. However, we hope that the Reviewer agrees that this is related but slightly beyond the scope of the paper. The evolution of the spectrum of the weights through training has been empirically analyzed in previous work aiming for “effective theories” of neural network training (e.g. Martin and Mahoney https://www.jmlr.org/papers/volume22/20-410/20-410.pdf). There, it is observed that the spectrum of correlation matrices of the weight can exhibit spikes, which can be modeled by spiked covariance models in RMT. **Implementation of $\mu$P** Our results are fully compatible with the existing open-source implementation of $\mu$P (https://github.com/microsoft/mup). We explicitly test this in a subset of experiments. Furthermore, the new rebuttal-time experiments of Adam + width scaling use this codebase. Also, we test networks up to very large widths and depths (up to 300 million parameters), which makes us confident that our implementation is correct. For the layer norm parameters, we follow the prescription on page 24 of Yang and Hu (https://arxiv.org/pdf/2203.03466). We will make the code available if/when the paper is accepted. We hope that this resolves these issues, and we are open to further discussion.

Authorsrebuttal2024-08-11

Response

As the end of the rebuttal period is approaching, we would like to once again thank the reviewer for their valuable feedback and time, and ask again if there are any other points that the reviewer would like to discuss. Otherwise, if we have settled all the open points (especially surrounding the main initial issues raised by the reviewer regarding the definition, which have been resolved and clarified with the reviewer’s valuable feedback) we kindly ask the reviewer to revise their original score.

Reviewer f2zq2024-08-13

After the discussion with the authors and agreeing on changes to be made to the paper, I increased my score and would recommend acceptance.

Reviewer f2zq2024-08-13

As an aside, I would still implore the authors do run a muP coordinate check on their implementation before writing up a camera-ready. I would feel much more comfortable having recommended the acceptance of this paper knowing this unit-test passed.

Authorsrebuttal2024-08-14

We deeply appreciate that the reviewer is now in favor of acceptance. To wrap it up, we will include the additional experiments on the hessian's eigenspectrum, and clarify the meaning of Super Consistency at all points in the paper according to the meaning and phrasing agreed here. We will also run a coord check experiment to verify that features updates are $\mathcal{\Theta}(1)$.

Reviewer 47kG2024-08-09

Thank you for the additional experiments and clarifications. I'm raising my score to 7, though I still think the following two points can be addressed further to improve the paper, but it could be too demanding in the short time frame of the rebuttal: 1. empirically show that super consistency scales to larger model. This is possible via parallel model sharding across more GPUs from the system level. It's also possible to use more efficient numerical algorithms other than power iteration. I'm curious if super consistency holds or even improves as a function of model scale. 2. For the point "more work has to be done in neural network optimization to establish (i.e. prove) how the optimal learning rate depends on the Hessian spectrum (together with other quantities)", I believe this is worth investigating further, but I know I am asking for too much for the scope of this paper.

Authorsrebuttal2024-08-12

We would like to thank the reviewer for the points they brought up, as well as for the score increase. Regarding the additional points: 1) We agree that it would be interesting to see how Super Consistency behaves at larger scales. In the final version of the paper, we will aim to scale up models further using the suggestions proposed by the reviewer. 2) We also believe that a very interesting future avenue would be to theoretically study the relationship between the optimal learning rate and the Hessian eigenvalue spectra. Understanding this dynamic could lead to better schedulers or even optimization algorithms that could speed up convergence in large models.

Reviewer 11Ly2024-08-11

Response

Thank you for the additional experiments and clarifications. Your response has addressed my concerns.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC