How Do the Architecture and Optimizer Affect Representation Learning? On the Training Dynamics of Representations in Deep Neural Networks

TL;DR

This paper investigates how the architecture and optimizer influence the evolution of representations in deep neural networks during training, revealing distinct training dynamics and layer behaviors in different models and optimization methods.

Contribution

It provides a detailed analysis of training dynamics in overparameterized DNNs, highlighting how architecture and optimizer choices affect representation evolution and layer synchronization.

Findings

Training phases are more distinguishable in SGD than Adam.

Memorization phases are clearer in SGD training.

Vision Transformers show synchronized layer dynamics unlike ResNet.

Abstract

In this paper, we elucidate how representations in deep neural networks (DNNs) evolve during training. Our focus is on overparameterized learning settings where the training continues much after the trained DNN starts to perfectly fit its training data. We examine the evolution of learned representations along the entire training process. We explore the representational similarity of DNN layers, each layer with respect to its own representations throughout the training process. For this, we use two similarity metrics: (1) The centered kernel alignment (CKA) similarity; (2) Similarity of decision regions of linear classifier probes that we train for the DNN layers. We visualize and analyze the decision regions of the DNN output and the layer probes during the DNN training to show how they geometrically evolve. Our extensive experiments discover training dynamics patterns that can emerge in layers depending on the relative layer-depth, architecture and optimizer. Among our findings: (i) The training phases, including those related to memorization, are more distinguishable in SGD training than in Adam training, and for Vision Transformer (ViT) than for ResNet; (ii) Unlike ResNet, the ViT layers have synchronized dynamics of representation learning.