Theoretical characterisation of the Gauss-Newton conditioning in Neural Networks

TL;DR

This work provides a theoretical characterization of Gauss-Newton conditioning in neural networks, establishing tight upper bounds on the GN matrix condition number for deep linear networks of arbitrary depth and width, and extending the analysis to two-layer ReLU networks. By expressing the GN spectrum through data covariance $\boldsymbol{\Sigma}$ and layer-wise weight conditioning, the authors show that the conditioning scales with $\kappa(\boldsymbol{\Sigma})$ and can be controlled via architectural choices such as skip connections and Batch Normalization. They further extend the analysis to non-linear activations in simple settings and demonstrate empirically that architectural components like residuals and BN can substantially improve conditioning, aligning with observed training dynamics. Overall, the paper links classic optimization conditioning to modern network design, offering principled guidance on width-depth scaling and normalization strategies to improve optimization speed and stability.

Abstract

The Gauss-Newton (GN) matrix plays an important role in machine learning, most evident in its use as a preconditioning matrix for a wide family of popular adaptive methods to speed up optimization. Besides, it can also provide key insights into the optimization landscape of neural networks. In the context of deep neural networks, understanding the GN matrix involves studying the interaction between different weight matrices as well as the dependencies introduced by the data, thus rendering its analysis challenging. In this work, we take a first step towards theoretically characterizing the conditioning of the GN matrix in neural networks. We establish tight bounds on the condition number of the GN in deep linear networks of arbitrary depth and width, which we also extend to two-layer ReLU networks. We expand the analysis to further architectural components, such as residual connections and convolutional layers. Finally, we empirically validate the bounds and uncover valuable insights into the influence of the analyzed architectural components.

Figures (23)

• Figure 1: Training loss (left) and condition number $\kappa$ of GN (right) for a ResNet20 trained on a subset of Cifar10 ($n=1000$) with different proportions of pruned weights. Weights were pruned layerwise by magnitude at initialization.
• Figure 2: a) Condition number at initialization under Kaiming normal initialization of GN $\kappa(\hat{{\bf G}}_O)$ and first upper bound derived in \ref{['lemma:UpperBoundCondNumHO']} and Eq.\ref{['eq:upper_bound_1_condnum_H_O_residual']} for whitened MNIST as a function of depth $L$ for different hidden layer widths $m$ for a Linear Network over 3 initializations. b) Scaling the width of the hidden layer proportionally to the depth leads to slower growth of the condition number (left) or improves the condition number with depth if the scaling factor is chosen sufficiently large (right).
• Figure 3: Comparison of derived upper bounds in Lemma \ref{['lemma:UpperBoundCondNumHO']} for the condition number at initialization for whitened MNIST over 20 runs. Note the logarithmic scaling of the $y$-axis.
• Figure 4: Adding skip connections between each layer for a general L-layer linear Neural Network.
• Figure 5: Comparison of condition number of GN $\kappa(\hat{{\bf G}}_O)$ between Linear Network and Residual Network with $\beta=1$ for whitened MNIST (left) and whitened Cifar-10 (right) at initialization using Kaiming normal initialization over three seeds. The upper bounds refer to the first upper bound in \ref{['lemma:UpperBoundCondNumHO']} and Eq. \ref{['eq:upper_bound_1_condnum_H_O_residual']}, respectively.

Theorems & Definitions (16)

• Remark R1: Difference between ${\bf H}_L$ and ${\bf G}_O$
• Lemma 1
• proof : Proof sketch
• Remark R2: Role of the data covariance
• Lemma 2
• Lemma 3
• Lemma 4
• Lemma 5
• Lemma 5
• Lemma 6: Corollary 5.35 in vershynin2010introduction