Dear Reviewer U3gK,

Thank you for the time you spent reading and reviewing our paper. Many of the questions you had related to the masses in particular, so let us briefly discuss how we think about them.

You are correct in that the masses absorb at least part of the question of determining optimal hyperparameters: it has been recognized by many practitioners that the learning rates for particular layers in a network (especially the embedding and the final layers) should be individually tuned, and tuning the masses plays this role in our framework. One should expect the optimal masses to differ for individual problems and architectures.

However, we still think the separation of “tune the masses” (aka relative learning rates of the layers) and “tune the global learning rate” is a useful way to subdivide the tuning problem. Many architectures take the form of an initial/embedding layer, then a number of residual layers, then a final layer. We found a useful scheme for allocating mass is to use the ratio 1 : M : 1 between the initial_layer : residual_body : final_layer. This means in particular if there are L residual layers, each one is given mass M/L. For this scheme, we reported the following experimental results in our paper (see Appendix D.6):

• For any given fixed M, the global learning rate exhibits hyperparameter transfer from a small to large network;
• The optimal M (in terms of lowest test/training loss) itself exhibits hyperparameter transfer.

Note that this scheme involves tuning a single mass M, rather than one mass for every layer.

Moreover, in contrast to earlier approaches, we concretely link the tuning of the masses to estimating the change in network output due to weight updates in each layer (Proposition 3), which gives a conceptual way to a priori reason about tuning the learning rates between separate parts of a network.

Together, hopefully this answers your first three questions. For the fourth question on initialization schemes, we want to clarify that there are two important issues when choosing an initializer:

1. what random matrix distribution do you use: orthogonal, Gaussian, uniform, etc.
2. given a choice of distribution, how do you scale it and how large are the resulting singular values?

Our belief is that the first question is not important so long as you carefully deal with the second question. In this work we use orthogonal initialization because we believe it is the conceptually simplest “base initializer” to subsequently scale according to question two (this is because the singular values of a random orthogonal matrix are all one). However, we want to emphasize that we believe that rigorously addressing this question is outside the scope of this work. See arXiv:2310.17813 for more discussion on orthogonal versus Gaussian init and arXiv:2011.14522 for more discussion on SP versus muP.

Thank you again for the time and effort you put into reviewing our paper!

Dear Reviewer FbvX,

We are grateful for your extremely thorough and helpful review! We’re very happy you found our paper interesting.

Before we go in-depth into your comments: many of them are related to the tightness of the bounds we prove on first/second derivatives, so let us first explain our picture of the situation. Almost every inequality in our paper is an instance of either:

1. The defining inequality $||W x|| \leq ||W|| \cdot ||x||$ for the spectral norm $||W||$ of a matrix $W$ (often $W$ is a gradient matrix, and $x$ is an activation vector);
2. The triangle inequality, where every term in the sum represents the change in network output due to a change in the weights of a particular module.

We view the question of whether either type of inequality is tight as being an "alignment" question: how aligned are the activation vectors to the singular value vectors of the gradient matrix for a single linear module, and how correlated are the various contributions to the change in total network output from all the different modules? We think these are very significant empirical questions about neural networks; of which there has been some investigation (e.g. arXiv:2310.17813) but further work is very much warranted. Based on your feedback and that of the other reviewers, we propose clearly highlighting this as an important avenue for future work in the “Limitations and Future Work” section.

At this point we also would like to contrast the approach to muP, which obtains \Omega(1) lower bounds in the infinite width (but crucially, constant batch size) limit. In the extreme case that the batch size is one, the gradient matrices are necessarily rank one, and inequalities of both type (1) and type (2) as above are automatically tight, and from this one can then deduce \Omega(1) estimates in the width >>> batch size regime (where one can make a low rank assumption on the gradient matrices). However, we question the relevance of such theoretical lower bounds obtained from this limiting case to real life neural network training. We believe an empirical study using realistic widths, depths and batch sizes would shed much more light on this issue.

Now, to answer your comments on the weaknesses:

1. A precise statement of what we termed "feature learning" would be "linearized change in the module output as a result of a weight update". We think this is a very useful concept that deserves an evocative name, and substantially informs how we think about the mass parameters. You’re right that this is different from other uses of the term e.g. in Tensor programs; we will heed your advice and better clarify the language in this section.
2. Thank you for some great suggestions on how to improve Section 2 in particular:

• We thought it was important to have less formal "motivation" section which outlined at a high level the relationships between the mathematical concepts in the paper, and certainly stressing the looseness of this is a good suggestion (so that the reader does not expect us to prove, e.g., that Lipschitz constants independent of network dimensions necessarily guarantee hyperparameter transfer).
• As for section 2.3, we wanted in this section to explicitly spell out how the modular norm could be actually used in real life optimization, so it’s a little bit different in its goals to section 2.1. Clarifying this would be a good suggestion.
• The tightness of equation 2.2 is partially tested in arXiv:2310.17813, but as above we believe more work should be done on this question.

To answer your questions:

1. In the experiments documented in the paper, we demonstrated that “normed SGD” was competitive with Adam for nanoGPT-like transformers. See Figures 1 in the main paper and Figures 9 and 10 in the appendix. Testing whether this remains true at a slightly larger scale is something we are actively working on currently (on a 160M parameter transformer model).
2. See the discussion above about tight bounds.
3. “Why is the fact that the norm of the weight updates in modular norm only depends on the learning rate in normed optimization (section 2.3) a sensible thing to do?” This is sensible because of Equation 3.1 in Definition 2: it implies that the learning rate will directly control the amount of “feature learning” measured in the module’s output norm.
4. We did not include a direct comparison to muP due to time and space constraints; while our work has a similar goal to muP, we do think the approach (going via non-asymptotic elementary inequalities rather than trying to identify infinite width/depth limits) is mostly orthogonal. We are not claiming our approach has better or worse real life performance than muP.
5. As you highlighted, we do believe the key question is "over the course of training". Based on your feedback, we will edit the "Limitations and Future Work" to highlight this (we teased some of this in the "Loss of well-normed-ness" section, but in hindsight this could be better foregrounded).

For your questions on limitations:

• See the above for a discussion about lower bound guarantees, including a discussion of the approach taken in muP.
• Correct, we only consider learning-rate transfer in this paper.

Thank you again for the time and effort you put into the thorough review of our paper!

Thank you for a very thorough response! I also very much appreciate all the changes the authors said they would make, and think they are a good idea.

I wanted to discuss a couple of points in the authors' response:

In the extreme case that the batch size is one, the gradient matrices are necessarily rank one, and inequalities of both type (1) and type (2) as above are automatically tight, and from this one can then deduce \Omega(1) estimates in the width >>> batch size regime (where one can make a low rank assumption on the gradient matrices).

Maybe I'm missing something, but I don't think this is true when looking at different inputs. Yes, this will be true when comparing the tightness of the bound on the same datapoint $x$ on which a gradient update was made, but the bounds will not necessarily be tight for a different datapoint $x^\prime$ – the activations for that datapoint might not be aligned to any extent with those of $x$. That's what makes the results about muP interesting and non-trivial in my opinion.

I would also push back on infinite model-size, fixed batch-width, not being a realistic limit. It seems to capture how people scale training in practice pretty well. If anything, I would take more of an issue with not considering the infinite training time limit, but I think being able to obtain asymptotic tightness results even in the finite training time limit is quite interesting.

“Why is the fact that the norm of the weight updates in modular norm only depends on the learning rate in normed optimization (section 2.3) a sensible thing to do?” This is sensible because of Equation 3.1 in Definition 2: it implies that the learning rate will directly control the amount of “feature learning” measured in the module’s output norm.

That makes sense. I think I would recommend putting this as a motivation (not formal, just colloquially explained in words) for normalising weight updates at the beginning of Section 2.3. I think at the moment Section 2.3 just comes a bit out of nowhere, and that bit of explanation makes it clear why normed optimisation is a desirable thing to do.

Dear Reviewer qUST,

We are grateful for your thorough and helpful review! First of all, we will heed your advice about restructuring the technical content. Second, regarding the result being lost in the literature, we are already proactively working with the community on larger-scale validations. For example, we are working with a software engineer at Midjourney on an open source project to port Modula into JAX (calling it “Modulax”) and they are planning to test it on large-scale diffusion models---including Midjourney v7 if it works well. And we are in contact with an engineer at Cerebras who reported good training stability with Modula on GPT-2 size models. In terms of what we can show you now, we will try to get the experiment on a 160M size model done and share the results with you before the end of the discussion period. Apologies for not having it done sooner.

Regarding your questions:

1. As you say, the tightness of equation 2.2 is partially tested in arXiv:2310.17813. We agree that a more thorough evaluation is needed. Doing this properly will require significant care and is probably worthy of a full paper. We propose clearly highlighting this as an important avenue for future work in the “Limitations and Future Work” section.
2. Thank you, this is a great and important point. We do have a clear, intuitive understanding of what Modula is doing “under the hood”. But due to space limitations and the mad dash to get the project done, we didn’t include this in the paper. To partially rectify this, we created open source docs and an FAQ that directly explain many of the mechanisms. We propose to add an intuitive explanation of the underlying mechanisms to the prospective camera ready. We will express this in clear, accessible language.
3. In our experience, we have never seen per-tensor gradient normalization being too aggressive. There is a body of work supporting this: consider the LARS and Fromage optimizers for older examples. See arXiv:2305.17212 for a more recent example. In addition, consider that Adam and recent extensions like Adam-mini can be interpreted as forms of per-tensor normalization, although people usually don’t think of them that way. And don’t forget Shampoo!

Thank you again for your time and effort. Again, we will try to get the suggested experiment to you by the end of the discussion period.

Dear Reviewer 3bTa,

We are grateful for your time spent reviewing our paper. We are sorry if the absence of the contribution list was disconcerting. We feel that our paper is chock-full of contributions since we are advancing a substantially novel perspective on deep learning optimization based on automatically and recursively metrizing the weight space of the neural network. To highlight a few key contributions, consider that:

• We propose the modular norm defined via a recursive algorithm (Definitions 3 and 4)
• We show that neural nets are Lipschitz (Proposition 2) and Lipschitz smooth (Proposition 5) in the modular norm

Establishing a solid, workable notion of Lipschitzness for neural networks is generally regarded as a major open problem by experts in optimization theory. See, for example, the survey “Optimization for deep learning: theory and algorithms” by Ruoyu Sun (arXiv:1912.08957) for an explanation of this. We are sorry that the presentation of the paper was not to your taste, but we hope that you will consider re-evaluating the paper in light of this clarification.

Regarding your other questions:

• Mass and sensitivity are held fixed from the start of training. Weight updates are normalized in the modular norm at each training step.
• The spectral perspective on training dynamics is already reconciled with muP and NTP in “A Spectral Condition for Feature Learning” by Yang et al (arXiv:2310.17813). See, for example, Figure 2 in that paper. We are generalizing that analysis to general architectures. We are not claiming to have better performance than muP.

Thank you again for your time and effort!

We really appreciate you engaging with us further.

"an actual application has to be found (either theoretical or practical)" We completely agree with you here. Actually, the application we are targeting in the paper is the definition of a single normalize function that can be applied to either Adam or SGD making the method exhibit good learning rate transfer. In contrast, in muP and all other approaches to this problem that we know of, different scaling rules are needed for different base optimizers. In other words, we believe that we have made the implementation of scalable training substantially easier across optimization algorithms. This application is what the experiments in Figures 1 and 4 target.

"even if the proposed method performs worse than muP, the experimental results would be interesting anyway" You're right, and we commit to including this comparison in the prospective camera ready. We'll try to get it done this weekend if we can.

"we do not know why ensuring Lipschitz smoothness in the modular norm should be useful for further research". You are right that we haven't found a killer application for this part yet. But we felt so excited about this result from a theoretical perspective, that we felt it was worth highlighting prominently in the main paper. We are hopeful that ourselves or other researchers are able to use automatic sharpness calculations productively in future work. We agree there is uncertainty here.

In conclusion With this paper, we resisted the urge to "just present one idea", which is the typical advice for ML conference papers. Instead, we threw everything we could at the problem and tried to write the most exciting paper we could based on that. As a result, there are many avenues left open. We intend to continue to work on this in future work, and we hope that maybe we could inspire other community members to also find these directions interesting.

If based on your reviewing you had suggestions about how we could restructure the paper further, we would love to hear them. Also, if you feel that we have addressed your concerns at least to some extent, we would appreciate it if you consider raising your score.

Best, Authors

Just to summarize the current reviewer positions for the AC:

• reviewer qUST liked the paper but felt that larger-scale experiments could increase its impact
• reviewer FbvX praised the work for both its originality and also its "fairly robust" and "very promising" experiments
• reviewer U3gK expressed concerns about the role of the mass parameters, which were largely assuaged when we pointed out that the mass of all the residual blocks can be tuned collectively
• reviewer 3bTa was thrown by the lack of a formal contribution list at the end of the introduction. However, we feel that the contribution list is essentially contained within the paper's abstract, and the NeurIPS conference has no requirement for a bulletpoint list of contributions at the end of the introduction.

Again, we are immensely grateful to all reviewers for their time and effort. Thank you!