Rapid training of deep neural networks without skip connections or normalization layers using Deep Kernel Shaping
\nameJames Martens \nameAndy Ballard \nameGuillaume Desjardins \nameGrzegorz Swirszcz \nameValentin Dalibard \addrDeepMind, London, UK \AND\nameJascha Sohl-Dickstein \nameSamuel S. Schoenholz \addrGoogle, Mountainview, USA
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 initialization-time 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 de
中文速览
深度神经网络难以训练的根本原因在于：随着层数加深，网络的"核函数（kernel function）"会退化，导致输出对输入的差异几乎不再敏感，从而让梯度下降变得极慢甚至失效。作者将这一问题用扩展版的Q/C映射分析框架加以形式化，并借助神经切线核（Neural Tangent Kernel, NTK）理论严格论证了退化核与训练缓慢、泛化差之间的必然联系。在此基础上，他们提出了"深度核塑形（Deep Kernel Shaping, DKS）"方法：通过精细的参数初始化、对激活函数做保持模型类不变的缩放与平移变换，以及少量结构调整，把网络初始化时刻的核函数"塑造"成良好的形态。实验结果表明，DKS使得没有归一化层的残差网络在ImageNet和CIFAR-10上能以接近标准ResNetV2的速度用SGD训练，配合K-FAC优化器甚至能让完全没有跳跃连接和归一化层的普通深网达到相当水平，且对逻辑sigmoid等传统上表现糟糕的激活函数同样有效。这项工作的重要意义在于：它为"网络结构设计"与"可训练性保障"的解耦提供了理论基础，也首次在ImageNet这一困难任务上证明无需跳跃连接和归一化层也能快速训练深网，从而为重新理解批归一化、残差连接等标准组件的作用机制提供了全新视角。
原文 arXiv:2110.01765;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2110.01765v1