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

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.