Rapid training of deep neural networks without skip connections or normalization layers using Deep Kernel Shaping

Abstract

Using an extended and formalized version of the Q/C map analysis of Poole et al. (2016), along with Neural Tangent Kernel theory, we identify the main pathologies present in deep networks that prevent them from training fast and generalizing to unseen data, and show how these can be avoided by carefully controlling the "shape" of the network's initializationtime kernel function. We then develop a method called Deep Kernel Shaping (DKS), which accomplishes this using a combination of precise parameter initialization, activation function transformations, and small architectural tweaks, all of which preserve the model class.In our experiments we show that DKS enables SGD training of residual networks without normalization layers on Imagenet and CIFAR-10 classification tasks at speeds comparable to standard ResNetV2 and Wide-ResNet models, with only a small decrease in generalization performance. And when using K-FAC as the optimizer, we achieve similar results for networks without skip connections. Our results apply for a large variety of activation functions, including those which traditionally perform very badly, such as the logistic sigmoid. In addition to DKS, we contribute a detailed analysis of skip connections, normalization layers, special activation functions like RELU and SELU, and various initialization schemes, explaining their effectiveness as alternative (and ultimately incomplete) ways of "shaping" the network's initialization-time kernel.

Learning from Randomly Initialized Neural Network Features

“…In summary, based on our experiments on different initializations and activation functions and our findings on ResNet-18, NNPK seems to hint at the choices that are known to work also well for fully trained models (Gotmare et al, 2018;Shah et al, 2016;Martens et al, 2021).…”

“…First, the NRF seems to improve by using skip connections for both initialization and activation function combinations. However, the performance of the trained model using the adjustments proposed in (Martens et al, 2021) seems to improve without the skip connections. This observation shows that there are cases for which NNPK may not reflect the performance of the final trained model (although the results may vary when using data augmentation and other types of regularizations).…”

“…There has been a number of efforts to replace or remove these two elements from ResNets (Zhang et al, 2019;Gaur et al, 2020;Bachlechner et al, 2021). The more recent adjustments proposed in (Martens et al, 2021) have shown further promise. This modification includes replacing the ReLU activation function with a scaled leaky ReLU and the use of delta orthogonal initialization for the convolutional kernels (and orthogonal initialization for the final dense layer).…”

“…Since the mean and scale values are initially set to 0 and 1, respectively, BatchNorm with default initialization will have no effect on the NNPK, but affects the performance of the trained model. We also examine whether NRF varies by removing skip connections and whether the changes proposed in (Martens et al, 2021) affect the results.…”

“…We create variants of the ResNet-18 model by removing BatchNorm and/or skip connections. In the first approach, we use the default He normal initialization with ReLU activation while in the second case, we use the delta orthogonal initialization proposed in (Martens et al, 2021) for the convolutional filters (and orthogonal initialization for the final dense layer). We also replace the ReLU activation with the scaled leaky ReLU proposed in (Martens et al, 2021).…”