[Paper Review] Freeze and Chaos for DNNs: an NTK view of Batch Normalization, Checkerboard and Boundary Effects.

Arthur Paul Jacot, Franck Gabriel+1 more|arXiv (Cornell University)|Jul 11, 2019

Generative Adversarial Networks and Image Synthesis17 references21 citations

TL;DR

This paper uses the Neural Tangent Kernel (NTK) to analyze deep neural networks (DNNs), revealing two regimes—freeze and chaos—governed by depth and nonlinearity. It shows that Batch Normalization and normalized nonlinearities avoid the freeze regime, which causes generator collapse in DCGANs, and proposes a layer-dependent learning rate and NTK-based parametrization to eliminate border artifacts and improve convergence in deep deconvolutional networks.

ABSTRACT

In this paper, we analyze a number of architectural features of Deep Neural Networks (DNNs), using the so-called Neural Tangent Kernel (NTK). The NTK describes the training trajectory and generalization of DNNs in the infinite-width limit. In this limit, we show that for (fully-connected) DNNs, as the depth grows, two regimes appear: freeze (also known as order), where the (scaled) NTK converges to a constant (slowing convergence), and chaos, where it converges to a Kronecker delta (limiting generalization). We show that when using the scaled ReLU as a nonlinearity, we naturally end up in the freeze. We show that Batch Normalization (BN) avoids the freeze regime by reducing the importance of the constant mode in the NTK. A similar effect is obtained by normalizing the nonlinearity which moves the network to the chaotic regime. We uncover the same freeze and chaos modes in Deep Deconvolutional Networks (DC-NNs). The freeze regime is characterized by checkerboard patterns in the image space in addition to the constant modes in input space. Finally, we introduce a new NTK-based parametrization to eliminate border artifacts and we propose a layer-dependent learning rate to improve the convergence of DC-NNs. We illustrate our findings by training DCGANs using our setup. When trained in the freeze regime, we see that the generator collapses to a checkerboard mode. We also demonstrate numerically that the generator collapse can be avoided and that good quality samples can be obtained, by tuning the nonlinearity to reach the chaos regime (without using batch normalization).

Motivation & Objective

• To understand the impact of architectural components like Batch Normalization and nonlinearity normalization on DNN training dynamics using the Neural Tangent Kernel (NTK).
• To identify and characterize the freeze and chaos regimes in deep fully-connected and deconvolutional networks via NTK analysis.
• To resolve generator collapse in DCGANs by avoiding the freeze regime through nonlinearity tuning or architectural modifications.
• To eliminate border artifacts in deep deconvolutional networks using a new NTK-based parametrization.
• To improve training convergence via a layer-dependent learning rate strategy grounded in NTK theory.

Proposed method

• Analyzes the infinite-width limit of fully-connected DNNs using the NTK to derive conditions under which the NTK converges to a constant (freeze) or Kronecker delta (chaos).
• Demonstrates that using scaled ReLU as nonlinearity leads naturally to the freeze regime due to dominant constant mode in the NTK.
• Shows that Batch Normalization reduces the influence of the constant mode in the NTK, effectively shifting the network to the chaotic regime.
• Applies the same NTK framework to Deep Deconvolutional Networks (DC-NNs), revealing checkerboard patterns in the image space as a signature of the freeze regime.
• Proposes a new NTK-based parametrization to eliminate border artifacts in DC-NNs by modifying the kernel structure.
• Introduces a layer-dependent learning rate that improves convergence by aligning with the NTK's depth-dependent behavior.

Experimental results

Research questions

• RQ1 What are the two distinct training regimes—freeze and chaos—in the infinite-width limit of deep fully-connected networks, and how do they emerge from the NTK?
• RQ2 How does Batch Normalization prevent the network from entering the freeze regime by modifying the NTK's constant mode?
• RQ3 Can normalizing the nonlinearity instead of using Batch Normalization also shift the network to the chaotic regime?
• RQ4 What role do checkerboard patterns play in the image space of deep deconvolutional networks, and how are they linked to the freeze regime?
• RQ5 Can border artifacts in deep deconvolutional networks be eliminated through a new NTK-based parametrization, and does a layer-dependent learning rate improve convergence?

Key findings

• The NTK of deep fully-connected networks converges to a constant in the freeze regime when using scaled ReLU, leading to slow convergence and poor generalization.
• Batch Normalization reduces the contribution of the constant mode in the NTK, effectively avoiding the freeze regime and enabling faster convergence.
• Normalizing the nonlinearity instead of using Batch Normalization also shifts the network to the chaotic regime, where the NTK converges to a Kronecker delta, improving generalization.
• Checkerboard patterns emerge in the image space of deep deconvolutional networks during the freeze regime, indicating mode collapse in generative models.
• The proposed NTK-based parametrization successfully eliminates border artifacts in deep deconvolutional networks, improving sample quality.
• A layer-dependent learning rate based on NTK theory significantly improves training convergence in deep deconvolutional networks, especially in the chaotic regime.