Modular Duality in Deep Learning

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!

First, we ran new experiments to address your questions about the long-term performance of duality-based optimizers. We created an anonymized GIF (link here) 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 for test accuracy. By the end, Adam and dualization both saturate around accuracy 60%, which is typical for an MLP on CIFAR-10.

Second, 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:

1. 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.
2. 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.
3. 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.
4. We demonstrate that dualized training algorithms automatically exhibit transferable learning rates.
5. 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.

In 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.

As for your comments on tightening the mathematical notation, thank you for them. We agree in most cases but not all of them:

• 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.
• We really like the reviewer’s suggestion of introducing the Jacobian-vector product. We will implement this idea as suggested.
• 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.
• 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.

As for your other questions:

• 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.
• 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.

We thank the reviewer for their time and effort reviewing for ICML.

First, 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.

Second, 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). 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!

With 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:

1. Proposing using Newton-Schulz iterations to do gradient orthogonalization
2. Proposing the idea of treating the polynomial coefficients in Newton-Schulz as tunable hyperparameters to accelerate the convergence

On 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.

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.

Given 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?

Regarding 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.

While 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.

We are also glad to answer your questions:

• 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.
• 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.
• 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.
• 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:

1. The neural network output is weight-lipschitz in the norm
2. The Lipschitz constant is non-dimensional (does not depend on e.g. width or depth)
3. The tightness of the Lipschitz guarantee is independent of network size

If 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.

Thank you again for your review. Again, we are immensely grateful for your time and effort reviewing for ICML.