Optimization and Generalization Guarantees for Weight Normalization

TL;DR

This paper analyzes optimization and generalization for deep WeightNorm networks with smooth activations. It derives a bound on the predictor Hessian $\| abla^2_\theta f(\theta; x)\|_2$ that scales with depth as $L^3$ and width as $m$, and inversely with the minimum weight norm, enabling gradient-descent convergence guarantees. It also establishes a uniform-convergence based generalization bound that is independent of width and scales as ${\cal O}(\sqrt{L}/\sqrt{n})$, using the Rademacher complexity of WeightNorm networks. Empirical results on CIFAR-10 and MNIST corroborate the theoretical links between normalization terms and training dynamics, illustrating practical impact of WeightNorm on optimization and generalization.

Abstract

Weight normalization (WeightNorm) is widely used in practice for the training of deep neural networks and modern deep learning libraries have built-in implementations of it. In this paper, we provide the first theoretical characterizations of both optimization and generalization of deep WeightNorm models with smooth activation functions. For optimization, from the form of the Hessian of the loss, we note that a small Hessian of the predictor leads to a tractable analysis. Thus, we bound the spectral norm of the Hessian of WeightNorm networks and show its dependence on the network width and weight normalization terms--the latter being unique to networks without WeightNorm. Then, we use this bound to establish training convergence guarantees under suitable assumptions for gradient decent. For generalization, we use WeightNorm to get a uniform convergence based generalization bound, which is independent from the width and depends sublinearly on the depth. Finally, we present experimental results which illustrate how the normalization terms and other quantities of theoretical interest relate to the training of WeightNorm networks.

Figures (2)

• Figure 1: Training neural networks for two different widths $m\in\{1024, 2048\}$ on the CIFAR-10 dataset, with two hidden layers $L=2$ of same width, with learning rate 0.001, and weights initialized independently from a uniform distribution $[-\frac{0.5}{\sqrt{m}}, \frac{0.5}{\sqrt{m}}]$. Each subfigure plots: (a): ${\min_{\substack{i\in[m]\\l\in[L]}}\left\lVert W^{(l)}_{i,t}\right\rVert_2}$; (b): $\left\lVert\nabla {\cal L}(\theta_t)\right\rVert_2^2/{\cal L}(\theta_t)$; where $t$ represents the number of iterations.
• Figure 2: Training neural networks for two different widths $m\in\{512, 1024\}$ on the MNIST dataset, with two hidden layers $L=2$ of same width, with learning rate 0.001 and the weights are initialized with a uniform distribution $[-\frac{5}{\sqrt{m}}, \frac{5}{\sqrt{m}}]$. Each subfigure plots: (a): ${\min_{\substack{i\in[m]\\l\in[L]}}\left\lVert W^{(l)}_{i,t}\right\rVert_2}$; (b): $\left\lVert\nabla {\cal L}(\theta_t)\right\rVert_2^2/{\cal L}(\theta_t)$; where $t$ represents the number of iterations.

Theorems & Definitions (42)

• Definition 3.1: Minimum weight vector norm
• Theorem 4.1: Hessian bound for WeightNorm
• Corollary 4.1: Hessian bound for WeightNorm under $\phi(0)=0$
• Remark 4.1: Comparison to networks without WeightNorm
• Remark 4.2: Dependence on the activation function
• Lemma 4.1: Predictor gradient bounds
• Corollary 4.2: Predictor gradient bounds under $\phi(0)=0$
• Remark 4.3: The Lipschitz constant is the same for networks of any depth
• Proposition 4.1: Empirical loss and empirical loss gradient bounds
• Corollary 4.3: Empirical loss and empirical loss gradient bounds under $\phi(0)=0$