Modular Duality in Deep Learning
Authors: Jeremy Bernstein
Laker Newhouse
Venue: ICML 2025 poster (2025)
Keywords: modular
duality
Newton-Schulz
Primary area: deep_learning->theory
TL;DR: Modular dualization maps gradients to weight space in neural nets, enabling fast and scalable training algorithms automatically optimized for different architectures. Successfully used by the community to speed up NanoGPT training.

Abstract: An old idea in optimization theory says that since the gradient is a dual vector it may not be subtracted from the weights without first being mapped to the primal space where the weights reside. We take this idea seriously in this paper and construct such a duality map for general neural networks. Our map, which we call modular dualization, forms a unifying theoretical basis for training algorithms that are a) fast and b) scalable. Modular dualization involves first assigning operator norms to layers based on the semantics of each layer, and then using these layerwise norms to recursively induce a duality map on the weight space of the full neural architecture. We derive GPU-friendly algorithms for dualizing Embed, Linear and Conv2D layers—the latter two methods are based on a Newton-Schulz iteration. We conclude with small experiments demonstrating the speed, scalability and novel numerical properties of duality-based optimizers. Our methods were used in the Muon optimizer, which recently set speed records for training NanoGPT and was scaled up to a 1.5 billion parameter transformer.

Modular Duality in Deep Learning
Jeremy Bernstein Laker Newhouse'
Abstract
weight LR weight.grad weight LR * bealize(weight.grae
type error! all good!
An old idea in optimization theory says that since the gradient is a dual vector it may not be subtracted from the weights without first being mapped the primal space where the weights reside. We take this idea seriously in this paper and construct such duality map for general neu- ral networks. Our map. which we call modular dualization, forms a unifying theoretical basis for training algorithms that are a) fast and b) scalable. Modular dualization involves first assigning op- erator norms layers based on the semantics of each layer, and then using these layerwise norms to recursively induce a duality map on the weight space of the full neural architecture. We derive GPU-friendly algorithms for dualizing Embed, Linear and Conv2D layers- the latter two meth- ods are based on a Newton-Schulz iteration. We conclude with small experiments demonstrating the speed, scalability and novel numerical prop- erties of duality- based optimizers. Our methods were used in the Muon optimizer, which recently set speed records for training NanoGPT and was scaled up to a 1.5 billion parameter transformer.
Why? The reason is that the loss function may not be equally smooth in all directions in weight space, and there is no rea- son for the sizes of different components of the raw gradient vector' to respect this heterogeneity. In other words, the ge- ometry of the loss function may be -on-istroppcc. Insisting type check should force the user to become cognizant of this issue and to find suitable duality map. A good dual- ity map should adjust the size and direction of the gradient to respect the smoothness structure of the loss function.
Duality maps on vector spaces are commonplace in physics and applied math. Examples include the musical isomor- phism differential geometry (Grosse, 2022), raising and lowering indices in general relativity (Carroll, 2019) and the bra-ket notation in quantum mechanics (Sakurai & Napoli- tano, 2020). Duality maps are also central to several opti- mization theories including mirror descent (Nemirovsky & Yudin, 1983), natural gradient descent (Amari, 2016) and steepest descent on normed space (Boyd & Vandenberghe, 2004). Despite the efforts of some prescient papers (Carlson et al., 2015b; Flynn, 2017), the latter kind of norm-based duality map is yet to puncture the deep learning mainstream.
We believe that duality is key theoretical concept that will help in building performant large-scale machine learning systems. To support this belief, we show in this paper that two important and seemingly disparate methods in contem- porary optimization research may be seen as approximations to a single duality map. These methods are maximal update parameterization (Yang & Hu, 2021, jPP), which is aimed at scalable training, and Shampoo (Shi et al., 2023), which is targeted at fast training. We show in Section 4. that both methods emerge as partial approximations to a single duality map induced by the RMS-RMS operator norm.
Introduction
This paper pursues a rigorous and first-principles theoretical framework for designing neural network training algorithms. We hope that such framework will facilitate the design of next generation of fast and scalable optimizers that are automatically tailored to different neural architectures.
While gradient descent is the workhorse of modern machine learning, the most vanilla form of the algorithm does not, in our view, pass a basic type check. For a gradient update to type check, we insist that the gradient must be passed through a duality map before being multiplied by a learning rate and applied to the weights:
'In this paper, we use the term "gradient" to mean the partial derivative of the loss function. This is in accordance with the terminology used in machine learning software libraries such as PyTorch (Paszke et al., 2019) and JAX (Bradbury et al., 2018). We chose this verbiage to help make the paper casily accessible to a broad machine learning audience. However, it is worth noting that this terminology is al odds with common mathematical parlance (Blondel & Roulet, 2024) where a "gradient" is a primal vector obtained by applying some form of duality map to the partial derivative of the loss function.
'IIT CSAIL, United States. Correspondence to: Jeremy Bem- stein <jbernstein edu. Laker Newhouse <lakem @ mit.edus.
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).

---

Modular Duality in Deep Learning
1.1 Summary of contributions
and they also propose a way to mix RMSprop-style pre- conditioning information (Tieleman & Hinton, 2012) into the weight updates. In contrast our work, the authors only derive duality maps for single linear layers, and these maps are then heuristically extended to all-layer updates. Nonethe- less, the authors achieve substantial wall clock speedups using variants of spectral descent to train small networks.
First contribution. We describe a procedure for construct- ing duality maps for general neural architectures. The pro- cedure, called modular dualization, works in three steps:
Step 1:
Operator norms are assigned to individual layers based on the input-output semantics of each layer;
Now, let us turn our attention duality structure gradient descent (Flynn, 2017). which constructs duality map on the full weight space of the neural architecture based on identifying a Finsler structure (Deimling, 1985) inherent to neural networks. Similar to modular dualization, Flynn (2's7)'s 's duality map works by assigning duality maps to each layer and then inducing a duality map on the full weight space. The substantial difference our approach is that Flynn (2017) leverages a weighted sum (L1 combination) of layerwise norms to construct his full duality map. This leads to optimization methods that only update single layer at each iteration, and the methods need to be heuristically extended to achieve all-layer updates. In contrast, we lever- age the modular norm (Large et al., 2024), which takes a weighted max (Li combination) of layerwise norms. In turn, our duality map leads directly to more conventional all all-layer optimizers.
Step 2:
Based on these operator norms, duality maps are constructed for individual layers:
Step 3:
Given the layerwise duality maps and the struc- ture of the neural architecture, a single duality map is recursively induced on the full weight space of the architecture.
Second contribution. We instantiate modular dualization for a rich family of neural wurchtectuurs--innteeee tional networks and transformers- by writing down duality maps for Linear, Embed and Conv2D layers. We also pro- vide GPU-friendly algorithms for computing these duality maps efficiently. One of these methods has already been applied in the Muon optimizer (Jordan et al., 2024b).
convolu-
Third contribution. We run two experiments. the first, we find that duality- optimizers are fast and scalable across width- see Figure 1. In the second, we see that duality-based training exhibits novel numerical properties: the weights move substantially further from their initial values than for non-dualized training-see Figure 2.
Another important difference between our work on mod- ular duality and prior work on duality structure gradient descent is that we fully "modularize" our theory-meanine that our construction explicitly recursive and as such it is easy to code up into a software package. In this re- gard, we are inspired by a line of work that attempts to build optimization algorithms that automatically adapt to the structure of general computation graphs. The earliest work we know of in this category is the PhD thesis of Grant (2004) on disciplined convex programming, which aims to infer the convexity properties of general functions by break- ing them up into subexpressions and applying composition theorems from convex analysis. More recent progress in this vein includes work universal ingirization-minnnizatooo algorithms (Streeter & Dillon, 2022; Streeter, 2023) and related papers on automatic majorization (Tran et al., 2015; Bernstein et al., 2023).
2 Related Work
This paper constructs duality map for general neural ar- chitectures. Our approach is based on assigning operator norms to individual network layers and using these layer- wise norms to recursively induce a duality map on the full neural architecture. The most closely related prior work is a series of papers on spectral descent (Carlson et al., 201 20;;;;;; 2016) and a paper on duality structure gradient descent (Flynn, 2017).
Spectral descent has been applied to restricted Boltzmann machines (Carlson et al., 2015a) and discrete graphical mod- els (Carlson et al., 2016). but let us focus on the more closely related paper on spectral descent for deep learning (Carlson et al., 2015b). In that paper, the authors propose assigning the Schatten-o0 norm (a.k.a. spectral norm) to individual linear layers. This assignment is based on the observation that neural networks admit natural majorization bounds in the Schatten-oo norm. The authors call the correspond- ing duality map for linear layers the "#-operator" name presumably inspired by the musical isomorphism (Grosse, 2022). The authors propose a cheap approximation to the #-operator based on sketching (Martinsson& & Tropp, 2020),
3 Theoretical Preliminaries
In this section, introduce duality maps, a means of con- structing duality maps based norms, and finally a norm called the modular norm that is well-suited to describe the geometry of general neural architectures.
3.1 Duality Maps
Given a vector space V, we say that a function f: V - R is a linear functional on V if f is linear. We define the
2

---

This restriction may seem absurd! After all, here the weight space W and its dual W are both just I", However, in- sisting upon this type check serves to remind us that the Modular Duality in Deep Learning
dual space V* to be the set of linear functionals on V. The dual space is itself a vector space provided that addition is
where, if the arg max is not unique, dualizee | returns any maximizer.
defined pointwise (f + g) (z) := f(x) + g(xr) and scalar multiplication is defined pointwise (af) == af(z) for any scalar a. By duality map we mean any function that sends a member of the dual vector space V to the primal vector space V. The function need not be an involution.
Given these definitions, minimizing the expression in the right- higtt-handann side of Equation (2) can be done using the follow- ing standard proposition, for which Bernstein & Newhouse (2024) provide a proof:
Proposition 1 (Steepest descent under a norm). For any vector g E IR" thought of as "the gradient", " any À > 0 thought of as "the sharpness", ", and any norm |-1| R," R with dual norm 1.1" and duality map dualize
Let C W - R denote the loss of a differentiable machine learning model with weight space W I The Taylor expansion of the loss at weight setting we E W is given by:
L(w Aw)- ((w) + Cww) Aw +
(1)
A |Awll 2
llel' A
Observe that, the first-order term, the gradient V ((w) acting as a linear functional: it is pairing with the weight vector Aur W in a linear way produce real number. As such, we shall say that the gradient belongs to the dual weight space: (w((w) E W", We shall forbid ourselves from directly subtracting a member of the dual weight space W* from the weight space W. If we would like to conduct a gradient descent update, then we had better find a duality map to send the gradient back to the primal space W.
arg min g Aw +
dualize
AwER"
In words: to find the minimizer of a linear term penalized by a squared norm, we need only evaluate the dual norm and a duality map. In this paper, we focus on constructing a duality map for the modular norm, which is a norm tailored to the weight space of general neural architectures. But first, we shall cover duality maps for more standard norms.
3.3 Basic Norms and Duality Maps
Many basic norms and duality maps already covered in prior work (Carlson et al., 2016; 201 ba;;; Flynn, 2017). For some warmup examples, the following duality maps for vector norms are standard:
curvature of the loss function may be highly heterogeneous. The next section will show one way to construct duality maps to account for this.
Example 1 (Duality map for the Euclidean norm). For a nonzero vector E Rd, we have dualize |-1, 9 = g/1gg/ For the zero vector, we take dualizo
3.2 Steepest Descent on a Normed Space
O.
Suppose that we have found a norm ||+|| W - IR and a sharpness parameter À > 0 that serve as a good model of the higher-order terms in the Taylor expansion of the loss function given in Equation (1):
Example 2 (Duality map for the infinity norm). For a vector g Rdd, we have dualizer, 9 sign(g), where the sign function is applied entrywise and we take sign(0) 0.
A 2
In neural networks, the weight spaces of individual layers tend to have matrix structure. And layers with the same shape weight matrix may have semantically different input and output spaces-think think embedding versus linear layers in a transformer. As such, we will need duality maps for different induced operator norms:
L(w + Aur E [(w)) + V vm((w) Aw + ||A
In other words, the norm provides a good characterization of the heterogeneity in curvature of the loss function. Then it makes sense to solve for a weight update Aw by minimizing the hand side of Equation (2). We will show that the minimizer can be expressed in terms of a dual norm and a duality map:
Definition 3 (Induced operator norm). Given matrix M Rdoa ddu and two normed vector spaces (R (l+lla) and (pddou I1-lee), the "a to 3" induced operator norm is:
Definition (Dual norm). Given a norm |I-|| R TR, the dual norm 1-IIt of vector g 12 is given by:
| 1 |x|aa
| 1a-88-
max 2r Rdin
max 9 t.
tER" [t|[="
For tensors, we define the duality map via
Definition 2 (Duality map based on a norm). Given a norm ||-II R TR, we consider the duality map:
dualizej.1 G := arg max (atten(G)) fTatten(T). [T|-1
dualize||1 == max g) t,
For linear layers, we will need the duality map for the RMS - RMS induced operator norm, This ends up as
-1
3

---

Modular Duality in Deep Learning
rescaled version of the spectral norm duality map from prior work (Carlson et al., 2015b; Flynn, 2017).
forward function is Lipschitz-contimuuus in the weights with constant 1 and in the inputs with constant M.sensitivity:
Example 3 (Duality map for the RMS- - RMS operator norm). For vector E IRd we define the RMS norm to be the normalized Euclidean norm: ||u||RsS = |/ell//vd. Given a matrix W E Rd Xdm, the RMS - RMS induced operator norm resolves to a rescaled spectral /ddn/dout x [[W],, w. where I1-1I= denotes the standard spectral norm. For a matrix G E Roou x thn with reduced singular value decomposition
Definition 5 (Well-normed module). Let M be a module on (X. y, W), where the input and output spaces have respec- five norms I-xx and I-lly. M is well-normed if for all imputs x E X, weights w W. weight perturbations Aw E W and input perturbations A E X, we have both:
norm: | WRRMS-RMS
V un M.fowwadd un, x) o Awl y < M.norm(Aur))
V forwwrdd w. x) o Ae) M.sensitivity x Axxxx
G = UEVT the corresponding duality map is given by
dualize| ||RMMs
-RMS G = v dout/ / din x UVT.
-RM
The o operator denotes summation over any shared tensor indices?. This definition of well-normed-neg can be used as a guiding principle in the design of library of atomic (i.e. handwritten) modules. First, norms should be assigned to the input and output space of each module based on the semantics of M.forward. Then a norm M.norm should be assigned to the module's 's weight space and number M.sensitivity should be chosen to make the module well- normed. Examples are given in Section 4.1.
And for embedding layers, we will need the duality map for the ( - RMS operator norm:
Example 4 (Duality map for the - RMS operator norm). Given a matrix W E IRdous Xdi, the & - RMS induced operator norm resolves the max RMS norm of For
max; [W)) matrix G E Rdom the corresponding duality map dualize lll, -RMS G simply normalizes each column of G to have unit RMS norm: col:(G) - (G/[[col((()))) for each = 1,...dda. din-
the columns: [W]e-nMs
Given such a library of well-normed atomic modules, a compound module built through any arbitrary sequence of module compositions and module concatenations is automat- ically well-normed (Large et al., 2024). And if the atomic modules are not only well-normed but are also smooth (Large et al., 2024, Definition 5), then there is an auto- matic procedure for computing sharpness coefficients for any compound module built from the library (Large et al., 2024, Appendix C). The relevant definition of module com- position is follows:
3.4 The Modular Norm
The modular norm (Large et al., 2024) is intended to help characterize the heterogeneous curvature of general neu- ral architectures. The construction first defines an abstract module type along with notion of what a good, or well- normed, module. Then combination rules are given for con- structing new well-normed modules from a library of exist- ing well-normed modules. So modules are a special case of combinator pattern from functional programming (Haskell Wiki Contributors, 2007) and are related to a monoidal cate- gory from category theory (Fong& & Spivak, 2019). Large et al. (2024) begin by defining an abstract module:
Definition 6 (Module composition). Consider module Mi with input, outpul and weight space (X,,3,, Wi) and mod- ule M2 with input, output and weight space (X2, V2, W2). M1 and M2 are composable if X2 Yr. Their composite module M M2 o M2 has input, oulput and weight space (2,, 2/2 Wi x W2) and attributes:
Definition (Module). Given input vector space X, output vector space Y and weight vector space W, module M is an object with the following four attributes:
(a) M.for w22) x))
Mg.forward( (w2, M₁ Mf.forward(w1,, x)))
(a) a function, M.forward: W x X + Y, which maps input and a weight vector to oulput;
(b) M.mass Mmamass M2-mass;
(c) M.sensitivity = M, M x My.sensitivity,
(b) number, M.mass 2 O, which is used to set the pro- portion feature learning that this module contributes to any supermodule;
(d) M.norm (w1, w2)) max(a, B), where:
1) a M2 My.sensitivity x
M.mass x normor(
Mi.mass
(c) a number, M.sensitivity > 0, which estimates the mod- ule's 's sensitivity to input perturbations;
2) 8
M.mass
x M2.norm(w22)
Ma.mass
(d) a norm over the weight space, M.norm : WW - R20 sometimes abbreviated to just I.IM-
and if M, .mss or Mn.mass is zero, the corresponding term in the max is set to zero.
"The expressions inside the norms the left-hand side are examples of "Jacobian vector products", or JVPs for short (Blondel & Roulet, 2024).
We shall care most about modules that are woll-normed (Large et al., 2024), which amounts to requiring that the
4

---

Modular Duality in Deep Learning
So the composite norm is taken to be a weighted max over the norms of the two sub-modules, where the weight space of the first module is coupled to the input sensitivity of the second module. The module masses provide freedom to tune the importance of each sub-module in the norm, and Large et al. (2024) prove that module mass provides precise control over the amount of feature learning that can happen in each sub-module.
outputs to evolve. Then we place norm on the weight space such that A is well-normed: this is typically the operator norm (Definition 3) induced by the input and output norms, Finally we are in position to solve for the duality map, which we shall call A. A.dualize. We now give some examples of this procedure for the basic layer types of Linear, Embed and Conv2D. The results are summarized in Table 1.
We start with the canonical example of an atomic module: Example 5 (The Linear module). The Linear module sends inputs from X = Ra to outputs in y = Rddatt The weight space is given by the matrix space W = Rdom Xdio, We endow the Linear module with attributes:
Concatenation is defined in similar way to composition:
Definition 7 (Module concatenation). Consider module M with input, output and weight space (2i, Wi) and mod- ule 22 with input, output and weight space (X2, V2, W2). We say that M1 and M2 are concatenatable if their inpur spaces match: Xi = X2. The tuple module M- = (M,, M2) has input, output and weight space (2 x 2/2 W1 x W2)
1. Linear.forward (W, x:) Wwr matrix-vector product;
2. Linear.sensitivity = 1;
and the following list of attributes:
3. Linear.mass μ4, where μ 0 is a hyperparameter;
4. (nnear.nom((W) = wwlnsss-n the RMS
(a) (rooward(((),, 12))x))
RMS induced operator norm.
(M,.orrward (1, 2),, Mg.forward(122.- (w2, x));
Since the Linear module is intended to map to and from vec- tors of roughly unit RMS norm, we place the RMS norm both the input and output space: 1-1|x ||- IRRS and I- = ||-lIRMs. Then Linear is well-normed if the imputs and weights belong to the unit balls {or E Red.. || & 1} and { E Rdow xdm, Linearr .orearr))rr(W)) < 1 Referring back to Section 3.3, the duality map corresponding to Linear.norm is then given by:
(b) M.mass MI. + Ma-mass;
M.sessiiiiity Mj Myssensitivity Mg-sensitivity;
(d) M.norm((w1,, (w1, W2)) max(a, 3), where:
1) a- Mi .mass
M. mass
M1.norm(w;
2) 3- =
M.mass
x M2.norm(ws) (w2)
My
and if mmsss or M2. mass is zero, the corresponding term in the max set to zero.
5. Linear. Annear.dualize(r) (G)
VV
dont din
x UVT, where the gradi-
ent G E iRdous Xdu has reduced SVD G = V.V
A shortcoming of the paper by Large et al. (2024) is that the power of the modular norm is not fully leveraged. In particular, the authors do modular normalization of training, where weight updates to certain modules are naively divided by their norm. This paper makes fuller use of the geometry of the modular norm by constructing the corresponding duality map, which we call modular dualization.
This single duality map recovers essential features of both maximal update parameterization (Yang & Hu, 2021, nP) and Shampoo (Gupta et al., 2018). In particular, the fac- tor of v dout/dn in dunalie recovers spectral update scaling (Yang et al., 2023) that leads to jaP (Initializing such that Linear.norm(W) (W- = 1 also recovers pP initialization scaling.) And the mapping G - UVT is equivalent to Shampoo without accumulation (Bernstein & Newhouse, 2024). As such, we believe duality maps may help reconcile different strands of deep learning research and provide a unifying basis for fast and scalable training algorithms.
4 Modular Dualization
In this section, we construct a duality map for general neu- ral architectures. Our strategy is to first write down duality maps for atomic modules, i.e. individual layers. We then extend to arbitrary compound modules, i.. full neural net- works, by showing how duality maps should pass through composition and concatenation.
The Embed module provides a useful counterpoint to the Linear module. The difference between the two modules stems from the fact that the input spaces of Embed and Linear have different semantics.
Example (The Embed module). The Embed module sends inputs from X' Re to outputs in y = Bduutt The weight space is given by the matrix space W Indo xdm We endow the Embed module with attributes:
4.1 Duality Maps for Atomic Modules
To construct a duality map for an atomic module A, the idea is to first fix norms on the input and output spaces that respect the semantics of A. forwwrdd. We should select norms that describe both how large we would like the inputs and outputs to be, and in what geometry we would like the
1. forwd W. :z)= -W, matrix-vector product;
Embed sinbivvsesiitivity = 1;
5

---

Modular Duality in Deep Learning
Module
Weight Space W
Module.norm
Module.dualize
Linear
T Rous x dia
W WIMS--RM
G G
dont UVT
dia
Embed
Rdow x din
W W
1 RMR
col,(G)
0oo,, RMM
Gol,
Conv2D
Root xdu x
W h² max
j=1 W RMMS-RMS
G..y ~
dans X V
dan
Table 1: Duality maps for three atomic modules: Linear, Embed, and Conv2D. These atomic modules are sufficient to build CNNS and transformers. In Linear. dualize, UEV denotes the reduced SVD of the gradient matrix G. In Conv2D. dualize, U2y E Vij denotes the reduced SVD of the slice of the gradient tensor G...j at kernel indices i.j. Section 5 provides GPU-friendly algorithms for computing these duality maps using a family of Newton-Schule iterations.
3. Embed.mast mass = μ4, where μ > 0 is hyperparameter;
Hont max max 13 IRMS. Then Conv2D is well-normed if the inputs and weights be- long to the unit balls {or E gW Hi xdn ; :| x & 1} and { {W E Rddou xdn xkxk Corv2
and l
Wons
4. Embed. W) || RMMS the l RMS induced operator norm.
< 13. Since
Embed is intended to map from one-hot vectors to vectors of roughly unit RMS norm, so we place the l norm on the input space and the RMS norm on the outpul spuce: 1.1xx = lle Il-lyy = ||-IRMs- Then Embed well-normed if the inputs and weights belong to the unit balls {x E Rdm : ||>| & 1} and Referring back to Section 3.3, the duality map for Embed. norm is:
the duality map for a max of norms decouples into one duality map per sub-norm, the duality map corresponding to Conv2D norm given by:
5. Conv2D.dualized G) does G...5 1 12 where G...g has reduced SVD UzJ Eig
dont x U., ViJ. din
RV
( E RRee xdi Embed.nero norm () < 1}.
4.2 Duality Maps for Bond Modules
5. Limbed.dualize( performs the mapping (G) (G) [col,(G) RMM for each column index j = =,,.., din-
Large et al. (2024) define another class of basic modules: bond modules. Bonds are handwritten modules without weights. An example of a bond is the ReLU nonlinearity. For a bond B, the weight space is the zero vector space W (0) and the modular norm B. the corresponding duality map is also B. B.dualize 0 - O. In a software package, one need not write norms or duality maps for bond modules.
Finally, we consider Conv2D module with a l l kernel. Conv2D has a more involved tensor structure than Linear and Embed. The calculations work by slicing up the weight tensor into a collection of 2.2 matrices.
0 - 0. As such,
Example 7 (The Conv2D module). The Conv2D module sends inputs from X' = W.... Hm. ddn outputs in y TR. Wont How x dout, We think of this mapping inpul image of width Win. height Hin and with din color channels to an output image of width Wout, height Hout and with dout color channels. The weight space given by the tensor space W- Redon ddax x xk where ls is the kernel size. We endow Conv2D with attributes:
4.3 Duality Maps for Compound Modules
Given two composable modules Mj and M2, the duality map for the composite M = M2 M1 is given by:
[M.maas dualizze))) My.mass
M.dualize(g 92))
xMy dualizeia))))
M.mass
Mg.sensitivity
M.mass
1. forward W.x) W ④ where @ denotes
Given two concatenatable modules M1 and M2, the duality map for the tuple M = (M1, is:
2D convolution;
2. Sonv2D.sensitivity 1:
My mass
Mg.mass
M. dualizzzlg.... .9.)
M, dddllize(gg.
MM. duduizzee)
M. mass
M.mass
3. Conv2D.mass 14, where 3 O is a hyperparameter;
The proofs of Section 4.3 follow in a straightforward manner from Definitions 6 and 7.
4. Conv2D nomm.ooomWWW. = 22 maxt W..ij|lMMs-RM the max RMS RMS norm over kernel indices.
We would like pixel intensities in the inputs and outputs to be order one and undergo order one change. We for- malize this by taking the input and output norms to be the spatial maximum of the RMS norms of all the color channel vectors: |
5
Fast Duality Maps
For modular dualization to be practically feasible, we need ways of computing duality maps quickly. Inspecting the duality maps listed in Table 1, we see that combed.dualics
Wi, Hin max W max,"",
1= |ayhh InMS
6

---

Modular Duality in Deep Learning
easy to implement since it just involves computing vec- tor norms of matrix columns. But Linear.dualize and Conv2D.dualize involve the projection:
ær closer and closer to +1. The final step is to realize that the effect of the matrix iteration is to apply this cubic f((x) to each singular value of X,. This shows that the spectral normalization Xo- G/ is stronger than what is required: we need only ensure that Xo has singular values no greater than v3 for the iteration to converge.
G = -UUVV - UV!
where VEVT is the reduced SVD of the matrix G. Since computing SVDs can be slow (Carlson et al., 2015b; Flynn, 2017), we discuss three approximations to this map via sketching, iterations for inverse matrix roots, and family of rectangular Newton-Schulz iterations.
This iteration has the advantage over sketching that it always works on all singular values, and since it does not compute inverse matrix roots the iteration is bell-behaved even on low-rank matrices.
Finally, there are in fact a family of degree 2n + 1 polyno- mial iterations of the form
5.1 Sketching
Sketching is randomized method (Martinsson & Tropp, 2020) that can be used to build low-rank approximations to the SVD. Carlson et al. (2015b) used sketching to pro- vide fast approximation to their #-operator. More recent papers have experimented with sketching in the context of Shampoo-type algorithms (Feinberg et al., 2023). A poten- tial downside of approximating Equation (3) via sketching is that randomized SVD methods usually try accurately approximate the largest singular values of matrix (Mar- tinsson& & Tropp, 2020, Section 11.2) while the value of Equation (3) may lie in its action on small singular values.
X = ax, + b(X X, +....... (X, )
for suitable a, b...., 2 instead of a, b = 3
2' 2. One should
choose coefficients a, b,,.,, 2 so that the univariate poly- nomial g(z) a. + bx. +... + 2x is a suitable approximation to sign(x). One may further accelerate the iteration by "tuning" the coefficients a, b...... 2 empirically.
6 Discussion
This paper develops the theory of modular duality and the procedure of modular dualization as means to construct duality maps for general neural architectures. Here, we comment on implications and connections.
5.2 Iterations for Inverse Matrix Roots
If G is a full rank matrix with reduced SVD UEVT,
then:
6.1 Neural Network Speedrunning
UV = (GGT) 1/2 G (G G) G) 1/4 = =GGTT) 10 G = G GG
We believe that the ideas in this paper can help in the design of faster training methods. In fact, based on work, a new NanoGPT training speed record was recently set using Newton-Schulz- duality map, packaged into an open- source optimizer called Muon (Jordan et al., 2024b).
This equation provides a route to approximating the map UEV UVT since one can compute inverse ma- trix roots such (GGT)-1/ via Newton iteration (Lakić, 1998). This is discussed in Chapter 7 of Higham (2008)'s 's book and also see Anil et al. (2020)'s 's paper. Care must be taken with inverses when the matrix G is ill-conditioned.
6.2 A Type System for Deep Learning
Part of the inspiration for this work is to build a fully-fledged type system for deep learning. We think that activation spaces should be typed by their intended norm and the intended size of activations in that norm. This informa- tion would help to construct well-normed modules (see Section 4.1). Modules should be typed according to Defi- nition 4. And, as suggested in the introduction, gradients should be explicitly typed as dual vectors. A duality map should flip the type of a dual vector to a primal vector. We plan to use the Modula deep learning package (Large et al., 2024) as a testbed for these ideas.
5.3 Rectangular Newton-Schulz Iteration
We developed a "rectangular Newton-Schulz iteration" for computing U by adapting Equation 5.22 in Higham (2008)'s 's book for computing the "matrix sign function". We later discovered this iteration has a long history (Kovarik, 1970; Björck& & Bowie, 1971). The method works by first normalizing the matrix G according to Xo- = G/||G]a ef- (or alternatively Xo G/||G| F) and then iterating:
1
X
X; 1 X,XXX X
2 2
6.3 Modular Duality: A Unifying Theoretical Framework for Fast and Scalable Training
then as 00, the sequence X, UVT To see this one can plot the univariate cubic function f(x) := zz } 3 and see that, for < x < v3, iterating this cubic will push
An important topic in contemporary optimization research is the design of fast and scalable training methods for neural
7

---

Modular Duality in Deep Learning
Adam (SP)
Adam (μP)
Dualization
Final training loss
22
1.0
109
26
10:1
24
Training accuracy
0.8
a
10"2
20
a
24
0.6
2
10-2
10-
Learning rate
109
10-1 109
0.4
a
a
a
a
a
a
a
a
Figure 1: Learning rate transfer with dualization. To test learning rate transfer across width, we train an MLP on CIFAR-10 for 20 epochs at range of widths and learning rates. We plot the final training loss and mark the best learning rate at each width with a red dot. Left: In standard parameterization (SP), Adam's optimal learning rate drifts to the left. Middle: Maximal update parameterization (Yang & Hu, 2021, μP) mostly corrects this drift. Right: Our duality- duality-based method has a fairly stable optimal learning rate and also reaches much lower loss. More experimental details are given in Appendix A.
0.2
With dualization
Without dualization
a
a
0.0
10-1
10
10a
Learning rate
Figure 2: Erasure of watermarked initial weights. It is commonly held that the weights stay close to initialization in very wide networks (Lee et al., 2019; Jesus et al., 2021). To visualize the change in weights, we "watermark" the hid- den layer weights of an MLP of width 1024 initialization by zeroing out matrix entries in the shape of the letter "a". We then train for 1000 steps CIFAR-IO, across ten differ- ent learning rates. For each run, we plot the final training accuracy along with an image of the learned weight matrix. Not only does dualized gradient descent reach higher train- ing accuracies than the non-dualized method, but dualized gradient descent also "erases" the watermark the highest stable learning rate, constituting substantial weight change. More experimental details are given in Appendix A.
networks. Two popular methods in this research space are maximal update parameterization (Yang & Hu, 2021, jPP), which allows increasing network width without changing the optimal learning rate, and Shampoo (Gupta et al., 2018), a variant of which (Shi et al., 2023) won a speed challenge at the AlgoPerf optimizer competition (Dahl et al., 2023).
We showed in Section 4.1 that essential features of both μP and Shampoo are recovered from the single duality map We think that, on basic theoretical level, uP and Shampoo should be viewed as partial approxima- tions to this duality map. This observation helps put μP and Shampoo on a consistent theoretical footing, orients the methods with respect to overlooked prior work on spectral descent (Carlson al., 2015b) and duality structure gradient descent (Flynn, 2017), and suggests new ways to generalize these methods arbitrary layer types and network architec- tures via the modular norm and modular dualization.
which all singular values have the same magnitude. For the case square weight matrix, the weight update is simply an orthogonal matrix, which acts as an isometry on inputs vectors meaning that feature learning happens trivially.
6.5 A Numerical Paradox: The Weights Don't Change!
Figure 1 shows that our duality-based optimizer is both scalable and fast: it transfers learning rate across width like IP,, and it reaches lower loss in the same number of steps.
Past work (Lee al., 2019; Jesus al., 2021) has pointed out an apparent paradox in deep learning: the weights seem to move a vanishing amount from initialization in the limit of large network width. This finding led to substantial inter- est in lincarized training dynamics (Jacot et al., 2018). Prior work attempted to resolve this paradox by showing that the weights move a roughly constant amount at any width when the change is measured in spectral norm (Yang et al., 2023). But duality maps lead to new story: Linear. duualie ramps up the stable rank of updates, causing the weights to move large width even in the Frobenius norm- provi-eeoveee the batch size is not too small. This result, shown in Figure 2, chal- lenges the belief that very wide neural networks cannot stray from their initialization; instead the numerical movement of the weights depends on the choice of optimizer.
6.4 On the Alignment of Activations and Updates
Recent work (Yang et al., 2023; Everett et al., 2024; Large et al., 2024) has singled out the following question as im- portant to the design of scalable deep learning systems: to what extent do gradient updates to neural network layers align with incoming activation vectors? This question is im- portant since it helps inform how large weight updates need to be to induce a certain amount of change in layer outputs. Duality maps such as Linear.dualize and Conv2D dualize may help simplify the answer to this question, since they project gradients to scaled orthognlt matrices for
8

---

Modular Duality in Deep Learning
7 Conclusion
Programming arXiv:2403.146000 2024. Cited on pages and 4.
This paper has proposed a recursive procedure called modu- lar dualization for building duality maps for general neural architectures. The procedure unifies past strands of opti- mization research Shampoo (Gupta et al., 2018) and μP (Yang & Hu, 2021). Duality-based optimizers have already led to significant wall-clock speedups in transformer train- ing ranging from 124M to 1.5B parameters (Jordan et al., 2024b). The rectangular Newton-Schulle iteration provides a GPU-friendly and numerically stable means of dualizing un- der the RMS RMS operator norm, while avoiding some of the downsides of sketching- based approaches (Carlson et al., 2015b). Overall, we hope that our theory of modu- lar duality provides a clarifying toolkit for the design and analysis of deep learning systems.
Boyd, and Vandenberghe, L. Convex Optimization. Cam- bridge University Press, 2004, Cited on page 1.
Bradbury, J., Frostig. R., Hawkins, P., Johnson, M. Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne S., and Zhang, Q. JAX: compos- able transformations of Python+NumPy programs, 2018, URL. http::/// /ttptt........g//g wwt.c/////aaao jttp//xxxmmmm/m Cited on page 1.
Carlson, D., Cevher, V., and Carin, L. Stochastic spectral descent for restricted Boltzmann machines. In Interna- tional Conference on Artificial Intelligence and Statistics, 2015a. Cited on pages 2 and 3.
Carlson, D., Hsich, Y.P.., Collins, E., Carin, L., and Cevher, V. Stochastic spectral descent for discrete graphical mod- els. Selected Topics in Signal Processing, 2016. Cited on pages 2 and 3.
Impact Statement
The methods developed in this paper may be used to make the training of machine learning systems more efficient, for either good or ill.
Carlson, D. E., Collins, E., Hsieh, Y.P.., Carin, L.,, and Cevher, V. Preconditioned spectral descent for deep learn- ing. In Neural Information Processing Systems, 2015b. Cited on pages 1, 2, 3, 4, 7, 8, and 9.
Acknowledgements
Many ideas in this paper, including Section 5.3, were devel- oped jointly with Tim Large before he left to work at tech company. We grateful to Phillip Isola for invaluable dis- cussions and for granting us the freedom to pursue this work. We also thank Jack Gallagher, Keller Jordan, Simo Ryu, Ro- gier Brussee, Tongzhou Wang, Victor Butoi, Jeffrey Cider and the anonymous reviewers for helpful conversations.
Carroll, S. M. Spacetime and Geometry: An Introduction to General Relativity. Cambridge University Press, 2019. Cited on page 1.
Dahl, G. E., Schneider, F., Nado, Z., Agarwal, N., Sastry, C. S., Hennig, P., Medapati, S., Eschenhagen, R., Kasim- beg, P., Suo, D., Bae, J., Gilmer, J., Peirson, A. L., Khan, B., Anil, R., Rabbat, M., Krishnan, S., Snider, D., Amid, E., Chen, K., Maddison, C. Vasudev, R., Badura, M., Garg, A., and Mattson, P. Benchmarking neural network training algorithms. arXiv:2306.07179 2023. Cited on page 8
References
Amari, S. Information Geometry and Its Applications. Springer, 2016. Cited on page 1.
Anil, R., Gupta, V., Koren, T., Regan, K., and Singer, Y. Scalable second order optimization for deep learning. arXiv:2002.09018, 2020. Cited on page 7.
Deimling, K. Nonlinear Functional Analysis. Springer Berlin, Heidelberg, 1985. Cited on page 2.
Everett, K. E., Xiao, L., Wortsman, M., Alemi, A. A., No- vak, R., Liu, P. J., Gur, I., Sohl-Dickstein, J., Kaelbling, L. P., Lee, J., and Pennington. J. Scaling exponents across parameterizations and optimizers. In International Con- ference on Machine Learning, 2024. Cited page 8.
Bernstein, J. and Newhouse, L. Old optimizer, new norm: An anthology. In Workshop on Optimization for Machine Learning, 2024. Cited on pages 3 and 5.
Bernstein, J., Mingard, C., Huang, K., Azizan, N., and Yue, Y. Automatic Gradient Descent: Deep Learning without Hyperparameters. arXiv:2304.0513 051187, 2023. Cited on page 2.
Feinberg, V., Chen, X., Sun, Y. J., Anil, R., and Hazan, E. Sketchy: Memory-efficient adaptive regularization with frequent directions. In Neural Information Processing Systems, 2023. Cited on page 7.
Björck, Å. and Bowie, C. An iterative algorithm for com- puting the best estimate of an orthogonal matrix. SIAM Journal on Numerical Analysis, 1971. Cited on page 7.
Flynn, T. The duality structure gradient descent algo- rithm: Analysis and applications to neural networks. 1XXi::7708.00523, 2017. Cited on pages 1, 2, 3, 4, 7,
Blondel, M. and Roulet, V. The Elements of Differentiable
and 8.

---

Modular Duality in Deep Learning
Fong, B. and Spivak, D. I. An Invitation to Applied Category Theory: Seven Sketches in Compositionality Cambridge University Press, 2019. Cited on page 4.
Large, T., Liu, Y., Huh, M., Bahng, H., Isola, P., and Bern- stein, J. Scalable optimization in the modular norm. In Neural Information Processing Systems, 2024. Cited on pages 2, 4, 5, 6, 7, and 8.
Grant, M. C. Disciplined Convex Programming. PhD dis- sertation, Stanford University, 2004. Cited on page
Lee, J., Xiao, L., Schoenholz, S., Bahri, Y., Novak, R., Sohl- Dickstein, J., and Pennington, Wide neural networks of any depth evolve linear models under gradient descent. In Neural Information Processing Systems, 2019. Cited on page 8.
Grosse, R. Metrics, Lecture 3 of CSC2541: Neural Net Training Dynamics, 2022. Cited on pages 1 and 2.
Gupta, V., Koren, T., and Singer, Y. Shampoo: Precondi- tioned stochastic tensor optimization. In International Conference Machine Learning. 2018. Cited on pages 5, 8, and 9.
Martinsson, P.-G. and Tropp. J. A. Randomized numer- ical linear algebra: Foundations and algorithms. Acta Numerica, 2020. Cited on pages 2 and 7.
Haskell Wiki Contributors. Combinator pattern. Haskell Wiki, 2007. URL https://w wikii.i.... wwwsseell.ooo Combinator_pen pa_tterr.. Cited on page 4.
Nemirovsky, A. S. and Yudin, D. B. Problem complexity and method efficiency in optimization. Wiley, 1983. Cited page 1.
Higham, N. J. Functions of Matrices. Society for Industrial and Applied Mathematics, 2008. Cited on page 7.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Köpf, A., Yang, E., Vito. Z., Rai- son, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L.., Bai, J., and Chintala, S. Pytorch: An imperative style, high-performance deep learning library. In Neural Information Processing Systems, 2019. Cited on page 1,
Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks. In Neural Information Processing Systems, 2018 8. Cited on page 8.
Jesus, R. J., Antunes, M. L.., da Costa, R. A., Dorogovtsev, S. N., Mendes, F. F,, and Aguiar, R. L. Effect of initial configuration of weights training and function of artificial neural networks. Mathematics, 2021. Cited on page 8.
Qiu, S., Potapczynski, A., Finzi, M., Goldblum, M., and Wilson, A. G. Compute better spent: Replacing dense lay- ers with structured matrices. In International Conference on Machine Learning, 2024. Cited on page 11.
Sakurai, J. and Napolitano, J. Modern Quantum Mechan- ics. Cambridge University Press, 2020. Cited page 1.
Jordan, K., Bernstein, J., Rappazzo, B., @fern- bear.bsky.social, Vlado, B., Jiacheng, Y., Cesista, F., Koszarsky, B., and @ @Grad62304977. modded- nanogpt: Speedrunning the nanogpt baseline, 2024a. URL https: /tttt/gittuu com/ KellerJordan/ modded-nanogee Cited on page 11.
Shi, H.-J M., Lee, T.-H., Iwasaki, S., Gallego-Posade J., Li, Z.,, Rangadurai, K., Mudigere, D., and Rabbat, M. A distributed data-parallel Py PyTccc implementation of the distributed Shampoo optimizer for training neural networks at-scale. arXiv:2309.06497, 2023. Cited on pages and 8.
Jordan, K., Jin, Y., Boza, V., You, J., Cecista, F., Newhouse, L., and Bernstein, J. Muon: An optimizer for hidden layers in neural networks, 2024b. URL htttps://wbb archive org/ http://eeee222000000444///11 //kel lelllerjoo joran.. github. 1 http//possss/mmono Cited on pages 2, 7, and 9.
Streeter, M. Universal majorization-minnnnizaooo algo- rithms. SxXiv:3308.00190 2023. Cited on page 2.
Streeter, M. and Dillon, V. Automatically bounding the Taylor remainder series: Tighter bounds and new applications. arXiv:2212. 14299, 2022. Cited page 2.
Kovarik, Z. Some iterative methods for improving orthonor- mality. SIAM Journal on Numerical Analysis, 1970. Cited on page 7.
Tieleman, T. and Hinton, G. RMSprop. Coursera: Neural Networks for Machine Learning, Lecture 6.5, 2012. Cited on page 2.
Krizhevsky, A. and Hinton, G. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009. Cited page 11.
Tran, D. T., Ono, N., and Vincent, E. Fast DNN train- ing based on auxiliary function technique. International Conference on Acoustics, Speech and Signal Processing, 2015. Cited on page
Lakić, S. On the computation of the matrix k-th root. Jour- nal of Applied Mathematics and Mechanics, 1998. Cited on page? 7.
10

---

Modular Duality in Deep Learning
Yang, G. and Hu, E. J. Tensor programs IV: Feature learn- ing in infinite-width neural networks. In International Conference Machine Learning, 2021. Cited on pages 5, 8, and 9.
with coefficients (2, (2,-1.5,0.5).
A.2 Erasure of Watermarked Initial Weights
We run experiments with hidden layer width 1024 and learn- ing rate ranging across 2 ..... 23, We train for 1000 steps with batch size 1024.
Yang, G., Simon, J. B., and Bernstein, J. A spectral condi- tion for feature learning. arXiv:2310.178/33 2023, Cited on pages 5 and 8.
The dualization uses ten steps of a quintic Newton-Schulz iteration with coefficients (3.0, -3.2, 1.2), followed by mul- tiplication by the dimensional constant. These quintic coef- ficients are different from above but lead to the same duality map UEV ~ /dout/ d x UV
A Experimental Details
The dataset for all experiments is CIFAR-10 (Krizhevsky & Hinton, 2009). We use the standard train and test splits with data augmentation.
Datasets,
To align the maximum stable learning rate between dual- ized and non-dualized training, we spectrally normalize the regular gradient descent update as g - g/1gll+.3. We also divide by 3 to match the scalar division that occurs in Linear. dualize due to the module masses for 3-layer MLP. This way both the dualized and non-dualized methods become unstable at the same learning rate, approximately
Architectures. The architecture for all experiments is 3-layer MLP with ReLU nonlinearity, one hidden layer, and no biases. Having three layers allows all three types of weight matrix shape to present: square, wide rectangular, and tall rectangular. In the learning rate transfer experiment, the width of the hidden layer varies from 32 4096. In the weight erasure experiment, the width is fixed 1024.
Ir 1.1, seen in Figure 2. Disabling spectral normaliza- tion does not change the qualitative findings.
Precision. All experiments the default precision float32. It is well demonstrated in the NanoGPT speedruns that Newton-Schulz iterations can run in bfloat16 (Jordan et al., 2024a). We believe is promising future direction to explore whether dualization leads to improved speed and scalability further reduced precisions.
A.1 Learning Rate Transfer
We run experiments with hidden layer widths 32, 64, 128, 256, 512, 1024, 2048, and 4096. For each width, we sweep between 10 and 20 different learning rates. We train for 20 epochs with batch size 128.
Adam (SP). The hyperparameters for Adam are the default ones in PyTorch, except for the sweep over the learning rate. Weight matrices are initialized according to the default PyTorch initialization (Kaiming uniform).
Adam (uP). In μP, the learning rate of each layer is equal to the global learning rate (which sweep) divided by that layer's input dimension din, and weight matrices are initialized as zero-mean Gaussians with standard deviation
v /min (di, dout)/ dz (Qiu et al., 2024). All other hyperparameters are the default ones in PyTorch.
Dualization. We use orthogonal weight initialization. Con- cretely, we create weight matrices with unit Gaussian entries and then iterate them through Newton-Schulle for 30 steps. Our duality-based optimizer in this experiment uses no mo- mentum. It passes the raw gradient through the duality map UEVT - /doue/ dn UVT, implemented via 5 steps of Newton-Schulee iteration and then multiplying by the dimen- sional constant. We use quintic Bewton-Schull iteration


Reviews:
[{"review_id": "ePrqjNOCyM", "reviewer": "Reviewer_NWmj", "summary": "This paper proposes a recipe for neural network design and optimization via \"modular dualization\". A module consists of a forward pass operation, \"mass\" and \"sensitivity\" parameters, and a *norm* associated with the weight space. This design allows for concatenation and composition of modules. The key insight is that the choice of \"norm\" encodes the desired semantics of the module, and thus affects the geometry of optimization over the weight space. The optimization direction over a given module's weight space is given by the dualization map, i.e. steepest descent with respect to the module's norm. A couple of key sample modules are provided, describing standard linear, embedding, and Conv2D layers, as well as (trivially) weight-less \"bond\" layers. The benefit of this perspective is demonstrated through the derivation of a new, highly performant optimizer Muon through the modular duality lens, and a simple derivation of the maximal update parameterization ($\\mu$P) rule.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer NWmj, we are sincerely grateful for your time and effort reviewing for ICML. We also really appreciate your thorough and positive review of our work.\n\nGiven your comments that *“the main ideas in this paper are appealingly simple to read and understand”* and also *“likely to have immediate impact on optimizer/model design”*, we wondered if you would be willing to champion our paper to the area chair?\n\nRegarding Reviewer GE3w’s review, we noticed that they compare our work to the optimizer anthology, which is a non-archival workshop paper. And even so, our work goes beyond the anthology by placing the ideas in a general and forward-looking theoretical framework, showing new experimental results such as on the novel numerical properties of dualized training methods, proposing the connection between Shampoo and muP, and establishing implications for deep learning software libraries.\n\nWhile Reviewer BSBb provides a careful and rigorous review, we feel that they unfairly characterize our work as incremental. While convergence rate analysis is certainly an important future direction, it was not a goal of our work. We ask that our work be evaluated as a piece of science accounting for its implications for unifying optimization theory, for building new kinds of neural network software libraries, for introducing new numerical linear algebra primitives to the deep learning optimization literature and for potential implications for deep learning number systems.\n\nWe are also glad to answer your questions:\n- **The role of the mass and sensitivity attributes.** If we compose two modules, the input sensitivity of the second module is used to re-scale the norm of the first module (part d of Definition 6). In turn, this means the duality map will calibrate the size of perturbations to the first module with regard to the input sensitivity of the second module. As for mass, this provides the user with control to manually re-scale the norms of certain modules in order to provide precise control over how much feature learning each submodule contributes to the overall network. The motivating application is to allow you to set the update size in the embedding layer in a transformer independent of how many residual blocks there are. See Section 3.3 in the modular norm paper for discussion of this.\n- **On the choice of the RMS-RMS norm for Linear modules.** We actually do not think that the RMS–RMS norm is necessarily always a good choice. The idea is that if you have RMS control on the inputs and you want RMS control on the outputs, then RMS–RMS control on the weight updates is a good idea. This seems to match behaviour in the hidden layers of transformers where best practice was already to RMS-normalize the activation spaces (e.g. LLaMa https://arxiv.org/abs/2302.13971). But, as you suggest, if your input or output data has different structure you might want to consider different norms such as L1 or L-infinity for two simple examples.\n- **On the intuitive benefits of boosting small singular values in the gradient.** We think the idea here is that the small singular values are not necessarily noise. If you inspect the singular value distribution of gradients, as done in say https://arxiv.org/abs/2310.17813, you notice that *most* of the gradient singular values are actually small compared to the max. From this perspective, it could seem wasteful to make effectively low rank gradient updates as you are not making use of a lot of signal in the gradient. Of course it’s possible that the *very tiny* singular values are still noise. This is an interesting question to explore further.\n- **On a general recipe for deriving maximal update schemes.** Actually, the purpose and construction of the modular norm is meant to provide a general recipe for obtaining feature learning for general architectures. The paper proposes that feature learning is obtained by scaling updates in a norm with three key properties:\n1. The neural network output is weight-lipschitz in the norm\n2. The Lipschitz constant is non-dimensional (does not depend on e.g. width or depth)\n3. The tightness of the Lipschitz guarantee is independent of network size\n\nIf you find a norm that achieves these properties, then it’s reasonable that using it to scale updates would confer precise and scale-independent control on the amount of feature learning. See Section 2.1 of the modular norm paper for informal discussion of this.\n\nThank you again for your review. Again, we are immensely grateful for your time and effort reviewing for ICML.", "claims_and_evidence": "The claims in this paper are well-supported by clear exposition and field-testing via Muon and $\\mu$P.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "I have verified the correctness of all the theoretical results in this paper.", "experimental_designs_or_analyses": "The numerical results are sensible and details are documented in Appendix A.", "relation_to_broader_scientific_literature": "This work seems to follow a very recent line of work that aims to design a rigorous modular framework for designing deep learning set-ups in a way that co-designs the architecture with the optimizer in mind. This is a valuable contribution to the community, as it has the potential to both unify many seemingly disparate threads (e.g.\\ $\\mu$P and Muon), as well as be the jumping off point for designing task-specific architectures/optimizers.", "essential_references_not_discussed": "Not that I'm aware of.", "other_strengths_and_weaknesses": "In addition to what's listed above, I think the main ideas in this paper are appealingly simple to read and understand. I think the well-normed module and modular norm ideas are likely to have immediate impact on optimizer/model design, especially since there seems to be revived interest in new deep learning optimizers that depart from the Adam family. The perspective of tying a layer's optimizer direction with its particular geometry makes a lot of sense, and may plausibly lead to more interpretable architecture behavior.\n\nI have a few minor questions not immediately answered in the paper:\n- What is the role of the mass and sensitivity attributes? Do these parameters affect the choice of norm depending on where in the architecture the module is?\n\n- Why is, e.g., RMS-RMS norm intuitively a good choice for a Linear module? As a possibly silly sanity check, if we apply a single Linear module on a linear least-squares problem, why should we expect/want the data *and* weight distribution to lie in unit balls, since the output $y = Mx$ can be made arbitrarily large or ill-conditioned. As a related note, would an \"optimal\" choice of norm depend on the data/activation distribution?\n\n- Why is boosting small singular values as in rectified Shampoo/Muon intuitively good? An immediate thought is that this would put the \"noise\" and \"signal\" directions of the weight update at the same magnitude. \n\n- It is shown $\\mu$P is recovered. Is there a general recipe for deriving maximal update / feature learning rules *given* a recipe of modules?", "other_comments_or_suggestions": "Minor comments:\n- Notation for spectral norm is a little confusing, since $\\|\\cdot\\|_\\ast$ is often used to denote the nuclear norm.\n\n- Different places use different symbols for scalar multiplication $\\times, \\ast$.", "questions_for_authors": "No critical questions; some clarification questions listed earlier.", "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": "8Zpc3qhso9", "reviewer": "Reviewer_GE3w", "summary": "The paper introduces a recursive procedure called modular dualization for constructing duality maps in general neural architectures. This method unifies two important optimization techniques—maximal update parameterization and Shampoo—by demonstrating that both are partial approximations of a single duality map induced by the RMS–RMS operator norm. The modular dualization procedure works by assigning operator norms to individual layers based on their input-output semantics, making the construction explicitly recursive and easy to implement in software packages. Essential features of both µP and Shampoo are recovered from the duality map `Linear.dualize`, placing these methods within a common theoretical framework. This unified approach has led to significant wall-clock speedups in training transformers ranging from 124 million to 1.5 billion parameters. Inspired by prior work on optimization algorithms that adapt to computation graph structures, the authors aim to provide a clarifying toolkit for the design and analysis of deep learning systems through their theory of modular duality.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and effort reviewing for ICML.\n\nFirst, we point out that the **optimizer anthology [3]** is a non-archival workshop paper, and therefore a conference submission on the same topic is in accordance with ICML policy. But even given this, our work goes beyond the anthology by placing the ideas in a general and forward-looking theoretical framework, showing new experimental results such as on the novel numerical properties of dualized training methods, proposing the unification of Shampoo and muP as approximations to a single duality map, and establishing implications for deep learning software libraries through the directly programmable structure of the duality maps.\n\nSecond, we are delighted the reviewer has characterized the **Shampoo algorithm** as *“gradient descent w.r.t. matrix operator norm”*—this characterization is actually the perspective proposed by the **optimizer anthology [3]**, which again is a non-archival workshop paper. In contrast, the original **Shampoo paper [1]** presents Shampoo as an approximation to full-matrix Adagrad—see, for example, Sections 1.1 and 1.2 of the Shampoo paper [1]. Even the full-matrix Adagrad perspective on Shampoo is controversial ([Xie et al 2025](https://arxiv.org/abs/2503.10537)). But, taken together, we see the reviewer’s characterization of Shampoo in this way as evidence of the appeal of the matrix norm and duality perspective!\n\nWith regard to Muon and Newton-Schulz, we are grateful for the reviewer bringing this up. We will clarify the language in the paper to make clear that on this axis we made two original contributions that were critical to the speed and success of Muon:\n1. Proposing using Newton-Schulz iterations to do gradient orthogonalization\n2. Proposing the idea of treating the polynomial coefficients in Newton-Schulz as tunable hyperparameters to accelerate the convergence\n\nOn top of these ideas, of course, Muon adds momentum and various systems innovations such as low-precision casts and a low overhead multi-GPU distributed implementation. We would be delighted to discuss any of these points further. Given these clarifications on the novelty of our contributions, we would be grateful if the reviewer would consider substantially increasing their score.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "No.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "Not applicable.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Overall the paper is well-written and it motivates well.\n\nThe theoretical novelty of this paper is quite limited. It looks like a concatenation of several previous works including steepest descent on a normed space [3], modular norm [2], and gradient descent w.r.t. matrix operator norm (Shampoo [1]). In other words, this paper looks more like a Systemization of Knowledge (SoK) paper instead of a standard conference paper. It is hard to tell which part of this paper is novel, i.e., not from previous paper. For instance, even Example 7 seems to be a simple extension of the linear modular based on the norm defined on line 313-314.\n\nRegarding the application contribution, please discuss the relation of this paper and Muon. The dualize function for the linear module is simply the dual norm of the operator norm of a matrix (with some rescaling factor), which is already introduced in Shampoo [1]. Muon takes the rectangular Newton-Schultz iteration. The author claims that Muon’s algorithm is based on the idea in this paper. However, it is simply an implementation method to calculate $UV^T$ without directly executing the SVD of a matrix. Claiming the credit of the invention of this implementation method is not well-supported. \n\n[1] Vineet Gupta, Tomer Koren, and Yoram Singer. “Shampoo: Preconditioned stochastic tensor optimization.” International Conference on Machine Learning. PMLR, 2018.\n\n[2] Large, T., Liu, Y., Huh, M., Bahng, H., Isola, P., and Bernstein, J. Scalable optimization in the modular norm. In Neural Information Processing Systems, 2024.\n\n[3] Jeremy Bernstein and Laker Newhouse. “Old optimizer, new norm: An anthology.” arXiv preprint arXiv:2409.20325 (2024).", "other_comments_or_suggestions": "The notation of RMS norm is a bit hard to understand. RMS seems to be the abbreviation of root mean square. Nevertheless, the RMS norm in the paper is defined to be a rescaled version of standard l2 norm with a scaling factor of $1/\\sqrt{d}$, which can be understood by taking the root mean square across all the dimension. It will be better to explain the reason to call this norm RMS norm or cite a paper which introduces RMS norm.", "questions_for_authors": "No.", "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": "RpqwIEnkby", "reviewer": "Reviewer_BSBb", "summary": "The contributions of this paper could be summarised as follows:\n* it combines the notion of dualization/steepest descent with the notion of a modular norm (a max-of-norms aggregation of norms tailored to each single module). This goes beyond previous works on steepest descent that only consider a l1-type aggregation. \n* it proposes a norm choice for the standard deep learning modules. The proposed choices are informally motivated, and lead to updates related to Shampoo and muP, two\nsuccessful techniques in deep learning.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer BSBb, thank you for your contributions to the conference. We are grateful for your constructive and thorough review of our paper. We hope to provide useful responses to your questions!\n\nFirst, we ran new experiments to address your questions about the long-term performance of duality-based optimizers.\nWe created an anonymized GIF ([link here](https://gifyu.com/image/bzygu)) so you can directly watch the training loss fall over 100 epochs on CIFAR-10 for Adam (SP), Adam (µP), and Dualization. The GIF shows that dualization has lower training loss than Adam at every epoch. For example, dualization reaches loss 1e-2 in epoch 17, while Adam reaches the same loss in epoch 56. And [here is another GIF](https://gifyu.com/image/bzyCY) for test accuracy. By the end, Adam and dualization both saturate around accuracy 60%, which is typical for an MLP on CIFAR-10.\n\nSecond, we really appreciate your questions about convergence analysis and nailing down norm choices. When evaluating our paper, we kindly ask that the reviewer consider a broad idea of what an important optimization paper might look like. We contend that our paper makes the following contributions of substantial importance:\n1. We build the first norm-based duality theory for deep learning that considers the tensor structure of the model and does not amount to updating one layer at a time. This is revitalizing interest in norm-based deep learning optimization theory, with exciting followup work based on our work involving ideas like linear minimization oracles, Frank-Wolfe analysis and trust region analyses.\n2. We introduce the Newton-Schulz orthogonalization primitive to the optimization literature. This is already having a substantial practical impact in industry and is inspiring followup research on experimenting with these methods in academia.\n3. We theoretically reconcile the Shampoo optimization algorithm with the maximal update parameterization. Anecdotally, these were both regarded as some of the hardest techniques to understand. We provide a new and easy way to unify these techniques that immediately suggests new ways to extend them.\n4. We demonstrate that dualized training algorithms automatically exhibit transferable learning rates.\n5. We also show that dualized algorithms have novel numerical properties. This is an important scientific contribution since it provides a direct counterexample to the idea that the weights don’t change in wide networks, which inspired a lot of NTK research. It may also have implications for computer number systems.\n\nIn short, we hope that you will take another look at our paper with an open mind. We agree that convergence analyses and an exhaustive experimental analysis of different norm choices are exciting directions, but they were not priorities for our paper.\n\nAs for your comments on tightening the mathematical notation, thank you for them. We agree in most cases but not all of them:\n- We agree that we glossed over the set-valued nature of “dualize” and how it should act on the zero input. We will clarify this as suggested.\n- We really like the reviewer’s suggestion of introducing the Jacobian-vector product. We will implement this idea as suggested.\n- Regarding your comment that we cannot add the gradient (in our parlance) to the weights, the reviewer has missed line 110 where we state that the weight space is the Euclidean space $W=R^n$, as is the case in deep learning. Therefore our presentation is sound. We will clarify this in the paper.\n- The reviewer correctly notes that many textbooks define gradients to live in primal space. But there is a gap between these definitions (e.g. Blondel/Roulet Proposition 2.4), which assume an inner product space, and deep learning where we lack a canonical inner product. Furthermore, in PyTorch/JAX, we usually call loss derivatives \"gradients\". Even more subtly, if we choose to equip the network with the dot product on flattened weight space, then the reviewer's gradient is equivalent to our paper's gradient! We glossed over these technicalities for accessibility, but propose adding an explanatory paragraph and welcome reviewer collaboration on this issue.\n\nAs for your other questions:\n- **watermark erasure experiment**. The method is SGD but with vanilla spectral normalization applied to the updates to match the learning rate scale to the dualized method (see Appendix A.2). Since the rank of the gradient is upper bounded by the batch size, batch size also limits the maximum possible stable rank of the dualized gradient, which is what drives watermark erasure.\n- **different Newton-Schulz iterations**. These experiments were simply run by different authors at different times. If you are interested to know more about Newton-Schulz, any coefficients that approximate sign(x) yield essentially the same duality map. The important practical consideration is the linear coefficient and the number of iterations, which set the inflation factor of small singular values.", "claims_and_evidence": "One shortcoming of the paper is the lack of theoretical foundation for the proposed framework: steepest descent has a general convergence theory (for example in the papers that are given as references, or in the book by Nesterov). Hence, the question arises how the choices for the module norms affect the rates and constants in this convergence theory. This would also be a possible way to motivate certain norm choices (for example, if the corresponding smoothness constants would be small). However, the paper in its current form makes no efforts in this direction.\n\nIt should be remarked that the paper also does not claim to make contributions in this direction.\n\nAnother (minor) shortcoming is that the motivation through a type system seems to not capture the situation, considering the standard mathematical formalization (see details below).", "methods_and_evaluation_criteria": "Evaluation criteria make sense (convergence and LR transferability).\n\nHowever, the choice of module norms have been already proposed in Large et al. 2024, and experimental or theoretical ablations/evidence for each particular choice is missing. For example, is the l1-RMS choice for embedding modules superior to using RMS-RMS?", "theoretical_claims": "There are no major proofs or theoretical claims that need to be checked.", "experimental_designs_or_analyses": "The experimental design seems sound.\n\nSome questions on experiments:\n\n* For Figure 1, how does it look after a larger number of epochs (20 epochs will not be sufficient for high accuracy)? Does the dualization approach still lead to lower loss, or is this effect only visible with a short training time?\n\n* In section 6.5 it is explained that the watermark erasure can be seen when the batch size is not too small: could you elaborate how the batch sizes comes into play? Also, is the method without dualization just SGD or Adam here?\n\n* Why are the iterations and coefficients for Newton-Schulz chosen differently in the two experiments?", "relation_to_broader_scientific_literature": "The main contribution of the paper seems to be to connect/motivate techniques that have been reported to improve training of deep learning models (e.g. muP and steepest descent methods). The amount of theoretical or empirical advancements in this paper itself however seems limited, especially given the similarity to Large et al. 2024.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "There are several issues with respect to correct mathematical notation in the paper (see below). These notation choices might have been made in order to keep it simple; however, the mathematical correctness is suffering from this choice:\n\n* The mapping \"dualize\" by its definition is a set-valued mapping (as the argmax is not unique in general). With the current definition (that reads as if dualize returns a single element of the space $W$), one quickly runs into formal issues: the equation in Proposition 1 is ill-defined, as the left-hand side is a set, but the right-hand side is a single element. \n* Example 1 is ill-defined if $g=0$. Again, in order to properly define this, dualize should be set valued, and then dualize(0) is the unit ball.\n* The lines 119-123 are strange: If the gradient is considered an element of the space $W^*$ as stated, then as long as $W^* \\neq W$ we *can not* add it to an element of $W$ by definition. However, this paragraph reads as if we would have a choice in doing so or not (\"we shall forbid ourselves\").\n* It should also be mentioned, that usually the gradient is defined as the element of $W$ that represents the linear mapping of the derivative (e.g. see Prop 2.4 in https://arxiv.org/pdf/2403.14606) via the Riesz representation theorem. Hence, in many textbooks the gradient is defined as an element of $W$, wheras the paper defines it as an element of $W^*$, which might lead to confusion. It would be beneficial to introduce the Jacobian-vector product (see Def 2.13 in  https://arxiv.org/pdf/2403.14606), and then motivate by considering the problem \n$$ \\arg \\min_{\\\\|w\\\\|\\leq 1} \\partial L(w)[\\Delta w]. $$\nIf the norm in the constraint is not the Euclidean norm, then the solution to this problem is not necessarily the negative gradient (scaled), and as a consequence we need to introduce the dualization mapping. \nMathematically speaking, as the gradient (as usually defined) is an element of the same space as the weight $w$, the motivation via a type system seems to not capture the situation (even though it is a useful metaphor).\n* The notation in lines 168-170 (right column) are not fully clear: what exactly is \"summation over any shared tensor indices\"? I think it would be much easier to introduce the Jacobian-vector product, then this quantity can be simply written as $\\partial_w M.forward(w,x)[\\Delta w]$.", "other_comments_or_suggestions": "Minor comments:\n\n* The multiplications is sometimes denoted as $\\times$ (Prop 1), sometimes as $*$ (Def 6). Please align these notations.\n* Section 3.4: \"are also smooth in an appropriate sense\". Can you provide a reference for this statement?\n* Definition 5 appears identically in Large et al., 2024. Please refer to it in the statement of Def. 5, to emphasize that this is not a new concept proposed in this paper.\n* Lines 431-436: how can papers from 2019 and 2021 inspire a paper from 2018?", "questions_for_authors": "For questions on experiments, see the according section.\n\nThe other questions follow from my comments above, but I will repeat some here:\n\n1) Why is the l1-RMS choice for embedding modules superior to using RMS-RMS? Did you run an ablation for this?\n\n2) Can the convergence theory of steepest descent give any insight into specific norm choices?", "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:

Meta-review: Dear Authors,

Thank you for your submission to ICML and for your contribution. Your paper proposes a novel, theoretically grounded framework for constructing duality maps tailored to neural architectures, and positions this framework as a unifying lens for interpreting and improving modern optimizers like Shampoo and maximal update parameterization (µP). The work presents a recursive construction of duality maps, practical algorithms for dualizing key layers, and early empirical results highlighting the potential of duality-based optimization.

The reviewers appreciated the clarity of exposition and the conceptual originality of the framework. Reviewer NWmj found the ideas compelling and potentially impactful for future optimizer/model co-design. Reviewers GE3w and BSBb initially raised concerns regarding the novelty relative to non-archival prior work, but ultimately acknowledged the distinctive theoretical contributions, unification of existing methods, and implications for practical optimization tools like Muon. Your thoughtful rebuttal, additional clarifications, and concrete follow-up experiments helped address their concerns and resulted in increased scores.

Given the strength of the theoretical contributions, the practical relevance demonstrated through recent adoption, and the overall positive trajectory of the discussion, I am pleased to recommend acceptance to ICML 2025. I encourage you to use the camera-ready version to further highlight the framework's implications for feature learning, norm choice, and optimizer design, and to clarify mathematical notation as discussed with reviewers.

Congratulations on your contribution.

Best regards,
AC
Decision: Accept (poster)