Demystifying Batch Normalization in ReLU Networks: Equivalent Convex Optimization Models and Implicit Regularization

Abstract

Batch Normalization (BN) is a commonly used technique to accelerate and stabilize training of deep neural networks. Despite its empirical success, a full theoretical understanding of BN is yet to be developed. In this work, we analyze BN through the lens of convex optimization. We introduce an analytic framework based on convex duality to obtain exact convex representations of weight-decay regularized ReLU networks with BN, which can be trained in polynomial-time. Our analyses also show that optimal layer weights can be obtained as simple closed-form formulas in the high-dimensional and/or overparameterized regimes. Furthermore, we find that Gradient Descent provides an algorithmic bias effect on the standard non-convex BN network, and we design an approach to explicitly encode this implicit regularization into the convex objective. Experiments with CIFAR image classification highlight the effectiveness of this explicit regularization for mimicking and substantially improving the performance of standard BN networks.

Accepted to ICLR 2022. First two authors contributed equally to this work; 36 pages, 13 figures

References in corpus (6)

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Prevalence of Neural Collapse during the terminal phase of deep learning training
• Rethinking the Usage of Batch Normalization and Dropout in the Training of Deep Neural Networks
• Revealing the Structure of Deep Neural Networks via Convex Duality
• Batch Normalization Provably Avoids Rank Collapse for Randomly Initialised Deep Networks

Cited by in corpus (6)

• Revealing the Structure of Deep Neural Networks via Convex Duality
• Convex Geometry and Duality of Over-parameterized Neural Networks
• Hidden Convexity of Wasserstein GANs: Interpretable Generative Models with Closed-Form Solutions
• The Convex Geometry of Backpropagation: Neural Network Gradient Flows Converge to Extreme Points of the Dual Convex Program
• Training Quantized Neural Networks to Global Optimality via Semidefinite Programming
• Implicit Convex Regularizers of CNN Architectures: Convex Optimization of Two- and Three-Layer Networks in Polynomial Time