Tracing Representation Progression: Analyzing and Enhancing Layer-Wise Similarity
Authors: Jiachen Jiang
Jinxin Zhou
Zhihui Zhu
Venue: ICLR 2025 Poster (2025)
Keywords: Representation Similarity
Saturation Event
Early Exit
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: Study how representations propagate across layers in transformers using sample-wise, layer-wise representation similarity; propose aligned training to promote early saturation events design multi-exit models with a single classifier

Abstract: Analyzing the similarity of internal representations within and across different models 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, a property that may approximately hold for residual networks. 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. This offers a justification for {\it saturation events}, where the model's top prediction remains unchanged across subsequent layers, indicating that the shallow layer has already learned the necessary knowledge. 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.

Published as a conference paper at ICLR 2025
TRACING REPRESENTATION PROGRESSION: ANALYZ- ING AND ENHANCING LAYER-WISE SIMILARITY
Jiachen Jiang, Jinxin Zhou & Zhihui Zhu Department of Computer Science and Engineering, The Ohio State University, {jiang. 2880, mhou.3820) zhu. 3440 3440]0osu.ee .du
ABSTRACT
Analyzing the similarity of internal representations within and across different models has been an important technique for understanding the behavior of deep neural networks. Most existing methods for analyzing the similarity between rep- resentations of high dimensions, such as those based on Centered Kernel Align- ment (CKA), rely on statistical properties of the representations for 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 simple sample-wise cosine similarity metric capable capturing the similarity and aligns with the complicated CKA. Our experimental results on common transformers reveal that representations across layers are pos- itively 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, property that may approximately hold for residual networks. We then show that an increase in representation similarity implies an increase predicted probability when directly applying the last-layer classifier any hidden layer representation. This offers justification for saturation events, where the model's top prediction remains unchanged across subsequent layers, in- dicating that the shallow layer has already learned the necessary knowledge. 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 per- formance with standard multi-exit architectures which consist of additional clas- sifiers 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 mod- els. We conduct experiments on both vision and NLP tasks to demonstrate the performance of the proposed aligned training.
INTRODUCTION
As one of the most significant breakthroughs deep neural network (DNN) architectures developed
in recent years, the transformer model (Vaswani et al., 2017) has driven recent advances in various vision and NLP tasks, such as vision transformer for image classification (Dosovitskiy et al., 2020) and image generation (Yu et al., 2021; Ramesh et al., 2021; Yu et al., 2022). BERT (Devlin, 2018), GPT (Radford al., 2019), and other various large language models (LLMss) (Zhao al., 2023) for natural language understanding and generation. It has been viewed as promising foundation model that can be adapted and extended to various applications and domains (Bommasani et al., 2021). Additionally, researchers have found that increasing the size (by stacking more layers and/or making them wider) can consistently improve performance, resulting in models of significant size (e.g., the 175B 1/5B-parameter GPT-3 and the 540B -arammeter PaLM). However, the inereencgg size has posed a significant challenge in studying and understanding exactly how these models solve tasks and in efficiently deploying them.

---

Published as a conference paper at ICLR 2025
1.
Shandard Akgindd
1.0
Stancard Algnod
3300
Standard Algres
9000
1003
02
200
(a) Cosine Similarity
(b) Layerwisc Accuracy (c) Cumulative Saturation Events
Figure 1: Illustration of the DeiT-S (Touvron al., 2021) (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 (Geva et al., 2022). 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.
Given the success of deep learning models, attributed to their ability to learn increasingly complex internal representations as they go deeper through their layers, promising direction for understand- ing these models is to study the hierarchical feature learning across layers. Recent work (Papyan et al., 2020; Fang et al., 2021: Zhu et al., 2021; Thrampoulidis et al., 2022; Tirer et al., 2023) have uncovered an intriguing phenomenon regarding the last-layer features and classifier of DNNs, called Neural Collapse (NC), across many different datasets and model architectures. Roughly speaking, NC refers to training phenomenon in which the last-layer features from the same class become nearly identical, while those from different classes become maximally linearly separable. Beyond the last-layer features, recent studies have also shown that deep classifiers progressively compress within-class features while enhancing the discrimination of between-clase features from shallow deep layers (He & Su, 2023; Rangamani et al., 2023; Wang et al., 2023). Another line of work attempts to compare the representations within and across DNNs. Various approaches have been proposed to quantify the representation similarity, such as the Canonical Correlation Anal- ysis (Thompson, 2000). Centered Kernel Alignment (CKA) (Korbblih et al., 2019), Orthogonal Procrustes Transformation (Hamilton et al., 2016) and Pointwise Normalized Kernel Alignment (PNKA) (Kolling et al., 2023). Representation similarity analysis has also been widely used in computational psychology and neuroscience as well (Edelman, 1998; Kriegeskorte et al., 2008).
As these approaches are designed to be invariant to certain transformations (such as orthogonal transformation) and can be applied to features with possibly different dimensions, they rely on statistical properties of the representations for entire training data. For instance, the NCC analysis captures the variance of the features from each class, while the widely used CKA for representation analysis replies on the inter-example structures; see Section 2 for the detailed definition of CKA. Consequently, it has been observed that CKA is sensitive to outliers and may give unexpected or counter-intuitiv results in certain situations (Davari et al., 2022). In this paper, are motivated by the following question: How are the representations of individual inputs progressively transformed from shallow deep layers?
Contribution In this work, we focus on transformer models, which have particular properties compared to other architectures: transformer contains identical blocks with residual connections in each block. Thus, the features in transformer model have the same dimension and may ex- hibit less rotation difference across layers. Motivated by this observation, we study the layer-wise representation similarity for transformer models a per-sample basis, allowing to directly ap- ply the last-layer classifier right after any hidden layer for classification or text generation tasks. This enables effortless multi-exit model that allows early exit during inference, thereby saving computation time. Our contributions can be summarized as follows.
Sample-wise cosine similarity captures representational similarity We introduce a straight- forward yet efficient sample-wise cosine similarity metric to examine the similarity of internal representations in transformers. Experiments show that the cosine similarity aligns with CKA, which is based on statistical properties of all the features, and is sufficient to reflect representa- tion similarity. In addition, as illustrated in Figure I (a), our experimental results (Standard) on common transformers show that representation similarity increases layers become closer. We
2

---

Published as a conference paper at ICLR 2025
provide theoretical justification for this phenomenon under the geodesic curve assumption for the learned transformer, a property that approximately holds for residual networks (Gai & Zhang, 2021), and hence for transformers.
Analysis for saturation events The similarity with last-layer representation suggests that the last-layer classifier can be directly applied right after any hidden layers for decision-making also known as the logit lens approach. Geva et al. (2022) discover a phenomenon called saturation events, where the model's final prediction becomes the top candidate in certain shallow layer and remains unchanged across all subsequent layers, indicating that the shallow layer has already learned the necessary knowledge. We show that an increase in representation similarity implies an increase in predicted probability across layers. This offers a justification for saturation events, stating that if a sample is correctly predicted the ě-th layer, will continue to be correctly predicted in subsequent layers, as the predicted probability increases progressively across layers. Aligned training for enhancing layer-wise representation similarity We propose an aligned training method to improve the effectiveness of shallow layers by increasing the similarity of internal representations between different layers. Motivated by the ACC phenomenon, where fea- tures from the last-hidden layers align with the common classifier, our aligned training approach deploys the common classifier to each layer and then minimizes the average of the cross-entropy losses from all the layers. As shown in Figure the aligned training can substantially enhance layer-wise representation similarity, thereby leading to more early saturation events and improving layer-wise accuracies. Consequently, the aligned training method can help identify the minimal number of layers needed by unleashing the power of shallow layers to transform features faster towards classifier across layers and push the redundancy behind.
Multi-exit models with single classifier Another important application of the proposed aligned training is to improve the inference efficiency of large models. During inference, the model al- lows early exit to save computation time. Previous works (Xin et al., 2020; Geng et al., 2021; Xin et al., 2021) design multi-exit models by introducing different classifiers to each layer, which may substantially increase the model size for a large number of classes, such ImageNet with 1000 classes and GPT3 (Brown et al., 2020) with a vocabulary of 50, 257 tokens. Instead of using separate classifiers for each exit, our multi-exit model employs common classifier, which to our knowledge is the first of its kind, maintaining the early exit capability and achieving performance on par with models that use multiple classifiers. We demonstrate the performance of the proposed aligned training method in both pretraining ViT and fine-tuning LLMs in NLP tasks, including fine-tuning BERT (Devlin, 2018) for text classification tasks on the General Language Under- standing Evaluation (GLUE) benchmark (Wang et al., 2018), and GPT2 (Radford et al., 2019) model for text generation on the Wikitext-103 dataset (Merity et al., 2016).
The MatFormer (Kudugunta al., 2023) introduces a nested structure into the Transformer by jointly training all submodels of different widths. In contrast, our method jointly trains submodels with varying layers. Exploring the potential integration of these approaches will be a focus of future research. While we mainly focus on transformer models, the proposed aligned training can be applied to other deep architectures, provided they have the same dimensions in each layer. Extending this approach to accommodate varied feature dimensions is the subject of future work.
2 MEASURING LAYER- WISE REPRESENTATIONAL SIMILARITY
In a transformer ff with L layers, the representations gradually evolve across layers, with the pro- gression from one layer (say f-th layer) to the next following residual update patterm:
(e+1) = (9) + focer (H (9).
(I)
where (!) E Rdx are input feature sequence of length s with hidden dimension of d. The residual block Loco Raxx + Rdxx describe the sequence-to-sequeece function mapping with parameters ((0) that mainly comprises two complementary stages of data transformation: the Multi- head Self-Attention across tokens and the MultiLayer Perceptron layer across features. Predictions are typically based on specific token of last layer representation H (T.) For instance, ViT uses the representation of class token [CLS] classify the image, while rego-regressive based lan- guage model (such as GPT) uses the representation of the previous tokens to predict the next word.
Consequently, we will focus on the feature (or representation) of this particular token in (() de- noted by (6) E IRd at the l-th layer. A linear classifier g(.) is applied to the last layer feature
3

---

Published as a conference paper at ICLR 2025
h (4) to make predications as' g(h( = arg max )h where l-ly denotes the j-th entry and W E R maps the dimensional features to K dimensional logits. We may also directly apply the last layer classifier g(.) to the hidden layer features h to make predictions via g(h!*) arg max, MSoftaxxa (Wh ) Given data samples S := {œR,,), where æ,, represents the i-th sample of class k with i E [n] = {l....,, n and & E [K], we define layer-wise accuracy as
(2) Ace's
( = k).. Lg (g
Kn
Existing work on measuring representational similarity Similarity analysis is widely applied in the literature, including research on learning dynamics (Morcos et al., 2018; Mehrer et al., 2018), effects of width and depth (Nguyen et al., 2020), differences between supervised and unsupervised models (Gwilliam & Shrivastava, 2022), robustness (Jones al., 2022; Nanda al., 2022), eval- uating knowledge distillation (Stanton et al., 2021), language representation (Kudugunta et al., 2019; Shi et al., 2022), and generalizability (McCoy al., 2019; Lee et al., 2022; Pagliardini et al., 2022). To enable measuring the similarity of features from from different architectures or layers that have different dimension, most existing methods for analyzing the similarity between representations of high dimensions, such as those based on Canonical Correlation Analysis (CCA) and widely used Centered Kernel Alignment (CKA) (Komblith et al., 2019) rely on statistical properties of the representations for set of data points. For instance, given input feature sequence z'
h
h K E _axNN N denoting the features of N = Kn training samples, the widely-
used CKA with a linear kernel quantifies similarities between features Z' and z"
CKA TTT=T z" (Z') z / (Z))
(z (1
(3)
CKA relies the similarity of inter-example structures since the gram matrix (Z') E RR- captures the pair-wise similarity of different samples, focusing on the consistency of relative po- sitions among features. Consequently, the CKA is invariant orthogonal transformations and isotropic scaling, and can be applied for the case where (2) and h (e") have different dimensions.
2.1 SAMPLE- WISE LAYER-WIL WISE REPRESENTATIONAL SIMILARITY IN TRANSFORMERS
In this work, we specifically focus on transformer architectures that obey the following particular properties of features across layers: (i) the features have the same dimension across layers since transformers are generally constructed by stacking identical blocks; (ii) the features may have no or less rotation ambiguity due to the residual connection (I). Based on these observations, for each sample Tk we propose to simply measure the cosine similarity of the corresponding feature vectors h (f) and h (f') layers and E' as?
COS 1 ( /
The above COsine Similarity (COS) measures the angle between feature vectors, providing clear geometric interpretation of feature alignment and similarity at the layer level. Unlike CKA (3), COS is not invariant to all transformations except for isotropic scaling. Furthermore, COS is computed for each individual sample and does not rely on inter-example structures. In the experiments, we compute the average COS over all the training samples.
To verify whether the proposed sample-wise COS is a good indicator of similarity structure within transformers, we train the DeiT-S model (Touvron et al., 2021) (a data-efficient vision transform) from scratch on both the CIFAR-10 and ImageNet- datasets. The feature dimension is set to 384 for both tasks across all layers. In Figure 2(a, b), we compute CKA and average cosine similarity between the features in each layer and the last layer. Additionally, we plot average COS between all pairs of layers and display the results heatmap in Figure 2(c, d). Based these results, make several observations.
'There is also a bias term in classification layer, but we omit it for simplicity of presentation. "hhe features in each layer are centered by reducing the global mean of all the samples.
4

---

Published as a conference paper at ICLR 2025
Myer
03
07
(a) CIFAR1O
(b) ImageNet1
(c) CIFARIO
(d) ImageNet|
Figure 2: Illustration of 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.
Observation 1: Simple-wise COS is sufficient and reflects CKA to measure layer-wise repre- sentation similarity We observe from Figure 2(a, b) that COS aligns with CKA and is sufficient to reflect the layer-wise representation similarity in transformer models. In other words, since CKA is invariant to orthogonal transformations while COS is not, there is little or no rotation difference among the features from different layers. This is attributed to the residual connections in (1), which create smoother transitions between layers and potentially lead to more stable feature representa- tions throughout the network. In the Appendix B.2, we design additional experiments on multi- layer perceptrons (MLPs) with and without skip connections to verify the effect of skip connections in climinating rotation ambiguity.
Progressively increasing layer-wise representation similarity. For models trained on small datasets, we observe ridge-to-plateal pattern in the heatmap plot in Figure 2(c): the initial layers undergo significant transformations, suggesting that lower layers rapidly refine the features to extract the most relevant information for classification tasks; in contrast, the higher layers exhibit a plateau in similarity scores, indicating that feature transformations stabilize and converge toward an optimal representation. This plateau also suggests redundancy in the higher layers, imply- ing that removing these redundant layers could improve efficiency without significantly sacrificing performance," 3 On the other hand, for models trained on large datasets, Figure 2(d) shows a con- sistent ridge pattern, characterized by rapid decay in similarity between adjacent layers. This indicates dynamic and continuous refinement process throughout the network. Nevertheless, for both Figure 2(c, d). we observe almost all nonnegative average cosine similarity between different layers, albeit the features are almost orthogonal when the layers are far apart. Moreover, in both Figure 2(c, d), we observe a progressive increase in representation similarity as the two layers get closer; across each row column, the cosine similarity increases as it approaches the diagonals. In appendix D, we observe similar phenomena on molti-modalityy models (CLIP).
Observation 2:
To understand this phenomenon, utilize the connection between residual network (ResNet) and dynamic system, viewing the residual update (1) as a discretization of a dynamic system (Haber & Ruthotto, 2017: Gai & Zhang, 2021). Specifically, Gai & Zhang (2021) proved that ResNet trained with weight decay attempts to learn the geodesic curve in the Wasserstein space. Since transformer is ResNet and weight decay is commonly applied in real-world training for example, DeiT' (Touvron et al., 2021) models are trained using the AdamW optimizer with a weight decay of 0.05-we use the following geodesic curve assumption (Wang et al., 2024).
Assumption 1. (Geodesic curve assumption) At the terminal phase of training, the transformer with weight decay has learned the geodesic curve in Wasserstein space P(TTe)), which is induced by the optimal transport map.
Based on this assumption, the following result shows monotonic increase in representation simi- larity as the layers get closer. Proofs are given in the Appendix A.
Theorem 1. (Representation similarity increases monotonically across layers) Under Assumption 1, for any layers l < lz < Eg, we have COS(h (f1) (fa)) 1g. < COS(h (t2) (fs)).
'e notice concurrent work (Men et al., 2024; Gromov et al., 2024: Jaiswal et al., 2024) that exploits representation similarity across layers for pruning redundant layers. For instance, Figure 2(c, d) show hidden layers ohey large similarity, indicating that some of the layers he skipped (Jaiswal al., 2024).
5

---

Published as a conference paper at ICLR 2025
Theorem 1 provides theoretical justification for the phenomenon observed in Figure 2 for in- stance, the similarity to the last layer features progressively increases from shallow to deep layers. However, note that in practice, the geodesic curve assumption may not hold precisely by prac- tical network, so some samples may not exhibit a strictly monotonic increase.
2.2 ANALYSIS FOR SATURATION EVENTS
We now use the layer-wise representation similarity to analyze the phenomenon of saturation events (Geva et al., 2022). To that goal, we first briefly introduce the neural collapse (AC) phe- nomenon (Papyan al., 2020) and its connection to layer-wise representation similarity.
Neural Collapse (NC) Roughly speaking, NCC concerns the terminal phase of training deep net- works and states that (i) within-class variable collapse (ACi)): i): the last-layer features from the same (L)
class become nearly identical, i.e., hgg h : E, kk,i (it) maximal distance (AC2)): those from different classes become maximally linearly separable, and (iii) self-duality (ACg):: the last- n
layer linear classifiers ww align with the class- features (L) To achieve this, deep classifiers progressively compress within-class features while enhancing the discrimination of between-class features from shallow to deep layers (He & Su, 2023; Rangamani et al., 2023; Wang et al., 2023). Our results show that transformer models achieve progressive compression and separation by pro- gressively aligning the features with the last-layer classifier from shallow to deep layers.
Aligument between layer-wise cosine similarity and accuracy. Motivated by the similarity be- tween features in shallow and deep layers, we apply the classifier to each hidden layer to obtain the layer-wise validation accuracy, which is plotted in Figure 2(a, b). We observe high correlation be- tween the layer-wise accuracy and the cosine similarity, with layer-wise accuracy also exhibiting a monotonic increase across layers. Our following result provides a justification for this phenomenon, demonstrating that an increase in representation similarity implies an increase in predicted probabil- ity across layers.
Theorem 2. (Predicted probability increases monotonically across layers) Assume that Theorem holds at ha i.e., COS(h (E+1) .h(.) > COS(N (L)) and that the last-layer features and classifers satisfy NC. Then, the predicted probability [SoftMax (W (d)l increases across layers:
[SoftMax Wh h 2,, )( [SoftMax(WW)
(4)
Saturation events The approach of applying last-layer classification to intermediate representation is also called logit lens (nostalgebraist). Recent studies (Belrose et al., 2023; Pal et al., 2023) use this method to decode hid- den states into probability distributions over the vocab- ulary, offering mechanistic interpretability of transform- ers. (Geva et al., 2022) discover phenomenon called saturation events, where the model's final predicted token becomes the top candidate in a certain shallow layer and remains unchanged across all subsequent layers. Specifi- cally, given an input sample æ the saturation layer Cx for a is defined as the smallest layer l such that
0.5
Saturate ECPARRR
COS (CRVAIZ1
DOS (NDO)
6.8
Kort of Sar
5.0
0.0
10
12
Figure 3: The DeiT-S models are trained on CIFARIO and Imagenet- K dataset from scratch. We measure the number of saturate event each layer and their average cosine similarity with last hidden states. More saturate events early layer indicates higher cosine similarity.
( # ... # gh ( =... (.)
Results in Appendix C. show saturate events also happen on recently developed LLMs such as LLaMA3 (Dubey et al., 2024). Our Theorem 2 offers a justification for satu- ration events, stating that if a sample is correctly predicted at the f-th layer, it will continue to be correctly predicted in subsequent layers, as the predicted probability increases progressively across layers. To further illustrate the relation between representation similarity and saturate events, we train a DeiT-S model on both CIFAR-10 10 and ImageNet- IK and plot the results in Figure 3. For the same model across different datasets, we observed that higher layer-wise representation similarity COS correlates with more early saturation events, suggesting COS is valuable metric for reflecting saturation events.

---

Published as a conference paper at ICLR 2025
Table 1: Comparison of the number of parameters across different architectures between multi-exit models with multiple classifiers (different classifiers at each layer) and a single classifier (ours). The last column shows that the percentage of saved parameters using single classifier
Models
Hidden Dim
of Classes
A Layers
Multiple Classifiers (#Params)
Single Classiber (#Paramss)
Faarm Saving
DefT S/Touvron 2021)
384
1,000
12
26.27M
22. 22.05M
16.07%
DeTT-B/
2021)
768
1,000
95.02M
86.57M
GPI-2(kadford al., 2019)
768
50,257
541.57M
117.35M
78,399
GPr-a al., 2020)
12,288
50,257
96
233,67B
175. 63.663
25.10%
LLAMA-2 (Touvron et al., 2023)
4,096
32,000
40
75.11B
70.35B
6.81%
3 ALIGNED TRAINING FOR ENHANCING REPRESENTATIONAL SIMILARITY
In the previous section, we observed that representations across layers within transformer models are positively correlated, resulting in saturation events when the last-layer classifier is directly ap- plied after any hidden layer for early prediction and enabling multi-exit model that shares single classifier. In this section, we propose an aligned training method to enhance the effectiveness of shallow layers by improving layer-wise feature similarity. This, in turn, promotes more early sat- uration events, determines the minimal effective depth, and enhances performance when used as multi-exit model. To the best of our knowledge, our work is the first to show that one common clas- sifier is sufficient for multi-exit models. Table 1 shows single classifier can significantly reduce the number of parameters and the computational complexity for multi-exit models, particularly for tasks with large number of classes and large feature dimensions. Examples include ImageNet, with 1000 classes, and LLMs, where the number of classes equals the vocabulary size, i.c., the number of all possible tokens for instance, the Llama-2 (Touvron et al., 2023) model has a vocabulary of 32, 000 tokens while the GPT3 (Brown et al., 2020) has 50, 257 tokens.
3.1 ALIGNED TRAINING FOR ENHANNING SHALLOW LAYER PERFORMANCE
The ability to capture the layer-wise similarity of representations for each sample enables us to develop efficient methods for enhancing this similarity during training. A first approach is to directly add the cosine similarity between (f) and h for all l I regularization term during the training process. However, as shown in the appendix (see Figure 14), this approach can only slightly improve layer-wise similarity and accuracy. We conjecture this is due to the imbalance between the cross-entropy loss and the cosine similarity. Instead, motivated by the self-duality between the class-mean features and the linear classifiers, as observed in the NC phenomenon, we propose a simple yet efficient method, named aligned training, to enhance the layer-wise wise similarity by jointly optimizing the following aligned loss that is the weighted average of the CE loss from all the layers
Caligned (z, 4) = Σ XcCce(Wh!! .y)
(S)
where y denotes the corresponding label for æ, d 0 is the weight for the 6-th layer. During the experiments, considering that shallow layers tend to have larger losses compared to deep layers, set the weight linearly increase with layers put more emphasis on the deeper layers", i.e.,
A = =2E/(I(T. + 1)). Roughly speaking, the aligned loss (5) introduces CE loss for intermediate layers and would encourage each layer features (2) to align with the common classifier W as implied by the ACC meenomenon-henee improving the representation similarity across layers.
Figure 4 displays the layer-wise representation similarity and accuracy by the proposed aligned training. We can observe that aligned training can significantly increase the layer-wise repre- sentation similarity and accuracy by aligning all the features the common classifier. Results in Appendix B.3 also show that aligned training can enhance progressive separation and compression from shallow to deep layers. To further illustrate the benefit of aligned training, define the notion of c-effective depth that modifies the notion exploited in Galanti et al. (2022) by replacing nearest neighbor classifier accuracy with our layer-wise accuracy in (2).
'Such a strategy of increasing weights is also employed in (Schuster et al., 2022). Uniformly weighting all layers (i.c., X₂ = 1/L.) may diminish the importance of the deeper layers. We provide an ablation study for the comparison of linear increasing weights and uniform weights in Appendix B.3.
7

---

Published as a conference paper at ICLR 2025
Ayen
1m0
byer
kep-
lews
tep-
Dot Lever
(a) Cos Similarity
(b) Layerwise Acc
(c) Aligned
(d) Multi-exit/classiiies
Figure 4: Comparison of ViT for ImageNet by standard training, proposed aligned training, and the multi-exii/classfffess,, in terms of (a) cosine similarity, (b) layer-wise testing accuracy and (c-d) cosine similarities between all pairs layers.
0.2
Loyes Stantad
Standars
Layers
Lapers Standars
02
Layars Standard
300
(a) Train Accuracy
(b) Training Curve
Figure 5: Comparison of standard training of 9, 12-layer DeiT-small model with aligned training of 12-layer model on CIFAR-16 10 in terms of (a) layer-wise train accuracy, and (b) convergence.
Definition 1. (c-effective depth). We define the c-effective depth d' f) of an I--layer transformer over dataset as the minimal layer such that Ace' (t) 2 c. Set de (S, I. if such
non-existent.
Minimal e-effective depth Denote the transformers learned by standard training and aligned train- ing as [standard and faligned, respectively. We observe that the aligned training yields model with a much smaller e-effective depth compared to standard training, i.e., d' (S, falignes) < d' dS, fsuadard). Specifically, Figure 5(a) displays the layer-wise training accuracy of three transformers with differ- ent number of layers L E [6, 9, 12) with standard training and 12-layer transformer with aligned training. For the standard training models, the train accuracy curves increase until the last layer with- out plateauing, even for the model with 12 layers, giving d' (SCIfARIO: fstandar)) = 12. While aligned training models unleash the power of shallow layers to transform features faster towards classifier across layers and push the redundancy behind, giving d' (SCIfARIO, Satigned) than e-effective depth of fstandard. On the other hand, the effective depth d' (S, faligned) increases with the complexity of the task, demonstrated by comparing the results from CIFARIO (Figure 5(a)) and ImageNet (Figure 4(b)). Thus, the effective depth, independent of the network's depth, can be leveraged to derive generalization bounds. This can be achieved by applying the approach from Galanti et al. (2022), which offers non-trivial estimates of generalization based on effective depth.
which is smaller
Models faligned with small effective depth also offer several advantages practical deployment. First, the layers beyond effective depth d' (ScIfArIo, (nlignea) are redundant, as they do not contribute to accuracy improvements and can be pruned, leading more efficient inference. Second, aligned training helps determine the minimal number of layers required for task. While more complex tasks typically demand more layers, identifying the exact number can be challenging. In standard training, multiple models of different sizes must be trained to determine the minimal layer count. In contrast, with aligned models, retraining is unnecessary -on can simply apply the last-layer linear classifier to intermediate layers. Third, models trained using aligned method only achieve slightly higher accuracy when truncated to 6 or 9 layers compared to models of the same size trained with standard methods (Figure 5(a)), but they also accelerate model convergence (Figure 5(b)) by providing immediate feedback to each layer, resulting in more effective parameter adjustments.
8

---

Published as a conference paper at ICLR 2025
120.0
1000
130.0
70030
20'
60
(a) DeiT on CIFARIO
(b) DeiT on ImageNet
(c) GPT2 on WikiText103
Figure 6: Illustration of the effect of aligned training versus standard training on cosine similarity and the number of saturation events across vision (a.b) and NIP (c) models. Aligned training leads to more early saturation events by increasing the cosine similarity with the last hidden states.
Saturation events and multi-exit model with a single classifier Another important application of the proposed aligned training is to improve the inference efficiency of large models. Figure 6 illustrates saturation events in both vision and NLP models (see the AlignedGPT setup in the ap- pendix). We observe that aligned training encourages earlier saturation events by increasing the cosine similarity with the final hidden states. This demonstrates that faligned has greater potential for supporting early exits, commonly referred to multi-exits in the literature (Xin et al., 2020; Geng et al., 2021; Xin et al., 2021). In previous work, multi-exit models typically use separate classifiers for each layer, which issue becomes more pronounced when dealing with a large number of classes, as dense linear classi- fiers require a substantial number of parameters. To our knowledge, by exploiting the representation similarity within the transformer, our proposed multi- exit model is the first to use single classifier for all the layers. In addition, when training multi-exit models, previous work (Geng et al., 2021) often uses additional KI. Ki--eieeeeecee terms to guide the logits of shallow layers by those of deep lay- ers. In contrast, by using a common classifier to align shallow representations with deep ones, our aligned training does not require KL-divergence or other such terms. For comparison, we implement the multi-exit training with multiple classifiers (Xin et al., 2021), denoted by "multi-exitCllassifeess and display the results in Figure 4. On one hand, we observe that the proposed aligned training with a shared classifeir achieves higher layer-wise wise representation similarity than liitti-xit/cllssiiiers.. On the other hand, the aligned training exhibits on-par performance as muuti-cxit/classiiiee in terms of layer-wise accuracy, which is remarkable as the former only uses a single classifier.
significantly increase the overall model size. As shown in Table this
To further show the performance of the proposed multi-exit mod- els with a single classifier, we allow exit at shallow layers if the confidence level (max of softmax logits) exceeds set thresh- old for each sample. On ImageNet dataset, Figure 7 displays the number of samples that exit at each layer. We observe that most samples exit at the last layer for standard training, while most exit early layers for aligned training. We then calculate the classification accuracy along with the ratio of speed improve- > Lxm"
Standars
100000
8000
2000
ment measured by Σi-1 ixmi where m² is the number of sam-
Figure 7: Number of samples exit at each layer.
ples that exit at the i-th layer of DeiT. The model trained with standard training achieves 80.28% accuracy, while the model trained with aligned training achieves 77.96% accuracy with 1.36x speedup, which compa- rable to those trained by multi-exit/Classifiers with 78.32 % accuracy and 1.42x speedup.
Effects on Transferability Concerns arise about whether aligning shallow layer features with deep layer features could diminish the transferability of shallow layers. To answer this question, we conduct experiments on (1) distribution shift and (2) task transfer. The results show that aligned training improves layer-wise accuracy for both pre-trained and downstream datasets while maintain- ing transferability. This indicates that the trained model can be effectively transferred without losing the universal patterns learned in shallow layers. Further details are provided in Appendix B.4,
3.2 APPLICATIONS ON LANGUAGE MODELS
We extend our aligned training approach to NLP tasks, demonstrating its effectiveness in fine-tuning I Language Models. For text classification tasks, we get AlignedBERT by finetuning a pretrained 12-

---

Published as a conference paper at ICLR 2025
40
10
16
Sunced
32
Sandeo
(a) SST-2
(b) MRPC
(c) QNLI
(d) MNLI
Figure 8: Layerwise accuracy for AliT and BERT and on SST-2, MRPC. QNLI and MNLI
datasets of the GLUE benchmark. (See more results on RTE and QQP in Figure 21).
Bandard
201
64
Lawis
LAyES
Lewe
(a) Prediction Accuracy
(b) Perplexity
(c) Coherence
(d) Diversity
Figure 9: Evaluation AlignedGPT and GPT2 on Wikitext-103 dataset in terms of (a) prediction
accuracy, (b) perplexity, (c) coherence, and (d) diversity. See Appendix C.2 for detailed definitions.
layer BERT Base model (Devlin, 2018) using aligned training method GLUE benchmark (Wang et al., 2018) tasks. This includes single-sented tSsks-SST-2, similarity and paraphrasing tasks- MRPC and QQP. as well natural language inference tasks-MNLI, QNLI and RTE. For compar- ison, we also finetune a BERT model using standard training. Then we evaluate the layerwise accuracy(or F-1 score) of both finetuned models using their last layer classifier. Figure 8 shows that AlignedBERT achieves better performance than the standard BERT across all layers.
For text generation task, we get AlignedGPT by finetuning pretrained 12-layer GPT2 model (Rad- ford et al., 2019) using aligned training method on Wikitext- 103 dataset (Merity et al., 2016). For comparison, we also finetune GPT2 model using standard training. Then we evaluate the two fine- tuned models from two perspectives(followw modeling quality, which assesses the intrinsic quality the model and measured by prediction accuracy and perplexity, and (2) generation quality, which measures the quality of the text produced by the model using coherence and diversity. Coherence is a measurement of relevance between prefix text and generated text, while diversity considers the recurrence of generation at varying n- gram levels. All evaluation metrics mentioned above can be found in Appendix C.2. Results show that AlignedGPT outperforms the standard one in prediction accuracy and exhibits lower perplexity across intermediate lFigureF (a,b)). Moreover, A AlignedGPT can also generate text with higher coherence and diversity using shallower layers(Figure (c,d)), which improves the the inference efficiency. More experimental setup and results can be found in Appendix C.2.
(Su et al., 2022; Su & Collier, 2023)), (1) language
4 CONCLUSION
In this paper, we demonstrate that a simple sample-wise wise cosine similarity metric effectively captures layer-wise representation similarity in transformer models, aligning with the more complex CKA metric. Our findings reveal that representations become more similar as layers get closer and show that increased representation similarity correlates with higher prediction accuracy, leading to satura- tion events where shallow layers can already make correct predictions. enhance this, we proposed an aligned training method that improves shallow layer effectiveness, resulting in more early satura- tion events and much higher layer-wish wise accuracies. Remarkably, when served as multi-exit models with common classifier, which to our knowledge is the first of its kind, they maintain the early exit capability and achieve performance on par with models that use multiple classifiers. Experiments on both vision and NL.P tasks demonstrate the performance of the proposed aligned training.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
We acknowledge support from NSF grants CCF-2240708, IIS-2312840, and IIS IIS-2402952, as well the ORAU Ralph E. Powe Junior Faculty Enhancement Award. We are grateful to Beidi Chen (Carnegie Mellon University) and Chong You (Google Research) for many valuable discussions and for helpful comments the manuscript.
REFERENCES
Nora Belrose, Zach Furman, Logan Smith, Danny Halawi, Igor Ostrovsky, Lev McKinney, Stella Biderman, and Jacob Steinhardt. Eliciting latent predictions from transformers with the tuned lens. arXiv preprint arXiv:2303.08122 2023.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportu- nities and risks of foundation models. arXiv preprint (0x.i::208..022: 2021.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve. Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European conference computer vision, pp. 213-229. Springer, 2020.
MohammadReza Davari, Stefan Horoi, Amine Natik, Guillaume Lajoie, Guy Wolf, and Fu- gene Belilovsky. Reliability of cka similarity measure in deep learning. arXiv preprint 1:Xiv:2210.1615:: 2022.
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.0480) 2018
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 11929, 2020.
Abhimanyu Dubey, Abhinav Jaubri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024.
Shimon Edelman. Representation is representation of similarities. Behavioral and brain sciences, 21(4):449 467, 1998.
Cong Fang, Hangfeng He, Qi Long, and Weijie J Su. Exploring deep neural networks via layer- peeled model: Minority collapse imbalanced training. Proceedings of the National Academy of Sciences, 118(43)xe21030111 2021.
Kuo Gai and Shihua Zhang. A mathematical principle of deep learning: Learn the geodesic curve in the wasserstein space. arXiv preprint arXiv:2102.092355 2021.
Tomer Galanti, Liane Galanti, and Ido Ben-Shaul, On the implicit bias towards minimal depth of deep neural networks. arXiv preprini arXiv: 2022.
Tianyu Gao, Xingcheng Yao, and Danqi Chen. Simcse: Simple contrastive learning of sentence embeddings. arXiv preprint arXiv:2104.0888 08821, 2021.
Shijie Geng, Peng Gao, Zuohui Fu, and Yongfeng Zhang. Romebert: Robust training multi-exit bert. arXiv preprint 10X::2111.0075: 2021.
Mor Geva, Avi Caciularu, Kevin Wang. and Yoav Goldberg. Transformer ford-forward layers build predictions by promoting concepts in the vocabulary space. In Proceedings of the 2022 Confer- ence on Empirical Methods in Natural Language Processing, pp. 30-45, 2022.
11

---

Published as a conference paper at ICLR 2025
Andrey Gromov, Kushal Tirumala, Hassan Shapourian, Paolo Glorioso, and Daniel A Roberts. The unreasonable ineffectiveness of the deeper layers. arXiv preprint arXiv:2403.17887, 2024.
Matthew Gwilliam and Abhinav Shrivastava. Beyond supervised vs. unsupervised: Representative benchmarking and analysis of image representation learning- Pruceedings of the IEEE/CVF Conference Computer Vision and Puttern Recognition, pp. 9642-9652, 2022.
Eldad Haber and Lars Ruthotto. Stable architectures for deep neural networks. Inverse problems, 34(1):014004, 2017.
William L Hamilton, Jure Leskovec, and Dan Jurafsky. Diachronic word embeddings reveal statis- tical laws of semantic change. arXiv preprint arXiv:160S.000% 2016.
Hangfeng He and Weijie J Su. A law of data separation in deep learning. Proceedings of the National Academy of Sciences, 120(3)):e222114411 2023.
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. arXiv preprint AXX4:199400075// 2019.
Ajay Jaiswal, Bodun Hu, Lu Yin, Yeonju Ro, Shiwei Liu, Tianlong Chen, and Aditya Akella. Fin- skipllm: A hidden gem for autoregressive decoding with adaptive feed forward skipping. arXiv 2024.
preprint arXiv:2404.0033000
Haydn T Jones, Jacob M Springer, Garrett T Kenyon, and Juston S Moore. If you' 'e trained one you' ve trained them all: inter-architecture similarity increases with robustness. Uncertainty in Artificial Intelligence, pp. 928-937. PMLR, 2022.
Camila Kolling, Till Speicher, Vedant Nanda, Mariya Toneva, and Krishna P Gummadi. Pointwise representational similarity. arXiv preprint arXiv:2305. 19294, 2023.
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In International conference on machine learning, pp. 3519- 3529. PMLR, 2019.
Nikolaus Kriegeskorte, Maricke Mur, and Peter A Bandettini. Representational similarity analysis- connecting the branches of systems neuroscience. Frontiers in systems neuroscience, 2:249, 2008
Sneha Kudugunta, Aditya Kusupati, Tim Dettmers. Kaifeng Chen, Inderjit Dhillon, Yulia Tsvetkov, Hannaneh Hajishirzi, Sham Kakade, Ali Farhadi, Prateek Jain, et al. Matformer: Nested trans- former for elastic inference. arXiv preprint 100.022310.07707, 2023.
Sneha Reddy Kudugunta, Ankur Bapna, Isaac Caswell, Naveen Arivazhagan, and Orhan Firat. In- vestigating multilingual nmt representations scale. arXiv preprint arXiv:190.022197, 2019.
Yoonho Lee, Huaxiu Yao, and Chelsea Finn. Diversify and disambiguate: -uttof-distributton ro- bustness via disagreement. In The Eleventh International Conference Learning Representa- fions, 2022.
Victor Weixin Liang, Yuhui Zhang, Yongchan Kwon, Serena Yeung, and James Y Zou. Mind the gap: Understanding the modality gap in multi-modal contrastive representation learning. Ad- vances in Neural Information Processing Systems, 35:17612- 2022.
Shangyun Lu, Bradley Nott, Aaron Olson, Alberto Todeschini, Hossein Vahabi, Yair Carmon, and Ludwig Schmidt. Harder or different? closer look at distribution shift in dataset reproduction. In ICML. Workshop on Uncertainty and Robustness Deep Learning, volume 5, pp. 15, 2020.
Thomas McCoy, Junghyun Min, and Tal Linzen. Berts of a feather do not generalize together: Large variability in generalization across models with similar test performance. arXiv preprint arXhv:1911.02969 2019.
Johannes Mehrer, Nikolaus Kriegeskorte, and Tim Kietzmann. Beware of the beginnings: inter- mediate and higherlevel representations in deep neural networks are strongly affected by weight initialization. In Conference Cognitive Computational Neuroscience, 2018.
12

---

Published as a conference paper at ICLR 2025
Xin Men, Mingyu Xu, Qingyu Zhang, Bingning Wang, Hongyu Lin, Yaojie Lu, Xianpei Han, and Weipeng Chen. Shortgpt: Layers in large language models are more redundant than you expect. arXiv preprint arXh::2403.3883, 2024.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture 2016.
models. arXiv preprint
Ari Morcos, Maithra Raghu, and Samy Bengio. Insights on representational similarity in neural networks with canonical correlation. Advances in neural information processing systems, 31, 2018.
Vedant Nanda, Till Speicher, Camila Kolling, John P Dickerson, Krishna Gummadi, and Adrian Weller. Measuring representational robustness of neural networks through shared invariances. In International Conference on Machine Learning, pp. 16368-16382. PMLR, 2022.
Thao Nguyen, Maithra Raghu, and Simon Kornblith. Do wide and deep networks learn the same things? uncovering how neural network representations vary with width and depth. arXiv preprint arXiv:2010.15327, 2020.
nostalgebraist. interpreting gpt: the logit lens. https://www. http://wwwwwwwwssss/eeoee com/posts: hwg.BB/wwpaaiorttrr/ieeerrrraaae -ht--tt--ttt--tttte lens.
Matteo Pagliardini, Martin Jaggi, François Fleuret, and Sai Praneeth Karimireddy, Agree to dis- agree: Diversity through disagreement for better transferability. arXiv preprint 20111:2224: 2022.
Koyena Pal, Jiuding Sun, Andrew Yuan, Byron C Wallace, and David Bau, Future lens: Anticipating subsequent tokens from single hidden state. arXiv preprint arXiv:2311.04897. 2023.
Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40): 24652-246633 2020.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference machine learning, pp- 8748-8763 8763. PMLR, 2021.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot toxt-to-image generation. In International conference on machine learning, pp. 8821-8831 Pmir, 2021.
Akshay Rangamani, Marius Lindegaard, Tomer Galanti, and Tomaso A Poggio. Feature learning in deep classifiers through intermediate neural collapse. In International Conference on Machine Learning, pp. 28729-28745. PMLR, 2023.
Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghani, Dara Bahri, Vinh Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. Advances in Neural Information Processing Systems,
2022.
Han Shi, Jiahui Gao, Hang Xu, Xiaodan Liang, Zhenguo Li, Lingpeng Kong, Stephen Lee, and James Kwok. Revisiting over-smoothing bert from the perspective of graph. arXiv preprint arXiv:2202.08625, 2022.
Samuel Stanton, Pavel Izmailov, Polina Kirichenko, Alexander A Alemi, and Andrew G Wilson Does knowledge distillation really work? Advances in Neural Information Processing Systems, 34:6906-6919 2021.
Yixuan Su and Nigel Collier. Contrastive search is what you need for neural text genera- tion. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https: / // //opereeview..oom net. for wtw.rmmmdd=ckkrr3333 fwI.
13

---

Published as a conference paper at ICLR 2025
Yixuan Su, Tian Lan, Yan Wang, Dani Yogatama, Lingpeng Kong, and Nigel Collier. A con- trastive framework for neural text generation. In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https: //pen ew net/ forum? d= 88a f.
Bruce Thompson. Canonical correlation analysis. 2000.
Christos Thrampoulidis, Ganesh Ramachandra Kini, Vala Vakilian, and Tina Behnia. Imbalance trouble: Revisiting neural-collapse geometry. Advances in Neural Information Processing Sys- tems, 35:27225-27233 2022.
Tom Tirer, Haoxiang Huang, and Jonathan Niles-Weed. Perturbation analysis of neural collapse. In International Conference on Machine Learning, pp. 34301-34329. PMLR, 2023.
Hugo Touvron, Matthicu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pp. 10347-10357. PMLR, 2021.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.09288. 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, łukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: 1804. 007461, 2018.
Peng Wang, Xiao Li, Can Yaras, Zhihui Zhu, Laura Balzano, Wei Hu, and Qing Qu. Understanding deep representation learning via layerwise feature compression and discrimination. arXiv preprint arXiv:2311.02960 2023.
Sicong Wang, Kuo Gai, and Shihua Zhang. Progressive feedforward collapse of resnet training. arXiv preprini arXiv:2405.00885 2024.
Xin, Raphael Tang. Jaejun Lee, Yaoliang Yu, and Jimmy Lin. Deebert: Dynamic early exiting for accelerating bert inference. arXiv preprint arXiv:2004. 22X2:2004411299, 2020.
Ji Xin, Raphael Tang, Yaoliang Yu, and Jimmy Lin. Berxit: Early exiting for bert with better fine- tuming and extension to regression. In Proceedings of the 16th conference of the European chapter of the association for computational linguisties: Main Volume, pp. 91-104, 2021.
Jiahui Yu, Xin Li, Jing Yu Koh, Han Zhang, Ruoming Pang. James Qin, Alexander Ku, Yuanzhong Xu, Jason Baldridge, and Yonghui Wu. Vector-quantized image modeling with improved vggan. arXiv preprint orXiv:2110004660 2021.
Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang, Burcu Karagol Ayan, al. Scaling autoregressive models for content- rich text-to-image generation. arXiv preprint arXiv:2206.10788) 2(3):5, 2022.
Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang, Zican Dong, al. A survey of large language models. arXiv preprint arXiv:2303. 18223, 2023.
Zhihui Zhu, Tianyu Ding, Jinxin Zhou, Xiao Li, Chong You, Jeremias Sulam, and Qing Qu. A ge- ometric analysis of neural collapse with unconstrained features. Advances in Neural Information Processing Systems, 34:29820-29834, 2021.
14

---

Published as a conference paper at ICLR 2025
Appendix
Notations and Organizations. The appendix provides theorem proofs as well as additional ex- perimental results both vision and language models. It is organized follows: First, present the proofs for the theorems in Appendix A. Next, we provide additional experiment results on vision models(Appende) B), language Appeniii C) and multi-modality models(Appenil) D).
A PROOF FOR THEOREMS
In this section, we provide proof for the theorem and 2.
A.1 PROOF FOR THEBREMS
As a result of the geodesic curve assumption, the forward propagation of transformer is along straight line in the feature space RT. Formally, the feature h [0, L] along the line hg
(E)
(1 h k
(0)
(L. In the following proof, we drop the subscript k, in h.
Let f. [0, 1]. Suppose the feature of the first layer is h and feature of the last layer is
(2) Then:
h (z) h sh (1)
The cosine of the angle between intermediate layer feature h (2) and last layer feature h (2) is,
C(x) cos(h (a), 10)) (h (a) 10
| |
Since [| is fixed, treat || constant. For simplicity, define N(x) (h( ,h and (xx) = [[)] Thus, we can get,
C(x) = (h (x) ())
= K N(x) D(z)
(8)
| (0) ||
where 1 |dlli is a positive constant. Note that
(xx) (h h (2) =1 (1- 2) (0) h 10)) + x 12
(xx) = |1 v - 21- 112 + 2x(1 (h (0) h()) + 2
To prove ((x) monotonically increase within E [0, 1]. we can take derivative of C(z) with respect to x. and show 4C(x) > 0 for x E [0, 1]. The derivative given by
dx
dC(x)
dN(x) D(x) dD(z) ((x)
K dx
(2
Assuming h (0) and h unit vectors(
11 - 1) and defining c = (1 (which
satisfies -1 c S 1), we have:
N(x) (1-r)c x c). +
D(x) = v(1 - r² 2x - x)c x
Taking derivative of both ((x) and D(x) with respect gives
dN (1 -c c) dx
dD (2x 1)(1 x.)
dr
D
15

---

Published as a conference paper at ICLR 2025
Then only need to check the sign of dN(m) dar
D dD(a) (a)) term,
(x)x) D(x) dD(x) N(x) (1-c) (2x 1)(1 (11-c)x+ +c)
dx
(
D ((1 - (z)x)) - (2x 1)(1 x) c) +c)
(2cx²
(1 + 3c)x + (1 + c))
D(x)
where x E [0, and c E [-1, 1]. Noticing that ((.) 2 0, we define
P(z) = -2r - (1 + 3c)x + (1+cc
(10)
We have P(0) 1 c [0, and P(I) 0. If P(x) is negative between [0, 1], for quadratic function, there is only one possibility: the axis of symmetry x 4c parabola opens upward:
1+3c is between 0 and 1, and the
1+3e 1
4e
c> c>0
(1I)
-
These conditions are contradictory, and no value of satisfies all of them simultaneously. Thus, P(x) 2 0 always holds for E [0, 1]. Consequently, C(x) cos(h!"),11)) increases mono- tonically for E [0, 1]. Thus, for any layers l le Es. the relationship cos(h (f1), ((a))) cos(h (f2), hh ((a) holds true.
A.2 PROOF FOR THEOREM 2
According to the COsine Similarity (COS) improvement, COS(H (E+1)
COS(h2 (6) (L) (1 COS(h (.) h ); by NCs. we have Assume the norm of (0) and h (e+1) are equal,
(E+1) (L),
By NCI, we have COS(h
COS
(f+1) w (1)) > COS
w (1)).
which is | (h ||b ((+1) h, so we have,
: h
(e+1)
h
(12)
Suppose that h (e+1)
h/ + Ah, so we have,
1 Ah 0
(13)
The logits for class le at layer f and layer & + 1 are denoted by 2 (E+1) : h respectively. They satisfy
w and (e+1)
24
2k
(e+1)
2/ 6%
(14)
where 6y = = Ah > For the class * lo logit,
(e+1)
(1) 2 + 6
(15)
where d, wf Ah. To prove that 8, 0, suppose there exist direction n that w n for all
k E [1, K]. So we have Ah dp + n and,
6, ! Ah "yw w
(16)
According to NC3, W form simplex ETF, meaning all weight vectors have unit norm and the same inner product between any two distinct vectors, i.e., for any i + k. wi we a - So have
6, adg. = K 5k <0
(17)
since & > 0 and K >
16

---

Published as a conference paper at ICLR 2025
The softmax output for class ls at layers l and l + 1 given by
e
/SoftMax(2(1111
Ek=1 eAk
K
e
Σ
K
E
!
ix
e 11+1)
E [1+1]
(orMMa(2((+211)))
k= eTh
(et1)
e+
((+1)
+ Ligh E
+E
+62
c
+ōx
Litk
e²i 0²;"+5, 46,
E eafe
For simplify, define r
then we have,
[SoftMax(z("))1 1
1+r &
[SoftMax( x ((+1)]]
e + reab
+ re(a-116z
Since a =
1
K < and 6% > 0, we have c(a-1)66 < 1, and hence
K-1
K-1
[SoftMax (z(+)))) > [SoftMax (2 (z)
(18)
which proves that
(e+1) [SoftMax Wh k,3 ) > [SoftMax (aaMMaxxhhhh0
Wh (0)
(19)
B ADDITIONAL EXPERIMENTS ON VISION MODELS
In this section, we first illustrate the box plots of cosine similarity validate in Appendix B. Sec- ondly, we validate that residual connections transformers resolve feature rotation ambiguity in Appendix B.2. Then, we describe the setup for the aligned training method in Appendix B.3, and demonstrate its effects on transferability in Appendix B.4. Finally, we apply the aligned training methods to the detection transformer in Appendix B.5.
Setup for Vision Experiments. We conduct experiments on both the CIFAR10 and ImageNet K datasets. The CIFAR10 dataset includes 60,000 color images in 10 classes, each measuring 32 32 pixels. ImageNetIK contains 1.2 million color images distributed in 1000 classes. To increase the diversity of our training data, we use a data augmentation strategy. This includes random crop and padding, random horizontal flip with probability of 0.5, and random rotation within 15 degrees. For optimization, we employ Adam with an initial learning rate 0. This rate decays according to the MultiStepLR at the 100th and 1 150th epochs, over total of 200 epochs. We set the weight decay at le-4. The global batch size for both datasets set at 256. For both vision and NLP tasks, we used 4 RTX A5000 GPUs with 24GB of memory each.
B. Box PLOTS OF SAMPLE- WISE COSINE SIMILARITY
There may be some rare samples with negative sample-wise cosine similarity between features from layers that far apart.
B.2 RESIDUAL CONNECTIONS ELIMINATE ROTATION AMBIGUITY
Section 2 demonstrates a consistent trend between COS and CKA. Additionally, when compute the cosine similarity of features from adjacent layers in Figure 11, most samples exhibit high similar- ity. These findings suggest that transformers do not have orthogonal transformations across layers. But why does this occur? In this section, we examine the role of skip connections in preventing orthogonal transformations.
17

---

Published as a conference paper at ICLR 2025
(a) CIFARIO
(b) INIK
Figure 10: Sample-wise cosine similarity of features from shallow layers and the last-hidden layer. The DeiT-S model is trained with standard training CIFAR-10 and ImageNet. It shows there are rare samples with negative sample-wiss cosine similarity.
Most transformer architectures include skip connections, which are added after the (i) self-attention layer and (ii) MLP layer. According to residual transformation (I), we obtain,
((+1) M.P(LN(MSA(LN((EE))) H())) MSA(LN(H())) H
( + f0co ( (5)
(a) COS of adjacent layers
(b) Norm Ratio
Figure 11: Cosine similarity of features from adjacent layers COS(N (f-1) () and norm ratios ||h (0) [1/l ( distributions. The DeiT-Small model is trained on Imagenet- 1K and evaluated on its validation dataset.
And the feature h is special token of (0) To investigate the effect of residual connections, we calculate the norm ratio ( f/llf (h (5) which is the ratio of the norm of skip connection (e) to the norm of the long branch face; (h (.). The results are displayed Figure 11. High norm ratios suggest that skip connections significantly influence the representational structure of VIT.
To provide further evidence that residual connections resolve the rotation ambiguity, we compared the MLP model with and without these connections and computed their COS and CKA values. For the MLP model without residual connections, as shown Figure 12(a), the CKA value is not con- sistent with accuracy and cosine similarity. A high CKA value might indicate significant similarity between features across layers, but does not necessarily correlate with high classification accuracy. This inconsistency primarily results from the fact that CKA does account for rotation the fea- ture space, suggesting that features could rotate without the residual connections. In contrast, for the MLP model with residual connections, as depicted in Figure 12(b), the CKA value aligns with layerwise accuracy, indicating that residual connections effectively eliminate the rotation ambiguity of features.
B.3 ALIGNED TRAINING DETAILS
Illustration of Train Once and Fit all devices. Figure 13 illustrate how aligned training support train once and fit all devices. After aligned training, one can directly fetch from shallow to deep layers of transformer according to the device computational resources and memory constrains.
18

---

Published as a conference paper at ICLR 2025
1.0
10
COA
Coulrs Samiaar
2.
Lapars
(a) Without residual connections
(b) With residual connections
Figure 12: Comparison layerwise accuracy, COS(COsine Similarity), and CKA (Centered Kernel Alignment) with the last layer of the 9-layer MLP models with and without residual connection the MNIST validation dataset. The models are trained from scratch using standard training. In the left figure, CKA fails to accurately reflect the change in layerwise accuracy for the MLP without residual connection. In the right figure, the presence of residual connection is the reason why CKA works well, it helps eliminate rotation ambiguity.
hg
Ljoin
hg
Aligned Training
Elastic Inference
Figure 13: Aligned training of transformer using joint CE loss of all layer features with common classifier and elastic inference for different memory constrains. Once the model is trained using the aligned method, it can all devices. Features from darker layers indicate better performance.
Alternative Approach for Enhancing Layer-wise Representation Similarity. In addition to us- ing aligned training loss enhance similarity, another method is to add the cosine similarity regularization term to the loss function.
Csim (h (0) h (1)
Σ 1 co (0,
And the total loss is the sum of this two term:
LCE (x. y) = Cen(Whi (,)) + Gaim (0) h (1)
where 3 > O is the regularization coefficient. According to Figure 14, the regularization term contributes minimally to the improvement of cosine similarity and layer-wise accuracy, compared to aligned training methods.
Using the CE-reg loss results in poor layerwise accuracy and lower cosine similarity compared to the aligned loss. The likely reason for this is an imbalance between the COS alignment objective and the primary classification objective. Our intuition that directly optimizing for high COS alignment may fail because the COS alignment loss primarily focuses aligning features across layers, without necessarily making the features discriminative enough for the classification task. In contrast, the cross-entropy (CE) loss directly optimizes for classification, and consequence, it naturally improves COS alignment. This suggests that while COS alignment important, it may not be sufficient its without the robust guidance provided by the CE loss.
Another approach involves adding the CKA term as a regularization term to the CE loss. However, this approach may not be effective and has several drawbacks. First, CKA is not always reliable in
19

---

Published as a conference paper at ICLR 2025
Sandad
Sandedd Aiyned
Zon
(a) Cos Similarity
(b) Layerwise Acc
Figure 14: Cosine similarity with last layer and layerwise accuracy of standard training using LCE( h y). standard-Res -eg training using CCE-reg (z, y) and Aligned training using Caligned (æ. y) The 12 layers DeiT model is trained ImageNetIK dataset. Regularization term helps little for improving the cosine similarity and layerwise accuracy. But our aligned training improves both a lot.
representing layer-wise similarity across all settings. particularly in scenarios with rotation ambigu- ity or when residual connections are absent. Second, it is computationally expensive, as it requires computing the Gram matrix evaluate relationships between features. Lastly. CKA might not per- form better than the COS regularization term and may yield similar results, falling short compared to the aligned loss. In transformers, both COS and CKA measure feature similarity and tend to exhibit similar trends. As shown in Figure 14, the COS regularization term contributes minimally to improving cosine similarity and layer-wise accuracy. Based on this, we infer that using CKA as regularization term would similarly have minimal impact on enhancing these metrics. Therefore, aligned training approaches may be more effective than relying solely on regularization terms.
Setup for Aligned Training in VIT. It's reported (Xin et al., 2021) that training only with this aligned loss would cause the performance drop in the last layer. So following (Xin et al., 2021), we choose the "alternating" training approach, which alternates objectives based on the iteration number. During odd-numeered iterations, we use the CE loss of the final layer Cev (Wr,, . y). For even-numbered iterations, the strategy involves using the aligned loss Latigned v).
Note that this training strategy, which uses a common classifier, no longer requires the KL- divergence term that is commonly used in mutii-exit/llsssffiir training. This is because the deep layers have been trained to capture the abstract and discriminative features of the input data, effec- tively serving as the teacher model. The KI.-divergence term is typically used to guide the shallower layers. However, when we use a common classifier, our aligned training method becomes latent knowledge self-distillation method. The shallow layers mimic align their feature representa- tions with those of the deep layers by aligning with the common classifier. As such, the deep layers, with their advanced feature representations, act as the teachers, while the shallow layers, in their quest to improve their feature extraction capabilities, assume the role of students. Therefore, the KL-divergence term is no longer necessary.
Linear Increasing Weight vs. Uniform Weight for Loss. In (5), we define the layer weights to increase linearly X = 2E/(L For the ablation study, we also consider a uniform weight- ing strategy, where all layers are assigned the same weight, i.e., X = 1/T. As shown in Figure 15, using uniform weights in the loss function tends improve the performance of shallow layers but degrade the final layer. This occurs because shallow layers, which typically learn general but less informative features, produce larger losses, while deeper layers achieve smaller losses. Uniformly weighting all layers disproportionately emphasizes shallow layers and diminishes the importance of deeper ones. contrast, linearly increasing the weights places greater emphasis on deeper layers, resulting in superior accuracy for the final layer. Therefore, linear increasing weight is selected instead of uniform weight for aligned training.
20

---

Published as a conference paper at ICLR 2025
0.8
0.6
Layerwise AcC
0.4
0.2
Uniform
Linear Increase Weight
0.0
3
11
Layers
Figure 15: Comparison of linear increasing weight vs. uniform weight for loss function design using Deit-Small on ImageNet-IK.
Aligned Training Enhance Neural Collapse. The aligned training method aligns features with the last-layer classifier from shallow to deep layers, enhancing the neural collapse (NC) phe- nomenon across layers. As shown in Figure 16, aligned training promotes progressive compression as features move closer to the last layer by noting that aligned results in lower NC1 across layers, where intermediate layers increasingly exhibit the stronger NC1 than standard model.
Algene
38
(a) CIFARIO
(b) ImageNet
Figure 16: Comparison of layerwise neural collapse between standard training and aligned training.
B.4 EFFECTS ON TRANSFER ABILITY
often claimed that shallow layers learn universal patterns while deep layers fit to class labels. Questions arise about whether the proposed aligned training approach is that aligning shallow layer features with deep layer features could cause the shallow layers to lose their transfer ability. To resolve this question, conduct two sets of experiments:
Distribution shift: we first train a DeiT on CIFARIO with standard training and align training, and then evaluate the layer-wise accuracy on CIFAR10.2 (Lu et al., 2020),
Transfer to different tasks: we first train a DeiT on ImageNet with standard training and align training, and then evaluate the layer-wise accuracy on CIFARIO by only fine-tune linear classi- fier, with the feature mapping fixed.
The results plotted in Figure 17. We observe that for both cases, the distribution shift and transferring to different tasks, layer-wise accuracy curves resemble those on the pre-trained datasets
WWithin-class variability collapse (Papyan et al., 2020: Zhu al., 2021) features (hr,}} from each layer computed as NC1 Tr(E D. where Ew 1 ΣK Σn, (br, ha), ha) captures the within-class covariance, Σg Σ (he h) - h) represents the between-clas covariance, represents the class-mean features, and hn represents the global mean of the features.
21

---

Published as a conference paper at ICLR 2025
1.0
30
3unde Lapordde
0.1
Sandar Algges
00
60
Lagers
Layers
(a) CIFARIO CIFARIO.2
ImageNct CTFAR10
Figure 17: The comparison of layer-wise accuracy between a standard model and an aligned model.
shown in Figure and Figure 5, demonstrating that aligned training not only improves layer-wise accuracy for the pre-trained datasets but also for the downstream datasets. In other words, the aligned training methods maintain transfcrability, ensuring that the trained model can effectively transferred.
B.5 ALIGNED TRAINING FOR DETECTION TRANSFORMER
In this section, we demonstrate that aligned training not only beneficial for classification tasks but also effective for other tasks such object detection. Following the De'Tr framework (Carion al., 2020), which employs an encoder-decoder architecture, the encoder extracts global image features, while the decoder predicts object classes and their bounding boxes using queries. Aligned training can be applied to the decoder, where predictions are made using intermediate features from its layers. This is achieved through auxiliary decoding losses (Carion et al., 2020). We evaluate the AP50 (average precision at 50% IoU) for predictions exiting from each decoding layer using aligned training and compare with predictions from the last layer of standard training, shown in Figure 18. As noted DeTr (Carion et al., 2020), the inclusion of aligned training losses critical for performance, and removing them results in significant drop in accuracy.
62
60
M150
56
54
Aligned Standard
2
Layers
Figure 18: Comparison of aligned training (using auxiliary decoding losses) and standard training using last layer loss only) of DeTr model.
C ADDITIONAL EXPERIMENTS ON LANGUAGE MODELS
In this section, first validate that the saturation events deseribed in Section 2.2 are consistently observed in LLaMA3, as shown in Appendix C.I. Then we show aligned training on BERT and GPT models in Appendix C.2.
22

---

Published as a conference paper at ICLR 2025
C.I SATURATE EVENTS LLAMA3
In this section, verify that saturation events, described in Section 2.2, also occur in large language models such as LLaMA3 (Dubey et al., 2024). As shown in Figure 19, we compare the saturation events and the 32-layer LLaMA3.1 8B model. Both models are given the same prompt to ensure a con- trolled comparison. The figure demonstrates how saturation events emerge across layers in models of varying sizes.
two variants of the LLaMA3 model: the 24-layer LLaMA3.2 3B model
Furthermore, as shown in Figure 20, we examine specific instance of saturation events in the 28-layer LLaMA3.2 3B model. Using the prompt "Simply put, the theory of relativity states," the model generates predictions over 24 decoding steps with greedy decoding. The figure displays the predicted words at each layer. using the last-layer classifier to obtain outputs. The color gradient represents the softmax logits of the last-layer predictions, with red indicating low probability (0) and blue indicating a high probability (1). Saturation events are observed when the prediction at given layer remains unchanged through subsequent layers until the final output. This visualization highlights the occurrence of saturation events during the progression of intermediate representations and their impact on the model's final predictions.
Llama-3.2-38
Llama-3.1
Count of Saturate Events
Layers
Saturation events for a 24-layer LLaMA3.2 model and a 32-layer LLaMA3.1 model using the same input prompt.
Figure 19;
C.2 ALIGNED TRAINING ON LANGUAGE MODELS
In this section, provide more details about the datasets and computational resources used. The datasets (GLUE, and Wikitext-1033 are publicly available under the MIT license. For NLP tasks, we used a single RTX A5000 GPUs with 24GB of memory. Then, will introduce the experimental setup for AlignedBERT and AlignedGPT models.
Setup in AlignedBERT The General Language Understanding Evaluation (GLUE) benchmark comprises nine tasks for assessing natural language understanding. In our AlignedBERT experi- ments on the GLUE dataset, we used a sequence length of 256. We employed Adam for optimiza- tion with an initial learning rate of 2e-5, and a batch size of 32. Each task underwent fine-tuning for three epochs. The WikiText- language modeling dataset consists of over 100 million tokens extracted from Wikipedia's verified good and featured articles. For AlignedGPT experiments on the WikiText- 103 dataset, we maintained the sequence length at 256 and used AdamW with an initial learning rate of 2e-5. In this case, we set the batch size to 8.
Setup in PPT. Our aligned training method can be used with any transformer-base lan- guage models. In this study, evaluated our method using the GPT-2 model. We finetune the GPT2 models using aligned training methods and then use intermediate layers GPT2 to generate the texts.
23

---

Published as a conference paper at ICLR 2025
11gh1
660
40L
10000
19960
19gh
speed
5gas
SAME
Scxt
NET
3001
D000
pood
Ioch
Beerr
1000
speed
Indext
Predent
Sye
500
BODCH
FACTAY
Figure 20: Specific instance of saturation events for I.LaMA3.2 3B model responding to the prompt: "Simply put, the theory of relativity states using greedy decoding. Predicted words are shown for each layer, with saturation events identified where predictions at layer & remain unchanged until the final layer.
Standard
Algned
Sandar
0.2
0:2
(a) RTE
(b) QQP
Figure 21: Layerwise accuracy for AlignedBERT and BERT and on RTE and QQP datasets of the GLUE benchmark.
Model and Baselines We finetune GPT-2 on the Wikitext- 103 dataset with the proposed objective Caligned for 40k training steps and generate the text continuation with nucleus sampling (Holtzman et al., 2019) with p 0.95 decoding methods. For the standard baseline model, we finetune model with CE loss LMLE- The model is finetuned using single 24G RTX A5000 GPU for 70 hours.
Evaluations Following (Su & Collier, 2023), evaluate the model from two perspectives: (I)
language modeling quality, assessing the inherent quality of the model, and (2) generation quality, measuring the quality of the text the model produces. In assessing language modeling quality, we calculate the prediction accuracy and perplexity of each layer. When evaluating generation quality, we measure the similarity between the prompt text and generated text using coherence. We employ generation repetition to gauge the diversity of the generated text. The metrics are defined as follows.
Prediction Accuracy The accuracy is computed the Wikitext-100 103 test set as,
D
Ace
ΣΣ
1 aarg max (o = 2z]
(20)
D
1=1
24

---

Published as a conference paper at ICLR 2025
where the D is the number of samples in the test dataset. Perppexity The perplexity computed on the test set of Wikitext-103. It's computed as the exponential of the test loss.
-Coherence Coherence measures the relevance between the prefix text and the generated text. We apply the advanced sentence embedding method, SimCSE (Gao et al., 2021), to mea- sure the semantic coherence or consistency between the prefix and the generated text. The coherence score is defined as follows,
Coherence = h2 ha/ / hall
(21)
where is the prefix text and à is the generated text and h SimCSE(#)) and ha. SimCSE(@). Higher coherence means more correlation to the given prompt.
Diversity Diversity measures the occurrence of generation at different n-gram levels. It is defined as:
unique n-grams(@))
Diversity = II
(22)
total n-grams [t]
n=2
A higher diversity score suggests fewer repeated words in the generated text.
D ADDITIONAL EXPERIMENTS ON MULTI-MODALITY MODELS
In this section, we demonstrate that the observation in Section 2.1 can also extend to multi-modall modiiimmddaltty models such the pre-trained CLIP models (Radford al., 2021). Given that the CLIP model comprises both a vision encoder and a text encoder, we evaluate the cosine similarity within each modality (vision or text) and across modalities (between the vision encoder and encoder). This comprehensive analysis highlights the robustness of the observed phenomena across diverse com- ponents of the CLIP architecture.
Specifically, given a pretrained CLIP model with vision encoder fvision and text encoder fex,, we extract the f-th layer vision feature i E 767 by taking the corresponding [CLS] token outputs of (-th layer from fvision; similarly, take the é-th layer vision feature g(0) E 1252 from the [EOS] token outputs of fext- Since CLIP projects both the vision features and text features to the same embedding space through vision projection matrix W vision E R 768x512 (followed by a layer normalization LN) and text projection matrix W text E R512x112 (also followed by layer normalization LN), we also apply these projection matrices to the hidden layer features obtain
LN(W vision 000 E 512 5
t
LN(W (0) 512
(23)
cOSiv"" v!o)
cosit!" el!l
0.9
COSiv!S esi
CLIP Similanty
COSit" v*
0.7
0.6
0.5
0.4
0.2
Layers
Figure 22: Illustaaiion of cosine similarity between the hidden layer the last layer for within and cross modalities in a 12-layer CLIP-B/32.
25

---

Published as a conference paper at ICLR 2025
14 0h
COS vr)
(b) COS(t f
(c) COS(
Figure 23:
Layer-wise cosine similarities within (a) vision encoder, (b) text encoder, and (c) cross vision and text encoder for pretrained 12-layers CLIP-B/32.
To measure within-modality cosine similarity, we calculate layer-wise similarity within the vision encoder as COS(v(*), v((") and within the text encoder as COS(t()) ((()) Similarly, for layer- wise cross- modos-my cosine similarity,
evaluate the relationships between the vision encoder and text encoder using COS(v t(e") and COS(t(O, o(e"). Figure 22 plots the layer-wise within- modality similarity and cross- cross-modality similarity by comparing the bidden-layer features to the last- layer features, while Figure plots all the pair-wise results. All experiments are conducted on the CIFAR10 validation dataset, where the text input to the text encoder is: "This is a photo of a (label)". From both figures, we observe clear pattern of progressively increasing layer-wise representational similarity in both the vision encoder and the text encoder. For cross-modalitty similarity, we note that the last-layer vision and text representations are not perfectly aligned (e.g.. COS(+(L), v(()) is not close to 1), phenomenon commonly referred to as the modality gap, which has been consistently observed across various multi-modal models (Liang al., 2022). Despite this, we also observe a progressive increase in layer-wise representation similarity across modalities. Together, these results highlight distinct trend of progressively increasing layer-wise representational similarity within and across modalities.
26


Reviews:
[{"review_id": "ACNvhfdbQ8", "reviewer": "Reviewer_ABue", "summary": "This work studies the feature similarity of neural networks and reveals that: (I) simple-wise cosine similarity can capture representation similarity; (ii) saturation events related with feature similarity and based on this observations, this work proposed a aligned training approach to enhance the representation thus benefit the performance and also the multi-exit inference approach.", "questions": "Please refer to the weaknesses part.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The manuscript is logically organized, including the observations and its applications.\n- Both empirical and theoretical justifications are provided.\n- The study covers both the vision and language domains.", "weaknesses": "- Line 342 “To the best of our knowledge, our work is the first to show that one common classifier is sufficient for multi-exit models.” This is not true, a lot of early exiting methods can do with a single classifier heads [1-3].\n- Line 240: “Progressively increasing layer-wise representation similarity”, thus observations might be different in other domain[2], is there any insights why autoregressive models seems not have progressively increasing layer-wise representation similarity?\n- Missing previous literature[4] which also studies layer-wise cosine similarity, yet in the language domain.\n\n[1] https://proceedings.neurips.cc/paper_files/paper/2022/file/6fac9e316a4ae75ea244ddcef1982c71-Paper-Conference.pdf\n\n[2] https://arxiv.org/pdf/2404.03865\n\n[3] https://arxiv.org/pdf/2403.03853\n\n[4] https://arxiv.org/pdf/2202.08625", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LoUQbVNAmu", "reviewer": "Reviewer_xxse", "summary": "This paper focuses on understanding the behavior of deep neural networks, particularly transformer models, by examining the similarity of internal representations across hidden layers. The authors introduce a sample-wise cosine similarity metric that aligns with more complex statistical methods and reveals increasing representation similarity as layers get closer. They provide a theoretical justification for this phenomenon under the geodesic curve assumption and demonstrate that enhanced representation similarity leads to increased predicted probability and earlier saturation events in model predictions. The paper proposes an aligned training method to improve shallow layer effectiveness, resulting in more early saturation events and higher layer-wise accuracies. Finally, the authors show that their approach enables multi-exit models with a single classifier, reducing parameter count and computational complexity while maintaining performance.", "questions": "Based on the weaknesses, I have some questions.\n\n1. Is enhancing similarity beneficial for all tasks besides the simple classification task?\n\n2. When enlarging the data/model size, will the method/analyses still works?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "In general, the proposed paper is well-written, the author provides detailed experiments, extensive theoretical analysis to analysis the feature pattern in Transformers. Based on these analyses, the author proposes a aligned training method for enhancing shallow layer performance. The proposed method achieves performance gain and speed boost on CV and NLP tasks.", "weaknesses": "In general, I think the proposed paper is well formulated and written. However, I still have some concerns about the paper:\n\n1. I don't think it is useful for enhancing similarity between all Transformer outputs (representations), it may help in simple tasks like image classification and sentence classification. But for complex tasks (like object detection/semantic segmentations), we may not want all representations to be similar. I think similar tasks may exist in NLP tasks (like parsing), Then the author should talk about the limitations or give more experiments to verify the effectiveness of the proposed method.\n\n2. In the paper, the author performs experiments on small datasets and small models, like DeiT trained on CIFAR10 and ImageNet. Also Bert/GPT2 on small NLP tasks. If the data increases, like a pre-trained CLIP/OpenCLIP on large datasets, will the findings/analyses be the same? Moreover, if the model becomes larger (like a LLM like Llama3), will the findings/analyses be the same? Does the saturation events still occur in those models? I'm curious about that.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yAPW9oPB0G", "reviewer": "Reviewer_aKPD", "summary": "This paper focuses on the analysis of layer-wise representations of Transformers. They demonstrates a series of beneficial observations, e.g., representations across layers are positively correlated. Meanwhile, they find the model's top prediction remains unchanged across subsequent layers. Following the observation, they further propose an aligned training strategy to improve the effectiveness of shallow layer, which is able to provide more early saturation events, minimal depth needed for the given task, multi-exit models.", "questions": "I have a few questions about this submission.\n1. According to Figure 6, does the proposed aligned training strategy decrease the performance more clearly in the last layers in large models? Please discuss the performance trade-offs between shallow and deep layers across different model sizes.\n2. With the common classifier, is there still a neural collapse phenomenon in the model? Is there any difference of this phenomenon across different layers? I'd like to see some discussions about this.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. By analyzing the similarity of representations among different layers, the paper demonstrates the possibility of early saturation events and shared classifier among different layers.\n2. They further presents a training strategy to improve effectiveness of shallow layers, such that they can enjoy more early saturation events, minimal depth, and so on\n3. Some analysis in this paper is insightful, e.g., the shadow layer is able to achieve approaching performance with the depth layer. This might inspire more efficient large model inference.\n4. The experiments are comprehensive.", "weaknesses": "1. a small weakness is that the previous approaches have observed this phenomenon that representations in the early layers can also achieve reasonable classifiers, though I think this is a tiny issue.\nPlease refer to the Questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AabWzBHiiG", "reviewer": "Reviewer_s3vt", "summary": "This paper studied the similarity of representations between the hidden layers of individual transformers and found that a simple cosine similarity metric can be used for similarity evaluation. Experimental results revealed that representations across layers are positively correlated and the authors introduce a multi-exit mechanism. The innovation lies in using the same classifier for different layers.", "questions": "1. Why use linearly increasing weights for aligned losses? Experiments for different choices on weights are preferred.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The deduction and experiments in the paper are relatively solid. The authors have made efforts to investigate the similarity of representations across different layers in transformers. \n\n2. The paper is well-written and easy to follow.", "weaknesses": "Using the same classifier for multi-exit is quite straightforward and its technical contribution seems limited. The experimental results seem not very promising compared with previous multi-exit methods. The authors may need to further enhance the novelty or provide more convincing evidence of the superiority of their approach compared with multi-exit/classifier to make the paper more acceptable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)