Summary
The authors propose the “modular norm”, a norm for deep learning that can be simply recursively composed, allowing for easily recursively computing (and hence controlling) the Lipschitz constant of the network and loss gradient.
They propose how to scale the gradient updates by the modular norm, and empirically demonstrate the effectiveness of such ‘normed optimization’ at achieving invariance of the optimal hyperparameters to scale.
Strengths
The paper proposes a very interesting idea: an architecture-agnostic norm for neural networks that gracefully scales in width and depth.
This seems like a very interesting take on what is currently being achieved through convoluted scaling of initialisations, learning rates, and residual branches following rules derived from asymptotic properties of these rules.
This paper has challenged how I think about sharpness and curvature of the loss landscape in a deep learning context, and the asymptotic properties thereof. I'll admit, I'm still digesting the take-aways, but I feel fairly confident many in NeurIPS community would benefit from reading this paper.
Lastly, the experimental results look on learning rate transfer are fairly robust, and look very promising.
Weaknesses
1. The use of the term ‘feature learning’ seems somewhat distinct in this paper from the way it's been used in the Tensor Programs (TP) literature. To the best of my understanding, in TP4, it refers to the lack of or presence of change in how each layer processes its inputs in the infinite width limit. In Proposition 3, the amount of “feature learning” seems to refer to a *bound* on how much the way a layer processes its inputs can change. I think it would have been great if the authors defined or front-loaded what they mean by ‘feature learning’ and ‘proportion of feature learning’ before Proposition 3. E.g., lines 128-131 were difficult to parse on a first read-through.
- Similarly, the controlling the ‘proportion of feature learning’ claims are teased again without clarifying what is meant by feature learning in lines 171-173. I think these lines could be cut to be honest.
2. I think the presentation could use a little bit of work. A lot of it is very good already. In particular, the notation, definitions and propositions read well, and the discussion of limitations and future work in section 5 was very useful and interesting. That being, I think section 2, which was setting up the motivation for what follows, left me a bit confused on my first read-through. I would have hoped to give concrete suggestions on how I'd like to see it changed, but that is a non-trivial task, so I'll point to what I think is confusing at the moment:
The authors mention two motivations for the modular norm in the abstract: a) graceful scaling in width and depth, and b) ease of constructing a Lipschitz-continuous network (in the modular norm). The elaboration on and setup of these two motivations than gets slightly jumbled, in my opinion. The introduction primarily focuses on a) – the graceful scaling. Then, section 2, to the best understanding, gives an argument for why achieving Lipschitz continuity with a tight constant that's invariant to the scaling dimension might result in the step-size being invariant across scale. As I understand it, this is a fairly loose argument, that's not being corroborated with further formalism or arguments later on.
I think, when first reading section 2.1, I would have appreciated:
1. Having a better sense of direction:
1. Making it clear that the link between graceful scaling of step-size and Lipschitz continuity is a loose motivating argument, that's absolutely useful, but will not be formalised later on.
2. Making it clearer that the goal of the following sections will be to present a modular norm in which it will be easy to specify networks with a Lipschitz-continuity that's invariant to a scaling dimension.
Then, I think the relationship before what section 2.3 lays out — achieving norm $\eta$ updates in modular norm — and section 2.1 could also have been made much clearer.
3. Some smaller points:
- On lines 81-83, it seems like bulletpoints (ii) and (iii) could be merged. They refer to the same desirable property, and one is just setting up the other.
- The sharpness bound in equation 2.1 could use a citation (or link to appendix derivation) for pedagogical purposes.
- It would have been great to define (or cite a reference for) what a matrix norm induced by input and output space norms is. I wasn't familiar with this term.
- The authors claim that equation (2.2) holds quite tightly. Do the authors have a reference?
Questions
1. The authors say “[scaling the updates in the right norm] in some instances, [...] enables training with a simpler optimizer—for example, training GPT with SGD rather than Adam—thus incurring a smaller memory footprint.” Do the authors have experiments or a reference to corroborate this? If stable training of large transformers was demonstrated in the modular norm, this would be an impressive feat.
2. The work aspirationally mentions graceful scaling in e.g. width/number of modules. If my understanding is correct, the many guarantees on how the properties of modules combine in the modular norm means that it would be easy to devise rules for setting e.g. the mass parameters so that the sharpness of the network is nicely bounded. However, the work doesn't actually show any results on the tightness of these bounds in any asymptotic setting? Without the bounds being tight (or not getting worse in appropriate scaling limits), why would we expect hyperparameters like the learning rate to transfer?
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?
4. Is there any reason the authors didn't compare to hyperparameter transfer in muP and its depth counterpart?
5. Empirically, are the Lipschitz bounds actually close-to-tight? Does the tightness persist throughout training?
Limitations
1. If I'm understanding things correctly, Proposition 3 does not guarantee the _presencence of feature learning_ in the modular norm, the same muP guarantees a change in the way a layer processes its inputs by at least $\Omega(1)$ as the width goes to infinity?
2. Only transfer of the step-size and mass allocation across scale is considered, and not of other training hyper-parameters (e.g. learning rate schedule).