Scalable Optimization in the Modular Norm
Authors: Tim Large
Yang Liu
Minyoung Huh
Hyojin Bahng
Phillip Isola
Jeremy Bernstein
Venue: NeurIPS 2024 poster (2024)
Keywords: scalable
optimization
modular
norm
normalization
modula
hyperparameter
transfer
architecture
aware
operator
module
tree
Primary area: optimization_for_deep_networks
TL;DR: Normalizing weight updates in a special norm leads to hyperparameter transfer; this norm can be inferred automatically from the computation graph.

Abstract: To improve performance in contemporary deep learning, one is interested in scaling up the neural network in terms of both the number and the size of the layers. When ramping up the width of a single layer, graceful scaling of training has been linked to the need to normalize the weights and their updates in the "natural norm" particular to that layer. In this paper, we significantly generalize this idea by defining the modular norm, which is the natural norm on the full weight space of any neural network architecture. The modular norm is defined recursively in tandem with the network architecture itself. We show that the modular norm has several promising applications. On the practical side, the modular norm can be used to normalize the updates of any base optimizer so that the learning rate becomes transferable across width and depth. This means that the user does not need to compute optimizer-specific scale factors in order to scale training. On the theoretical side, we show that for any neural network built from "well-behaved" atomic modules, the gradient of the network is Lipschitz-continuous in the modular norm, with the Lipschitz constant admitting a simple recursive formula. This characterization opens the door to porting standard ideas in optimization theory over to deep learning. We have created a Python package called Modula that automatically normalizes weight updates in the modular norm of the architecture. Both the Modula package and code for our experiments are provided in the supplementary material.

Scalable Optimization in the Modular Norm
Tim Large* Columbia University
Yang Liu Lawrence Livermore National Lab
Minyoung Huh MIT CSAIL
Hyojin Bahng MIT CSAIL
Phillip Isola MIT CSAIL
Jeremy Bernstein MIT CSAIL
Abstract
To improve performance in contemporary deep learning, one is interested in scaling up the neural network in terms of both the number and the size of the layers. When ramping up the width of single layer, graceful scaling of training has been linked to the need to normalize the weights and their updates in the "natural norm" particular to that layer. In this paper, we significantly generalize this idea by defining the modular norm, which is the natural norm on the full weight space of any neural network architecture. The modular norm is defined recursively in tandem with the network architecture itself. We show that the modular norm has several promising applications. On the practical side, the modular norm can be used to normalize the updates of any base optimizer so that the learning rate becomes transferable across width and depth. This means that the user does not need to compute optimizer-specffic scale factors in order to scale training. On the theoretical side, we show that for any neural network built from "well-behaved" atomic modules, the gradient of the network is Lipschitz-comtinuuus in the modular norm, with the Lipschitz constant admitting a simple recursive formula. This characterization opens the door porting standard ideas in optimization theory over to deep learning. We have created a Python package called Modula that automatically normalizes weight updates in the modular norm of the architecture. The package is available via pip install modula a with source code here.
1 Introduction
Given the practical impact of deep learning systems trained the largest scale, there is need for training algorithms that scale gracefully: without instability and-if possible possible-without without manual tuning. However, current best practices for training have developed somewhat organically and do not live on a bedrock of sound numerical analysis. For example, while the Adam optimizer [1] ubiquitous in the field, errors have been found in its proof of convergence [2], and empirically Adam has been found to scale poorly as either the width [3] or the depth [4] of the network is ramped up.
To remedy this situation, patchwork of learning rate correction factors have recently been proposed [3-6]. The general idea to retrofit a base optimizer such as Adam or SGD with special correction factors intended to render the optimizer's optimal learning rate invariant to scale. But this situation is not ideal: the correction factors are reportedly difficult to use. Lingle [7] suggests that this may be due to their "higher implementation complexity, many variations, complex theoretical background". What's more, the correction factors are optimizer-speciffi meaning that if one switches to a different optimizer one must either look up recalculate separate set of correction factors.
The goal of this paper is to simplify matters. We show that both Adam and SGD can be made to scale gracefully with width and depth by simply normalizing their updates in a special norm associated with
* denotes equal contribution. Correspondence to (jbernstein, 1o
minhuh @mii.. edu.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

width
number of blocks
Adam
SGD
2
26
2
2
29
2
2
2
23
2
26
2
normed Adam
normed SGD
train loss
10
100
10-1
10
27 2
29
210
22
23
24
25
26
learning rate
width
number of blocks
Figure 1: Learning rate transfer in the modular norm. We train GPT with context length 128 WebText for 10k steps on OpenWehText Left: Learning rate sweeps for normed Adam (Adam with updates
normalized in the modular norm) with three transformer blocks and varying width. The optimal learning rate (marked by red dots) transfers well across scales. Mid-left: The same, but varying the number of blocks width 128. Mid-right: Comparing normed versus unnormed Adam and SGD fixed learning rate and varying width. For each method, we tune the learning rate at the scale marked by the dotted line. The normed methods scale better. Right: The same, but scaling number of blocks.
the network architecture see Figure 1, package called Modula that constructs this norm automatically and in tandem with the architecture.
We call this norm the modular norm, and provide a Python
The modular norm is constructed recursively, leveraging the module tree perspective on neural architectures. It enough to define how the modular operations: composition and concatenation. We show how other basic operations on modules, such as addition and mallaa-muliipliaatoon. can be implemented through composition and concatenation. And then higher-order structures, such as residual networks, can be built using these basic operations.
propagates through only two elementary
Beyond its practical relevance, the modular norm may also prove usefuu to theoreticians. Various optimization ceaimization-theoettiic quantities are accessible and efficiently calculable in the modular norm. For instance, show that the gradient of any neural network built from "well-behaved" atomic modules Lipschitz-comee continuous in the modular norm of the architecture. This opens the door to porting several more-or-less textbook optimization theory analyses [8] over the world of deep learning.
1.1 Related work
Metrization It is by now well-known that deep networks do not easily or naturally admit Lipschitz- continuity or smoothness guarantees in the Euclidean norm [9-13]. Researchers have attempted to address this problem: for instance, Bernstein et al. [12] propose distance function called deep relative trust, which combines Frobenius norms across network layers. However, deep relative trust only constructed for the multilayer perceptron and, when used to normalize updates, its employment of the Frobenius norm precludes good width scaling. In contrast, Yang et al. [14] equip individual layers with the RMS RMS operator norm, finding this to enable good width scaling. Researchers have also looked at building neural net distance functions outside the context of scalability [15-17].
Asymptotics The metrization-base approach to scaling developed in this paper contrasts with the tradition of asymptotic scaling analyses- the study of infinite width and depth limits-mort common in the deep learning theory literature [3-5, 18, 19]. These asymptotic analyses follow an old observation of Neal [20] that interesting properties of the neural network function space exactly calculable in the infinite width limit and at initialization. This tradition has continued with asymptotic studies of the neural tangent kernel [21] as well as infinite depth limits [4, 5, 22]. However, there is increasing recognition of the limits of these limits, with researchers now often trying to relax limiting results [23-25]. And ultimately, from a practitioner's perspective, these results can be difficult to make sense of [7]. In contrast, our framework eschews any kind of limiting or probabilistic analysis. As consequence, we believe our framework is simpler, more easily relatable to basic mathematical concepts, and ultimately more relevant to what one may encounter in, say, Porch [26] program.

---

Majorization In recent work, Streeter and Dillon [27] propose a universal majorize- algo- rithm [28]: a method that automatically computes and minimizes majorizer for any computational graph. Despite its generality, current downsides to the method include its overhead, which can be 2x per step [29], as well as the risk that use of full majorization may be overly pessimistic. Indeed, Cho and Shin [30] find that an optimization approach leveraging second-order information converges significantly faster than a majorization- -insiireed approach. Related ideas appear in [31, 32].
Descent in Normed Spaces
We define the modular norm in $3. This section is intended to prime the reader for what is to come. In this section, and the rest of the document, the diamond operator denotes tensor contraction.
2.1 What's in norm?
Suppose that we wish to use gradient descent to minimize a loss function L W + R over a weight space W RR.. What properties of the loss L and weight space W would we desire for this to be sensible? Three such properties are:
(i) the loss function is differentiable, meaning that the gradient map V W W exists;
(ii) the weight space W carries norm ||-| W > E, which need not be the Euclidean norm;
(iii) the loss is Lipschitz smooth in the norm |1-1], with sharpness constant >> O, meaning that:
L(w + Aw) S L(w) + vw Aw + 2 ?1
(2.1)
Under these conditions, the weight update given by Aw = arg min [v. o Aw + $ || |2] is guaranteed to reduce the loss. The particular norm ||-Il influences the direction of this weight update, while the sharpness constant À influences the size of the update.
In deep learning, we would ideally like the optimal step-size to remain invariant as we scale, say, the width and the depth of the network. Thus, fundamental problem is to design a norm such that, first, Inequality (2. 1) actually holds (and is not hopelessly lax), and second, the corresponding sharpness constant is invariant to the relevant architectural dimensions. If the norm is chosen poorly, the practitioner may end up having to re-tune the step size as the network is scaled up. In this paper, we design& norm for neural networks that meets these requirements: the modular norm.
2.2 Preview of the modular norm
The weight space of a deep neural network is a Cartesian product W = Wi x x... x Wr. where W3 is the weight space layer k. Yang al. [14] consider the problem of metrizing individual layers. For instance. if layer A is a linear layer with weight space Wi = do don xdm, then they equip this layer with the RMS-RMS operator norm, ||-||RMS-RMA RMS. This is the matrix norm induced by equipping the input and output space of the layer with the root-mean-sqpe -quare (RMS) vector norm, |a=liRss := L a for x E Rdd. The advantage of this non-standard matrix norm is that allows one to estimate the amount of feature change induced by gradient update. In other words, the inequality
| x #|IMMS S AW ||Ms|
RMMSRMS
[
(2.2)
turns out to hold quite tightly when A AW is gradient update and a corresponding layer input. This is because gradient updates to a layer are (sums of) outer products that align with layer inputs.
Once we know how metrize individual layers, natural question is: can we combine layer-wise norms to produce a norm on the full weight space W = II W/ of the network? Naïvely, there are many ways do this: one could take any positive linear combination of the layer-wise norms (L combination), the square root of any combination of the squared layer-wise norms (L² combination), and on. But we want the norm be useful by the criteria of $2.1. To this end, we propose the modular norm ll,,,w, which ends up as a max (L combination) of scaled layer-wise norms l|-l
[[w),,... wt. llw == max (81 |/ ll,,,
.., sr. wt, llwz))
(2.3)
The positive scalar constants , 8t, are determined by both the architecture of the network and set of user-specified "mass" parameters. The precise construction of the modular norm, working
3

---

operations for building modules
X E L
module M
attributes:
ye y
Add
M.for
M.mass
E 7
M.sensitivity
2
M
M
M
M
M2
M.norm
composition
concatenation
module sum M, + M2
Figure 2: Modules and trees of modules. module an object that maps an input and a weight vector to output. Left: In addition to the standard forward function, our modules are endowed with two numbers- mass and sensitivity-ann via the binary operations of composition and concatenation. We provide rules for composing and concatenating all module attributes. Right: Compound modules are binary trees, where the leaves modules and the internal nodes compose and concatenate their children. Here we illustrate sum of modules, which leverages a special utility module Add see Table for more on this.
norm. Middle: New compound modules are built
recursively over the module tree of the network, is given in $3; there, we also explain how the modular norm satisfies the criteria $2.1. and the role played by the mass parameters. For now, let us explain what good the modular norm yields in practice.
2.3 Normed optimization
The main practical use the modular norm is to normalize weight updates. With reference to Equation (2.3), we define the following operation weight updates Aw (A.,.-... Awy) W:
W
AWL
(ormalize(Aw) =
(2.4)
Awi W
Awzllwe
Provided none the Aws, are zero, then normalize(Aw)) is unit vector in the modular norm. We propose using normalize as wrapper, along with an explicit learning rate schedule, for any base optimizer such as Adam or SGD. The resulting normed optimizer is thus made architecture-aware the normalize function. In pseudo-code- actual Modula code-this amounts to:
delta_w w optim(v.grad() net nww.alimmeeeeeeeee (delta_w) W estep)) delta_w
get update from base optimizer # normalize update in the modular norm apply update with learning rate eta
We find this wrapper to significantly improve the scalability of the base optimizer. It renders the optimal learning rate roughly invariant to width and depth, with seemingly no cost accuracy. In some instances, it enables training with a simpler optimizer-for example, training GPT with SGD rather than Adam-thus incurring smaller memory footprint.
Normalization in the modular norm essentially forces individual layers to learn at specified, regulated rates. We view this balancing learning across the network; individual layer learn too fast and destabilize training. This balance is determined by the architecture, along with user-specified parameters that provide precise control over the relative learning speed in different submodules.
For variety of experiments with normed optimization, see $4 and Appendix D. But first, we detail the construction of the modular norm along with its core properties.
3 Constructing the Modular Norm
Our strategy is first define the abstract notion a module, which includes a norm as an attribute. We depict this concept in Figure 2. Then, by providing rules for composing and concatenating modules, we recursively define a norm for any module built via an arbitrary sequence compositions and concatenations: the modular norm!

---

3.1 Modules
module is a re-usable, composable object useful for building complicated neural networks. Our definition of a module augments the PyTorch module [26] with two real numbers and norm: Definition (Module). Given input vector space X, output vector space y and weight vector space W, module M is object with the following four attributes:
(a) a function, M.forward : W x + Vv which maps an input and a weight vector to an output-we often abbreviate this attribute to just M = M.forward;
(b) number, M.mass O, which will turn out set the proportion of feature learning that this module contributes to any supermodule;
(c) number, M.sensitivity 2 0, which estimates the module's sensitivity input pertarbations:
(d) a norm over the weight space, M.norm W R2,, sometimes abbreviated to just II-IIM-
Before we say more about the intended roles of these attributes, let us mention the three kinds of modules that we will care about practice:
ii) atomic modules, whose attributes are hand-declared, and have weights. Examples include linear modules, embedding modules, and convolution modules.
(ii) bond modules, whose attributes are hand-declared, but have no weights. Formally, their weight space is the zero vector space W 0. An example is the ReLU non-linearity module.
(iii) compound modules, built out of other modules, with automatically inferred attributes.
Note that the space of objects that type-check as a module by Definition I is vast. Since we need to hand-declare atomic and bond modules in order build interesting compound modules, we should have an idea of what makes for a "good" module. Simply put, a module is good when its attributes are predictive of its behaviour. To formalize this idea, we say that module is well-normed if its forward function, sensitivity, and norm satisfy the following two relationships:
Definition 2 (Well-normed), Let M be a module on (X, (X W), where the input and output spaces have respective norms ||-I1x and || Il-lby. M is well-normed if for all inputs x E X and weights w W:
V m W. f.rwwwrrd un, .x) Awlly S M.norm(Aw))
for all Aw E W;
(3.1)
V x ( w. x) o Aælly < M.sensitivity * ||
for all A E X.
(3.2)
well-normee--nss. means that the norm function and sensitivity are good match for the forward function. The first inequality says that a well-normed module is Lipschitz-continuuss over its weight space with constant one. The second inequality says that well-normed module is Lipschitz- continuous over its input space with constant M.sensitivity. In practice, we will be interested in well-normed modules where these inequalities hold fairly tightly, since then M.sensitivity and M.norm will let us estimate the sensitivity of the module to input and weight perturbations. Appendix B provides many examples of well-normed atomic and bond modules.
The remaining attribute M.mass will turn out to control the proportion of feature learning that a module contributes to any compound module in which it participates. We formalize this concept in $3.3. But before that, we need to understand how to build compound modules.
3.2 Compound modules: Building new modules from old
We consider building new modules from old ones via the binary operations of composition and concatenation, illustrated in Figure 2. Composition is denoted via the serial combination M₂ o M1. and concatenation via the parallel combination (M1, M2), alternatively referred to as a module tuple. These simple binary combinations will let us build basic algebraic operations on modules (Table as well as complex neural network architectures. We start by defining module composition:
Definition 3 (Module composition). Consider module M1 with input, output and weight space (X,,)) Wi) and module M2 with input, output and weight space (X2, W2). Mj and M2 are composable if X₂ = Yr. Their composite M M2 o Mj lives on (X,, Wi x W2) with attributes:
(a) M.forward (2), x)) Mg.forward w2, M₁ Mf.foraard(wj,r
5

---

Operation
Shorthand
Definition
Modula Expression
module addition
Mi + M
Add o (M1, M2)
M_1 M_2
scalar multiplication
a M
Mula M
a M
iterated composition
M
M M4-1 with M == Identity
M
Table 1: Arithmetic with modules. Composition and concatenation let us define an extended arithmetic modules. The utility modules Add, Mul, and Identity are defined in Appendix B.2.
(b) M.mass = M1.mass + Mg.mass;
(c) M.sensitivity = M₁ sensitiviiy * My.sensitivity;
(d) Annom(((1,)) u22) given by:
M.mass
M.mass
max M2 Mg.sensitivity *
nor))))(w)),
* M2.norm(wz))
Mj.mass
Ma.mass
where if M mass M2. is zero, the corresponding term in the max is set to zero.
At this stage, we make two comments about this definition. First, in the definition of the composite norm, notice that the norm of the first module couples with the sensitivity of the second module. This reflects the fact that the output of the first module is fed into the second module and not vice versa. Second, observe that the masses of the submodules are involved in setting the balance of the composite norm. Before we further motivate this definition, let us first define module concatenation:
Definition 4 (Module concatenation). Consider module M1 with input, output and weight space (1,,31 Wi) and module M2 with input, output and weight space (X2, Y2, W2). We say that M1 and M2 concatenatable if their input spaces match: X1 X2. The tuple M (M1, M2) has input, output and weight space (X1,3h x Y/2 Wi x W2) and attributes:
(a) (()),, w2), )²) (M, (w1, .)), Mo.forward(ww,, a));
(b) M.mass MI + massa
(c) M.sensitivity = M + Myssensitivity;
(d) M. M.norm(w,, w2) given by:
M.mass
M.mass
* *ornorr(w..). M2.mass -Mg.nrr((wg)
max
Mj.mass
where .mmss or Ma.mass is zero, the corresponding term in the max is to zero.
Concatenation is simpler than composition in the sense that neither module is fed through the other, and therefore, sensitivity does not appear in the concatenated norm. To further motivate these definitions, observe that two basic and desirable properties follow as immediate consequences:
Proposition 1 (Composition and concatenation associative). If modules M1, M2, M3 succes- sively composable, then M3o(M2oM;) equals (M3oM2)oM) in all altributes. modules My, M2, M3 mutually concatenatable, then ((M1, M2), Ms) equals (M1, (M2, M3)) in all attribtee..
Proposition 2 (Composition and concatenation preserve well-normedness). If modules M1 and M2 are well-normed and composable, then their composite M2 o M₁ is also well-normed. If modules M and M, are well-normed and concatenatable, then their tuple (M₂, M2) is also well-normed with respect to the L' combination norm on the output space: [I(.,))llx x = |1-1l,, + +1-llyz.-
The proofs follow directly from the definitions and the chain rule. Proposition implies that one may build complicated compound modules without worrying in which order successive combinations taken. Proposition implies that complicated compounds automatically inherit Lipschitz guarantees.
Taken together, Definitions 3 and 4 define the modular norm M.norm of any compound module M.

---

ResMLP on CIFARIO
GPT OpenWebText
1.8
blocks
mass
width=128 hlocks=3
1.7
0.25
0.5
1.6
1
tran
1.5
mass=0.25
mass=0.5
32
5
1.8
width=512 blocks=6
1.7
1.6
7
6
1.6
5
1.5
1.5
mass=
mass=2
101
100
10
10
10
10
10
learning rate
blocks
learning rate
Figure 3: Exploring mass allocation. We tune the total mass of the hidden layers, training with normed Adam. Left group: Learning rate sweeps for ResMLP CIFAR-IO, for varying depth and mass. The bottom right subplot reports the best train loss at each mass and depth. Mass 0.5 was best all depths. Right group: Learning rate sweeps for GPT on Open DpenWeeText.. for varying mass. Both optimal mass and learning rate transferred from the small model (top) to the large model (bottom).
3.3 Mass allocation in compound modules
Suppose we wish to train network with input layer, output layer, and L. blocks between:
Network OutputLayer HiddenLayers InputLayer
(3.3)
OutputLayer o Block* InputLayer.
(3.4)
Then how much learning should happen in the output layer, compared to the blocks, compared to the input layer? And what if we scale the number of blocks L-do want relatively less Icarning to occur in the network's extremities? Or do we want the input and output layers to learn non-trivially in the L - limit? Since answering these questions is difficult priori, introduced the mass parameter to allow a user to set the proportional contribution each module has toward learning:
Proposition 3 (Feature learning is apportioned by mass). Consider compound module M derived in any fashion from well-normed modules MI...., M Given weight setting w = (W).... Wr), where we denote the weights of module MR. let us perturb w by Aw (Awi..... AWL). If we decompose the linearized change in the output of module M into one contribution per sub-module:
M(w, x) Aw Vw, M(w, x) Aw + + V M(w, x) AWL.
(3.5)
then the kth term in this decomposition satisfies:
M
V wg M(w x) Awally <
* M.norm(Aw).
(3.6)
M.mass
In words: module mass provides the flexibility needed to build complicated compound modules involving many sub-modules, while maintaining precise control over how much learning any sub- module can contribute to the overall compound. Proposition is proved in Appendix E.
In practice, obtained the best training performance by maintaining constant amount of learning in the input and output layers even as the number of blocks scaled (Figure 6). In other words, seems be good idea to assign OutputLayer.. mass Hidden Layers. masss. InputLayer. .masaaaassssssss mass in proportion 1 m : 1, where m is independent of the number of blocks The exact mass of the hidden layers needs to be tuned rates in the input and output layers in μP [18]; this tuning can be done on small model prior to scaling (Figure 3). We further discuss mass allocation in Appendix D.6.
new architecture -just
needs tune separate learning
3.4 Smoothness in the modular norm
In this section, we study the second derivatives of module using the modular norm measuring stick. Let us start by defining the notion of sharpness that we will consider:
7

---

Definition 5 (Module sharpness). Let M be module on V. W),, where the input and output spaces have respective norms I|-Ilx and ll-lly. We say that M (a, 3.7))sharp for constants a, 6,9,7 7 2 O if, at all inputs a E X and weights w E W, the second derivatives of M are bounded as:
lww o V M(w,
x) o Awolo - a Aw MM
A
for all Aw,, Au E WW
(3.7)
AAw V wa M(w, x) Am y S / A M
|
for all Au E W and An E X:;
(3.8)
| V ua M(w, Ae) y S 7 A x
||
for all Ax, At Ex.
(3.9)
While one may ultimately be interested in the sharpness of module with respect to weight perturba- tions, Definition also tracks sharpness with respect input perturbations. In fact, tracking this extra information is essential for propagating sharpness bounds up the module tree. Appendix C details the procedure for automatically calculating the sharpness constants of compound module starting from the sharpness constants of all its submodules; see Propositions 8 and 9 for the specific formulae. Here we highlight one major corollary of these formulae, proved in Appendix for a specific choice of block multipliers, the sharpness constant of residual network is independent of depth:
Proposition 4. Suppose M is a well-normed, (a, 3. () -hharpp module on (X.X, W) with unit sensi- tivity. Define the depth L residual module Res (M) via the module arithmetic of Table I as:
Res (M) == ( * Identity + 1 * M)²
(3.10)
Then this residual module Resy, (M) is in fact (a + B 3.3+ 3 3. 7) sharpp independent of depth L.
For optimization purposes, one may be more interested in the sharpness of the loss function rather than the sharpness of the neural network. Fortunately, it possible to convert sharpness bounds on modules into sharpness bounds on loss functions, provided a little is known about the error measure:
Proposition 5 (Loss functions are smooth in the modular norm). Let M be a module on (X, W) and let E y x T + R measure the error between module output and target in target space T. The loss L : W- + R records the module's 's average error on data distribution 'D over X x T
((w) == Ex,t~D E(M(w, t).
(3.11)
Suppose that the error measure is a --Lipshhitz and r-smooth in the module outpua, the sense that: for all Ay E y and t E Ti
VV t) Ayl S Ayllyy
(3.12)
Ay V2 (y t) AاI Ayاly االاا
for all Ay, Aÿ E y and t T.
(3.13)
If the module M is well-normed and (a, 9, -sharp, then the loss function L satisfies the following three inequalities at all weight settings w E W and for all weight perturbations Aw, Aw W:
(i) Aw O uw L o Aw S (oa + T) | |A@
V
(ii) V W Aw) V Cw IM S (aa + T) Awi|M
where II. is the dual norm of ||-||M;
(iii) [[w + Aw) - (C(w) + Vo Aw]l < ( + +) //
The proof is given in Appendix E, and we present estimates for a and 7 for common error measures in Appendix C.4. Notice that inequalities (i), (ii) and (iii) are the standard inequalities of smooth optimization [8]. albeit expressed in the modular norm. In fact, (i) implies (ii) implies (iii). In words, inequality (ii) says that the gradient of the loss is Lipschitz-continuel
in the modular norm. The Lipschitz constant depends on the module only through the module's ''s first sharpness coefficient a.
4 Experiments
Our experiments aimed test the scalability of training with normed versions of Adam and SGD: whether one can tune the learning rate on a small model, and expect the learning rate to remain close to optimal on models of much larger width and depth. In addition to the learning rate, normed optimization in Modula requires a mass parameter to apportion feature learning between the input, output and hidden layers; we also tested the sensitivity of this parameter, whether it affects learning rate transfer, and to what extent the optimal mass itself transfers across width and depth.
8

---

ResMIP on CIFAR-10
ResNet on CIFAR-10
Adam
train loss
1.6
1.0
normed Adam
1.4
0.5
SGD
normed SGD
25
2
29
23
2
2
27
29
2
2
width
blocks
width
blocks
Figure 4: Learning rate transfer on CIFAR-10. We tume the learning rate on a small model-an at the scale marked by the dotted line-and test the performance models of increasing width and depth at this fixed learning rate. We find that normed Adam and SGD scale better than their unnormed counterparts on both ResMLPs and ResNets. See Figure for the same experiment on GPT.
All SGD experiments were done with momentum / = 0.9, and all Adam experiments used 3 = 0.9 and B/2 = 0.99. No weight decay was used in any experiment. Every experiment was done with a linear decay learning rate schedule. As for initialization, we used orthogonal initialization for Linear and Conv2D modules, and Gaussian weights projected module. This was to ensure all modules were well- well-normed at initialization. Precise versions of our architectures are described in Appendices B.5 and B.7. We compare with nanoGPT using standard initialization in Appendix D.4 to make sure our changes recover standard performance. We actually found unnormed Adam using our GPT architecture transferred learning rate better than in nanoGPT.
unit norm ball for our Embed
We found that normed optimization, with both Adam and SGD as the base optimizer, allows for successful learning rate transfer across width and depth for GPT training on OpenWehTPe (Figure 1), as well as ResMLP and ResNet training on CIFAR-I0 (Figure 4). We present expanded results in Appendix D.5, including results on test loss. We reproduce the standard finding that train and test loss are remarkably simillar in large language model pretraining. As for mass allocation, Figure 3 shows that optimal mass transfers with depth for training ResMLP on CIFAR-IO with normed Adam, and also that both mass and learning rate transfer quite well from a smaller GPT on Open WebText to larger one. We detail more experiments on mass allocation in Appendix D.6.
5 Discussion: Limitations and Future Work
This paper was influenced by four main streams of work: first, the Tensor Programs series, starting at TP-IV [3, 4, 18]; second, the papers universal majorize-minimizz algorithms [27, 28]; third, work on deep network metrization [12, 14, 31]; and fourth, the open source deep learning ecosystem [26, 33, 34] including the PyTorch module tree and Karpathy's YouTube video autograd [35]. We have distilled and synthesized key ideas from these sources, creating a framework that we believe to be simpler than Tensor Programs, computationally lighter than universal majorization majorization-miiiiizzaooo more general than prior work on metrization and more scalable than the PyTorch module tree. We have packaged these ideas into a (soon-to-be open-source library called Modula. Inevitably, Modula has limitations. We highlight some of them here, along with associated avenues for future work.
Loss of well-normed-ness, We have emphasized well-normed-ness (Definition 2) as an important criterion in module design. We show Appendix B. that, for example, the Linear module is well-normed when its weights lie within a spectral norm ball. In our experiments, we initialize all weights so that all modules are well-normed, but we do not enforce this property throughout training. Future work could explore regularization as means to enforce well-normed-nee throughout training, with the hope of attaining better scalability or improved generalization.
Overhead of normalization. As discussed in Appendix A.3, we implement normalization for Linear and Conv2D modules using two steps of online power iteration. While online power iteration an established and fast primitive in deep learning in fact, coming from the GAN literature [36]-ir it does add a modest overhead to training time, as discussed in Appendix A.4. We think it may possible to mitigate this overhead by constructing atomic modules with more exotic operator norms. For example, if one equips feature vectors with the 0 norm rather than the RMS norm, then the induced _ matrix norm is cheaper to compute than the RMS-RM operator norm. In fact, _

---

operator normalization has the convenient feature that it decouples over matrix rows, making it more local than spectral normalization and, dare-we-say, more biologically plausible.
Automatic step-size selection. Beyond scalability, recent work has explored the question of automatic learning rate selection [31, 37--9],, with the Prodigy optimizer [37] serving popular example. We tested the Adam version of Prodigy and found it performs well at small scales, essentially working by an implicit form line search. However, Prodigy will always break large enough widths, since requires a lower bound (dg) on Adam's 's initial learning rate; Yang et al. [3] showed that no such lower bound exists. We believe this issue could be fixed by rebuilding Prodigy on top of Modula. More broadly, we think that designing line search methods in a properly-norm space is good idea.
Acknowledgements
We are grateful to Chris Mingard, Virgile Richard and Evan Kiely for useful discussions early project. Tongzhou Wang and Jyo Pari provided helpful feedback on the writing and figures.
The work was supported by Packard Fellowship and Sloan Research Fellowship PI, by the MIT- IBM Watson AI Lab, by ONR MURI grant N00014-22-1-2740 and the MIT Quest for Intelligence. TL was supported by Simons Junior Fellowship. This work was performed under the auspices of the U.S. Department of Energy by Lawrence Livermore National Laboratory under Contract DE-AC25-07NA27344
We are grateful to the anonymous reviewers for their helpful and constructive feedback on this manuscript. Sadly, have not had time to integrate much of their feedback into this camera-ready version of the paper. We will integrate the feedback into the aiv version of the paper.
Contribution Statement
All authors were involved in project conception and discussions, which were initiated by JB. TL and JB developed the theory. MH and YL made core experimental observations. YL, MH, JB, and HB ran experiments. TL and JB did most of the writing, while JB, MH and YI. made the figures. PI contributed guidance and helpful feedback throughout the course of the project. JB wrote the Modula package with help from MH.
10

---

References
[1] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2015. Cited on page 1.
[2] Sashank J. Reddi, Satyen Kale and Sanjiv Kumar. On the convergence of Adam and beyond. In International Conference on Learning Representations, 2018. Cited on page 1.
[3] Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu al. Tuning large neural networks via zero-shot hyperparameter transfer. In Neural Information Processing Systems, 2021. Cited on pages 1, 2, 9, and 10.
[4] Greg Yang. Dingli Yu, Chen Zhu and Soufiane Hayou. Tensor programs VI: Feature learning in infinite depth neural networks. In International Conference on Learning Representations, 2024.
Cited on pages 1, 2, 9, and 21.
[5] Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin and Cengiz Pehlevan. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit. In International Conference on Learning Representations, 2024. Cited on pages and 2.
[6] Samy Jelassi, Boris Hanin, Ziwei Ji, Sashank Reddi, Srinadh Bhojanapalli et al. Depth dependence of yPP learning rates in ReLU MIPs. arXiv:2305.07810, 2023. Cited page 1.
[7] Lucas Lingle. A large-scale exploration of j--transfer. arXiv:2404.05728, 2024. Cited pages and 2.
[8] Hamza Fawzi. Topics in convex optimisation. University of Cambridge, Lent 2023. Lecture 3. Cited on pages 2 and 8.
[9] Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 2021. Cited on page 2.
[10] Haochuan Li, Jian Qian, Yi Tian, Alexander Rakhlin and Ali Jadbabaie. Convex and non-convex optimization under generalized smoothness. In Neural Information Processing Systems, 2023. Cited on page2. 2.
[11] Jingzhao Zhang, Tianxing He, Suvrit Sra and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. In International Conference on Learning Representations, 2020. Cited on page 2.
[12] Jeremy Bernstein, Arash Vahdat, Yisong Yue and Ming-Yu Liu. On the distance between two neural networks and the stability of learning. In Neural Information Processing Systems, 2020. Cited on pages? 2 and
[13] Michael Vernon Nelson. Gradient conditioning in deep neural networks. Master's thesis, Brigham Young University. 2022. Cited on page 2.
[14] Greg Yang. James B. Simon and Jeremy Bernstein. A spectral condition for feature learning. arXiv:2310.17813. 2023. Cited on pages 3, and 16.
[15] Nikita Dhawan, Sicong Huang, Juhan Bae and Roger Grosse. Efficient parametric approxi- mations of neural network function space distance. In International Conference on Machine Learning, 2023. Cited on page 2.
[16] Ari Benjamin, David Rolnick and Konrad Kording. Measuring and regularizing networks in function space. In International Conference on Learning Representations, 2019, Cited on page
[17] Behnam Neyshabur, Ruslan Salakhutdinov and Nathan Srebro. Path-SGD: Path normaloedd optimization in deep neural networks. Neural Information Processing Systems, 2015. Cited on page2 2.
[18] Greg Yang and J. Edward Hu. Tensor programs IV: Feature learning infinite-width neural networks. In International Conference on Machine Learning, 2021. Cited pages 2, 7, and 9.

---

[19] Jachoon Lee, Jascha Sohl-Dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz al. Deep neural networks as Gaussian processes. In International Conference on Learning Represeniations, 2018. Cited on page 2.
[20] Radford M. Neal. Bayesian Leurning for Neural Networks. Ph.D. thesis, Department of Computer Science, University of Toronto, 1994. Cited on page 2.
[21] Arthur Jacot, Franck Gabriel and Clement Hongler. Neural tangent kemnl:: Convergence and generalization in neural networks. In Neural Information Processing Systems, 2018. Cited on
page 2.
[22] Mufan Bill Li, Mihai Nica and Daniel M. Roy. The neural covariance SDE: Shaped infinite depth-and and-width networks initialization. In Advances in Neural Information Processing Systems, 2022. Cited on page 2.
[23] Aitor Lewkowyez, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein and Guy Gur-Ari. The large learning rate phase of deep learning, 2021. Cited on page
[24] Daniel A. Roberts, Sho Yaida and Boris Hanin. The Principles of Deep Learning Theory. Cambridge University Press, 2022. Cited on page 2.
[25] Chaoyue Liu, Libin Zhu and Mikhail Belkin. On the linearity of large non-linear models: When and why the tangent kernel is constant. Neural Information Processing Systems, 2020. Cited on page
[26] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury al. PyTorch: An imperative style, perf-peeformance deep learning library. In Neural Information Processing Systems, 2019. Cited on pages 2, 5, and
[27] Matthew J. Streeter and Joshua V. Dillon. Automatically bounding the Taylor remainder series: Tighter bounds and new applications. arXiv:2212.114299 2022. Cited pages 3 and 9.
(28) Matthew J. Streeter. Universal majorization-miiinizzzaoonn algorithms. arX1v:2308..0//99 2023. Cited on pages 3 and 9.
[29] Matthew Streeter. Beyond automatic differentiation, 2023, URL blog/beyond http://eyoonnuttttoottt-deeeeeee nuttomatic--diffeeetiiiie
Cited on page 3.
[30] Namhoon Cho and Hyo-Sang Shin. Automatic optimisation of normalised neural networks. 1067 2023. Cited on page 3.
[31] Jeremy Bernstein, Chris Mingard, Kevin Huang, Navid Azizan and Yisong Yue. Automatic gradient descent: Deep learning without hyperparameters. arXiv:2304.05187, 2023. Cited on pages 3, 9, 10, and 16.
[32] Dung T. Tran, Nobutaka Ono and Emmanuel Vincent. Fast DNN training based on auxiliary function technique. International Conference on Acoustics, Speech and Signal Processing, 2015. Cited on page 3.
[33] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary et JAX: composable transformations of Python+NumPP programs, 2018. URL http:///ttt/ com/ google/jax. Cited on page 9.
[34] Rami AI-Rfou, Guillaume Alain, Amjad Almahairi, Christof Angermueller, Dzmitry Bah- danau al. Theano: A Python framework for fast computation of mathematical expressions. 16Xiv:1605.02668 2016. Cited on page 9.
[35] Andrej Karpathy. The spelled-out intro to neural networks and backpropagation: Building micrograd, 2018. URL https: http://wwwwwwtbbeeee//// youtube. com/watch?v=VNi Bwwpwjj33Stttt -S1tkkuo. Cited on page
9.
[36] Takeru Miyato, Toshiki Kataoka, Masanori Koyama and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In International Conference on Learning Representations, 2018. Cited on page 9.
12

---

[37] Konstantin Mishchenko and Aaron Defazio. Prodigy: An expeditiously adaptive parameter-free learner. 2rX1::230600601,,, 2024, Cited on page 10.
[38] Aaron Defazio and Konstantin Mishchenko. Learning-rate- learning by D-adaptation. In International Conference on Machine Learning, 2023. Cited on page 10.
[39] Maor Ivgi, Oliver Hinder and Yair Carmon. DoG is SGD's best friend: A parameter-free dynamic step size schedule. In International Conference on Machine Learning, 2023. Cited on page 10.
[40] Kaiming He, Zhang, Shaoqing Ren and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. International Conference on Computer Vision, 2015. Cited on page 19.
[41] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (GELUs). arXiv:/606.0844 2016. Cited on page 19.
[42] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra and Christopher Ré. FlashAttention: Fast and memorr--eficieen exact attention with 1O-awareness. arXiv:2205.14135, 2022. Cited on page 22.
[43] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei et al. Language models are unsupervised multitask learners. Technical report, OpenAl, 2019. Cited on pages 23 and 27.
[44] Andrej Karpathy. nanoGPT code repository, 2022. URL https: / watauu..com/karpatyyy.. nanoGPT. Cited on pages 23, 27, and 28.
[45] Kaiming He, X. Zhang, Shaoqing Ren and Jian Sun. Deep residual learning for image recogni- tion. Computer Vision and Pattern Recognition, 2015. Cited on page 27.
[46] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009. Cited on page 27.
[47] Andrej Karpathy. Tiny Shakespeare. https: http:///hhggggaceeeeeeee http://ggnngfaac..../aaaaaaaaaaa cottee/etttttttttstrrtttttt////a siny_shakespeare, 2022. Cited on page 27.
[48] Ronen Eldan and Yuanzhi Li. TinyStories: How small can language models be and still speak coherent English? arXii:2305.07759, 2023. Cited on page 27.
[49] Aaron Gokaslan and Vanya Cohen. OpenWebText corpus. http://Skkylonn07.gg/h .itthu..
io DoeDeenNeeTTeetCopuss, 2019. Cited on page 27.
13

---

Contents of the Appendices
Appendix A The Modula Package
15
A.I
The Vector class
15
A.2
The Module class
15
A.3 Normalization in Modula
15
A.4 Overhead
16
Appendix B Module and Network Design
17
B.1 Atomic modules
17
B.2 Bond modules
18
B.3 Module broadcasting
19
B.4
Mass taring
20
B.5
Compound modules and neural networks
................................
20
B.6
Case study I: Attention
................................
22
B.7 Case study II: GPT
................................
23
Appendix C
More on Smoothness and Sharpness
24
C.I Underlying every estimate: The Gauss-Newton decomposition
24
C.2 Sharpness under composition and concatenation
24
C.3 Sharpness under module broadcasting
24
C,4 Smoothness estimates for common measures
25
Appendix D
Experimental Details
27
D.I Datasets
27
D.2 Architectures
27
D.3 Hardware
28
D.4
Comparing to standard nanoGPT architecture
................................
28
D.5 Full sweeps
29
D.6 Mass allocation
29
D.7 Context length
30
Full sweep results
30
Appendix E
Proofs
35
14

---

normalizes to have unit modular norm regularizes in-place Appendix A
The Modula Package
We created Python package called Modula that realizes our module framework in code. Modula supplements PyTorch's Tensor class with two new classes: Vector and Module.
A.1 The Vector class
The Vector class is used to store the weights of module. It allows for basic algebraic operations to performed on module weights without needing to write for loops over lists of tensors. For example, if v and as v_1 v_2 for the vector sum, or v_1 v_2 for the elementwise product. Internally, a Vector stores list of tensors and implements operations using efficient PyTorch foreach primitives.
are vectors with the same sub-structure then may write expressions such
A.2 The Module class
The most significant aspect of the Modula package is the Module class. A Module must have six attributes: two float attributes, namely mass and sensitivity. And four methods:
forward
Vector, x: Tensor) Tensor
returns an output tensor # randomly samples a weight vector
initialize() Vector
normalize (w: Vector)
regularizeww: Vector, strength:
float)
The norm of module is not specifically implemented, instead we use the normalize method which is how the norm is directly used optimization.
We refer to modules with hand-specified attributes as bonds they have no weights and atoms they have weights. Modules formed by combining existing modules called compounds. Modula automatically constructs the attributes of compound modules. We provide reference implementations for many common modules-set Appendix B. We equip atoms with their natural operator norm, and compute spectral norms via online power iteration. Reference modules may be imported as follows:
from modula. bond from modula. atom from modula. compound
import Identity, ReLU, Abs, FunctionalAttention import Linear, Embed, Conv2D import ResMLP, ResCNN, Attention, GPT
To make building compounds easier, Modula overloads the following operations modules:
M_2 M_1
# composes module M_2 with module M_1 # acts tuple module in any further composition returns the module sum multiplies module by scalar returns the Lth iterate of module
(M_1, M_2)
M_1 M_2
** L
For example, the compound
(L-) Identity() 1/L M()
builds an L-layer residual network from base module M. Comparing with Equation (3.10), we see that Modula expressions closely resemble their mathematical counterparts.
Finally, all modules come with a convenience method tare float), which resets the module mass to m, with default m=1.
A.3 Normalization in Modula
We normalize any base optimizer in the modular norm using the following pattern:
15

---

# get update from base optimizer # normalize update in the modular norm # apply update to weights delta_w optim(v.grad() net www.aoomlizeeeee (delta_w) (drmalize(delta__))
- lr delta_w
Computation of net normalize (delta_w) requires an efficient estimation of the spectral matrix norm, in the last two dimensions, of the constituent lensors of delta_w: this can be done very quickly to reasonable accuracy using power iteration. We implement this by storing running estimate of the top singular vector for each constituent tensor of delta_w. At initialization, is sampled Gaussian, and each time we normalize weight update, the previous update's ' estimated singular vector is used as the starting value for the power iteration. This enables us to use just two steps of power iteration weight update. Indeed, for any base optimizer with momentum, successive weight updates should be fairly close; for training without momentum more steps of power iteration may be required.
A.4
Overhead
To test the overhead of normalization in the modular norm, we trained width 64 ResMLP with 8 blocks and block-depti for 10k steps on the CIFAR- 10 dataset. We repeated the experiment with and without normalization, and in each case with three different random seeds. Without normalization, the training took 101 seconds, and with normalization the training took 124 seconds. So in this experiment, the overhead of modular normalization was around 23%.
We note that the of the Modula package is free write atomic modules with cheaper or more efficient normalize functions. For instance, the Frobenius norm can be used as a proxy for the spectral norm whenever the weight updates have low stable rank [14, 31]. And we note in $5 that one could explore more exotic norms such as the - operator norm, which is cheaper to compute than the standard spectral norm. Beyond these suggestions, one could explore CUDA-level optimizations to spectral norm computation, which is something that we have not explored.
16

---

Module M
M.forward
M.mass
M.sensitivity
M.norm
Linear
W, + dout Wæ
W + ||
Embed
E, 2. + vd Ex
E + max; |E 12
Conv2D
C
Lus C C
C + max;
||C.is]l++
Three atomic modules. These are the three atoms implemented in Modula Modula-enough to build ResNet and GPT networks. By including explicit dimensional scale factors in the forward functions, we are able to use the standard spectral norm ||-I|= and Euclidean norm ||-112, rather than their rescaled versions. da, and dout denote the input and output dimension of the Linear module. d denotes the embedding dimension of the Embed module. K denotes the kernel size of a Conv2D module with dout output channels and din input channels. @ denotes convolution.
Table 2:
Appendix B Module and Network Design
In this appendix, we list the basic, hand-declared modules that serve as building blocks for more complicated neural networks. Then go to show how these modules may be combined to yield interesting neural networks. This includes discussion of module broadcasting (Appendix B.3) and mass taring (Appendix B.4). The appendix culminates with case studies attention (Appendix B.6) and transformers (Appendix B.7).
B.1 Atomic modules
An atomic module or atom for short is module with nonzero mass and nonzero parameter space, whose attributes are specifically declared rather than derived. Setting an atom's 's mass to zero has the effect of freezing its weights under normed optimization.
Atom (Linear). For positive integers dout and d the linear module Linear(dout din) corresponds the standard linear layer with dia input features and dout, output features. module, it has input space X' = Rd.,, output space Rdaut and weights W = Rd x the space of dout x din matrices.
Its four attributes (forward function, mass, sensitivity, norm) given in Table 2. Note the presence of the / //utt/ dnn factor in the forward function: this convention means that can work with the standard 2 operator norm ||-||~ rather than the RMS-RMS operator norm.
Writing f (dout forwwad its derivative and second derivative at (W, ax) are given by:
V o (AW, Ax) / dout/dnn ((AW) + W (Aæ)),
(B.1)
AW Ax) V (AW AE / dout/dn dm (AW) + (AW) (Aw)
(B.2)
from which conclude that Linear(d (dour, dia) is well-normed, using the RMS norms on its input and output, so long as its arguments satisfy:
| |l, ||]||RM S 1.
(B.3)
These conditions will be automatically satisfied for many neural networks under orthogonal initial- ization of the weights, and especially if a linear module is immediately preceded by something like LayerNorm module. Moreover, orthogonal initialization guarantees that the well-normed inequality
I Ax ||Rss < ||RM
(B.4)
holds tightly in nearly-square matrices at initialization, which is important for getting good signal propagation through the whole network.
Moreover, inspection of the second derivative formula above shows it is always (0, 1, -hsarpp with respect to the RMS norms on the input and output spaces.
17

---

Atom 2 (Embed). For positive integers n and d, the embedding module (mbed d) corresponds to n class, token, or positional embeddings in a d-dimensional embedding space. As module, has input space IR,", output space I& and weights W = Rdxx the space of d x n matrices.
Its attributes are listed in Table 2.
This is at first sight similar to the linear module, the key difference being that in applications we expect the inputs of Embed(n, d) to be one-hot vectors; as such we consider its input space to carry the L -orm..
As for the linear module, Embed(n, d) is well-normed and (0, 1. 1.0rppharp with respect to the 1 norm on the input space R and the RMS norm on the output space Rd.
Atom 3 (Conv2D). For positive integers dout, din, K as well as H, W, the 2D-convolution module Conv2D (dout, din: K) corresponds are the number of channels for the input and output respectively (we suppress optional stride and padding arguments here for simplicity). Its input space is X' IR xW its output space is y = oou XH and its weights are W = IR dout x dia X K.X. x
convolutional layer with a K x K kernel; din and dout
Its attributes are listed in Table 2.
In fact, one could alternatively build Cdnv2D(dout- dia K) starting from K2 different Linear(dout, din.) modules (of mass 1/K each) and concatenating them, and composing with (parameter-less) convolution module. As such, Conv2D is nerl-normed and (0, 1, Oh-rpr... However, in our Modula package, we choose to explicitly declare Conv2D so as to take advantage of Pytorch's efficient implementation of convolution; the presentation here reflects this.
B.2
Bond modules
A bond module or bond is a module with zero mass and zero parameter space. They are the "glue" between the atomic modules, needed to construct complex neural networks.
Note that we need not specify a weight space, or mass or norm arguments for a bond module. Moreover, when discussing whether a bond module is (a, B, r)sharp,,, the inequalities for and are vacuous; thus for bond modules we will abbreviate this notion to y-sharp.
To begin, we need two bond modules that are essentially "utility", they are crucial for defining basic secondary module operations. These modules are also "type polymorphic" the sense that they work with any underlying vector space.
Bond (Add). For any vector space Y' the adder module Add has inputs y x y and outputs y. has forward function
Add. fod.forwwrdd (y1, U2) + y1 Y2
(B.5)
and sensitivity 1. Its significance is that it allows for concatenable modules to be auded:
M1 + M2 = Add o (M1, M2).
(B.6)
For any norm ||-lly on the vector space y, Add is well-normed with respect to the L combination norm (1 y2) byxy 11y Ily + || its input space. Furthermore, Add is O-sharp.
Bond 2 (Mulx). ). For any normed vector space )' and real number À the scalar multiplier module Mul, has inputs y and outputs V. Its forward function is:
forward y + x
(B.7)
and its sensitivity is [A]. Its significance is that it allows for scalar multiplication of modules:
X M Mul, M.
(B.8)
It is well-normed with respect to any norm V, and O-sharp. When À = 1, we call this the identity module Identity = Mulj Note that * Identity = Mul, for any X.
The remaining bond modules are used explicitly as non-linearities in neural networks. Bond (Abs). In any dimension d, the absolute value bond module Abs has inputs and outputs Rd.
forward function
Abs.forward (#1...., xx) + (|r1b ..., (sal)
(B.9)
and sensitivity 1. It well-normed for any norm 1dd
18

---

Bond 4 (ReLU and ReU), In any dimension d, we define the "rectified linear unit" bond module ReLU have input space X' R., output space y R forward function
ReLU.forrwrd (21,..., Ta) (0, xz xz)) ......dd
(B.10)
and sensitivity 1/v2. For this choice of sensitivity, ReLU is not well-normed with input space X set to the full IR However, it is well-normed if the imput space is, informally, set of dense vectors with balanced signs. For illustration, ReLU is rigorously well-normed with respect to the input space
X {signt for E Rd with # positive entries = # negative entries).
(B.11)
and RMS norm inputs and ouputs. For more this design decision, [40]. We also define ScaledReLU := /2 ReLU be the unit sensitivity counterpart to ReLU.
Bond 5 (GELU and ScaledGeLU). The "Gaussian error linear unit" bond module GELU [41] is essentially a smoothed version of ReLU. In any dimension d, GELU has inputs X' C Rd outputs ) = Rd and forward function
GELU.forward (I1,..., Ta) + (I,
(B.12)
where ((x) 1 /2x
/2 dl is the cumulative distribution function of the standard Gaussian.
GELU is well-normed in the same sense as ReLU. We similarly set ScaledGeLU = v * GELU.
Bond (MeanSubtract). For any dimension d, the mean subtraction module MeanSubtract has inputs and outputs Rd. It centers its input to have mean zero. The forward function is given by:
MeanSubtract.forwart (Z1...., z.) + (x1 ...... Zd +
(B.13)
and has sensitivity 1. It is well-normed, and since it is linear mapping, it is O-sharp. Bond 7 (RMSDivide). For any dimension d, the RMS division bond module RMSDivide has inputs and outputs Rdd. It normalizes its input to have unit RMS norm. The forward function is given by:
vdx
RMSDivide.forward x H
x
(B.14)
||I||RMS ||0|||2
and has sensitivity While is not automatically well-normed, as long as its inputs have ||a [allRss PV 1, the required inequality is not very far off. Similarly, it is approximately 1-sharp.
Bond (LayerNorm). For any positive integer d, the layer normalization bond module LayerNorm has inputs and outputs Rd, and is just defined as the composition of modules
LayerNorm RMSDivide MeanSubtract.
(B.15)
As with RMSDivide, it approximately well-normed and approximately 1-sharp.
B.3 Module broadcasting
Let us briefly discuss a supplementary module operation, which we refer to as module broadcasting, Definition 6. Suppose is module with inputs X, outputs y and weights W. Then for any Z of M is the module (th) with the same weight space W. mass, sensitivity and norm as M, but inputs the Cartesian power Xh X X and outputs yh = = x... Y, and forward function
the
(21,..,..., - (M.forward (w,
..., (MM....... .1),
M.forward( xn))
(B.16)
Since this is not defining module with a new set of weights, will usually just refer the broadcast module by the same name M, and consider this as just an extension of its forward function.
For example, this allows us define the action of linear modules Linear(dout, din) on inputs æ d to give outputs y Exdoa, where e is the context length parameter for transformer (see Appendix B.6, Appendix B.7, where it is also crucial for the construction of muldi-h attention). Additionally, one can view the basic Abs. ReLU and GELU modules being broadcasts of the usual one-variable functions to take inputs and outputs in Rdd
Let us briefly note:
19

---

Proposition 6. If M is well-normed, then so is any broadeast of M taking X to y", as long as the norms on X and yhh are taken to be either the "mean L norms
(7) (1 .... [
[[ (ir1......6)lll.
(B.17)
(y........ ynllly (1 +... 11.m
(B.18)
for 1 S S do; when oo this is just the max norm. In the case that M is bond module (so
W 0, any scalar multiple of the mean L norm can be used (including the standard L norm).
The situation for sharpness is bit more complicated; we discuss this in C.3.
B.4
Mass taring
In order make working with the mass parameter of modules bit easier, let us introduce an auxiliary operation:
Definition 7 (Tare). For any module M and positive real number Mnew, the module tare(M, mnew) has the exact same inputs, outputs and weights as M; the same forward function, the same sensitivity and the same norm; but has mass
tare M, Mmew) mass =
-aew-
(B.19)
This resets the of M. If M is a compound module, one could also reset the masses of all its submodules, by taking tare (M,, mnew Mass computation graph for tare meew). M.mass
for every submodule M₁, "reconstruct" the
This way, one can build complex modules starting from atomic modules with unit masses, and then using tare later to reset their masses to desired quantities for better feature learning with normed descent in Proposition 3.
B.5 Compound modules and neural networks
Composition, concatenation and the secondary operations of addition, scalar multiplication and iter- ated concatenation allow us to build a wide variety of neural networks which thus come automatically endowed with the modular norm.
Deep neural networks are typically built as long series of compositions. Let us introduce some terminology:
Definition 8 (Blocks and deep networks). A deep neural network is a module M formed by composition
= OutputLayer Block,, .... o Block: o InputLaaer
(B.20)
where InputLayer, Block Block,, OutputLayer are modules; the number of blocks L > is the depth of the network.
Typically, each of Block)..... Block,, will be copies of the same module (allowing them to take different weight values, of course), that the network can be written as iterated composition
OutputLayer Block InputLayer.
(B.21)
InputLayer, Block, OutputLayer be principle be any module likes, but usually InputLayer is often some form of embedding module, and OutputLayer is usually a linear module.
As for the form of Block, we found the following design principle to be quite useful in practice:
Arrange so that each Block has unit sensitivity.
This ensures that the sensitivity of the whole network stays bounded as I, -> oo (this will also be the case if we ensure that Block. sensitivity = 1 + O(1L), but unit sensitivity has the advantage that the modular norm becomes very explicit). With this in mind:
20

---

Compound (Residual network). Suppose that M is module of unit sensitivity whose inputs and outputs are the same space X. For any L 2 consider the residual block
Block L-1 L * Identity + 1/4 M
(B.22)
and write Res (M) Block* This is of unit sensitivity, well-normed if M is, and moreover by Proposition 4 is sharp with O(1) sharpness if M is.
A general residual network with residue M is any neural network of the form
OutputLayer o Res, (M) o InputLayer.
(B.23)
In practice, we will want to apply one more operation: we will want to tare the mass of the residual blocks. To this end, the residual network with residue M, depth and total block mass m i
OutputLayer o tare(Res((M),, m) InputLayer.
(B.24)
Let us give two basic example of residual networks. Compound 2 (ResMIPP.. This is simple residual variation the multi-layer perceptron. For width 2 1, consider the unit sensitivity module
M(d d) MeanSubtract Abs Linear(d, RMSDivide.
(B.25)
This particular order of operations is inspired by a reecent paper of Yang et al. [4].
We invite the reader to compare this to something like ReLU o Linear(d, d) o LayerNorm: three core operations are being performed (but in a different order in both cases): the inputs are being normalized; the inputs are being centered; and the inputs are passed through a nonlinearity that mutates just the negative coordinates.
The ResMLP network has as its residue an iterated composition of M(d), where the number of copies of M(d) in each residue is called the block depth and denoted B. It also has just linear initial and final modules. Thus the ResMLP network of depth L, width d, block depth B and total block mass m is
ResMLP = Linear(dout d) tare Fare(Res (M(d) m) Linear dm)
(B.26)
where din is the number of features of the data, and dout the desired number of output features the network.
Usually we suggest taking = 1 or 2, and m 1.
Compound 3 (ResNet). This is version ResNet for image classification tasks. For a width d 2 1 and kernel size K, consider similarly above the unit sensitivity module
M(d, K)- = MeanSubtract e Abs o Conv2D d, K) RMSDivide.
(B.27)
As in the ResMLP, the ResNet network is residual network with its residue an iterated compo- sition of B copies M(d, K) where is the block depth. Its initial and final modules are given by
InputLayer = Conv2D(d, Cn. K) OutputLayer Linear(dout. drotal) AvgPool
(B.28) (B.29)
where AugPool is an additional bond module implementing adaptative average pooling. Here, Gin the number of channel dimensions of the input image, diotal = d H W is the total dimension of the hidden representation, and dout is the desired number of output features (note that in Modula we include an additional dummy module Flatten to change the tensor shape before passing through the final layer). The ResNet network of depth L, width d, block depth B, kernel size K and total block mass m is thus:
ResNet = OutputLayer tare(Res;, (M(dR K) m) InputLayer.
(B.30)
As defaults, suggest taking B = 2, K 3 and m 20.
21

---

B.6 Case study I: Attention
Let us now focus on the construction of single multi-headed attention module in this framework. The attention module should have. as both inputs and outputs, A' = Rxx where & is the context length and is the embedding dimension. The attention module itself will depend on three additional dimensional arguments:
h, the number of heads;
dq. the key/query dimension;
dy, the value dimension;
as well an x matrix mask, which we usually take to be either
i j
maskij
(B.31)
-o0
otherwise
for causal attention, and mask = for non-casual attention.
The core of the attention module is a bond module which we call functional attention. Bond (FuncAttention). Take positive integers E, do. dy and mask matrix mask. The correspond- ing functional attention is the bond module of unit sensitivity, inputs Rexdq x Bexda Rexdy, outputs y = Rixdy, and forward function
gk dQ
iuncAAttention.forwarr k. = softmax
mask
v.
(B.32)
Moreover, we set FuncAttention.seesiiivii = 1.
In theory, one could try break up attention further into constituent more basic modules (such as scaled dot product, softmax, etc), but keeping FuncAttention as the basic unit one to leverage efficient implementations of attention such as FlashAttention [42].
In fact, perhaps surprising result that with the above do scaling of the dot product, estimate the sensitivity and sharpness of FuncAttention. This relies on giving norms for the input and output spaces; these norms chosen to be
(q k. l) = |] + ||M + ||o||ooRMS,
|| =
|o
(B.33)
where ||- RMS is the infinity-RMS-nu on iexd defined from the standard root-mean-squan norm [|-IIRMs on Id by
|| |a max :=1.....
||att ||RMS-
(B.34)
Proposition 7. Over the space of imputs q, k, with each || ||&|| RMMS |[ S 1, the functional attention module FuncAttention is well-normed, and moreover is sharp with sharpness constant 7 =3
The proof is given in Appendix We thus choose to adopt a - scaling in our imple- mentation attention- a rigorous bound as above is not possible for -calingg, for instance.
da
We can then immediately define single head of attention.
Compound (Single-headed attention). For positive integers E, d, do: dy and a choice mask, take four instances of the linear module, for the query, key, value and exit parameters:
Query Linear(do, d)
(B.35)
Key (mnear(d,,40 d)
(B.36)
Value Linear (d,
(B.37)
Exit Linear(d, (d, dv)
(B.38)
which by broadcasting we consider to have inputs shape exd The single-headed attention Attention module is then the composition
Attention = Exit
3 FuncAttention o (Query, Key, Value).
(B.39)
The scalar multiplication factor of & ensures that Attention has unit sensitivity.
22

---

For multiple heads of attention, we simply take advantage of module broadcasting (Definition 6):
Compound (Multi-headed attention). For positive integers e, d, h, de, dv and choice of mask, take four instances of the linear module:
Query Linear (h do. d)
(B.40)
Key Linear (h (4, d)
(B.41)
Value Linear (h dv, d)
(B.42)
Exit Linear h dv)
(B.43)
which by broadcasting we consider to have inputs of shape Rexd
The multi-headed attention
MultHHeadAttention module then the composition:
MultiHeadAttention = Exit FuncAttention (M) (Query, Key, Value)
(B.44)
where FuncAttention is broadcast over the heads dimension. Note that in Modula, we do this by creating dummy bond modules called AddHeads and RemoveHeads to reshape the tensors and create/remove the explicit head dimension.
As in the single-headed case, the scalar multiplication factor of ensures unit sensitivity.
B.7 Case study II: GPT
Let us now build an auto-regressive transformer similar to GPT-2 [43] or nanoGPT [44] in this framework. Fix positive integers E, d, h, do. dy (usually h divides d and do dv d/h). In addition to Compound 5 from the earlier, consider the 2-layer MI.P:
MLP = Linear(d, 4d) /2 GELU o Linear(4d,
(B.45)
where we are using the scalar correction so that GELU has unit sensitivity, and using module broadcasting so that it take inputs and outputs gexd Fix a depth L 2 1, and consider the following two modules, whose input and output spaces are Réxd;
BlockMLP == 2L-1 Identity + 1
2E MLP LayerNorm.
(B.46)
2L
BlockAtto == 2L-1 2L
Identity + 2L
1
MultiHeadAttention o LayerNorm.
(B.47)
where LayerNorma refers to taking LayerNorm in the embedding dimension (i.e. the rows of matrices in f distinct from normalizing all x d coordinates together). This alternately be thought of as just taking the usual LayerNorm on Rd and broadcasting it to take inputs and outputs x
Suppose that N is the number of tokens. For the initial module, take two embeddings of the N tokens and context positions
Embedbok Embed(N, d),
Embed Embed
(B.48)
and form the mass one, sensitivity one module
InputLayer = tare( Embedrok + & Embed 1).
(B.49)
The final module is just
OutputLayer Linear (N, d) LayerNorm
(B.50)
The depth L 1, width d, total block mass > 0 GPT module is thus
GPT OutputLayer tBroc(Blo Blockam)) m) InputLayer.
(B.51)
We suggest, as default value, total block mass of 5.
23

---

Appendix C More on Smoothness and Sharpness
C.1 Underlying every estimate: The Gauss-Newton decomposition
All our estimates of sharpness for compound modules, as well as the smoothness estimate Proposi- tion 5 for loss functions, depend on an application of the chain rule compute second derivatives which in the optimization context is sometimes called the Gauss-Newton decomposition.
Namely, if f de Rd and Rd = g f is computed by
Rd then the second derivative of their composite
e (V (Vooo) Vgolo
(C.1)
for any v, w RRd for short
2²h(s-) = vig(Vf())) + Vg(3ff()))
(C.2)
Indeed, this amounts to simply the following expression for partial derivatives:
8hh
8gg 05k 0. Oy dx; dz;
Og 8 fx
E
+ Σ
(C.3)
dr,ōr,
Tyk dx,0r,
C.2 Sharpness under composition and concatenation
Here, we state the two formulae for computing the sharpness of a composition and a concatenation of two modules. The proofs are given in Appendix E.
Proposition 8 (Sharpness under composition). Suppose that M2 and M1 are well-normed, compos- able modules that are respectively (a12, $2, 72) and (0n,,) 11,77 7 --sharp. Under the shorthand Mass and μ =\t M2 sensitiviv,,,, the composite M2 o M1 is (a, ---sparp for:
that pk My .msss Ma
a prz ua + +2202 + PIPP23 + P2
PiT
(C.4)
3 p1 + /1P2/2 + ng P172:
(C.5)
7 127 + 1272.
(C.6)
Proposition (Sharpness under concatenation). Suppose that M1 and M2 are well-normed, concate- natable moduley that are respectively 31 )sharp and (az, B2, 72)-sharp. Under the shorthand that Pk
Ma.mass mass My ...masst,,maass and μ4 = M2 sensitviiv the tuple (M1, M2) is (a, 7)-sharp for:
a pιan + 1502.
(C.7)
8 p1 + P222,
(C.8)
71 + 7
(C.9)
Taken together, Propositions 8 and 9 specify recursive procedure for computing the sharpness of any compound module that is built from a set of well-normed modules known sharpness.
Remark 1. These two sets formulas are actually associative, as the reader may verify using their favorite computer algebra package. This means, for instance, that if M1, M2, M2 are successively composable, well-normed and each (ax, 3k, 7 -s)-sharp. then the two sets of sharpness estimates coming from applying the above formulas for M3 (M2 M,) and (Ma M2) M1 actually coincide.
C.3 Sharpness under module broadcasting
Suppose M a well-normed module with inputs X, outputs y and weights W, and suppose moreover that it is (a, B, -hhhpp The broadcast module (A) has the same weights, mass, sensitivity and norm, but takes 2 to )ah.
By Proposition 6, ((h) is well-normed, long as the norms on x and th are taken be
[[ S * (||=1 1ix + + || ) 1/p
(C.10)
II(y,.... yn)llyn S (1/371 Bb + + 11:) 1/p
(C.11)
24

---

with the for 1 S I bo; unless M is bond module (and thus weight-less), we must take S = h h11, otherwise S can be any positive scalar.
A natural question is whether M(A) is also sharp, and if so what its sharpness constants are, with respect to these norms. More or less the same proof as for Proposition 6 shows that the a and B bounds for sharpness depends subtly on the chosen p. We highlight three cases where one can say something interesting.
always true, with the same a, B. The 7 bound is trickier however, and
Case 1. p o0, 1. For the 1 norm, have that ( is (ar, B, -sharpp with the same a; 8 by a more or less immediate proof.
Case 2. p < 00,
For the "standard" Io-norms, we have that M(A) is same a, 3, % The proof is direct, using the inequality
+ + 2 s (of + +...
(C.12)
for any positive reals Ii...... very pessimistic sharpness estimates for large h.
..... Tn. however this a very weak inequality and so leads to ......,,............... Th:
Case 3. p 2 = 1/Vh. This is the "RMS norm" case. As in Case 2, one could use very weak inequality to obtain the pessimistic result that M is (a, 8, vh 7)-shprp. However, one could also make the following observation: if h is large, and 21.... 24 are sampled from any normal distribution N(,, )²), then
(H
1/2
v3 v3 (oi+.................. x
(C.13)
In particular, this justifies the statement that "for large h, the broadeast module M approximately (a, B, /3 * >)-shprp" While in actual deep learning contexts, the assumption that Zi...., Th are sampled from normal distribution may not be valid, two sides of Equation (C.13) to stay O(1) h co, and even if the "V3 rule" is insufficient, the effective sharpness of the broadcast module should not blow up as + oc.
should still expect the ratio between the
C.4 Smoothness estimates for common error measures
Suppose é: y on together with sharpness of a neural network imply smoothness of the corresponding average error loss function. The precise estimates are that & is L-Lisschitz and 7-smooth in the module output, in the sense that:
R is an error measure. In Proposition 5, showed that smoothness estimates
vyl(y,t) Ayl S a |laylyy
for all Ay E ) and E T; for all Ay, Aü y andteT.
(C.14)
Ay V (v. Ayl S | | الaلly
(C.15)
We now present estimates on a and 7 for square and cross-entropy error. Both estimates will be in terms of the value of the average loss function L itselff rather than being truly global over the entire output space Y. Thus, to apply them to real learning problems, one should measure the average loss at initialization, and use this for estimates for a and т; we are implicitly making the assumption that under gradient descent the loss decreases.
Square error
Consider square error for a d-class classification problem. Thus, y R2d and T = (I. Consider the RMS norm V, and define the error function
fyy t) = 2d vi .
+ (ve
vd)² +...
for y. E y x T.
(C.16)
(the slightly non-standard scalings are due to the choice of RMS norm on y). The first and second partial derivatives of are given by
ai Dy:
828
1
(y. 1) _(11
5% VaVa)
Dy. (y. 1) 5
(C.17)
25

---

The desired constants 7 can then be computed as maxima:
de
8
max
Σ
Ziz
max
θy,ūyyj 2,25
(C.18)
||
By.
||*|||Ms
which from the above formulas amounts exactly
T=1. To translate this into bound for the average loss function L, note that square root is a concave function. Thus if we have outputs yı ,I......3 yB with true classes ......, tp Jensen's 's inequality yields
v V(y,t),
(C.19)
1
E
Vely,, to) S V B Σ (y) vD
(C.20)
B
allowing us use VL as our estimate for Proposition 5.
Cross-entropy error
Consider cross-entropy error for a d-class classification problem. Thus, y Rd and T 1,...,d. .d. For y Rd and t E T, write
ev
pe(y) Eee
(C.21)
and consider the error function
f(y, t) = - log(p:()))
(C.22)
The first and second partial derivatives of l are given by
828 dy:U5
ay: (y, 1) = Sic,
(y,t) = 6tpPt PiPj-
(C.23)
Consider again the RMS norm on y. An estimate can thus be computed as
1/2
max Σ de ||=||
[_ 8.)) S Vdlve
By: 2, V
(C.24)
using the basic fact that if P..... Pd are non-negative numbers that sum to 1, then
p +.... (p: 1)2 + + - log((t))
(C.25)
(Indeed, for fixed Pt, the left hand side is maximized p1 Pe and other p; = 0; one then easily checked that 2(p - 1)² < - log(p) for all 0 < < 1.)
A similar concavity argument to the square error case then enables us to use a = vd * VC as the first derivative bound for average cross-entropy loss.
The second derivative bound depends on more subtle information geometry. Indeed, r can be computed to be
T ==
(C.26)
where À is the largest eigenvalue of the matrix diag(p) pp" It is possible for this eigenvalue to be quite large (for instance, if p1 = = 1/2 and all other = 0, then X = 1/2). However, the average eigenvalue is
<1-Σ
d2 <
(C.27)
we presumed that, in the course of gradient descent optimizing the weights of a module M, the output perturbations VM Aur are only generically aligned with the eigenvectors of diag(p) pp". then we could use the "effective" smoothness bound
Perhaps this is a dubious assumption however. A more conservative, but perhaps still dubious, assumption comes from assuming that the logits y have roughly (O, 1) entries at least this could be more or less true at initialization. In this case, the largest eigenvalue À is with high probability bounded
A A 1/vd justifying "approximate" smoothness bound of vd.
(C.28)
26

---

Adam nanoGPT
Adam modulaGPT
normed Adam modulaGPT
width
32
64
Test loss
128
1024
6
blocks
2
test loss
4
8
16
32
2-
64
10
10-2
10
10
10
10
100
learning rate
Figure 5: Comparing to a standard transformer implementation. Since we used our own well- normed GPT implementation for the experiments in this paper (here referred modulaGPT) wanted to check its performance was on par with a standard nanoGPT impiementation. These plots show learning rate sweeps for varying width and depth for Adam nanoGPT, well as Adam and normed Adam on modulaGPT. Even without normed updates, the architectural changes and orthogonal initialization used in Modula seem to already improve transfer compared to nanoGPT.
Appendix D Experimental Details
D.1 Datasets
All experiments with ResMLP and ResNet [45] are done with the CIFAR-10 [46] image dataset with standard train and test splits. For data augmentation the training use random crop, random horizontal flip and PyTorch AutoAugment.
For the GPT [43] transformer experiments, compared three different datasets:
(a) The Shakespeare corpus, using character-level tokens [47];
(b) The TinyStories database [48] using sub-word level tokenization;
(c) OpenWebText using sub-word level tokenization [49].
No data augmentation was used on the language data. We used data splitting code from [44].
D.2 Architectures
Full details of the ResMLP, ResNet and GPT architectures we used are detailed Appendices B.5 and Appendix B.7. In every experiment, we used:
(a) cross-entropy loss with no weight decay;
(b) block depth B = for ResMLP and ResNet;
(c) kernel size K for ResNet:
h heads for GPT, with query and value dimensions dg dv d/h where is the embedding dimension (width);
(c) context length 128 for GPT, except for the experiment in Appendix D.7.
27

---

free mass
mass=0.25
mass=0.5
mass=
mass=2
depth
train loss
1.8
4
1.6
16
32
10
10
10-1 10
10
10
10
10
10
10
learning rate
We train a ResMLP with width 64 and 2 layers per block on CIFAR-10. In the first sub-plot titled "free mass", we set every atomic module to have unit mass, so that depth is scaled the masses of the imput and output layer become insignificant relative to the total mass of the hidden layers. In the other four subplots, we tare the total mass of the hidden layers the value indicated the subplot title. As can be seen. the taring strategy seems to work much better than the free mass strategy. So, at least in this experiment, it is good to keep constant fraction of learning in the input and output layers even as depth scaled.
Figure 6: Comparing mass allocation strategies.
D.3 Hardware
All experiments were run on NVIDIA GPUs using (roaat32-precision We used a combination of TITAN-RTH RTX-3090, V100, Ada6000, and H100 devices. Each data point in the experiments takes up to hours, depending on the computing device used. We ran over 1000 training runs in total.
D.4 Comparing to standard nanoGPT architecture
Our implementation GPT in Modula has certain differences from off-the-shelf architectures such as nanoGPT [44]. We would summarize the overall changes to transformer architecture and training following three points:
(I) the mathematical architecture has slightly different coefficients;
(II) initialize weight matrices to be orthogonal rather than Gaussian;
(III)) train using normalized weight updates.
The architectural choices we made were entirely informed by the desire for the network to be well- normed and have unit sensitivity: in particular this means that the network enjoys favorable signal propagation properties. In the language of modules, these architectural changes can be summarized as:
(a) Each residual block in architecture is of the form
21-1 2L
Identity + zz Block
(D.I)
where Block = BlockMLP or BlockAtt.- compared to Identity + VL
Block suggested for
nanoGPT;
(b) We use scaled dot product attention with 1 scaling, rather than
de
vdo
(c) The forward function our Linear and Embed modules (see Appendix B. includes scale factors / din and vd respectively.
(d) We use several additional scalar multiplications to keep the network unit sensitivity:
Each Attention module (B.44) has a scalar factor t
Each MLP module (B.45) has scalar factor of /2:
The token and position embeddings (B.49) have scalar factor of 2-
In Figure we ran a comparison of the performance of the standard (unnormed) Adam optimizer trained on OpenWeeText with:
28

---

Shakespeare
Tiny TinyStories
Open WebText
6
20
3
15
test test loss
4
10
2-
5
2
10
100
10
10
10
10
learning rate
Figure 7: Mass and learning rate sweeps across datasets of increasing diffiiulty. small GPT architecture of width 128 and 3 transformer blocks trained on the Shakespeare, TinyStores and WeenWe datasets. We varied the learning rate as well as the total mass of the blocks. Optimal mass and learning rate well from the much smaller Shakespeare dataset.
transfer reasonably well from TinyStories to OpenWebText, and less
1. the nanoGPT architecture with Gaussian initialization;
2. our implementation of GPT with orthogonal initialization.
We found that even without using the normed optimizer, our implementation with orthogonal initial- ization transferred learning rate better. We suggest that even the base Adam optimizer benefits from the above architectural changes.
D.5 Full sweeps
In Figures 12, the end of this Appendix, we report on full learning rate sweep experiments, across width and depth, for GPT on OpenWebText and TinyStories, and ResMLP, ResNet CIFAR- 10.
We consistently find that the normed Adam optimizer matches or outperforms unnormed Adam in both test and training loss, all the while exhibiting significantly better transfer across width. The difference in depth transfer is less stark, however we posit that, in part, unnormed Adam is already benefiting from architectural changes we made to improve depth scaling.
Notice too that normed SGD consistently significantly outperforms ordinary SGD, often coming close matching the performance of Adam. We would like highlight this, since SGD has significantly lower memory requirement than Adam, and does not require any tuning of 32.
D.6
Mass allocation
novel feature of our normed optimization framework is the need to choose a mass parameter for each atomic module. In the context of networks of the form
Network OutputLayer HiddenLayers InputLayer
(D.2)
where HiddenLayers = Block We typically do this by assuming that InputLayer, OutputLayer have mass 1, and by hand resetting the mass of HiddenLayers to be a fixed total mass 0, by calling tare Lardddiiden,,,ers, m)).
In this Appendix, explore some different aspects the choice of
First, we tested whether or not calling tare necessary in the first place. Not using tare would leave the "free mass" of HiddenLayers.mass mase = L Block mass; accordingly as grows large, the feature learning allotment (see Proposition 3) for InputLayer and OutputLayer would grow smaller. Indeed, as the reader can see in Figure 6, this "free mass" arrangement for a ResMLP network on CIFAR-10, 10, allowing the mass of HiddenLayer transfer with depth must fix mass.
grow with I, is very undesirable, and for good learning rate
29

---

8
context length
7
32
64
test loss
6
128
256
5
512
4
1024
3
10
100
learning rate
Figure 8: Context length transfer. We trained GPTs of various context lengths using normed Adam. As can seen, learning rate transferred quite well across context length.
The mass is thus left tunable parameter. We then tested the transferability of mass tuning. Specifically, we wanted to know:
1. whether one can tune m on network of small width/depth, and expect that same m be close to optimal on a larger network;
2. whether learning rate transfer across width/depth is itself dependent on selecting a good mass m
3. how sensitive the tuning for m is: if there is a broad range of acceptable masses, certain precise values lead to big improvements train or test loss.
Figures and answer Question 1 aabve in the affirmative, in the context of ResMLP CIFAR-10 and GPT on OpenWetText Moreover, in the context of ResMLP on CIFAR-10, they give an answer of Question 2 and Question 3: learning rate transfer occurs at range of values of m.
Figure address Question the context of transformers, on three different datasets. Across all three datasets, mass in the region ~ 5 to 10 is reasonable.
D.7 Context length
Additionally, we also tested the dependence the optimal learning rate for GPT training on Open- WebText on the context length; the results are in Figure 8 Interestingly, we report good transfer of the optimal learning rate from small contexts long contexts.
D.8 Full sweep results
The next four pages of the Appendix list results of our full learning rate sweeps over width/depth for GPT on TexTee and TinyStories, and ResMLP, ResNet CIFAR-10.
30

---

SGD
normed SGD
Adam
normed Adam
7.0
width
32
train loss
6.0
64
128
256
5.0
512
1024
4.0
7.0
depth
train loss
6.0
16
5.0
32
64
4.0
10
10a
10-1 1
10
10
10
step size
step size
step size
step size
SGD
normed SGD
Adam
normed Adam
7.0
width
32
6.0
64
test loss
128
256
5.0
512
1024
4.0
7.0
depth
6.0
test loss
5.0
32
64
4.0
10
100
10
10
10
3
10
10
100
step size
step size
step size
step size
Figure 9: Learning rate transfer for GPT on WeenWeeText. Training is done for 10k steps, at batch size 128, with SGD, Adam, and their normed versions. The total block mass for normed SGD/Adam is m = Width scaling experiments are done at fixed depth and depth scaling experiments are done at fixed width 128.
31

---

SGD
normed SGI
Adam
normed Adam
width
5.0
32
train loss
4.0
64
128
256
512
2.0
1024
100
depth
5.0
train loss
3.0
16
32
2.0
64
1
10"
10a
10
10
10a
10
2
10
10"
10a
step size
step size
step size
step size
SGD
normed SGD
Adam
normed Adam
width
5.0
32
64
test loss
4.0
128
3.0
256
512
2.0
1024
depth
5.0
test loss
4.0
3.0
32
2.0
64
m807
mm
10
100
10A
10
10
10
10
10 2
10 10
10
step size
step size
step size
step size
Figure 10: Learning rate transfer for GPT on TinyStories. Training is done for 10k steps, batch size 128, with SGD, Adam, and their normed versions. The total block mass for normed SGD/Adam m Width scaling experiments are done at fixed depth 3, and depth scaling experiments are done at fixed width 128.
32

---

SGD
normed SGD
Adam
normed Adam
width
32
1.8
train loss
64
128
1.6
256
512
1.4
1024
depth
1.8
train loss
1.6
16
1.4
32
64
10
10a
10
step size
step size
step size
step size
SGD
normed SGD
Adam
normed Adam
width
32
1.8
64
test loss
128
1.6
256
512
1.4
1024
depth
1.8
test loss
1.6
1.4
32
64
10
100
10A
10
10 3
10
10
step size
step size
step size
step size
Figure 11: Learning rate transfer for ResMLP. ResMLP architectures on CIFAR-10 are trained
for 10k steps, at batch size 128, with SGD, Adam, and their normed versions. The total block mass for normed SGD/Adam m Width scaling experiments are done at fixed depth 3, and depth scaling experiments are done at fixed width 128.
33

---

SGD
normed SGD
Adam
normed Adam
1.5
width
32
train loss
1.0
64
128
256
0.5
512
1024
1.5
depth
train loss
1.0
16
0.5
32
64
10P
10a
10 100 10a
$10 210
$10 10a
step size
step size
step size
step size
SGD
normed SGD
Adam
normed Adam
1.5
width
32
64
test loss
1.0
128
256
0.5
512
1024
1.5
depth
test loss
1.0
0.5
32
64
10
10"
10 100 10
10 10 210
10 100 10' step size
step size
step size
step size
Figure 12: Learning rate transfer for ResNet. ResNet architectures on CIFAR-10 are trained for
10k steps, at batch size 128, with SGD, Adam, and their normed versions. The total block mass for normed SGD/Adam is m 20. Width scaling experiments are done at fixed depth 3, and depth scaling experiments are done at fixed width 128.
34

---

Appendix E Proofs
Proposition 3: Feature learning is apportioned by mass
To prove Proposition 3, it suffices to induct on the construction of a compound module by composition and concatenation, with the atomic modules (where the inequality is just part of well- normed-ness) as the base case.
Indeed, suppose either M M2 M1 or M (M1, M2). Suppose wy is weight for one of the atomic modules of M, and write m for the mass of this atomic module. Then wy must be weight of either M M2; the inductive assumption is that
m. M, Awxxl
m
|
(E.
Ms.mass
where = 1 or 2 accordingly.
Case 1. M = M2 М₁ and wy is a weight of M₁. From the chain rule we then must have:
V M A
VV, o w. M1 Awwll
(E.2)
v My.sensitivity * V M1 o Aww
by well-normed-ne
(E.3)
< Mg.sensitivity * Mi.mass
m
AwilM
by assumption
(E.4)
|
(E.5)
M.mass
where the last line is by the definition of the norm of module composition.
Case M2 M, and wa is weight of M2. The chain rule is needed in this case, and we proceed
straight from the inductive assumption:
V mn M Aww
|| w. M2
Awrk
(E.6)
m
|
(E.7)
Ma-mass
m
A
IM-
(E.8)
M.mass
Case M (M1, M2). Given the symmetric roles of M1, M2. without loss of generality assume wk is a weight of M1. Then,
V mn M Awall
|| m, Mi Awxl
(E.9)
Aw [M;
(E.
Ma.mass
m
Aw IM-
(E.11)
M.mass
This completes the proof.
Proposition 4: Sharpness of residual networks
Suppose M a well-normed module unit sensitivity on (X, X, W) and (a, B, 7)-sharp. Then, by Proposition for any L 2 1, the module * M is well-normed, sensitivity 7. and (Lar B, --harp
Identity is also well-normed, sensitivity -_ب and (0, 0, 0)-sharp. In particular, The module -_1 the sum
Mfes _ Identity + + MM well-normed, unit sensitivity, and (Lo, 3, 1 -hharp: has the same mass as the original module.
(E.
We induct on the statement for 2, ... that M (ax, 8k, -h)-sharp where
L
2(1 + 2 +.... + (k 1)) B + (12 2 + +...4 (k: 1)2)
aa-
(E.13)
k
k2
LA
(1+2 + + (k 1))
3x B
(F.14)
Lk
k
L 7
(E.15)
35

---

The base case is clearly true, and given the statement for M which has exactly & times the mass as is (0 S641, Tk+1 1-+trppharppp by applying Proposition 8 with
Mres. we see that MK+1 = Mres M res
P1 k+1
and P2 = k+1' where
k2
(k + 1)2 La 2k
k2
ac+1 =
(k 1)2 a& +
(Is + 1)2
L(k 1)2?
(E.16)
k
k
Bk+1 =
Bk + k+1 B + L(k 11)
(E.17)
k + 1
1
78+1 = 7/k T. 7.
(E.
which yields the induction. Setting k L, observe that 1 1 ... +L 1) = L - 1) and + 2 +....4
. (L - 1)2
L 1) - 1), giving
- 1
L(L - 1)(2L 1)
= a +
A 3
S a 5 3
(E.19)
L
673
1
3, - B +
7
S 3 27
(E.20)
2T.
1L =
(E.21)
which proves the result.
Proposition 5: Smoothness in the modular norm
To establish the first inequality, start by applying the Gauss-Newton decomposition (C.1) of the Hessian, followed by the fact that the error is o-Lisschitz and 7-smooth, followed by the well-normedness and (a, 3...,,.pspeeeess of the module M:
Aw V www L Aw
(E.22)
Em V (Aw V Mo. Aw) (V Mo V V
(v Ae)
(E.23)
F. yDD [o Au V www M Awily
M Awlly
V Aw by]
(F.24)
< (oa + +) | M |
(E.25)
The second inequality follows from the first via the fundamental theorem of calculus:
Ww + Aw) - V ((w) = max
Vw + Aww [wL(w) о Aw
(E.26)
|A =1
( Aww V C(w + tdw)
max
Awl dt (E.27)
|| d =1
(aa + r) ||ww|| A / di
< max
(E.28)
A
= (oa + T) |
(E.29)
The third inequality follows from the second by again applying the fundamental theorem of calculus, followed by the Cauchy-Schwarc inequality:
[(w + Aw) - [L(w) + Vw o Aw]]
(E.30)
LwC(w +dw) [ww(w) Aw df
(E.31)
/ [|VoL(w iAw) v.n liM || dt
(F.32)
< (aa + +T) ||A lM
dd
(E.33)
= 3 (oa + 7) |uu 2-
(E.34)
This completes the proof.
36

---

Proposition 6: Broadcast modules are well-normed
Suppose M is module with inputs X, outputs Y and weights W, broadcast take X to ) We take norms on these spaces to be
(1 , mb)ll S (1|æ) 11xx + ... + || /
(B.35)
[ ........ vnlllan S (1lv: [| + +... 11 [13))
1/p
(E.36)
where = unless M is a bond module. Write μ = M.sensitivity. Then, for perturbations in the weight direction, which only occur if M is not a bond module:
1/P
M(w, Ti,.... 2I.......... 2a) o Awllyn
Σ
M(w, xz) e Awlly
(E.37)
S ||wl|M
applying well-normed-ness.
(E.38)
For perturbations in the input direction, we have:
1/8
| V vn M (AæI...... Axa)llsn S
Σ
Va, Mo Mo
(E.39)
1/P
SS* *
En 113
(E.40)
|1(aæ Axta)lo"
(E.41)
which proves the proposition.
Proposition 7: Sensitivity of attention
We prove that the functional attention module FuncAttention of Bond is well-normed and of unit sensitivity.
Recall we use the following norms on the inputs X' -Rexda Rexxo y' = Rexdv;
Rxdy and outputs
(1(q, k. v)llx || || + ||
||ally || |oolss
(F.42)
We will also make use of the - -pperatoo norm for l x l matrices, which we write as
Σ Byyl
oo = max
(E.43)
== .......
observe that for B e Ex and x E Igdd we have
B a RMS I
| oo ||||||/oloo |ooRRs
(E.44)
Writing FuncAttention.forwarr for short, recall that
F(q, k, v) softmax( 2 qk. + M)v
(E.45)
where M is the mask (our proof will apply equally for the standard causal mask and also the non-causal M 0).
We will prove that any (q. v) satisfying ||q|] |||| || S for any (Aq, Ak, Av) we have
|V k v) o (Aq, Ak, Au)lly s ||(gg, Ak, Av)(x
(E.46)
For short, write A = softmax gkt M) for the attention matrix and its derivative as
A4 = (g.k) softmax( 1 qk + M) (Aq. Ak).
(E.47)
da
37

---

Now, the derivative of F splits into two terms
V (Aq. Ak, Av) = A(Av) + (AA)e.
(B.48)
To complete the proof, we claim that coc--
and ||AA S IlggllooMMs + ||
(F.49)
The calculation of the norm of A follows by definition from its construction by softmax. For the calculation of the norm of A, direct calculation yields that
A = पे An Ag. ky Aimk km) + do Ai (qi, Akj 2mm Aim
(E.50)
where we are writing q, = qi= and so on. Taking absolute values, applying the Bauchy-Schware inequality and summing over j we have
A Ae |IRMS (E) [k; Em km (RRS)
(E.51)
+ ||
|RMS (E, [E,A;;1AAA
Em Aim A km [RMS).
(E.52)
We now use the following inequality: given any non-ne non-negative reals Pi..... Pe which sum to 1, and : Fy any vectors 21,...;3
in inner product space with norm ||-||, we have by Jensen's inequality
Σ,P,l 1|a
Empm Lm S Eopbllag - EmPma 112)
(E.53)
E,ppllz3 112 1|E 211²))
(B.54)
(E ||=,,
(E.55)
< max ||
(B.56)
Applying the matrix AA, we thus have
Lg & [Aq.IlRMsss max |ky ||MMS + |RRMS max || |RMS.
(E.57)
3
Taking the max over this shows the 0 -orratoo---ommm of A most
booMSs RMS + RMS R
(E.58)
which, since [] sRMs
|k
RMR S 1, completes the proof.
Proposition Sharpness of functional attention
In this section, we estimate the second derivative of the forward function F of functional attention at (g, k. v) in perturbation directions (Ag, Ak, Av) and (Aq. Ak, Aö):
2 = (Aq, Ak, Aq) o 2 F o (Aq, Ak, Av).
(E.59)
We will prove that functional attention is s-sharp where in fact 7 = 3; this amounts to proving that
2 < 3|| [Qq, Ak,
Ak, Aö)
(E.60)
We continue with all the notation of the previous section. Moreover, to simplify the calculation, we suppress all factors of dg (indeed, one absorb them rescaled inner product (,, +)). We also, in addition to the shorthand æ; = T. for x matrices æ, we adopt the shorthand for an é matrix B and a f x d matrix æ, and any i,j = 1....,.. C::
B. B := x2 Em B Bmm
(E.61)
We note three crucial inequalities regarding [B .... any & E matrix B with non-negative entries whose rows sum to I, and & x matrices x, y:
Σ, Bi
B, æ| < max |
(E.62)
Σ, B.y
BB x و|1 11 < max |z,1|;
(E.63)
B; [B, x) ||I| v الl (max || (r)]l)((axx ||
(E.64)
38

---

All three inequalities follow from standard expectation/variance inequalities for random variables on the finite {I....,{} with distributions given by ........ By.
With these conventions, the expression for A A is thus
A (ag [A,
+ Auya. [A, Ak))))
(E.65)
Let us also write
A4 V (q.b) softmax(
do qk
M) o (Aq, AE)
(E.66)
A = Aij(A [A, kliiis Asy(:: [A, Ak)
(E.67)
well
A = (Áq. A V (Aq. Ak).
(E.68)
In these terms, the second derivative A F just
F (Av) (AA)(AT) + A)v.
(E.69)
From the estimates of the previous section, we have
[| [(AA) RMS S Al DoRMS + | (aRMS) Av RRS
(E.70)
A| AD) RoRMS - (1) A bORMS + | A
ccRMS
(E.71)
so our task is to estimate the L -p=-opprrtoornnrmm of A Thus, we calculate A=A:
Ag = AAq [A, Akja)
(B.72)
+ An A, Ak)iss)
(H.73)
+ A A (Aqi Alis)
(E.74)
AA (gin [A, Aklist
(B.75)
Aij(Aq: Emm (A A im kim)
(E.76)
Aij(g; -E,, (AA) Akim)
(E.77)
We estimate the - -pooprrrtor--ormm of these six terms one by one. The first (E.72), (E.73) are the simplest, using inequality (E.62):
max maxE;/AAy(aa [A, AA S max A,
Aq
A,
Akjul
(E.78)
I max
max
Akkl
(E.79)
=l|
RMS
A
cRRMS
(E.80)
max E, [A, AM S ||
RMS
Ak
coRMS
likewise.
(E.81)
For the term (E.74), we have
AAi [A (Aij [A, + Aiy(q: [A, (A,)))))) (Aqi, [A, (Ayy)
(E.82)
Take absolute values, sum over j, and apply Cauchy- and incqualities (E.63),(E.64):
21 b)))) s Ai
(I Al [A, k) || ||q.1] [A, k) [I,, Akjy
(E.83)
< A q: ||l| max ||, || + || |lll Aq. (max | (mxx Ak I]).
(F.84)
Taking the max over i and applying ||g|| | RO | 1:
max E, (Aq, [A, ()))) Ag bcRMS Aq boRMS [Aqll AkllooRMS- (E.85)
The term (E.75) is similar:
max E, A (qi, [A,AK))] < | (oollosMMs/// Aqll collo + ||Ak| sORMS Ak OOOMS
(E.86)
39

---

For term (E.76), observe that
max 52m (AA) im kmm
S |
co
Ak
RMS
(F.87)
A RMS
A ooRMS
(E.88)
and so by Cauchy- Seecha- and the fact that the rows of sum to
i Ey|Ai1 (Aq: -m (AA) km)) S max Aq Em (A A km
(E.89)
RMS Aq RMSS + AgloooRMs
Ak
moRMS
(F.90)
By a similar argument, for term (E.77) we have:
(g -2m (AA) im Amm)
A k: sRRMS
A OORMS + ||
(E.91)
Thus, we have an estimate on the ~ -peratoonnnrm A
| | oo- I 2 2| + 3 A
|| + 3||
+|| |||
(E.92)
where all the norms on the right hand side are ||-looRMS-
Adding this together with (E.70) and (E.71), we obtain (all norms being |/:||=RMs:
| 2
A A + 3 | + || |||| + 2|| |||
(E.93)
+ | A + || | + || | + ||А
(E.94)
<(
|| + || + || + ||]
(E,95)
which is the desired result.
Proposition Sharpness under composition
Suppose M = M2 o M1 where M1, M2 are well-normed modules on respectively (Nx, Vv. Wa) and
moreover (ax, Ae, 7ik )-sharp for & = 1, 2. If pı mass for k = 1, 2, note that by the definition of M.mass, the modular norm on the composite M, we have for any Aw (Aw1, Aw2) E W1 W2:
| [lM] S 1 |/wwll and ||Aw2||M; S p2 A
(E.96)
We must prove that is (a, 8, sharp where:
a | pia + p222 + 2 n PIP222 + PiT.
(E.97)
3 p1B1 + 1IP232 + "1 P172:
(E.98)
7 12211 + uł72.
(E.99)
Turning to the second derivative of M(.,), -), we prove the first Inequality (F.97). The Gauss-Newton decomposition (C.1) for any Au (Awi, Aw2) and Air Awn, Aur yields
Aw V Au =VM, o (Aw) V Aw,)
(E.100)
(Awy, VM, Aw;)
v (Awz, VM, Aw,)
(E.
Applying the well-normed and sharpness inequalities, the norm of the first (E.100) of these terms is bounded by
p2 Aw o V o Awi Ib S 1200 Aw [|,, || l|M,
(E. 102)
<
M.
(E3)
The second term (F.) breaks into four separate terms:
Awr o 2 wo M2 Aw
(E.104)
+ +VM AW)) 2
M2 Awz
(E.105)
+ +Aw2 o V2 We M2 VM o Aw
(E.106)
(VM, Aw))
2 2
M2 o (VM, Aw)
(E.
40

---

In particular, applying the well-normed and sharpness inequalities, this is bounded by
Org
A
W22
20
M2
(E.108)
+
Aw
A
22
Ma
(E.109)
+-32
||Mg
A
M1
(E.110)
+72
Awi
||M,
Awi
||M
(E.111)
which is less than
(păa2 +
Pip22 2 + p || Nkw/lll
(E.112)
which completes the proof of Inequality (C.4).
Inequalities (F.98) and (F.99) simpler. For the first of these, note have
Aw o V o Aæ = =VM (Awi V o Ar)
(E.113)
+ (Aws, VM, o Awi) o V (VM, Ax).
(E.114)
Term (E.113) bounded by P2 Auw M Amlly
/2/ Aw IM, ||
(E.115)
$ p1 | |
(E.116)
Term (E.114) breaks into two separate terms
Awy M2 (VM1 Ar) + (VM1 o Aws) o 2 M2 (VM1 Az).
(E.117)
In particular this bounded by
8ell Awy M1 Ax 1.x,
Aw MM, /1llAx/xxx $ (2i22 #i p1 | |M
Ax
(E.118)
which completes the proof of Inequality (E.98). Finally, for (F.99), we have
Ax ViMo =V & AT)
(E.119)
(VM, o
(VM, o AT).
(E. 120)
Term (E.119) is bounded by
P2l|Ax o V Azly, S / ||
(E.121)
while Term (E.120) is bounded by
VM o A VM, A S 1
|
(E.122)
which together give Inequality (E.99).
Proposition 9: Sharpness under concatenation
Suppose = (M1, M2) where M1, M2 are well-normed modules on respectively (Kx, Vs, Wh) and moreover (cry, S&, Tik )sharp for 1, 2. pa Ma.s for k = 1,2, as in the previous proof we M.mmas
have for any Aw (Aw1, Awy) Wi x Wy:
|Aw Ilm, S J. P1 ||w/| and ||Awz|lm, I <p2ll S2llAww|M We must prove that M is (a, (e,8,,)-shary where
(E.
pian + p5a2.
(E.124)
p1 + P232.
(E.
7 71 72.
(E.126)
Now, for the first these identities, we have for Aw = (Awi, Aw2) and Aw (Awi, , Aw2):
AAw 2 Aw/byx
||(w V o A Aw, Aw, V
x
(E. 127)
| V o Aw y + | o V o Aww)
(E.128)
< on Aw R2, + 02 A W2 1 < (pia + 2²02) A &
(E. 129)
(F.130)
which shows = ppa1 + 2322. The expressions for 8, follow similarly.
41

---

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: Our main claim is that normalizing Adam and SGD updates in the modular norm leads to good learning rate transfer across width and depth. We believe this claim is supported by the experiments in our paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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 NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
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 of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations in the discussion section ($5).
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 a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-specfiiatiinn 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 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?
42

---

Answer: [Yes]
Justification: We state theoretical results as formal propositions and provide their proofs in Appendix E.
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 to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
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 whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: See Appendix for an overview of our code, Appendix B for the detailed network architectures and Appendix D for the parameters of our experiments. In addition, we provide the source code for our experiments and the Modula package.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer 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 is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a 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 an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
43

---

5. Open access to data and code
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:
We make use of standard datasets and provide our code in the supplemental materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/
www.mmisssioiiiinonioeeiii for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for 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. aumbiiicguddee//oddduubbiii Srooeeubmiissoooo///
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. 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 to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: 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 to understand the results?
Answer: [Yes]
Justification:
See Appendix D for the full details of our experiments. Also, see our code.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to 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: [No]
Justification: Due to computational resource constraints, we opted to run large number of experiments check that trends hold across several distinct architectures and datasets, rather than running repeats on each individual experiment. Each hyperparameter sweep involves the order of 100 training runs, and we are working under academic resource limits. We believe that the fact the reported trends hold across varied experimental settings supports the significance of our results.
Guidelines:
The answer NA means that the paper does not include experiments.
44

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-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 not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We report on this in Appendix D.3.
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 compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: /ttppp:://neuue http://rrwrp..cpppueeaeeee/teeea
Answer: [Yes]
Justification:
We believe that no ethics guidelines were violated.
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., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: No societal impact of the work was discussed. Potentially the work could have a positive impact in terms of reducing carbon emissions caused by sweeping hyperparameters for large-scale models.
45

---

Guidelines:
The answer NA means that there 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.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We are not releasing any new datasets or pre-trained models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make 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: [Yes]
Justification: We only use public and open-source resources. We have cited these works. Licenses were not provided from the original source.
Guidelines:
The answer NA means that the paper does not 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 (c.g., CC-BY 4.0) should be included for each asset.
46

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com dataeets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors 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: [Yes]
Justification:
We provide code and instructions on how to use the new modules that we
define.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
We did crowdsource and we did human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We did not use human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
47

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
48


Reviews:
[{"review_id": "mzW72kfHrs", "reviewer": "Reviewer_3bTa", "summary": "This paper introduces the *modular norm*, which is a norm designed to be adapted to neural network (NN) optimization. Specifically, an optimization process involving a gradient computed according to this norm scales with the size of the NN to be optimized. This paper provides the algorithm of computation of the modular norm, along with a small set of training experiments.", "questions": "How does the proposed optimization process compare to optimization with NTK/muP parameterization?", "limitations": "Lack of experimental validation, lack of clear list of contributions.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "# Originality\n\nCreating a norm making the optimization process scale with the size of the trained NN is original.\n\n# Clarity\n\nOverall, the general idea is easy to understand.\n\n# Significance\n\nFinding the scaling laws linking the hyperparameters (learning rates, initialization scale, penalty, etc.) to the architecture of a NN is a major field of research in NN optimization.\n\n# Quality\n\nLines 87--92, the authors state an interestng motivation of their work:\n\n> In machine learning, we we would ideally like [...] that meets these requirements: the *modular norm*.\n\nOverall, the motivation and the idea of the modular norm are appealing.", "weaknesses": "# Clarity\n\nA major issue of the paper is the absence of list of contributions. In other words, the authors do not make any general claim about their results. The list of contributions *must* appear somewhere (ideally in the beginning), so that one would have a formal basis to evaluate the paper (are the claims significant enough? are they well justified? etc.). \n\nSection 4 (experiments): one would expect a description of the implementation of the modular norm in an optimizer. As such, I do not understand what is computed (mass? sensitivity? norm?), when (at each step? epoch? ...) what is used in the optimizer and how.\n\nMajor typographical issue make the main text hard to read: some notations overlap with each other (line 148, lines 162--163, etc.).\n\n# Significance\n\nNo list of contributions is provided, so it is difficult to evaluate the significance of the paper.\nBesides, the set of experiments is very limited.\n\n# Quality\n\nThe contribution of this paper is unclear for the reasons aforementioned. \n\nMoreover, given the context, the motivation and the related works, one would expect a comparision with similar \"scaling\" techniques (muP and NTK parameterizations, for instance. It is well-known that standard training (with the usual parameterization) diverges as the number of \"blocks\" or the \"width\" tends to infinity (Fig. 4), so it is not surprising that the proposed method performs better in this situation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer 3bTa,\n\nWe 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:\n- We propose the modular norm defined via a recursive algorithm (Definitions 3 and 4)\n- We show that neural nets are Lipschitz (Proposition 2) and Lipschitz smooth (Proposition 5) in the modular norm\n\nEstablishing 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.\n\nRegarding your other questions:\n- Mass and sensitivity are held fixed from the start of training. Weight updates are normalized in the modular norm at each training step.\n- 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.\n\nThank you again for your time and effort!", "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": "0akcAgbvxT", "reviewer": "Reviewer_qUST", "summary": "The authors propose a new normalization strategy for deep models rooted in the introduction of a new framework and on feature learning considerations. The authors provide a few experimental examples as motivation, and then start introducing their framework. They formally define what a module is and its norm and then present a few results on module composition and relation to gradient/hessian-dependent quantities. The authors showcase that, on some toy experiments, their \"modula\" package yields improvements over vanilla SGD, closing the gap to Adam.", "questions": "1) Can you show experimentally that equation 2.2 holds tight? Many results in your paper are based on inequalities. And yes, I am aware of G. Yang's discussion of the spectral norm - yours is similar. Yet I think further motivation is needed to believe your inequalities fully.\n\n2) Can you trace back your findings to some framework-independent considerations? I think it is very important to show - even in a small architecture, what actually your normalization is doing. Tensor programs is similarly unsatisfying in the sense that for many researchers, having a program doing the normalization for you is a bit fishy. I think you should outline your contributions also in a language that everyone is using, so to make clear to researchers what is that you believe is needed -- precisely -- for SGD to work.\n\n3) One thing I am afraid of is that sometimes normalization can be a bit aggressive. Did you ever found this to be the case? I am thinking (for instance) on some findings about \"edge of stability\" : in the paper, the authors show that if you normalize the learning rate with the biggest hessian eigenvalue, then things can go south. Here, you are normalizing blocks (modules) though, so it might be different. Probably is.", "limitations": "Experiments, see above.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The paper is well-written, presentation is formal, notation is pleasant. Illustrations are very well done and in general the whole work is very thought-through. I also like the idea: not modifying the optimizer or the architecture, but normalizing stuff before it is fed into the optimizer. This is very flexible and I hope the authors can scale this up on bigger models.", "weaknesses": "There are 2 weaknesses I think are hurting the paper, but both are solvable.\n\n1) While the paper reads well, I found a lot of distance between the discussion of SGD in formula 2.1 and the results of Proposition 5. In between, there is little discussion of SGD. When reading, I felt a bit lost into the definitions and formalism - seemed the connection was just used as motivation but one has to to through 4 pages to get back at it. I think the authors should put proposition 5 sooner as a main result, and actually as a motivation for introducing the modular norm.\n\n2) This is solvable but hurts the contribution a lot: experiments. Resnet experiments are ok - but where we really need the thing to shine is transformers since there SGD has a huge gap. I think the current status is lacking a bigger transformer architecture: the context length used is 128 for GPT - very small. I know that if you have resource limitations, this can be demanding, but I am afraid your idea would be lost in the literature if you don't provide further evidence. If I have to make a suggestion: pythia 160M (https://github.com/EleutherAI/pythia/blob/main/models/160M/pythia-160m.yml) is a good model: probably takes < 0.5 day on a single GPU (train e.g. on slimpajama). If you feel very confident, try 24 layers. I think you only need to run this with SGD lr = [1e-3, 1e-2, 1e-1, 1], use the standard Adam parameters reported in the link above, and then run module on SGD. I would keep the context to 1k or 2k. I think if this works many people will pay attention.\n\nI am happy to raise to accept if you can provide some experimental evidence in the rebuttal.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer qUST,\n\n\nWe 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.\n\nRegarding your questions:\n1. 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.\n2. 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.\n3. 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!\n\nThank 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.", "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": "erzPlRky4l", "reviewer": "Reviewer_FbvX", "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.\nThey 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.", "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.\n2. 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?\n3. 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? \n4. Is there any reason the authors didn't compare to hyperparameter transfer in muP and its depth counterpart?\n5. 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?\n2. 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).", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The paper proposes a very interesting idea: an architecture-agnostic norm for neural networks that gracefully scales in width and depth.\nThis 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.\nThis 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.\n\nLastly, 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.\n\t- 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.\n2. 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:\n\nThe 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.\n\nI think, when first reading section 2.1, I would have appreciated:\n\n1. Having a better sense of direction:\n\t1. 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.\n\t2. 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.\n\t\nThen, 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.\n\n3. Some smaller points:\n- 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.\n- The sharpness bound in equation 2.1 could use a citation (or link to appendix derivation) for pedagogical purposes.\n- 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.\n- The authors claim that equation (2.2) holds quite tightly. Do the authors have a reference?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer FbvX,\n\nWe are grateful for your extremely thorough and helpful review! We’re very happy you found our paper interesting.\n\nBefore 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:\n1. 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);\n2. 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.\n\nWe 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.\n\nAt 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.\n\nNow, to answer your comments on the weaknesses:\n1. 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.\n2. Thank you for some great suggestions on how to improve Section 2 in particular:\n- 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).\n- 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.\n- 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.\n\nTo answer your questions:\n1. 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).\n2. See the discussion above about tight bounds.\n3. “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.\n4. 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.\n5. 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).\n\nFor your questions on limitations:\n- See the above for a discussion about lower bound guarantees, including a discussion of the approach taken in muP.\n- Correct, we only consider learning-rate transfer in this paper.\n\nThank you again for the time and effort you put into the thorough review of our paper!", "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": "anaUJciQpb", "reviewer": "Reviewer_U3gK", "summary": "The authors are tackling the difficult problem of try to scale the parameters of a network so that different sizes of network have similar optimisation properties.  This would greatly aid in hyper-parameter tuning.  They do this by introducing a new norm which is defined for the whole network.  They show experimental results providing evidence that their approach is successful.", "questions": "1, Is there a principled way to determine the masses.\n2. Are the masses for a convolutional layer in VGG the same as that for ResNet or AlexNet\n3. Do the masses depend on the problem\n4. How does the scalability compare to the commonly used initialisation schemes for weights?", "limitations": "I see no issues here.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This is a hard problem that is worthy of study.  The approach is mathematically rigorous in so far as it introduces a new norm.  The norm is designed to capture features of the network that are important in scaling and this is backed up by experimental results.", "weaknesses": "The approach feels a bit ad hoc, particularly the introduction of \"masses\".  It feels a little bit like the problem of adjusting hyper-parameters has been pushed onto determining masses for modules.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer U3gK,\n\nThank 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.\n\nYou 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.\n\nHowever, 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):\n- For any given fixed M, the global learning rate exhibits hyperparameter transfer from a small to large network;\n- The optimal M (in terms of lowest test/training loss) itself exhibits hyperparameter transfer.\n\nNote that this scheme involves tuning a single mass M, rather than one mass for every layer.\n\nMoreover, 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.\n\nTogether, 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:\n1. what random matrix distribution do you use: orthogonal, Gaussian, uniform, etc.\n2. given a choice of distribution, how do you scale it and how large are the resulting singular values?\n\nOur 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.\n\n\nThank you again for the time and effort you put into reviewing our paper!", "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:
To all Reviewers,

We are sincerely grateful for your contributions to the conference and for your feedback on our work. Reviewer FbvX commented that the paper “has challenged how I think about sharpness and curvature of the loss landscape in a deep learning context” and that they “feel fairly confident many in NeurIPS community would benefit from reading this paper”. We found this feedback really encouraging—thank you! We will also listen to and try to address all the reviewer feedback—see our individual responses below.

Best,

Authors

Meta-review: This manuscript proposes a modular norm that can be used within the training process of a neural networks. There are some limitations and deficiencies of the manuscript in its current form, particularly aspects of the presentation, a lack of certain comparative "baselines," and incomplete evaluation of the proposed formalism. However, there was sufficient excitement from several reviewers about the novel ideas presented within the manuscript to counterbalance those issues—the problem the manuscript is working to tackle is undoubtedly important. While there is still clearly work to be done, it is reasonable to consider the present manuscript a sufficient first step in an interesting direction. The authors should make the relevant adjustments to the manuscript that arose during the discussion phase and address the specific weaknesses raised about aspects of the presentation (note that as discussed, this need not be a contribution list but additional clarity in intro about the positioning and advancements of the work would be appreciated).
Decision: Accept (poster)