Automatic Gradient Descent: Deep Learning without Hyperparameters

Introduction

Automatic differentiation has contributed to the rapid pace of innovation in the field of deep learning. Software packages such as PyTorch (Paszke et al., 2019) and Theano (Al-Rfou et al., 2016) have advanced a programming paradigm where the user (1) defines a neural network architecture by composing differentiable operators and (2) supplies training data. The package then automatically computes the gradient of the error on the training data via recursive application of the chain rule. At this point, the user must become involved again by (3) selecting one of numerous optimisation algorithms and (4) manually tuning its hyperparameters: in particular, the initial learning rate and the learning rate decay schedule (Goodfellow et al., 2016).

But manually tuning hyperparameters is irksome. An abundance of hyperparameters makes it difficult to rank the performance of different deep learning algorithms (Lucic et al., 2017; Schmidt et al., 2021) and difficult to reproduce results in the literature (Henderson et al., 2018). Hyperparameters confound our efforts to build a scientific understanding of generalisation in deep learning (Jiang et al., 2020; Farhang et al., 2022). And, when training neural networks at the largest scale, in pursuit of stronger forms of artificial intelligence, hyperparameter grid search can rack up millions of dollars in compute costs (Sharir et al., 2020).

Are hyperparameters just a fact of life? The thesis of this paper is that no: they are not. Deep learning involves fitting a known function to known data via minimising a known objective. If we could characterise these components both individually and in how they interact, then—in principle—there should be no leftover degrees of freedom to be tuned (Orabona & Cutkosky, 2020). Taking this idea and running with it leads to automatic gradient descent (AGD): a neural network optimiser without any hyperparameters. AGD is complementary to automatic differentiation and could help to automate general machine learning workflows.

Two existing tools are central to our derivation, and it is their novel combination that presents the main theoretical contribution of this paper. First, a classic tool from convex analysis known as the Bregman divergence (Bregman, 1967; Dhillon & Tropp, 2008) is used to characterise how the neural network interacts with the loss function. And second, a tool called deep relative trust (Bernstein et al., 2020) is used to characterise the highly non-linear interaction between the weights and the network output. With these tools in hand, we can apply the majorise-minimise meta-algorithm (Lange, 2016) to derive an optimiser explicitly tailored to deep network objective functions. To summarise, the derivation of AGD follows three main steps:

Functional expansion. We use a Bregman divergence to express the linearisation error of the objective function $\mathcal{L}({\bm{w}})$ in terms of the functional perturbation $\Delta{\bm{f}}$ to the network ${\bm{f}}$.

Architectural perturbation bounds. We use deep relative trust to relate the size and structure of the weight perturbation $\Delta{\bm{w}}$ to the size of the induced functional perturbation $\Delta{\bm{f}}$.

Majorise-minimise. We substitute deep relative trust into the Bregman divergence to obtain an explicitly architecture-dependent majorisation. Minimising with respect to $\Delta{\bm{w}}$ yields an optimiser.

This paper derives automatic gradient descent (AGD) by applying the majorise-minimise meta-algorithm to deep network objective functions. AGD trains all tested network architectures without hyperparameters, and scales to deep networks such as ResNet-50 and large datasets such as ImageNet. AGD trains out-of-the-box even when Adam and SGD fail to train with their default hyperparameters.

1 Related work

First-order optimisers leverage the first-order Taylor expansion of the objective function $\mathcal{L}({\bm{w}})$—in particular, the gradient $\nabla_{\bm{w}}\mathcal{L}({\bm{w}})$. Theoretical treatments include mirror descent (Nemirovsky & Yudin, 1983), natural gradient descent (Amari, 1998) and the Gauss-Newton method (Björck, 1996). These methods have been explored in the context of deep learning (Pascanu & Bengio, 2014; Azizan & Hassibi, 2019; Sun et al., 2022). First-order methods are amenable to deep learning since the gradient of the objective is available via recursive application of the chain rule—a.k.a. error back-propagation (Rumelhart et al., 1986).

Second-order optimisers leverage the second-order Taylor expansion of the objective function $\mathcal{L}({\bm{w}})$—in particular, the gradient $\nabla_{\bm{w}}\mathcal{L}({\bm{w}})$ and Hessian $\nabla^{2}_{\bm{w}}\mathcal{L}({\bm{w}})$. Examples include Newton’s method (Nocedal & Wright, 1999) and cubic-regularised Newton’s method (Nesterov & Polyak, 2006). Naïvely, second-order methods are less amenable to deep learning since the cost of the relevant Hessian computations is prohibitive at high dimension. That being said, efforts have been made to circumvent this issue (Agarwal et al., 2017).

The majorise-minimise meta-algorithm (Lange, 2016) is an algorithmic pattern that can be used to derive optimisers. To apply the meta-algorithm, one must first derive an upper bound on the objective which matches the objective up to $k$th-order in its Taylor series for some integer $k$. This majorisation can then be minimised as a proxy for reducing the original objective. Figure 2 illustrates the meta-algorithm for $k=1$.

The Lipschitz smoothness assumption—a global constraint on the eigenvalues of the Hessian—is often used to derive and analyse neural network optimisers (Agarwal et al., 2016). But this assumption has been questioned (Zhang et al., 2020) and evidence has even been found for the reverse relationship, where the Hessian spectrum is highly sensitive to the choice of optimiser (Cohen et al., 2021).

These considerations motivate the development of theory that is more explicitly tailored to neural architecture. For instance, Bernstein et al. (2020) used an architectural perturbation bound termed deep relative trust to characterise the neural network optimisation landscape as a function of network depth. Similarly, Yang & Hu (2021) sought to understand the role of width, leading to their maximal update parameterisation. Tables 1 and 2 provide some points of comparison between automatic gradient descent and these and other frameworks.

2 Preliminaries

The Manhattan norm $\|{\,\cdot\,}\|_{1}$ of a vector ${\bm{v}}$ is defined by $\|{{\bm{v}}}\|_{1}\vcentcolon=\sum_{i}|{{\bm{v}}_{i}}|$.

The Euclidean norm $\|{\,\cdot\,}\|_{2}$ of a vector ${\bm{v}}$ is defined by $\smash{\|{{\bm{v}}}\|_{2}\vcentcolon=\sqrt{\sum_{i}{\bm{v}}_{i}^{2}}}$.

The infinity norm $\|{\,\cdot\,}\|_{\infty}$ of a vector ${\bm{v}}$ is defined by $\|{{\bm{v}}}\|_{\infty}\vcentcolon=\max_{i}|{{\bm{v}}_{i}}|$.

The singular value decomposition allows us to measure the size of a matrix in two different ways:

The Frobenius norm $\|{\,\cdot\,}\|_{F}$ of a matrix ${\bm{M}}$ is given by $\|{{\bm{M}}}\|_{F}\vcentcolon=\sqrt{\sum_{i}\sigma_{i}({\bm{M}})^{2}}$.

The operator norm $\|{\,\cdot\,}\|_{*}$ of a matrix ${\bm{M}}$ is given by $\|{{\bm{M}}}\|_{*}\vcentcolon=\max_{i}\sigma_{i}({\bm{M}})$.

While the operator norm $\|{{\bm{M}}}\|_{*}$ reports the largest singular value, the quantity $\|{{\bm{M}}}\|_{F}/\sqrt{\min(m,n)}$ reports the root mean square singular value. Finally, we will need to understand two aspects of matrix conditioning:

The rank of a matrix counts the number of non-zero singular values.

The stable rank of a matrix ${\bm{M}}$ is defined by $\operatorname{rank_{stable}}{\bm{M}}\vcentcolon=\|{{\bm{M}}}\|_{F}^{2}/\|{{\bm{M}}}\|_{*}^{2}$.

This section develops a framework for applying the majorise-minimise meta-algorithm to generic optimisation problems in machine learning. In particular, the novel technique of functional expansion is introduced. Section 3 will apply this technique to deep neural networks. All proofs are supplied in Appendix A.

Given a machine learning model and a set of training data, our objective is to minimise the error of the model, averaged over the training data. Formally, we would like to minimise the following function:

We have adopted a compact notation where the dependence of ${\bm{f}}({\bm{x}};{\bm{w}})$ on ${\bm{w}}$ is at times suppressed. The perturbation hierarchies of a generic machine learning model and a deep neural network are visualised in Figures 2 and 3, respectively. The linearisation error of the objective perturbation $\Delta\mathcal{L}$ decomposes as:

In words: the linearisation error of the objective decomposes into two terms. The first depends on the linearisation error of the machine learning model and the second the loss. This decomposition relies on nothing but differentiability. For a convex loss, the second term may be interpreted as a Bregman divergence:

A Bregman divergence is just the linearisation error of a convex function. Two important examples are:

Our methods may be applied to other convex losses by calculating or bounding their Bregman divergence.

2 Functional expansion and functional majorisation

Before continuing, we make one simplifying assumption. Observe that the first term on the right-hand side of Proposition 1 is a high-dimensional inner product between two vectors. Since there is no clear reason why these two vectors should be aligned, let us assume that their inner product is zero:

While it is possible to work without this assumption (Bernstein, 2022), we found that its inclusion simplifies the analysis and in practice did not lead to a discernible weakening of the resulting algorithm. In any case, this assumption is considerably milder than the common assumption in the literature (Pascanu & Bengio, 2014; Lee et al., 2019) that the model linearisation error is itself zero: $\left[\Delta{\bm{f}}({\bm{x}})-\nabla_{\bm{w}}{\bm{f}}({\bm{x}})\Delta{\bm{w}}\right]=0$.

Armed with Proposition 1 and Assumption 1, we are ready to introduce functional expansion and majorisation:

So the perturbed objective $\mathcal{L}({\bm{w}}+\Delta{\bm{w}})$ may be written as the sum of its first-order Taylor expansion with a Bregman divergence in the model outputs averaged over the training set. It is straightforward to specialise this result to different losses by substituting in their Bregman divergence:

Under Assumption 1, for cross-entropy loss, if ${\bm{y}}^{\top}\bm{1}=1$:

When the functional perturbation is reasonably “spread out”, we would expect $\|{\Delta{\bm{f}}({\bm{x}})}\|_{\infty}^{2}\approx\|{\Delta{\bm{f}}({\bm{x}})}\|_{2}^{2}/d_{L}$. In this setting, the functional majorisation of cross-entropy loss agrees with the functional expansion of mean squared error to second order. While the paper derives automatic gradient descent for the square loss, this observation justifies its application to cross-entropy loss, as in the case of the ImageNet experiments.

3 Recovering existing frameworks

We briefly observe that three existing optimisation frameworks may be recovered efficiently from Theorem 1:

Substituting the linearised functional perturbation $\Delta{\bm{f}}({\bm{x}})\approx\nabla_{\bm{w}}{\bm{f}}({\bm{x}})\Delta{\bm{w}}$ into Corollary 1 and minimising with respect to $\Delta{\bm{w}}$ is the starting point for the Gauss-Newton method.

Majorise-Minimise for Deep Learning Problems

Substituting the linearised functional perturbation $\Delta{\bm{f}}({\bm{x}})\approx\nabla_{\bm{w}}{\bm{f}}({\bm{x}})\Delta{\bm{w}}$ into Corollary 2 and minimising with respect to $\Delta{\bm{w}}$ is the starting point for natural gradient descent.

In this section, we will focus our efforts on deriving an optimiser for deep fully-connected networks trained with square loss. The derivation for cross-entropy loss is analogous. Proofs are relegated to Appendix A.

For $\eta>0$, the data $({\bm{x}},{\bm{y}})$, weights ${\bm{W}}_{k}$ and updates $\Delta{\bm{W}}_{k}$ should obey:

While results can be derived without adopting Prescription 1, the scalings substantially simplify our formulae. One reason for this is that, under Prescription 1, we have the telescoping property that $\prod_{k=1}^{L}\|{{\bm{W}}_{k}}\|_{*}=\sqrt{d_{L}/d_{0}}$. For a concrete example of how this helps, consider the following bound on the norm of the network outputs:

The output norm of a fully-connected network ${\bm{f}}$ obeys the following bound:

So, under Prescription 1, the bound is simple. Furthermore, the scaling of the update with a single parameter $\eta$ reduces the problem of solving for an optimiser to a single parameter problem. To see how this might make life easier, consider the following lemma that relates weight perturbations to functional perturbations:

When adjusting the weights ${\bm{w}}=({\bm{W}}_{1},...,{\bm{W}}_{L})$ of a fully-connected network ${\bm{f}}$ by $\Delta{\bm{w}}=(\Delta{\bm{W}}_{1},...,\Delta{\bm{W}}_{L})$, the induced functional perturbation $\Delta{\bm{f}}({\bm{x}})\vcentcolon={\bm{f}}({\bm{x}};{\bm{w}}+\Delta{\bm{w}})-{\bm{f}}({\bm{x}};{\bm{w}})$ obeys:

So, under Prescription 1, the single parameter $\eta$ directly controls the size of functional perturbations.

In terms of enforcing Prescription 1 in practice, the norms of the data $({\bm{x}},{\bm{y}})$ may be set via pre-processing, the norm of the update $\Delta{\bm{W}}_{k}$ may be set via the optimisation algorithm and the norm of the weight matrix ${\bm{W}}_{k}$ may be set by the choice of initialisation. While, yes, $\|{{\bm{W}}_{k}}\|_{*}$ may drift during training, the amount that this can happen is limited by Weyl (1912)’s inequality for singular values. In particular, after one step the perturbed operator norm $\|{{\bm{W}}_{k}+\Delta{\bm{W}}_{K}}\|_{*}$ is sandwiched like $(1-\eta/L)\cdot\|{{\bm{W}}_{k}}\|_{*}\leq\|{{\bm{W}}_{k}+\Delta{\bm{W}}_{K}}\|_{*}\leq(1+\eta/L)\cdot\|{{\bm{W}}_{k}}\|_{*}$.

With both functional majorisation and deep relative trust in hand, we can majorise the deep network objective:

For an FCN with square loss, under Assumption 1 and Prescription 1:

Observe that the majorisation only depends on the magnitude of the scalar $\eta$ and on some notion of angle $\operatorname{tr}\Delta{\bm{W}}_{k}^{\top}\nabla_{{\bm{W}}_{k}}\mathcal{L}/\|{\Delta{\bm{W}}_{k}}\|_{*}$ between the perturbation matrix $\Delta{\bm{W}}_{k}$ and the gradient matrix $\nabla_{{\bm{W}}_{k}}\mathcal{L}$. To derive an optimiser, we would now like to minimise this majorisation with respect to $\eta$ and this angle. First, let us introduce one additional assumption and one additional definition:

The gradient satisfies $\operatorname{rank_{stable}}\nabla_{{\bm{W}}_{k}}\mathcal{L}=1$ at all layers $k=1,...,L$.

At a weight setting ${\bm{w}}$, the gradient summary $G$ is given by:

The gradient summary is a weighted average of gradient norms over layers. It can be thought of as a way to measure the size of the gradient while accounting for the fact that the weight matrices at different layers may be on different scales. This is related to the concept of the gradient scale coefficient of Philipp et al. (2017).

We now have everything we need to derive automatic gradient descent via the majorise-minimise principle:

For a deep fully-connected network, under Assumptions 1 and 2 and Prescription 1, the majorisation of square loss given in Lemma 5 is minimised by setting:

We present pseudocode for this theorem in Algorithm 1, and a PyTorch implementation in Appendix B. Via a simple derivation based on clear algorithmic principles, automatic gradient descent unifies various heuristic and theoretical ideas that have appeared in the literature:

Relative updates. The update is scaled relative to the norm of the weight matrix to which it is applied—assuming the weight matrices are scaled according to Prescription 1. Such a scaling was proposed by You et al. (2017) and further explored by Carbonnelle & Vleeschouwer (2019) and Bernstein et al. (2020). There is evidence that such relative synaptic updates may occur in neuroscience (Loewenstein et al., 2011).

Depth scaling. Scaling the perturbation strength like $1/L$ for networks of depth $L$ was proposed on theoretical grounds by Bernstein et al. (2020) based on analysis via deep relative trust.

Width scaling. The dimensional factors of $d_{k}$ and $d_{k-1}$ that appear closely relate to the maximal update parameterisation of Yang & Hu (2021) designed to ensure hyperparameter transfer across network width.

Gradient clipping. The logarithmic dependence of the update on the gradient summary may be seen as an automatic form of adaptive gradient clipping (Brock et al., 2021)—a technique which clips the gradient once its magnitude surpasses a certain threshold set by a hyperparameter.

2 Convergence analysis

This section presents theoretical convergence rates for automatic gradient descent. While the spirit of the analysis is standard in optimisation theory, the details may still prove interesting for their detailed characterisation of the optimisation properties of deep networks. For instance, we propose a novel Polyak-Łojasiewicz inequality tailored to the operator structure of deep networks. We begin with two observations:

For square loss, the objective is bounded as follows:

For square loss, the norm of the gradient at layer $k$ is bounded as follows:

These results help us prove that automatic gradient descent converges to a point where the gradient vanishes:

Consider a fully-connected network trained by automatic gradient descent (Theorem 2) and square loss for $T$ iterations. Let $G_{t}$ denote the gradient summary (Definition 11) at step $t\leq T$. Under Assumptions 1 and 2 and Prescription 1, AGD converges at the following rate:

This lemma can be converted into a convergence rate to a global minimum with one additional assumption:

For some $\alpha>0$, the gradient norm is lower bounded by:

This lower bound mirrors the structure of the upper bound in Lemma 7. The parameter $\alpha$ captures how much of the gradient is attenuated by small singular values in the weights and by deactivated $\operatorname{relu}$ units. While Polyak-Łojasiewicz inequalities are common in the literature (Liu et al., 2022), our assumption is novel in that it pays attention to the operator structure of the network. Assumption 3 leads to the following theorem:

For automatic gradient descent (Theorem 2) in the same setting as Lemma 8 but with the addition of Assumption 3, the mean squared error objective at step $T$ obeys:

3 Experiments

The goal of our experiments was twofold. First, we wanted to test automatic gradient descent (AGD, Algorithm 1) on a broad variety of networks architectures and datasets to check that it actually works. In particular, we tested AGD on fully-connected networks (FCNs, Definition 10), and both VGG-style (Simonyan & Zisserman, 2015) and ResNet-style (He et al., 2015) convolutional neural networks on the CIFAR-10, CIFAR-100 (Krizhevsky, 2009) and ImageNet (Deng et al., 2009, ILSVRC2012) datasets with standard data augmentation. And second, to see what AGD may have to offer beyond the status quo, we wanted to compare AGD to tuned Adam and SGD baselines, as well as Adam and SGD run with their default hyperparameters.

To get AGD working with convolutional layers, we adopted a per-submatrix normalisation scheme. Specifically, for a convolutional tensor with filters of size $\mathtt{k_{x}}\times\mathtt{k_{y}}$, we implemented the normalisation separately for each of the $\mathtt{k_{x}}\times\mathtt{k_{y}}$ submatrices of dimension $\mathtt{channels_{in}}\times\mathtt{channels_{out}}$. Since AGD does not yet support biases or affine parameters in batchnorm, we disabled these parameters in all architectures. To at least adhere to Prescription 1 at initialisation, AGD draws initial weight matrices uniform semi-orthogonal and re-scaled by a factor of $\sqrt{\mathtt{fan\_in}/\mathtt{fan\_out}}$. Adam and SGD baselines used the PyTorch default initialisation. A PyTorch implementation of AGD reflecting these details is given in Appendix B. All experiments use square loss except ImageNet which used cross-entropy loss. Cross-entropy loss has been found to be superior to square loss for datasets with a large number of classes (Demirkaya et al., 2020; Hui & Belkin, 2021).

Our experimental results are spread across five figures:

Figure 1 presents some highlights of our results: First, AGD can train networks that Adam and SGD with default hyperparameters cannot. Second, for ResNet-18 on CIFAR-10, AGD attained performance comparable to the best-tuned performance of Adam and SGD. And third, AGD scales up to ImageNet.

Figure 4 displays the breadth of our experiments: from training a 16-layer fully-connected network on CIFAR-10 to training ResNet-50 on ImageNet. Adam’s learning rate was tuned over the logarithmic grid $\{10^{-5},10^{-4},...,10^{-1}\}$ while for ImageNet we used a default learning rate of 0.1 for SGD without any manual decay. AGD and Adam performed almost equally well on the depth-16 width-512 fully-connected network: 52.7% test accuracy for AGD compared to 53.5% for Adam. For ResNet-18 on CIFAR-10, Adam attained 92.9% test accuracy compared to AGD’s 91.2%. On this benchmark, a fully-tuned SGD with learning rate schedule, weight decay, cross-entropy loss and bias and affine parameters can attain 93.0% test accuracy (Liu, 2017). For VGG-16 on CIFAR-100, AGD achieved 67.4% test accuracy compared to Adam’s 69.7%. Finally, on ImageNet AGD achieved a top-1 test accuracy of 65.5% after 350 epochs.

Figure 5 compares AGD to Adam and SGD for training an eight-layer fully-connected network of width 256. Adam and SGD’s learning rates were tuned over the logarithmic grid $\{10^{-5},10^{-4},...,10^{-1}\}$. Adam’s optimal learning rate of $10^{-4}$ was three orders of magnitude smaller than SGD’s optimal learning rate of $10^{-1}$. SGD did not attain as low of an objective value as Adam or AGD.

Figure 6 shows that AGD can train FCNs with width ranging from 64 to 2048 and depth from 2 to 32 and Figure 7 shows that AGD successfully trains a four-layer FCN at varying mini-batch size: from 32 to 4096.

This paper has proposed a new framework for deriving optimisation algorithms for non-convex composite objective functions, which are particularly prevalent in the field of machine learning and the subfield of deep learning. What we have proposed is truly a framework: it can be applied to a new loss function by writing down its Bregman divergence, or a new machine learning model by writing down its architectural perturbation bound. The framework is properly placed in the context of existing frameworks such as the majorise-minimise meta-algorithm, mirror descent and natural gradient descent.

Recent papers have proposed a paradigm of hyperparameter transfer where a small network is tuned and the resulting hyperparameters are transferred to a larger network (Yang et al., 2021; Bernstein, 2022). The methods and results in this paper suggest a stronger paradigm of hyperparameter elimination: by detailed analysis of the structure and interactions between different components of a machine learning system, we may hope—if not to outright outlaw hyperparameters—at least to reduce their abundance and opacity.

The main product of this research is automatic gradient descent (AGD), with pseudocode given in Algorithm 1 and PyTorch code given in Appendix B. We have found AGD to be genuinely useful, and believe that it may complement automatic differentiation in helping to automate general machine learning workflows.

The analysis leading to automatic gradient descent is elementary: we leverage basic concepts in linear algebra such as matrix and vector norms, and use simple bounds such as the triangle inequality for vector–vector sums, and the operator norm bound for matrix–vector products. The analysis is non-asymptotic: it does not rely on taking dimensions to infinity, and deterministic: it does not involve random matrix theory. We believe that the accessibility of the analysis could make this paper a good starting point for future developments.

Here we list some promising avenues for theoretical and practical research. We are exploring some of these ideas in our development codebase: https://github.com/C1510/agd_exp.

Stochastic optimisation. Automatic gradient descent is derived in the full-batch optimisation setting, but the algorithm is evaluated experimentally in the mini-batch setting. It would be interesting to try to extend our theoretical and practical methods to more faithfully address stochastic optimisation.

More architectures. Automatic gradient descent is derived for fully-connected networks and extended heuristically to convolutional networks. We are curious to extend the methods to more varied architectures such as transformers (Vaswani et al., 2017) and architectural components such as biases. Since most neural networks resemble fully-connected networks in the sense that they are all just deep compound operators, we expect much of the structure of automatic gradient descent as presented to carry through.

Regularisation. The present paper deals purely with the optimisation structure of deep neural networks, and little thought is given to either generalisation or regularisation. Future work could look at both theoretical and practical regularisation schemes for automatic gradient descent. It would be interesting to try to do this without introducing hyperparameters, although we suspect that when it comes to regularisation at least one hyperparameter may become necessary.

Acceleration. We have found in some preliminary experiments that slightly increasing the update size of automatic gradient descent with a gain hyperparameter, or introducing a momentum hyperparameter, can lead to faster convergence. We emphasise that no experiment in this paper used such hyperparameters. Still, these observations may provide a valuable starting point for improving AGD in future work.

Operator perturbation theory. Part of the inspiration for this paper was the idea of applying operator perturbation theory to deep learning. While perturbation theory is well-studied in the context of linear operators (Weyl, 1912; Kato, 1966; Stewart, 2006), in deep learning we are concerned with non-linear compound operators. It may be interesting to try to further extend results in perturbation theory to deep neural networks. One could imagine cataloging the perturbation structure of different neural network building blocks, and using a result similar to deep relative trust (Lemma 4) to describe how they compound.

Acknowledgments

The authors are grateful to MIT SuperCloud, Oxford Hydra, NVIDIA and Virgile Richard for providing GPUs. Thanks are due to Greg Yang and Jamie Simon for helpful discussions. A paper with Greg and Jamie is in preparation to explain the relationship between muP (Yang & Hu, 2021) and the operator norm.

Appendix A Proofs

Here are the proofs for the theoretical results in the main text.

First, since $\sum_{i}{\bm{y}}_{i}=1$, cross-entropy loss may be re-written:

The linear term $-{\bm{f}}({\bm{x}})^{\top}{\bm{y}}$ does not contribute to the linearisation error and may be neglected. Therefore:

To establish the inequality, let $\otimes$ denote the outer product and define $p\vcentcolon=\operatorname{softmax}(f({\bm{x}}))$. Then we have:

where we have used that $p\otimes p$ is positive definite and then applied Hölder’s inequality with $\|{p}\|_{1}=1$.

The result follows by substituting Assumption 1 into Proposition 1 and applying Definition 9.

Combine Lemma 1 with Theorem 1 to obtain the result.

Combine Lemma 2 with Theorem 1 to obtain the result.

For any vector ${\bm{v}}$ and matrix ${\bm{M}}$ with compatible dimensions, we have that $\|{{\bm{M}}{\bm{v}}}\|_{2}\leq\|{{\bm{M}}}\|_{*}\cdot\|{{\bm{v}}}\|_{2}$ and $\|{\operatorname{relu}{\bm{v}}}\|_{2}\leq\|{{\bm{v}}}\|_{2}$. The lemma follows by applying these results recursively over the depth of the network.

We proceed by induction. First, consider a network with $L=1$ layers: ${\bm{f}}({\bm{x}})={\bm{W}}_{1}{\bm{x}}$. Observe that $\|{\Delta{\bm{f}}({\bm{x}})}\|_{2}=\|{\Delta{\bm{W}}_{1}{\bm{x}}}\|_{2}\leq\|{\Delta{\bm{W}}_{1}}\|_{*}\cdot\|{{\bm{x}}}\|_{2}$ as required. Next, assume that the result holds for a network ${\bm{g}}({\bm{x}})$ with $L-1$ layers and consider adding a layer to obtain ${\bm{f}}({\bm{x}})={\bm{W}}_{L}\circ\operatorname{relu}{}\circ{\bm{g}}({\bm{x}})$. Then:

where the inequality follows by applying the triangle inequality, the operator norm bound, the fact that $\operatorname{relu}{}$ is one-Lipschitz, and a further application of the triangle inequality. But by the inductive hypothesis and Lemma 3, the right-hand side is bounded by:

The induction is complete. To further bound this result under Prescription 1, observe that the product $\left[\prod_{k=1}^{L}\|{{\bm{W}}_{k}}\|_{*}\right]\times\|{{\bm{x}}}\|_{2}$ telescopes to just $\sqrt{d_{L}}$, while the other product satisfies:

Combining these observations yields the result.

Substitute Lemma 4 into Corollary 1 and decompose $\nabla_{\bm{w}}\mathcal{L}({\bm{w}})^{\top}\Delta{\bm{w}}=\sum_{k=1}^{L}\operatorname{tr}(\Delta{\bm{W}}_{k}^{\top}\nabla_{{\bm{W}}_{k}}\mathcal{L})$. The result follows by realising that under Prescription 1, the perturbations satisfy $\|{\Delta{\bm{W}}_{k}}\|_{*}=\sqrt{d_{k}/d_{k-1}}\cdot\frac{\eta}{L}$.

The inner product $\operatorname{tr}\frac{\Delta{\bm{W}}_{k}^{\top}\nabla_{{\bm{W}}_{k}}\mathcal{L}}{\|{\Delta{\bm{W}}_{k}}\|_{*}}$ that appears in Lemma 5 is most negative when the perturbation $\Delta{\bm{W}}_{k}$ satisfies $\Delta{\bm{W}}_{k}/\|{\Delta{\bm{W}}_{k}}\|_{*}=-\nabla_{{\bm{W}}_{k}}\mathcal{L}/\|{\nabla_{{\bm{W}}_{k}}\mathcal{L}}\|_{*}$. Substituting this result back into Lemma 5 yields:

Under Assumption 2, we have that $\|{\nabla_{{\bm{W}}_{k}}\mathcal{L}}\|_{F}^{2}/\|{\nabla_{{\bm{W}}_{k}}\mathcal{L}}\|_{*}=\|{\nabla_{{\bm{W}}_{k}}\mathcal{L}}\|_{F}$ and so this inequality simplifies to:

Taking the derivative of the right-hand side with respect to $\eta$ and setting it to zero yields $(\exp\eta-1)\exp\eta=G$. Applying the quadratic formula and retaining the positive solution yields $\exp\eta=\tfrac{1}{2}(1+\sqrt{1+4G})$. Combining this with the relation that $\Delta{\bm{W}}_{k}/\|{\Delta{\bm{W}}_{k}}\|_{*}=-\nabla_{{\bm{W}}_{k}}\mathcal{L}/\|{\nabla_{{\bm{W}}_{k}}\mathcal{L}}\|_{*}$ and applying that $\|{\Delta{\bm{W}}_{k}}\|_{*}=\sqrt{d_{k}/d_{k-1}}\cdot\frac{\eta}{L}$ by Prescription 1 yields the result.

The result follows by the following chain of inequalities:

where the second inequality holds under Prescription 1.

By the chain rule, the gradient of mean square error objective may be written:

where $\otimes$ denotes the outer product and ${\bm{D}}_{k}$ denotes a diagonal matrix whose entries are one when $\operatorname{relu}$ is active and zero when $\operatorname{relu}$ is inactive. Since the operator norm $\|{{\bm{D}}_{k}}\|_{*}=1$, we have that the Frobenius norm $\|{\nabla_{{\bm{W}}_{k}}\mathcal{L}({\bm{w}})}\|_{F}$ is bounded from above by:

In the above argument, the first inequality follows by recursive application of the operator norm upper bound, and the second inequality follows from the Cauchy-Schwarz inequality. The right-hand side simplifies under Prescription 1, and we may apply Lemma 6 to obtain:

Theorem 2 prescribes that $\exp\eta=\tfrac{1}{2}(1+\sqrt{1+4G})$, and so \eta=\log\big{(}1+\frac{\sqrt{1+4G}-1}{2}\big{)}. We begin by proving some useful auxiliary bounds. By Lemma 7 and Prescription 1, the gradient summary is bounded by:

The fact that the gradient summary $G$ is less than two is important because, for $x\leq 1$, we have that $\log(1+x)\geq x\log 2$. In turn, this implies that since $G<2$, we have that $\eta=\log\frac{1+\sqrt{1+4G}}{2}\geq\frac{\sqrt{1+4G}-1}{2}\log 2$. It will also be important to know that for $G<2$, we have that $\tfrac{1}{2}\cdot G\leq\tfrac{\sqrt{1+4G}-1}{2}\leq G$.

With these bounds in hand, the analysis becomes fairly standard. By an intermediate step in the proof of Theorem 2, the change in objective across a single step is bounded by:

where the second and third inequalities follow by our auxiliary bounds. Letting $G_{t}$ denote the gradient summary at step $t$, averaging this bound over time steps and applying the telescoping property yields:

where the final inequality follows by Lemma 6 and the fact that $\mathcal{L}({\bm{w}}_{T})\geq 0$.

By Assumption 3, the gradient summary at time step $t$ must satisfy $G_{t}\geq\alpha\times\sqrt{2\cdot\mathcal{L}({\bm{w}}_{t})}$. Therefore the objective at time step $t$ is bounded by $\mathcal{L}({\bm{w}}_{t})\leq G_{t}^{2}/(2\alpha^{2})$. Combining with Lemma 8 then yields that:

Appendix B PyTorch Implementation

The following code implements automatic gradient descent in PyTorch (Paszke et al., 2019). We include a single gain hyperparameter which controls the update size and may be increased from its default value of 1.0 to slightly accelerate training. We emphasise that all the results reported in the paper used a gain of unity.