Don’t Judge a Language Model by Its Last Layer: Contrastive Learning with Layer-Wise Attention Pooling (2022)

Abstract: Recent pre-trained language models (PLMs) achieved great success on many natural language processing tasks through learning linguistic features and contextualized sentence representation. Since attributes captured in stacked layers of PLMs are not clearly identified, straightforward approaches such as embedding the last layer are commonly preferred to derive sentence representations from PLMs. This paper introduces the attention-based pooling strategy, which enables the model to preserve layer-wise signals captured in each layer and learn digested linguistic features for downstream tasks. The contrastive learning objective can adapt the layer-wise attention pooling to both unsupervised and supervised manners. It results in regularizing the anisotropic space of pre-trained embeddings and being more uniform. We evaluate our model on standard semantic textual similarity (STS) and semantic search tasks. As a result, our method improved the performance of the base contrastive learned BERT _{base} and variants.

## introduction pre-trained language models (plms) @xcite @xcite have shown competitive performance on many natural language processing (nlp) tasks. also, contrastive learning using the plms shows the highest performance in sentence representation. contrastive learning is to learn effective representations by staying semantically close sample pairs together while dissimilar ones are far apart @xcite . in general, plms use either [cls] tokens in the last layer, av g which is the average representation of tokens in the last layer @xcite , or av g f l which is the average * these authors contributed equally. † these authors are corresponding authors. figure 2: spearsman's correlation score depending on the pooling methods of plms for each domain representation of tokens in the first and last layers @xcite , to pool out sentence representation from word representations. however, since language models show performance gaps by domain when trained on different objectives, the fixed pooling strategy has limitations in performance improvement. figure 1 and 2 show the spearman's correlation score of each layer or pooling method in plms. we evaluated the test set of the standard semantic textual similarity (sts) dataset @xcite @xcite @xcite @xcite . the comparison of performance when pooling each layer shown in figure 1 indicates that using only a specific layer for pooling is insufficient. other layers other than the last layer may contain substantial information for sentence representation. for example, for the sts benchmark (sts-b) task @xcite , bert base with [cls] embedding scored the highest at the fourth layer (48.66%), which is about 20% higher than the last layer. figure 2 shows that simply pooling from more layers impedes the performance by comparing models pooled from the first and last layer and the last layer. in addition, there is no consistent tendency to compare effectiveness for a given layer between [cls] pooling and average pooling. motivated by this point, we designed the attention networks and task-agnostic pooling methods to assign more weights to spots that need more focus in the layer and lead to representation vector optimization. our proposed method outperforms previously fixed pooling strategies in contrastive learning. in addition, contrastive learning models with layer-wise attention pooling show a higher semantic search performance with the same parameters. in summary, the contributions of this paper are as follows: ## method in this section, we present a layer-wise pooling strategy based on attention mechanisms to improve the quality of sentence representations from language models. in addition, we describe the process of applying the proposed pooling strategy to be leveraged on three contrastive learning schemes.