Towards Training Without Depth Limits: Batch Normalization Without
  Gradient Explosion
Alexandru Meterez Amir Joudaki Francesco Orabona Alexander Immer Gunnar Rätsch Hadi Daneshmand
Normalization layers are one of the key building blocks for deep neural
networks. Several theoretical studies have shown that batch normalization
improves the signal propagation, by avoiding the representations from becoming
collinear across the layers. However, results on mean-field theory of batch
normalization also conclude that this benefit comes at the expense of exploding
gradients in depth. Motivated by these two aspects of batch normalization, in
this study we pose the following question: "Can a batch-normalized network keep
the optimal signal propagation properties, but avoid exploding gradients?" We
answer this question in the affirmative by giving a particular construction of
an Multi-Layer Perceptron (MLP) with linear activations and batch-normalization
that provably has bounded gradients at any depth. Based on Weingarten calculus,
we develop a rigorous and non-asymptotic theory for this constructed MLP that
gives a precise characterization of forward signal propagation, while proving
that gradients remain bounded for linearly independent input samples, which
holds in most practical settings. Inspired by our theory, we also design an
activation shaping scheme that empirically achieves the same properties for
certain non-linear activations.