Exact, Tractable Gauss-Newton Optimization in Deep Reversible Architectures Reveal Poor Generalization

TL;DR

It is shown for the first time that exact Gauss-Newton (GN) updates take on a tractable form in a class of deep reversible architectures that are sufficiently expressive to be meaningfully applied to common benchmark datasets.

Abstract

Second-order optimization has been shown to accelerate the training of deep neural networks in many applications, often yielding faster progress per iteration on the training loss compared to first-order optimizers. However, the generalization properties of second-order methods are still being debated. Theoretical investigations have proved difficult to carry out outside the tractable settings of heavily simplified model classes -- thus, the relevance of existing theories to practical deep learning applications remains unclear. Similarly, empirical studies in large-scale models and real datasets are significantly confounded by the necessity to approximate second-order updates in practice. It is often unclear whether the observed generalization behaviour arises specifically from the second-order nature of the parameter updates, or instead reflects the specific structured (e.g.\ Kronecker) approximations used or any damping-based interpolation towards first-order updates. Here, we show for the first time that exact Gauss-Newton (GN) updates take on a tractable form in a class of deep reversible architectures that are sufficiently expressive to be meaningfully applied to common benchmark datasets. We exploit this novel setting to study the training and generalization properties of the GN optimizer. We find that exact GN generalizes poorly. In the mini-batch training setting, this manifests as rapidly saturating progress even on the \emph{training} loss, with parameter updates found to overfit each mini-batchatch without producing the features that would support generalization to other mini-batches. We show that our experiments run in the ``lazy'' regime, in which the neural tangent kernel (NTK) changes very little during the course of training. This behaviour is associated with having no significant changes in neural representations, explaining the lack of generalization.

Figures (15)

• Figure 1: Training loss and accuracy on (a) MNIST and (b) CIFAR-10 in a full-batch scenario where each dataset is trimmed to a fixed subset of $n = 1024$ images. GN converges much faster than Adam and SGD.
• Figure 2: Training loss, test loss, and test accuracy on (a) MNIST and (b) CIFAR-10 in a mini-batch scenario. GN does not exhibit the same properties observed in the full-batch setting. In fact, Adam reaches lower training and test loss.
• Figure 3: Percentage change in training loss after each update. GN decreases the loss for the current mini-batch more than SGD and Adam early in training.
• Figure 4: NTK and CKA similarity evolution across training for GN, Adam and SGD. Top two panels include (a) the rate of change of the NTK and (b) the NTK similarity during training with respect to initialization. Bottom three panels, along the same axis, include the CKA similarities for the (c) first, (d) middle and (e) last block with respect to their initial values.
• Figure 5: Experiments on CIFAR using a model without Inverted Bottleneck (Full-batch on the left, mini-batch on center and right). While the theoretical guarantees do not hold in this setting, the results follow the same trend observed in Figure \ref{['fig:mini_batch']}.

Theorems & Definitions (3)

• Theorem 4.3
• Proposition 4.4
• proof