Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Tracing Representation Progression: Analyzing and Enhancing Layer-Wise Similarity

Jiachen Jiang, Jinxin Zhou, Zhihui Zhu

TL;DR

The paper addresses understanding how internal representations evolve across transformer layers and how to leverage this progression for efficiency. It introduces a per-sample cosine similarity measure to quantify layer-wise similarity and establishes a geodesic-curve-based justification for monotonic similarity increase, while showing that the last-layer classifier can be effectively applied to intermediate layers. Building on this, it proposes an aligned training objective that jointly optimizes multiple layer classifications to boost representation similarity, leading to earlier saturation, higher layer-wise accuracies, and a single-classifier multi-exit framework with substantial parameter savings. The approach is validated across vision and language tasks, demonstrating faster inference and strong performance, and is shown to preserve transferability across distributions and tasks.

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.

Tracing Representation Progression: Analyzing and Enhancing Layer-Wise Similarity

TL;DR

The paper addresses understanding how internal representations evolve across transformer layers and how to leverage this progression for efficiency. It introduces a per-sample cosine similarity measure to quantify layer-wise similarity and establishes a geodesic-curve-based justification for monotonic similarity increase, while showing that the last-layer classifier can be effectively applied to intermediate layers. Building on this, it proposes an aligned training objective that jointly optimizes multiple layer classifications to boost representation similarity, leading to earlier saturation, higher layer-wise accuracies, and a single-classifier multi-exit framework with substantial parameter savings. The approach is validated across vision and language tasks, demonstrating faster inference and strong performance, and is shown to preserve transferability across distributions and tasks.

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.
Paper Structure (37 sections, 2 theorems, 31 equations, 23 figures, 1 table)

This paper contains 37 sections, 2 theorems, 31 equations, 23 figures, 1 table.

Table of Contents

  1. Introduction
  2. Contribution
  3. Measuring Layer-wise Representational Similarity
  4. Existing work on measuring representational similarity
  5. Sample-wise layer-wise representational similarity in transformers
  6. Analysis for Saturation Events
  7. Neural Collapse ($\mathcal{NC}$)
  8. Saturation events
  9. Aligned Training for Enhancing Representational Similarity
  10. Aligned Training for Enhancing Shallow Layer Performance
  11. Minimal $\epsilon$-effective depth
  12. Saturation events and multi-exit model with a single classifier
  13. Effects on Transferability
  14. Applications on Language Models
  15. Conclusion
  16. ...and 22 more sections

Key Result

Theorem 2.1

(Representation similarity increases monotonically across layers) Under assum:geodesic, for any layers $\ell_1 < \ell_2 < \ell_3$, we have $\mathop{\mathrm{COS}}\nolimits(\boldsymbol{h}_{k,i}^{(\ell_1)}, \boldsymbol{h}_{k,i}^{(\ell_3)}) < \mathop{\mathrm{COS}}\nolimits(\boldsymbol{h}_{k,i}^{(\ell_2)

Figures (23)

  • Figure 1: Illustration of the DeiT-S touvron2021training (pretrained on ImageNet) fine-tuned on CIFAR-10 with standard method and the proposed aligned training in terms of (a) cosine similarity of features from shallow and the last-hidden layer, (b) layer-wise testing accuracies by applying the last-layer classifier to each layer, as well as (c) cumulative saturation events geva2022transformer. We observe that our proposed aligned training can substantially enhance layer-wise representation similarity, thereby improving layer-wise accuracies and promoting more early saturate events.
  • Figure 2: Illustration of a DeiT-S model trained with standard training on CIFAR-10 and ImageNet in terms of (a-b) similarity of features from shallow layers and the last-hidden layer measured by CKA and COS, as well as layerwise validation accuracy, and (c-d) cosine similarities between all pairs of layers. For both datasets, cosine similarity can reflect the trend of layerwise accuracy.
  • Figure 3: The DeiT-S models are trained on CIFAR10 and ImageNet-1K datasets from scratch. We measure the number of saturation events at each layer and their average cosine similarity with the final hidden states. More saturation events in early layers indicate higher cosine similarity.
  • Figure 4: Comparison of ViT for ImageNet by standard training, proposed aligned training, and the multi-exit/classifiers, in terms of (a) cosine similarity, (b) layer-wise testing accuracy and (c-d) cosine similarities between all pairs of layers.
  • Figure 5: Comparison of standard training of 6, 9, 12-layer DeiT-small model with aligned training of 12-layer model on CIFAR-10 in terms of (a) layer-wise train accuracy, and (b) convergence.
  • ...and 18 more figures

Theorems & Definitions (3)

  • Theorem 2.1
  • Theorem 2.2
  • Definition 3.1