[Paper Review] Are All Layers Created Equal?

Chiyuan Zhang, Samy Bengio+1 more|arXiv (Cornell University)|Feb 6, 2019

Adversarial Robustness in Machine Learning49 references110 citations

TL;DR

The paper empirically shows that layers in deep networks are heterogeneous: some are robust to post-training re-initialization/re-randomization, while others are critical, with robustness varying by architecture and task.

ABSTRACT

Understanding deep neural networks is a major research objective with notable experimental and theoretical attention in recent years. The practical success of excessively large networks underscores the need for better theoretical analyses and justifications. In this paper we focus on layer-wise functional structure and behavior in overparameterized deep models. To do so, we study empirically the layers' robustness to post-training re-initialization and re-randomization of the parameters. We provide experimental results which give evidence for the heterogeneity of layers. Morally, layers of large deep neural networks can be categorized as either "robust" or "critical". Resetting the robust layers to their initial values does not result in adverse decline in performance. In many cases, robust layers hardly change throughout training. In contrast, re-initializing critical layers vastly degrades the performance of the network with test error essentially dropping to random guesses. Our study provides further evidence that mere parameter counting or norm calculations are too coarse in studying generalization of deep models, and "flatness" and robustness analysis of trained models need to be examined while taking into account the respective network architectures.

Motivation & Objective

• Investigate whether individual layers play distinct roles in trained deep networks.
• Assess layer robustness to post-training re-initialization and re-randomization across architectures and tasks.
• Understand how layer robustness relates to network capacity, task difficulty, and generalization.

Proposed method

• Train various architectures (FCNs, VGG, ResNet, transformers, ViTs, MLP-Mixers) on MNIST, CIFAR-10, ImageNet, and LM1B.
• After training, perform re-initialization or re-randomization of individual layers and measure test/test-perplexity performance.
• Use checkpointing to re-initialize layers to their initial values or re-sample from initialization distributions without retraining.
• Analyze layer-wise weight distances to initialization using 2-norm and infinity-norm.
• Examine joint robustness by re-initializing/re-randomizing groups of layers and test the effects.

Experimental results

Research questions

• RQ1 Do layers within deep networks have heterogeneous roles, being either robust or critical to performance?
• RQ2 How does layer robustness vary across architectures (FCN, VGG, ResNet, Transformers) and datasets (MNIST, CIFAR-10, ImageNet, LM1B)?
• RQ3 Can robustness properties of layers explain generalization and the effect of overparameterization?
• RQ4 What is the impact of joint layer perturbations on network performance, and can constraints like freezing or removing layers mitigate loss?
• RQ5 Are robustness patterns observed in vision models generalizable to convolution-free architectures (ViTs, MLP-Mixers) and language models?

Key findings

• Re-randomizing any layer typically collapses performance to random guessing, showing non-robustness of that layer.
• Re-initializing the bottom layer tends to be the most detrimental, while higher layers often remain robust to re-initialization.
• Layer robustness correlates with network capacity and task difficulty: wider networks tend to have more robust top layers, and harder tasks keep more layers as sensitive.
• ResNets distribute critical layers throughout depth, unlike FCNs/VGGs where bottom layers are more critical.
• Groups of layers can be jointly robust or fragile, and architectural constraints (freezing/removing layers) affect the observed robustness.
• The robustness pattern generalizes to language models and convolution-free architectures (ViTs, MLP-Mixers), with higher layers often more robust, though specifics vary by component (e.g., layernorms in transformers).