SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lachlan E. MacDonald*
Mathematical Institute for Data Science
lemacdonald@protonmail.com
Jack Valmadre
Australian Institute for Machine Learning  
University of Adelaide
Hemanth Saratchandran
Simon Lucey
Australian Institute for Machine Learning
We introduce a general theoretical framework, designed for the study of gradient optimisation of deep neural networks, that encompasses ubiquitous architecture choices including batch normalisation, weight normalisation and skip connections. Our framework determines the curvature and regularity properties of multilayer loss landscapes in terms of their constituent layers, thereby elucidating the roles played by normalisation layers and skip connections in globalising these properties. We then demonstrate the utility of this framework in two respects. First, we give the only proof of which we are aware that a class of deep neural networks can be trained using gradient descent to global optima even when such optima only exist at infinity, as is the case for the cross-entropy cost. Second, we identify a novel causal mechanism by which skip connections accelerate training, which we verify predictively with ResNets on MNIST, CIFAR10, CIFAR100 and ImageNet.
Deep, overparameterised neural networks are efficiently trainable to global optima using simple first order methods. That this is true is immensely surprising from a theoretical perspective: modern datasets and deep neural network architectures are so complex and irregular that they are essentially opaque from the perspective of classical (convex) optimisation theory. A recent surge in inspired theoretical works [22, 14, 13, 1, 52, 51, 36, 28, 29, 33, 32] has elucidated this phenomenon, showing linear convergence of gradient descent on certain classes of neural networks, with certain cost functions, to global optima. The formal principles underlying these works are identical. By taking width sufficiently large, one guarantees uniform bounds on curvature (via a Lipschitz gradients-type property) and regularity (via a Polyak-Łojasiewicz-type inequality) in a neighbourhood of initialisation. Convergence to a global optimum in that neighbourhood then follows from a well-known chain of estimates [23].
Despite significant progress, the theory of deep learning optimisation extant in the literature presents at least three significant shortcomings:
*Most done while at the Australian Institute for Machine Learning, University of Adelaide
1. It lacks a formal framework in which to compare common practical architecture choices. Indeed, none of the aforementioned works consider the impact of ubiquitous (weight/batch) normalisation layers. Moreover, where common architectural modifications such as skip connections are studied, it is unclear exactly what impact they have on optimisation. For instance, while in [13] it is shown that skip connections enable convergence with width polynomial in the number of layers, as compared with exponential width for chain networks, in [1] polynomial width is shown to be sufficient for convergence of both chain and residual networks.
2. It lacks theoretical flexibility. The consistent use of uniform curvature and regularity bounds are insufficiently flexible to enable optimisation guarantees too far away from initialisation, where the local, uniform bounds used in previous theory no longer hold. In particular, proving globally optimal convergence for deep neural nets with the cross-entropy cost was (until now) an open problem [5].
3. It lacks practical utility. Although it is presently unreasonable to demand quantitatively predictive bounds on practical performance, existing optimisation theory has been largely unable to inform architecture design even qualitatively. This is in part due to the first item, since practical architectures typically differ substantially from those considered for theoretical purposes.
Our purpose in this article is to take a step in addressing these shortcomings. Specifically:
1. We provide a formal framework, inspired by [45], for the study of multilayer optimisation. Our framework is sufficiently general to include all commonly used neural network layers, and contains formal results relating the curvature and regularity properties of multilayer loss landscapes to those of their constituent layers. As instances, we prove novel results on the global curvature and regularity properties enabled by normalisation layers and skip connections respectively, in contrast to the local bounds provided in previous work.
2. Using these novel, global bounds, we identify a class of weight-normalised residual networks for which, given a linear independence assumption on the data, gradient descent can be provably trained to a global optimum arbitrarily far away from initialisation. From a regularity perspective, our analysis is strictly more flexible than the uniform analysis considered in previous works, and in particular solves the open problem of proving global optimality for the training of deep nets with the cross-entropy cost.
3. Using our theoretical insight that skip connections aid loss regularity, we conduct a systematic empirical analysis of singular value distributions of layer Jacobians for practical layers. We are thereby able to predict that simple modifications to the classic ResNet architecture [20] will improve training speed. We verify our predictions on MNIST, CIFAR10, CIFAR100 and ImageNet.
In this section we give a summary of the principles underlying recent theoretical advances in neural network optimisation. We discuss related works after this summary for greater clarity.
2.1 Smoothness and the PL-inequality
Gradient descent on a possibly non-convex function  $\ell : \mathbb{R}^p \to \mathbb{R}_{\geq 0}$  can be guaranteed to converge to a global optimum by insisting that  $\ell$  have Lipschitz gradients and satisfy the Polyak-Lojasiewicz inequality. We recall these well-known properties here for convenience.
Definition 2.1. Let  $\beta >0$ . A continuously differentiable function  $\ell :\mathbb{R}^p\to \mathbb{R}$  is said to have  $\beta$ -Lipschitz gradients, or is said to be  $\beta$ -smooth over a set  $S\subset \mathbb{R}^p$  if the vector field  $\nabla \ell :\mathbb{R}^p\rightarrow \mathbb{R}^p$  is  $\beta$ -Lipschitz. If  $S$  is convex,  $\ell$  having  $\beta$ -Lipschitz gradients implies that the inequality
$$
\ell \left(\theta_ {2}\right) - \ell \left(\theta_ {1}\right) \leq \nabla \ell \left(\theta_ {1}\right) ^ {T} \left(\theta_ {2} - \theta_ {1}\right) + \frac {\beta}{2} \| \theta_ {2} - \theta_ {1} \| ^ {2} \tag {1}
$$
holds for all  $\theta_1, \theta_2 \in S$ .
The  $\beta$ -smoothness of  $\ell$  over  $S$  can be thought of as a uniform bound on the curvature of  $\ell$  over  $S$ : if  $\ell$  is twice continuously differentiable, then it has Lipschitz gradients over any compact set  $K$  with (possibly loose) Lipschitz constant given by
$$
\beta := \sup  _ {\theta \in K} \| D ^ {2} \ell (\theta) \|, \tag {2}
$$
where  $D^2\ell$  is the Hessian and  $\| \cdot \|$  denotes any matrix norm.
Definition 2.2. Let  $\mu >0$ . A differentiable function  $\ell :\mathbb{R}^p\to \mathbb{R}_{\geq 0}$  is said to satisfy the  $\mu$ -Polyak-Lojasiewicz inequality, or is said to be  $\mu$ -PL over a set  $S\subset \mathbb{R}^p$  if
$$
\left\| \nabla \ell (\theta) \right\| ^ {2} \geq \mu \left(\ell (\theta) - \inf  _ {\theta^ {\prime} \in S} \ell \left(\theta^ {\prime}\right)\right) \tag {3}
$$
for all  $\theta \in S$
The PL condition on  $\ell$  over  $S$  is a uniform guarantee of regularity, which implies that all critical points of  $\ell$  over  $S$  are  $S$ -global minima; however, such a function need not be convex. Synthesising these definitions leads easily to the following result (cf. Theorem 1 of [23]).
Theorem 2.3. Let  $\ell : \mathbb{R}^p \to \mathbb{R}_{\geq 0}$  be a continuously differentiable function that is  $\beta$ -smooth and  $\mu$ -PL over a convex set  $S$ . Suppose that  $\theta_0 \in S$  and let  $\{\theta_t\}_{t=0}^{\infty}$  be the trajectory taken by gradient descent, with step size  $\eta < 2\beta^{-1}$ , starting at  $\theta_0$ . If  $\{\theta_t\}_{t=0}^{\infty} \subset S$ , then  $\ell(\theta_t)$  converges to an  $S$ -global minimum  $\ell^*$  at a linear rate:
$$
\ell \left(\theta_ {t}\right) - \ell^ {*} \leq \left(1 - \mu \eta \left(1 - \frac {\beta \eta}{2}\right)\right) ^ {t} \left(\ell \left(\theta_ {0}\right) - \ell^ {*}\right) \tag {4}
$$
for all  $t\in \mathbb{N}$
Essentially, while the Lipschitz constant of the gradients controls whether or not gradient descent with a given step size can be guaranteed to decrease the loss at each step, the PL constant determines by how much the loss will decrease. These ideas can be applied to the optimisation of deep neural nets as follows.
2.2 Application to model optimisation
The above theory can be applied to parameterised models in the following fashion. Let  $f: \mathbb{R}^p \times \mathbb{R}^{d_0} \to \mathbb{R}^{d_L}$  be a differentiable,  $\mathbb{R}^p$ -parameterised family of functions  $\mathbb{R}^{d_0} \to \mathbb{R}^{d_L}$  (in later sections,  $L$  will denote the number of layers of a deep neural network). Given  $N$  training data  $\{(x_i, y_i)\}_{i=1}^N \subset \mathbb{R}^{d_0} \times \mathbb{R}^{d_L}$ , let  $F: \mathbb{R}^p \to \mathbb{R}^{d_L \times N}$  be the corresponding parameter-function map defined by
$$
F (\theta) _ {i} := f (\theta , x _ {i}). \tag {5}
$$
Any differentiable cost function  $c: \mathbb{R}^{d_L} \times \mathbb{R}^{d_L} \to \mathbb{R}_{\geq 0}$ , convex in the first variable, extends to a differentiable, convex function  $\gamma: \mathbb{R}^{d_L \times N} \to \mathbb{R}_{\geq 0}$  defined by
$$
\gamma \left(\left(z _ {i}\right) _ {i = 1} ^ {N}\right) := \frac {1}{N} \sum_ {i = 1} ^ {N} c \left(z _ {i}, y _ {i}\right), \tag {6}
$$
and one is then concerned with the optimisation of the composite  $\ell \coloneqq \gamma \circ F:\mathbb{R}^p\to \mathbb{R}_{\geq 0}$  via gradient descent.
To apply Theorem 2.3, one needs to determine the smoothness and regularity properties of  $\ell$ . By the chain rule, the former can be determined given sufficient conditions on the derivatives  $D\gamma \circ F$  and  $DF$  (cf. Lemma 2 of [45]). The latter can be bounded by Lemma 3 of [45], which we recall below and prove in the appendix for the reader's convenience.
Theorem 2.4. Let  $S \subset \mathbb{R}^p$  be a set. Suppose that  $\gamma: \mathbb{R}^{d_L \times N} \to \mathbb{R}_{\geq 0}$  is  $\mu$ -PL over  $F(S)$  with minimum  $\gamma_S^*$ . Let  $\lambda(DF(\theta))$  denote the smallest eigenvalue of  $DF(\theta)DF(\theta)^T$ . Then
$$
\left\| \nabla \ell (\theta) \right\| ^ {2} \geq \mu \lambda (D F (\theta)) \left(\ell (\theta) - \gamma_ {S} ^ {*}\right) \tag {7}
$$
Note that Theorem 2.4 is vacuous ( $\lambda(DF(\theta)) = 0$  for all  $\theta$ ) unless in the overparameterised regime ( $p \geq d_L N$ ). Even in this regime, however, Theorem 2.4 does not imply that  $\ell$  is PL unless  $\lambda(\theta)$  can be uniformly lower bounded by a positive constant over  $S$ . Although universally utilised in previous literature, such a uniform lower bound will not be possible in our global analysis, and our convergence theorem does not follow from Theorem 2.3, in contrast to previous work. Our theorem requires additional argumentation, which we believe may be of independent utility.
Convergence theorems for deep linear networks with the square cost are considered in [4, 2]. In [22], it is proved that the tangent kernel of a multi-layer perceptron (MLP) becomes approximately constant over all of parameter space as width goes to infinity, and is positive-definite for certain data distributions, which by Theorem 2.4 implies that all critical points are global minima. Strictly speaking, however, [22] does not prove convergence of gradient descent: the authors consider only gradient flow, and leave Lipschitz concerns untouched. The papers [14, 13, 1, 52, 51, 36] prove that overparameterized neural nets of varying architectures can be optimised to global minima close to initialisation by assuming sufficient width of several layers. While [1] does consider the cross-entropy cost, convergence to a global optimum is not proved: it is instead shown that perfect classification accuracy can be achieved close to initialisation during training. Improvements on these works have been made in [33, 32, 7], wherein large width is required of only a single layer.
It is identified in [28] that linearity of the final layer is key in establishing the approximate constancy of the tangent kernel for wide networks that was used in [14, 13]. By making explicit the implicit use of the PL condition present in previous works [14, 13, 1, 52, 51, 36], [29] proves a convergence theorem even with nonlinear output layers. The theory explicated in [29] is formalised and generalised in [45]. A key weakness of all of the works mentioned thus far (bar the purely formal [45]) is that their hypotheses imply that optimisation trajectories are always close to initialisation. Without this, there is no obvious way to guarantee the PL-inequality along the optimisation trajectory, and hence no way to guarantee one does not converge to a suboptimal critical point. However such training is not possible with the cross-entropy cost, whose global minima only exist at infinity. There is also evidence to suggest that such training must be avoided for state-of-the-art test performance [8, 15, 26]. In contrast, our theory gives convergence guarantees even for trajectories that travel arbitrarily far from initialisation, and is the only work of which we are aware that can make this claim.
Among the tools that make our theory work are skip connections [20] and weight normalisation [41]. The smoothness properties of normalisation schemes have previously been studied [42, 40], however they only give pointwise estimates comparing normalised to non-normalised layers, and do not provide a global analysis of Lipschitz properties as we do. The regularising effect of skip connections on the loss landscape has previously been studied in [35], however this study is not tightly linked to optimisation theory. Skip connections have also been shown to enable the interpretation of a neural network as coordinate transformations of data manifolds [19]. Mean field analyses of skip connections have been conducted in [49, 30] which necessitate large width; our own analysis does not. A similarly general framework to that which we supply is given in [47, 48]; while both encapsulate all presently used architectures, that of [47, 48] is designed for the study of infinite-width tangent kernels, while ours is designed specifically for optimisation theory. Our empirical singular value analysis of skip connections complements existing theoretical work using random matrix theory [18, 37, 39, 34, 16]. These works have not yet considered the shifting effect of skip connections on layer Jacobians that we observe empirically.
Our theory also links nicely to the intuitive notions of gradient propagation [20] and dynamical isometry already present in the literature. In tying Jacobian singular values rigorously to loss regularity in the sense of the Polyak-Lojasiewicz inequality, our theory provides a new link between dynamical isometry and optimisation theory [43, 38, 46]: specifically dynamical isometry ensures better PL conditioning and therefore faster and more reliable convergence to global minima. In linking this productive section of the literature to optimisation theory, our work may open up new possibilities for convergence proofs in the optimisation theory of deep networks. We leave further exploration of this topic to future work.
Due to this relationship with the notion of dynamical isometry, our work also provides optimisation-theoretic support for the empirical analyses of [24, 6] that study the importance of layerwise dynamical isometry for trainability and neural architecture search [27, 44, 25]. Recent work on deep kernel shaping shows via careful tuning of initialisation and activation functions that while skip connections and normalisation layers may be sufficient for good trainability, they are not necessary [31, 50]. Other recent work has also shown benefits to inference performance by removing skip connections from the trained model using a parameter transformation [11] or by removing them from the model altogether and incorporating them only into the optimiser [12].
Finally, a line of work has recently emerged on training in the more realistic, large learning rate regime known as "edge of stability" [9, 3, 10]. This intriguing line of work diverges from ours, and its integration into the framework we present is a promising future research direction.
4 Formal framework
In this section we will define our theoretical framework. We will use  $\mathbb{R}^{n\times m}$  to denote the space of  $n\times m$  matrices, and vectorise rows first. We use  $\mathrm{Id}_m$  to denote the  $m\times m$  identity matrix, and  $1_{n}$  to denote the vector of ones in  $\mathbb{R}^n$ . We use  $\otimes$  to denote the Kronecker (tensor) product of matrices, and given a vector  $v\in \mathbb{R}^n$  we use  $\mathrm{diag}(v)$  to denote the  $n\times n$  matrix whose leading diagonal is  $v$ . Given a matrix  $A\in \mathbb{R}^{n\times m}$ , and a seminorm  $\| \cdot \|$  on  $\mathbb{R}^m$ ,  $\| A\|_{row}$  will be used to denote the vector of  $\| \cdot \|$  -seminorms of each row of  $A$ . The smallest singular value of  $A$  is denoted  $\sigma (A)$ , and the smallest eigenvalue of  $AA^T$  denoted  $\lambda (A)$ . Full proofs of all of our results can be found in the appendix.
Our theory is derived from the following formal generalisation of a deep neural network.
Definition 4.1. By a multilayer parameterised system (MPS) we mean a family  $\{f_l : \mathbb{R}^{p_l} \times \mathbb{R}^{d_{l-1} \times N} \to \mathbb{R}^{d_l \times N}\}_{l=1}^L$  of functions. Given a data matrix  $X \in \mathbb{R}^{d_0 \times N}$ , we denote by  $F : \mathbb{R}^{\sum_{i=1}^L p_i} \to \mathbb{R}^{d_L \times N}$  the parameter-function map<sup>2</sup> defined by
$$
F (\vec {\theta}) := f _ {L} \left(\theta_ {L}\right) \circ \dots \circ f _ {1} \left(\theta_ {1}\right) (X), \tag {8}
$$
for  $\vec{\theta} = (\theta_{1},\dots ,\theta_{L})^{T}\in \mathbb{R}^{\sum_{i = 1}^{L}p_{i}}$
Definition 4.1 is sufficiently general to encompass all presently used neural network layers, but we will assume without further comment from here on in that all layers are continuously differentiable. Before we give examples, we record the following result giving the form of the derivative of the parameter-function map, which follows easily from the chain rule. It will play a key role in the analysis of the following subsections.
Proposition 4.2. Let  $\{f_l : \mathbb{R}^{p_l} \times \mathbb{R}^{d_{l-1} \times N} \to \mathbb{R}^{d_l \times N}\}_{l=1}^L$  be a MPS, and  $X \in \mathbb{R}^{d_0 \times N}$  a data matrix. For  $1 \leq l \leq L$ , denote the derivatives of  $f_l$  with respect to the  $\mathbb{R}^{p_l}$  and  $\mathbb{R}^{d_{l-1}}$  variables by  $Df_l$  and  $Jf_l$  respectively, and let  $f_{<l}(\vec{\theta}, X)$  denote the composite
$$
f _ {<   l} (\vec {\theta}, X) := f _ {l - 1} \left(\theta_ {l - 1}\right) \circ \dots \circ f _ {1} \left(\theta_ {1}\right) (X). \tag {9}
$$
The derivative  $DF$  of the associated parameter-function map is given by
$$
D F (\vec {\theta}) = \left(D _ {\theta_ {1}} F (\vec {\theta}), \dots , D _ {\theta_ {L}} F (\vec {\theta})\right), \tag {10}
$$
where for  $1 \leq l < L$ ,  $D_{\theta_l}F(\vec{\theta})$  is given by the formula
$$
\left(\prod_ {j = l + 1} ^ {L} J f _ {j} \left(\theta_ {j}, f _ {<   j} (\vec {\theta}, X)\right)\right) D f _ {l} \left(\theta_ {l}, f _ {<   l} (\vec {\theta}, X)\right), \tag {11}
$$
with the product taken so that indices are arranged in descending order from left to right.
All common differentiable neural network layers fit into this framework; we record some examples in detail in the appendix. We will see in the next section that insofar as one wishes to guarantee global smoothness, the usual parameterisation of affine layers is poor, although this defect can be ameliorated to differing extents by normalisation strategies.
4.1 Smoothness
In this subsection, we give sufficient conditions for the derivative of the parameter-function map of a MPS to be bounded and Lipschitz. We are thereby able to give sufficient conditions for any associated loss function to have Lipschitz gradients on its sublevel sets. We begin with a formal proposition that describes the Lipschitz properties of the derivative of the parameter-function map of a MPS in terms of those of its constituent layers.
$^2\mathbb{R}^{d_L\times N}$  is canonically isomorphic to the space of  $\mathbb{R}^{d_L}$ -valued functions on an  $N$ -point set.
Proposition 4.3. Let  $\{f_l : \mathbb{R}^{p_l} \times \mathbb{R}^{d_{l-1} \times N} \to \mathbb{R}^{d_l \times N}\}_{l=1}^L$ , and let  $\{S_l \subset \mathbb{R}^{p_l}\}_{l=1}^L$  be subsets of the parameter spaces. Suppose that for each bounded set  $B_l \subset \mathbb{R}^{d_{l-1} \times N}$ , the maps  $f_l, Df_l$  and  $Jf_l$  are all bounded and Lipschitz on  $S_l \times B_l$ . Then for any data matrix  $X \in \mathbb{R}^{d_0 \times N}$ , the derivative  $DF$  of the associated parameter-function map  $F : \mathbb{R}^p \to \mathbb{R}^{d_L \times N}$  is bounded and Lipschitz on  $S := \prod_{j=1}^L S_j \subset \mathbb{R}^p$ .
Proposition 4.3 has the following immediate corollary, whose proof follows easily from Lemma B.1.
Corollary 4.4. Let  $c: \mathbb{R}^{d_L} \to \mathbb{R}$  be any cost function whose gradient is bounded and Lipschitz on its sublevel sets  $Z_{\alpha} := \{z \in \mathbb{R}^{d_L}: c(z) \leq \alpha\}$ . If  $\{f_l: \mathbb{R}^{p_l} \times \mathbb{R}^{d_{l-1} \times N} \to \mathbb{R}^{d_l \times N}\}_{l=1}^L$  is any MPS satisfying the hypotheses of Proposition 4.3 and  $X \in \mathbb{R}^{d_0 \times N}$ , then the associated loss function  $\ell := \gamma \circ F$  has Lipschitz gradients over  $S := \prod_{l=1}^L S_l$ .
The most ubiquitous cost functions presently in use (the mean square error and cross entropy functions) satisfy the hypotheses of Corollary 4.4. We now turn to an analysis of common layer types in deep neural networks and indicate to what extent they satisfy the hypotheses of Proposition 4.3.
Theorem 4.5. Fix  $\epsilon >0$ . The following layers satisfy the hypotheses of Proposition 4.3 over all of parameter space.
1. Continuously differentiable nonlinearities.
2. Bias-free  $\epsilon$ -weight-normalised or  $\epsilon$ -entry-normalised affine layers<sup>3</sup>.
3. Any residual block whose branch is a composite of any of the above layer types.
Consequently, the loss function of any neural network composed of layers as above, trained with a cost function satisfying the hypotheses of Corollary 4.4, has globally Lipschitz gradients along any sublevel set.
The proof we give of Theorem 4.5 also considers composites bn  $\circ$  aff of batch norm layers with affine layers. Such composites satisfy the hypotheses of Proposition 4.3 only over sets in data-space  $\mathbb{R}^{d\times N}$  which consist of matrices with nondegenerate covariance. Since such sets are generic (probability 1 with respect to any probability measure that is absolutely continuous with respect to Lebesgue measure), batch norm layers satisfy the hypotheses of Proposition 4.3 with high probability over random initialisation.
Theorem 4.5 says that normalisation of parameters enables global analysis of the loss, while standard affine layers, due to their unboundedness, are well-suited to analysis only over bounded sets in parameter space.
4.2 Regularity
Having examined the Lipschitz properties of MPS and given examples of layers with global Lipschitz properties, let us now do the same for the regularity properties of MPS. Formally one has the following simple result.
Proposition 4.6. Let  $\{f_l:\mathbb{R}^{pl}\times \mathbb{R}^{d_{l - 1}\times N}\to \mathbb{R}^{d_l\times N}\}_{i = 1}^L$  be a MPS and  $X\in \mathbb{R}^{d_0\times N}$  a data matrix. Then