A Comprehensive and Modularized Statistical Framework for Gradient Norm Equality in Deep Neural Networks
This work addresses the problem of gradient explosion and vanishing for deep learning practitioners by providing more practical tools and insights, though it is incremental in building on existing concepts like Gradient Norm Equality.
The paper tackles the challenge of evaluating gradient norm stability in modern deep neural networks by proposing Block Dynamical Isometry, a metric requiring weaker assumptions, and a modularized statistical framework based on free probability to handle complex network connections and components. It applies this to analyze and improve initialization, normalization, and network structures, resulting in a new normalization technique that is 30% faster than batch normalization without accuracy loss.
In recent years, plenty of metrics have been proposed to identify networks that are free of gradient explosion and vanishing. However, due to the diversity of network components and complex serial-parallel hybrid connections in modern DNNs, the evaluation of existing metrics usually requires strong assumptions, complex statistical analysis, or has limited application fields, which constraints their spread in the community. In this paper, inspired by the Gradient Norm Equality and dynamical isometry, we first propose a novel metric called Block Dynamical Isometry, which measures the change of gradient norm in individual block. Because our Block Dynamical Isometry is norm-based, its evaluation needs weaker assumptions compared with the original dynamical isometry. To mitigate the challenging derivation, we propose a highly modularized statistical framework based on free probability. Our framework includes several key theorems to handle complex serial-parallel hybrid connections and a library to cover the diversity of network components. Besides, several sufficient prerequisites are provided. Powered by our metric and framework, we analyze extensive initialization, normalization, and network structures. We find that Gradient Norm Equality is a universal philosophy behind them. Then, we improve some existing methods based on our analysis, including an activation function selection strategy for initialization techniques, a new configuration for weight normalization, and a depth-aware way to derive coefficients in SeLU. Moreover, we propose a novel normalization technique named second moment normalization, which is theoretically 30% faster than batch normalization without accuracy loss. Last but not least, our conclusions and methods are evidenced by extensive experiments on multiple models over CIFAR10 and ImageNet.