Fine-Tuned Transformers Show Clusters of Similar Representations Across Layers (2021)

Abstract: Despite the success of fine-tuning pretrained language encoders like BERT for downstream natural language understanding (NLU) tasks, it is still poorly understood how neural networks change after fine-tuning. In this work, we use centered kernel alignment (CKA), a method for comparing learned representations, to measure the similarity of representations in task-tuned models across layers. In experiments across twelve NLU tasks, we discover a consistent block diagonal structure in the similarity of representations within fine-tuned RoBERTa and ALBERT models, with strong similarity within clusters of earlier and later layers, but not between them. The similarity of later layer representations implies that later layers only marginally contribute to task performance, and we verify in experiments that the top few layers of fine-tuned Transformers can be discarded without hurting performance, even with no further tuning.

## introduction fine-tuning pretrained language encoders such as bert @xcite and its successors @xcite @xcite ) has proven to be highly successful, attaining state-of-the-art performance on many language tasks, but how do these models internally represent task-specific knowledge? in this work, we study how learned representations change through fine-tuning by studying the similarity of representations between layers of untuned and task-tuned models. we use centered kernel alignment (cka; @xcite to measure representation similarity and conduct extensive experiments across three pretrained encoders and twelve language understanding tasks. we discover a consistent, block diagonal structure (figure 1c , d and albert models, where early layer representations and later layer representations form two distinct clusters, with high intra-cluster and low inter-cluster similarity. given the strong representation similarity of later model layers, we hypothesize that many of the later layers only marginally contribute to task performance. we show in experiments that the later layers of task-tuned roberta and albert can indeed be discarded with minimal impact to performance, even without any further fine-tuning. ## experimental setup models for the majority of our experiments, we consider three commonly used language-encoding models: roberta @xcite , albert @xcite and electra @xcite . because of the large number of exper-iments being performed, we use roberta base , albert largev2 and electra base rather than the largest available versions of these models. tasks we use the tasks included in the glue benchmark @xcite excluding the datapoor wnli, namely: cola @xcite , mnli @xcite , mrpc @xcite , qnli @xcite , qqp, foot_0 rte @xcite , sst-2 @xcite , and sts-b @xcite . we include four additional tasks to cover a more diverse set of task formats and difficulties: boolq @xcite and yelp review polarity @xcite classification tasks, and hellaswag @xcite and cosmosqa @xcite multiple-choice tasks. ## representation similarity with cka to analyze how learned representations change via fine-tuning, we use centered kernel alignment (cka; @xcite to measure representation similarity. cka is invariant to both orthogonal transformation and isotropic scaling of the compared representations, making it ideal for measuring the similarity of neural network representations, and has applied to bert-type models in prior work @xcite . given two sets of representations @xmath0, d 2 the hidden dimensions, cka computes a similarity score between 0 and 1, where a higher score indicates greater similarity. further details on cka are provided in appendix a. using cka, we can compare the similarity of representations between different layers of the same model or even different models. for our analysis, we use the representations of the cls token, i.e. the token whose final layer representation is fed to the task output head. 2 we compute cka over the validation examples of each task. to provide intuition for cka scores, we first show in figure 1 an example of the comparison formats using albert fine-tuned on rte. orig-orig the top left plot shows the similarity of representations across the layers of the untuned albert model on rte inputs. adjacent layers have high similarity scores, only gradually decreasing as more distant layers are compared. ft-orig we show layers of the task-tuned model on the y-axis and untuned model on the x-axis. the cls representations of the later layers in the task-tuned model appear highly dissimilar to any of the untuned model: in other words, the representations differ starkly from those used for albert's masked language modeling (mlm) and sentence order prediction (sop) pretraining. this coheres with prior work showing that representations of later layers are most likely to change during fine-tuning @xcite . ft-ft next, we compare layers within a single fine-tuned model. we observe a block-diagonal structure in the representation similarities-two distinct clusters of earlier (approx. first 10) and later (approx. last 14) layers that have high inter-cluster but low intra-cluster similarity. when considered together with ft-orig, we can infer that the earlier layer representations resemble those used for pretraining, whereas the later layers encode a representation suitable for tackling the task. the high internal similarity between the top few layers and the sharp block diagonal structure of the similarity matrix imply that the representations starkly differ. ## truncating fine-tuned models to test our hypothesis that the later layers of tuned task-models only marginally contribute to task performance, we propose a simple experiment where we feed the representations from an intermediate layer directly to the task output head, effectively discarding the later layers. we refer to these as truncated models. we test three different configurations: (a) untuned , where we feed intermediate representations from a fine-tuned model to the tuned task output head without any further finetuning, (b) tuned , where we fine-tune only the output head, and (c) tunedorig , where we use representations from the base model (not fine-tuned on the task), but we fine-tune the output head. performance of the untuned trunated models indicates the extent to which an intermediate representation can be directly substituted for the final layer's representation; the tuned and tunedorig models provide an upper-bound of performance using the cls representation of a given layer of a finetuned and non-fine-tuned encoder respectively. our results are shown in figure 3 . for roberta and albert, we find that the untuned truncated models perform comparably to the tuned truncated and full fine-tuned models 3 at the later layers. for instance, the top 4 layers of the roberta for yelp polarity model can be discarded with no further tuning and minimal impact to performance (95.5 vs 96.1). on the other hand, tunedorig models perform very poorly compared to the tuned models across all layers, showing that task-tuned intermediate representations are crucial for good performance, even when fine-tuning the output head. for albert, which shares parameters between layers, a larger fraction of layers can be discarded with minimal impact to performance for both un-tuned and tuned truncated models. on the other hand, we do not find a similar pattern in electra models. the untuned truncated models perform extremely poorly when discarding almost any number of layers, and even the tuned truncated models quickly drop in performance with even one or two layers discarded. these results are consistent with our cka analyses that showed that the learned and task-tuned representations for electra do not share the same structure as those of roberta and albert. we speculate that this differences stems from the different pretraining objectives-replaced token detection is a binary prediction problem, whereas masked language modeling involves predicting a distribution over a large number of tokens-leading to differences in learned representations that propagate even to fine-tuned models. we leave further is equivalent to a regular fine-tuned model. investigation these differences to future work. ## related work while cka @xcite was initially proposed as an interpretability method for computer vision models, it has more recently seen application to nlp models. wu et al. (2020) applied cka to pretrained transformers models such as bert and gpt-2, focusing on cross-model comparison-our analysis builds on their findings, with greater focus on layer-wise comparisons and implications for fine-tuning and discarding layers. sridhar and sarah (2020) use cka to measure the impact of a proposed model architecture change on the learned representations. voita et al. (2019) and merchant et al. (2020) apply similar representation similarity analyses to transformers, with the latter also investigating freezing and dropping layers from models. more broadly, significant work has been done on better understanding and interpreting the capabilities of bert-type models-rogers et al. ( 2020 ) offers a thorough survey of this line of work. of particular relevance to our work: work on model probing @xcite @xcite has studied the extent to syntactic and semantic features are represented at different layers of bert-type models. our results on model truncation also cohere with existing work on early exit in bert models @xcite , wherein models are explicitly fine-tuned to dynamically skip the later layers of a bert encoder and directly to the output head, often to reduce inference times of models. our results somewhat differ as we show that models can also be truncated or exited early without any explicit tuning. it has also been shown in the computer vision domain that models with residual networks work akin to an ensemble of deep and shallow models @xcite . ## conclusion we show a consistent pattern to the structure of representation similarity in task-tuned roberta and albert models, with strong representation similarity within clusters of earlier and later layers, but not between them. we further show that the later layers of task-tuned roberta and al-bert models can often be discarded without hurting task performance, verifying that the later layers of these models truly have similar representations. however, we find that electra models exhibit starkly different properties from the other two mod-els, which prompts further investigation into how and why these models differ.