Modular Duality in Deep Learning

Introduction

In this paper, we pursue a rigorous and first-principles theoretical framework for designing neural network training algorithms. We hope that building such a framework will facilitate the design of a 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:

Why? The reason is that the loss function may not be equally smooth in all directions in weight space, and there is no reason for the sizes of different components of the raw gradient vector to respect this heterogeneity. In other words, the geometry of the loss function may be non-isotropic. Insisting on a type check should force the user to become cognizant of this issue and to find a suitable duality map. A good duality 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 isomorphism in differential geometry (Grosse, 2022), raising and lowering indices in general relativity (Carroll, 2019) and the bra-ket notation in quantum mechanics (Sakurai & Napolitano, 2020). Duality maps are also central to several optimization theories including mirror descent (Nemirovsky & Yudin, 1983), natural gradient descent (Amari, 2016) and steepest descent on a normed space (Boyd & Vandenberghe, 2004). Despite the efforts of some prescient papers (Carlson et al., 2015b; Flynn, 2017), the latter kind of duality map involving normed vector spaces is yet to puncture the deep learning mainstream.

We believe that duality is a 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 contemporary optimization research may be seen as approximations to a single duality map. These methods are maximal update parameterization (Yang & Hu, 2021), which is aimed at scalable training, and Shampoo (Shi et al., 2023), which is targeted at fast training. We show in Section 4.1 that both methods emerge as partial approximations to a single duality map induced by the RMS–RMS operator norm.

The main contribution of this paper is to describe a procedure for constructing duality maps for general neural architectures. The procedure, which we call modular dualization, works in three steps:

Operator norms are assigned to individual layers based on the input-output semantics of each layer;

Based on these operator norms, duality maps are constructed for individual layers;

Given the layerwise duality maps and the structure of the neural architecture, a single duality map is recursively induced on the full weight space of the architecture.

To instantiate this procedure for a rich family of neural architectures—including convolutional networks and transformers—we write down duality maps for $\mathsf{Linear}$, $\mathsf{Embed}$ and $\mathsf{Conv2D}$ layers. We also provide novel, GPU-friendly algorithms for computing these duality maps. Overall, we hope that modular dualization will help in the principled design of the machine learning systems of the future.

Related Work

This paper constructs a duality map for general neural architectures. Our approach is based on assigning operator norms to individual network layers and using these layerwise 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., 2015a; b; 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 models (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-$\infty$ 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-$\infty$ norm. The authors call the corresponding duality map for linear layers the “#-operator”—a name presumably inspired by the musical isomorphism (Grosse, 2022). The authors propose a cheap approximation to the #-operator based on sketching (Martinsson & Tropp, 2020), and they also propose a way to mix RMSprop-style pre-conditioning information (Tieleman & Hinton, 2012) into the weight updates. In contrast to our work, the authors only derive duality maps for single linear layers, and these maps are then heuristically extended to all-layer updates. Nonetheless, the authors achieve substantial wallclock speedups using variants of spectral descent to train small networks.

Now, let us turn our attention to duality structure gradient descent (Flynn, 2017), which constructs a 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 (2017)’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 to our approach is that Flynn (2017) leverages a weighted sum ($L_{1}$ combination) of layerwise norms to construct his full duality map. This leads to optimization methods that only update a single layer at each iteration, and the methods need to be heuristically extended to achieve all-layer updates. In contrast, we leverage the modular norm (Large et al., 2024), which takes a weighted max ($L_{\infty}$ combination) of layerwise norms. In turn, our duality map leads directly to more conventional all-layer optimizers.

Another important difference between our work on modular duality and prior work on duality structure gradient descent is that we fully “modularize” our theory—meaning that our construction is explicitly recursive—and as such it is easy to code up into a software package. In this regard, 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 breaking them up into subexpressions and applying composition theorems from convex analysis. More recent progress in this vein includes work on universal majorization-minimization algorithms (Streeter & Dillon, 2022; Streeter, 2023) and related papers on automatic majorization (Tran et al., 2015; Bernstein et al., 2023).

Theoretical Preliminaries

In this section, we introduce duality maps, a means of constructing duality maps based on norms, and finally a norm called the modular norm that is well-suited to describe the geometry of general neural architectures.

Observe that, in the first-order term, the gradient $\nabla_{\bm{w}}\mathcal{L}({\bm{w}})$ is acting as a linear functional: it is pairing with the weight vector $\Delta{\bm{w}}\in\mathcal{W}$ in a linear way to produce a real number. As such, we shall say that the gradient belongs to the dual weight space: $\nabla_{\bm{w}}\mathcal{L}({\bm{w}})\in\mathcal{W}^{*}$. We shall forbid ourselves from directly subtracting a member of the dual weight space $\mathcal{W}^{*}$ from the weight space $\mathcal{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 $\mathcal{W}$.

2 Steepest Descent on a Normed Space

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 $\Delta{\bm{w}}$ by minimizing the right-hand side of Equation 4. We will show that the minimizer can be expressed in terms of a dual norm and a duality map:

where, if the $\operatorname*{arg\,max}$ is not unique, $\operatorname{dualize}_{\|{\cdot}\|}$ returns any maximizer.

Given these definitions, minimizing the expression in the right-hand side of Equation 4 can be done using the following standard proposition, for which Bernstein & Newhouse (2024) provide a proof:

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 defined on general neural architectures. The next section reviews duality maps for more standard norms.

3 Basic Norms and Duality Maps

Many basic norms and duality maps are already covered in prior work (Carlson et al., 2016; 2015a; 2015b; Flynn, 2017). For some warmup examples, the following duality maps for vector norms are standard:

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 embedding versus linear layers in a transformer. As such, we will need duality maps for different induced operator norms:

4 The Modular Norm

The modular norm (Large et al., 2024) is intended to help characterize the heterogeneous curvature of general neural architectures. The construction first defines an abstract module type along with a notion of what is a good, or well-normed, module. Then combination rules are given for constructing new well-normed modules from a library of existing well-normed modules. So modules are a special case of combinator pattern from functional programming (Haskell Wiki Contributors, 2007). Modules are also related to the monoidal category from category theory (Fong & Spivak, 2019). We begin by defining the abstract notion of a module:

Given input vector space $\mathcal{X}$, output vector space $\mathcal{Y}$ and weight vector space $\mathcal{W}$, a module $\mathsf{M}$ is an object with the following four attributes:

a function, $\mathsf{M}\mathsf{.forward}:\mathcal{W}\times\mathcal{X}\to\mathcal{Y}$, which maps an input and a weight vector to an output;

a number, $\mathsf{M}\mathsf{.mass}\geq 0$, which is used to set the proportion of feature learning that this module contributes to any supermodule;

a number, $\mathsf{M}\mathsf{.sensitivity}\geq 0$, which estimates the module’s sensitivity to input perturbations;

We shall care most about modules that are well-normed, which amounts to requiring that the forward function is Lipschitz-continuous in the weights with constant 1 and in the inputs with constant $\mathsf{M}\mathsf{.sensitivity}$:

Let $\mathsf{M}$ be a module on $(\mathcal{X},\mathcal{Y},\mathcal{W})$, where the input and output spaces have respective norms $\|{\cdot}\|_{\mathcal{X}}$ and $\|{\cdot}\|_{\mathcal{Y}}$. $\mathsf{M}$ is well-normed if for all inputs ${\bm{x}}\in\mathcal{X}$ and weights ${\bm{w}}\in\mathcal{W}$:

The $\diamond$ operator denotes summation over any shared tensor indices. This definition of well-normed-ness can be used as a guiding principle in the design of a 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 $\mathsf{M}\mathsf{.forward}$. Then a norm $\mathsf{M}\mathsf{.norm}$ should be assigned to the module’s weight space and a number $\mathsf{M}\mathsf{.sensitivity}$ should be chosen to make the module well-normed. Examples are given in Section 4.1.

Given such a library of well-normed atomic modules, a compound module built through any arbitrary sequence of module compositions and module concatenations is automatically well-normed (Large et al., 2024). And if the atomic modules in the library are not only well-normed but are also smooth in an appropriate sense, then Large et al. (2024) give an automatic procedure for computing sharpness coefficients for any compound module built from the library. The relevant definition of module composition is as follows:

Consider module $\mathsf{M}_{1}$ with input, output and weight space $(\mathcal{X}_{1},\mathcal{Y}_{1},\mathcal{W}_{1})$ and module $\mathsf{M}_{2}$ with input, output and weight space $(\mathcal{X}_{2},\mathcal{Y}_{2},\mathcal{W}_{2})$. $\mathsf{M}_{1}$ and $\mathsf{M}_{2}$ are composable if $\mathcal{X}_{2}=\mathcal{Y}_{1}$. Their composite module $\mathsf{M}=\mathsf{M}_{2}\circ\mathsf{M}_{1}$ has input, output and weight space $(\mathcal{X}_{1},\mathcal{Y}_{2},\mathcal{W}_{1}\times\mathcal{W}_{2})$ and attributes:

$\mathsf{M}\mathsf{.forward}(({\bm{w}}_{1},{\bm{w}}_{2}),{\bm{x}}))=\mathsf{M}_{2}\mathsf{.forward}({\bm{w}}_{2},\mathsf{M}_{1}\mathsf{.forward}({\bm{w}}_{1},{\bm{x}}))$;

$\mathsf{M}\mathsf{.mass}=\mathsf{M}_{1}\mathsf{.mass}+\mathsf{M}_{2}\mathsf{.mass}$;

$\mathsf{M}\mathsf{.sensitivity}=\mathsf{M}_{1}\mathsf{.sensitivity}*\mathsf{M}_{2}\mathsf{.sensitivity}$;

$\mathsf{M}\mathsf{.norm}(({\bm{w}}_{1},{\bm{w}}_{2}))$ given by:

where if $\mathsf{M}_{1}\mathsf{.mass}$ or $\mathsf{M}_{2}\mathsf{.mass}$ is zero, the corresponding term in the $\max$ is set to zero.

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.

Module concatenation is defined in a similar way to module composition:

Consider module $\mathsf{M}_{1}$ with input, output and weight space $(\mathcal{X}_{1},\mathcal{Y}_{1},\mathcal{W}_{1})$ and module $\mathsf{M}_{2}$ with input, output and weight space $(\mathcal{X}_{2},\mathcal{Y}_{2},\mathcal{W}_{2})$. We say that $\mathsf{M}_{1}$ and $\mathsf{M}_{2}$ are concatenatable if their input spaces match: $\mathcal{X}_{1}=\mathcal{X}_{2}$. The tuple module $\mathsf{M}=(\mathsf{M}_{1},\mathsf{M}_{2})$ has input, output and weight space $(\mathcal{X}_{1},\mathcal{Y}_{1}\times\mathcal{Y}_{2},\mathcal{W}_{1}\times\mathcal{W}_{2})$ and the following list of attributes:

$\mathsf{M}\mathsf{.forward}(({\bm{w}}_{1},{\bm{w}}_{2}),{\bm{x}}))=(\mathsf{M}_{1}\mathsf{.forward}({\bm{w}}_{1},{\bm{x}}),\mathsf{M}_{2}\mathsf{.forward}({\bm{w}}_{2},{\bm{x}}))$;

$\mathsf{M}\mathsf{.mass}=\mathsf{M}_{1}\mathsf{.mass}+\mathsf{M}_{2}\mathsf{.mass}$;

$\mathsf{M}\mathsf{.sensitivity}=\mathsf{M}_{1}\mathsf{.sensitivity}+\mathsf{M}_{2}\mathsf{.sensitivity}$;

$\mathsf{M}\mathsf{.norm}({\bm{w}}_{1},{\bm{w}}_{2})$ given by:

where if $\mathsf{M}_{1}\mathsf{.mass}$ or $\mathsf{M}_{2}\mathsf{.mass}$ is zero, the corresponding term in the $\max$ is set to zero.

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 modules are sometimes just naïvely divided by their norm. In this paper we make fuller use of the geometry implied by the modular norm by constructing the corresponding duality map, which we call modular dualization.

Modular Dualization

In this section, we construct a duality map for general neural 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.e. full neural networks, by showing how duality maps should pass through composition and concatenation.

To construct a duality map for an atomic module $\mathsf{A}$, the idea is to first fix norms on the input and output spaces that respect the semantics of $\mathsf{A}\mathsf{.forward}$. 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 outputs to evolve. Then we place a norm on the weight space such that $\mathsf{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 $\mathsf{A}\mathsf{.dualize}$. We now give some examples of this procedure for the basic layer types of $\mathsf{Linear}$, $\mathsf{Embed}$ and $\mathsf{Conv2D}$. The results are summarized in Table 1.

We start with the canonical example of an atomic module:

$\mathsf{Linear}\mathsf{.forward}({\bm{W}},{\bm{x}})={\bm{W}}{\bm{x}}$, the matrix-vector product;

$\mathsf{Linear}\mathsf{.sensitivity}=1$;

$\mathsf{Linear}\mathsf{.mass}=\mu$, where $\mu\geq 0$ is a hyperparameter;

The $\mathsf{Embed}$ module provides a useful counterpoint to the $\mathsf{Linear}$ module. The difference between the two modules stems from the fact that the input spaces of $\mathsf{Embed}$ and $\mathsf{Linear}$ have different semantics.

$\mathsf{Embed}\mathsf{.forward}({\bm{W}},{\bm{x}})={\bm{W}}{\bm{x}}$, the matrix-vector product;

$\mathsf{Embed}\mathsf{.mass}=\mu$, where $\mu\geq 0$ is a hyperparameter;

Finally, we consider a $\mathsf{Conv2D}$ module with a $k\times k$ kernel. $\mathsf{Conv2D}$ has a more involved tensor structure than $\mathsf{Linear}$ and $\mathsf{Embed}$. The calculations work by slicing up the weight tensor into a collection of $k^{2}$ matrices.

$\mathsf{Conv2D}\mathsf{.forward}({\bm{W}},{\bm{x}})={\bm{W}}\circledast{\bm{x}}$, where $\circledast$ denotes 2D convolution;

$\mathsf{Conv2D}\mathsf{.sensitivity}=1$;

$\mathsf{Conv2D}\mathsf{.mass}=\mu$, where $\mu\geq 0$ is a hyperparameter;

2 Duality Maps for Bond Modules

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 $\mathsf{ReLU}$ nonlinearity. For a bond $\mathsf{B}$, the weight space is the zero vector space $\mathcal{W}=\{0\}$ and the modular norm $\mathsf{B}\mathsf{.norm}=0\mapsto 0$. As such, the corresponding duality map is also $\mathsf{B}\mathsf{.dualize}=0\mapsto 0$. In a software package, one need not write norms or duality maps for bond modules.

3 Duality Maps for Compound Modules

First, given two composable modules $\mathsf{M}_{1}$ and $\mathsf{M}_{2}$, the duality map for the composite $\mathsf{M}=\mathsf{M}_{2}\circ\mathsf{M}_{1}$ is given by:

And second, given two concatenatable modules $\mathsf{M}_{1}$ and $\mathsf{M}_{2}$, the duality map for the tuple $\mathsf{M}=(\mathsf{M}_{1},\mathsf{M}_{2})$ is:

The proofs of Equations 11 and 12 follow in a straightforward manner from Definitions 6 and 7.

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 $\mathsf{Embed}\mathsf{.dualize}$ is easy to implement since it just involves computing vector norms of matrix columns. But $\mathsf{Linear}\mathsf{.dualize}$ and $\mathsf{Conv2D}\mathsf{.dualize}$ involve the projection:

where ${\bm{U}}{\bm{\Sigma}}{\bm{V}}^{\top}$ is the reduced SVD of the matrix ${\bm{G}}$. Since computing SVDs can be slow (Carlson et al., 2015b; Flynn, 2017), here we discuss three fast approximations to Equation 13 via sketching, iterations for inverse matrix roots, and a new family of rectangular Newton-Schulz iterations that we propose. Which method works best may depend on the condition number of the matrix ${\bm{G}}$ or the available computational resources.

Sketching is a randomized method (Martinsson & Tropp, 2020) that can be used to build low-rank approximations to the SVD. Carlson et al. (2015b) already used sketching to provide a fast approximation to their $\#$-operator. More recent papers have experimented with sketching in the context of Shampoo-type algorithms (Feinberg et al., 2023). A potential downside of approximating Equation 13 via sketching is that randomized SVD methods usually try to accurately approximate the largest singular values of a matrix (Martinsson & Tropp, 2020, Section 11.2) while the value of Equation 13 may lie in its action on the small singular values.

2 Iterations for Inverse Matrix Roots

Given a full rank matrix ${\bm{G}}$ with reduced SVD ${\bm{U}}{\bm{\Sigma}}{\bm{V}}^{\top}$, we have that:

This provides a route to approximating Equation 13 since one can compute inverse matrix roots such as $({\bm{G}}{\bm{G}}^{\top})^{-\nicefrac{{1}}{{2}}}$ via Newton iteration (Lakić, 1998). This is discussed in Chapter 7 of Higham (2008)’s book and also see Anil et al. (2020)’s paper. Care must be taken with inverses whenever the matrix ${\bm{G}}$ is ill-conditioned.

3 Rectangular Newton-Schulz Iteration

The iteration in Equation 15 has the advantage over sketching that it always works on all singular values, and since the iteration does not compute inverse matrix roots it is well-behaved even on low-rank matrices.

Finally, there are in fact a family of degree $2n+1$ polynomial iterations of the form

for suitable $a,b,c,...,z$ that could be used instead of Equation 15. One should choose coefficients $a,b,c,...,z$ so that the univariate polynomial $g(x)=a\cdot x+b\cdot x^{3}+c\cdot x^{5}+...+z\cdot x^{2n+1}$ is a suitable approximation to $\operatorname{sign}(x)$. One may try to further accelerate the iteration by “tuning” the coefficients $a,b,c,...,z$ empirically.

We came up with Equation 15 by inspecting Equation 5.22 in Higham (2008)’s book, which provides a related iteration for computing the “matrix sign function” for square matrices. We developed the graphical understanding ourselves and used this as the basis for proposing the higher-order polynomial iterations.

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.

Part of the inspiration for this work is the idea of building 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 information would help in the construction of well-normed modules (see Section 4.1). Modules should be typed according to Definition 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.

2 Neural Network Speedrunning

We believe that the ideas in this paper can help in the design of faster training methods. In fact, a new NanoGPT training speed record was recently set (Jordan, 2024) using our rectangular Newton-Schulz iteration. We communicated the iteration to Keller Jordan through our workshop paper (Bernstein & Newhouse, 2024).

3 Modular Duality: A Unifying Theoretical Framework for Fast and Scalable Training

An important topic in contemporary optimization research is the design of fast and scalable training methods for neural networks. In fact, the theme of the Optimization for Machine Learning workshop at this year’s NeurIPS conference is “scaling up optimization” (OPT, 2024). Two popular methods in this research space are maximal update parameterization (Yang & Hu, 2021, $\mu$P), which allows for 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 inaugural AlgoPerf optimization competition (Dahl et al., 2023).

We showed in Section 4.1 that essential features of both $\mu$P and Shampoo are recovered from the single duality map $\mathsf{Linear}\mathsf{.dualize}$. We think that, on a basic theoretical level, $\mu$P and Shampoo should be viewed as partial approximations to this duality map. This observation helps put $\mu$P and Shampoo on a consistent theoretical footing, orients the methods with respect to overlooked prior work on spectral descent (Carlson et al., 2015b) and duality structure gradient descent (Flynn, 2017), and suggests new ways to generalize these methods to arbitrary layer types and network architectures via the modular norm and modular dualization.

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 important 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 important 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 $\mathsf{Linear}\mathsf{.dualize}$ and $\mathsf{Conv2D}\mathsf{.dualize}$ may help simplify the answer to this question, since they project gradients to scaled semi-orthogonal matrices for which all singular values have the same magnitude.

5 A Numerical Paradox: The Weights Don’t Change!

Past work (Lee et al., 2019; Jesus et 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 has motivated a substantial amount of work on linearized training dynamics (Jacot et al., 2018). We attempted to resolve this paradox in prior work 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 change the story again: $\mathsf{Linear}\mathsf{.dualize}$ ramps up the stable rank of updates, so the weights should move a non-trivial relative amount at large width even in the Frobenius norm—provided the batch size is not too small.

Conclusion

Acknowledgements

Many ideas in this paper were developed jointly with Tim Large before he left to work at a tech company. We are grateful to Phillip Isola for invaluable discussions. We also thank Jack Gallagher, Keller Jordan, Simo Ryu, Rogier Brussee, Tongzhou Wang, Victor Butoi, Jeffrey Cider and Volkan Cevher for helpful conversations.

References