Tracing Representation Progression: Analyzing and Enhancing Layer-Wise Similarity

TL;DR

This paper investigates layer-wise representation similarity in transformer models, showing that simple cosine similarity aligns with complex metrics, and proposes an aligned training method to enhance early layer effectiveness and interpretability.

Contribution

It introduces a simple cosine similarity measure for analyzing transformer representations, provides theoretical insights, and proposes an aligned training approach to improve shallow layer performance.

Findings

Cosine similarity aligns with CKA in capturing layer-wise similarity.

Representation similarity increases as layers get closer, correlating with model confidence.

Aligned training enhances early layer accuracy and enables effective multi-exit models.

Abstract

Analyzing the similarity of internal representations has been an important technique for understanding the behavior of deep neural networks. Most existing methods for analyzing the similarity between representations of high dimensions, such as those based on Centered Kernel Alignment (CKA), rely on statistical properties of the representations for a set of data points. In this paper, we focus on transformer models and study the similarity of representations between the hidden layers of individual transformers. In this context, we show that a simple sample-wise cosine similarity metric is capable of capturing the similarity and aligns with the complicated CKA. Our experimental results on common transformers reveal that representations across layers are positively correlated, with similarity increasing when layers get closer. We provide a theoretical justification for this phenomenon under the geodesic curve assumption for the learned transformer. We then show that an increase in representation similarity implies an increase in predicted probability when directly applying the last-layer classifier to any hidden layer representation. We then propose an aligned training method to improve the effectiveness of shallow layer by enhancing the similarity between internal representations, with trained models that enjoy the following properties: (1) more early saturation events, (2) layer-wise accuracies monotonically increase and reveal the minimal depth needed for the given task, (3) when served as multi-exit models, they achieve on-par performance with standard multi-exit architectures which consist of additional classifiers designed for early exiting in shallow layers. To our knowledge, our work is the first to show that one common classifier is sufficient for multi-exit models. We conduct experiments on both vision and NLP tasks to demonstrate the performance of the proposed aligned training.