Trans-Encoder: Unsupervised sentence-pair modelling through self- and mutual-distillations

In NLP , a large volume of tasks involve pairwise comparison between two sequences ( e.g. , sentence similarity and paraphrase identification ) . Predominantly , two formulations are used for sentence-pair tasks : bi-encoders and cross-encoders . Bi-encoders produce fixed-dimensional sentence representations and are computationally efficient , however , they usually underperform cross-encoders . Crossencoders can leverage their attention heads to exploit inter-sentence interactions for better performance but they require task finetuning and are computationally more expensive . In this paper , we present a completely unsupervised sentence-pair model termed as TRANS-ENCODER that combines the two learning paradigms into an iterative joint framework to simultaneously learn enhanced bi- and crossencoders . Specifically , on top of a pre-trained language model ( PLM ) , we start with converting it to an unsupervised bi-encoder , and then alternate between the bi- and cross-encoder task formulations . In each alternation , one task formulation will produce pseudo-labels which are used as learning signals for the other task formulation . We then propose an extension to conduct such self-distillation approach on multiple PLMs in parallel and use the average of their pseudo-labels for mutual-distillation . TRANS-ENCODER creates , to the best of our knowledge , the first completely unsupervised cross-encoder and also a state-of-the-art unsupervised bi-encoder for sentence similarity . Both the bi-encoder and cross-encoder formulations of TRANSENCODER outperform recently proposed state-of-the-art unsupervised sentence encoders such as Mirror-BERT ( Liu et al. , 2021 ) and SimCSE ( Gao et al. , 2021 ) by up to 5 % on the sentence similarity benchmarks . 1 INTRODUCTION . Comparing pairwise sentences is fundamental to a wide spectrum of tasks in NLP such as information retrieval ( IR ) , natural language inference ( NLI ) , semantic textual similarity ( STS ) and clustering . Two general architectures usually used for sentence-pair modelling are bi-encoders and cross-encoders . In a cross-encoder , two sequences are concatenated and sent into the model ( usually deep Transformers like BERT/RoBERTa ; Devlin et al . 2019 ; Liu et al . 2019 ) in one pass . The attention heads of Transformers could directly model the inter-sentence interactions and output a classification/relevance score . However , a cross-encoder needs to recompute the encoding for different combinations of sentences in each unique sequence pair , resulting in a heavy computational overhead . It is thus impractical in tasks like IR and clustering where massive pairwise sentence comparisons are involved . Also , task finetuning is always required for converting PLMs to cross-encoders . By contrast , in a bi-encoder , each sequence is encoded separately and mapped to a common embedding space for similarity comparison . The encoded sentences can be cached and reused . A bi-encoder is thus much more efficient . Also , the output of a bi-encoder can be used off-the-shelf as sentence embeddings for other downstream tasks . That said , it is well-known that in supervised learning , bi-encoders underperform cross-encoders ( Humeau et al. , 2020 ; Thakur et al. , 2021 ) since the former could not explicitly model the interactions between sentences but could only compare them in the embedding space in a post hoc manner . In this work , we ask the question : can we leverage the advantages of both bi- and cross-encoders and bootstrap knowledge from them in an unsupervised manner ? Our proposed TRANS-ENCODER addresses this question with the following intuition : As a starting point , we can use bi-encoder representations to tune a cross-encoder . With more powerful inter-sentence modelling , the crossencoder should resurface more knowledge from the PLMs than the bi-encoder given the same data . In turn , the more powerful cross-encoder can distil its knowledge back to the bi-encoder . We can repeat this cycle to iteratively bootstrap from both the bi- and cross-encoders . 2 TRANS-ENCODER . The general idea of TRANS-ENCODER is simple yet extremely effective . In §2.1 , we first transform an off-the-shelf PLM to a strong bi-encoder , serving as an initialisation point . Then , the bi-encoder produces pseudo-labels and the PLM subsequently learns from these pseudo-labels in a cross-encoder manner ( §2.2 ) . Consecutively , the cross-encoder further produces more accurate pseudo-labels for bi-encoder learning ( §2.3 ) . This self-distillation process is visualised in Fig . 1 . Then in , §2.4 , we propose a further extension called mutual-distillation that stabilises the training process and boosts the encoder performance even more . 2.1 TRANSFORM PLMS INTO EFFECTIVE BI-ENCODERS . Off-the-shelf PLMs are unsatisfactory bi-encoders.1 To have a reasonably good starting point , we leverage a simple contrastive tuning procedure to transform existing PLMs to bi-encoders . This approach is concurrently proposed in both ( Liu et al. , 2021 ) and ( Gao et al. , 2021 ) . Let f ( · ) denote the encoder model ( a Transformer-based model ) ; X be a batch of randomly sampled raw sentences . For any sentence xi ∈ X , we send it to the encoder twice to create two views of the same data point : f ( xi ) and f ( xi ) . Note that the two encodings slightly differ due to the dropout layers in the encoder.2 We then use the standard InfoNCE loss ( Oord et al. , 2018 ) to cluster the positive pairs together and push away the negative pairs in the mini-batch : LinfoNCE = − |X |∑ i=1 log exp ( cos ( f ( xi ) , f ( xi ) ) /τ ) ∑ xj∈Ni exp ( cos ( f ( xi ) , f ( xj ) ) /τ ) . ( 1 ) τ denotes a temperature parameter ; Ni denotes all negatives of xi within the current data batch , which includes all xj , xj where i ̸= j ( i.e. , |Ni| = 2|X | − 2 ) . Intuitively , the numerator is the similarity of the self-duplicated pair ( the positive example ) and the denominator is the sum of the similarities 1In this work , a bi-encoder is the same as a universal text encoder . When embedding a pair of sentences under the bi-encoder architecture , the same encoder is used ( i.e. , the two-branches of bi-encoder share weights ) . 2Mirror-BERT has the option to also randomly mask a span in xi and using drophead ( Zhou et al. , 2020 ) to replace dropout as hidden states augmentation . between xi and all other negative strings besides xi . The loss encourages the positive pairs to be relatively close compared to the negative ones . In experiments , we use checkpoints released by SimCSE and Mirror-BERT . But in principle , any techniques could be used here as long as the method serves the purpose of transforming BERT to an effective bi-encoder and does not require additional labelled data ( see §6 ) . 2.2 SELF-DISTILLATION : BI- TO CROSS-ENCODER . After obtaining a sufficiently good bi-encoder , we leverage it to label sentence pairs sampled from the task of interest . Specifically , for a given sentence-pair ( sent1 , sent2 ) , we input them to the bi-encoder separately and get two embeddings ( we use the embedding of [ CLS ] from the last layer ) . The cosine similarity between them is regarded as their relevance score . In this way we have constructed a self-labelled sentence-pair scoring dataset in the format of ( sent1 , sent2 , score ) . We then employ the same model architecture to learn from these score , but with a cross-encoder formulation . The cross-encoder weights are initialised from the original PLM.3 For the sentence-pair ( sent1 , sent2 ) , we concatenate them to produce “ [ CLS ] sent1 [ SEP ] sent2 [ SEP ] ” and input it to the cross-encoder . A linear layer ( newly initialised ) then map the sequence ’ s encoding ( embedding of the [ CLS ] token ) to a scalar . The learning objective of the cross-encoder is minimising the KL divergence between its predictions and the self-labelled scores from the bi-encoder . This is equivalent to optimising the ( soft ) binary cross-entropy ( BCE ) loss : LBCE = − 1 N N∑ n=1 ( yn · log ( σ ( xn ) ) + ( 1− yn ) · log ( 1− σ ( xn ) ) ) ( 2 ) where N is the data-batch size ; σ ( · ) is the sigmoid activation ; xn is the prediction of the cross-encoder ; yn is the self-labelled ground-truth score from the bi-encoder . Note that while the cross-encoder is essentially learning from the data produced by itself ( in a bi-encoder form ) , usually , it outperforms the original bi-encoder on held-out data . The cross-encoder directly discovers the similarity between two sentences through its attention heads , finding more accurate cues to justify the relevance score . The ability of discovering such cues could then generalise to unseen data , resulting in stronger sentence-pair scoring capability than the original bi-encoder . From a knowledge distillation perspective , we can view the bi- and cross-encoder as the teacher and student respectively . In this case the student outperforms the teacher , not because of stronger model capacity , but smarter task formulation . By leveraging this simple yet powerful observation , we are able to design a learning scheme that iteratively boosts the performance of both bi- and cross-encoder . 2.3 SELF-DISTILLATION : CROSS- TO BI-ENCODER . With the more powerful cross-encoder at hand , a natural next step is distilling the extra gained knowledge back to a bi-encoder form , which is more useful for downstream tasks . Besides , and more importantly , a better bi-encoder could produce even more self-labelled data for cross-encoder learning . In this way we could repeat §2.2 and §2.3 , continually bootstrapping the encoder performance . We create the self-labelled sentence-scoring dataset in the same way as §2.2 except that the crossencoder is used for producing the relevance score . The bi-encoder is initialised with the weights after SimCSE training from §2.1.4 For every sentence pair , two sentence embeddings are produced separately by the bi-encoder . The cosine similarity between the two embeddings are regarded as their predictions of the relevance score . The aim is to regress the predictions to the self-labelled scores by the cross-encoder . We use a mean square error ( MSE ) loss : LMSE = − 1 N N∑ n=1 ( xn − yn ) 2 ( 3 ) 3We have a design choice to make here , we could either ( 1 ) directly use the weights of the bi-encoder or ( 2 ) use the weights of the original PLM . We will show in experiments ( §5.2 ) that ( 2 ) is slightly better than ( 1 ) . 4Again , we have the choice of using the cross-encoder weights from §2.2 ( the extra linear layer is disregarded ) . We will discuss more in experiments . where N is the batch size ; xn is the cosine similarity between a sentence pair ; yn is the self-labelled ground-truth . In experiments , we will show that this resulting bi-encoder is more powerful than the initial bi-encoder . Sometimes , the bi-encoder will even outperform its teacher ( i.e. , the cross-encoder ) . Choice of loss functions : maintaining student-teacher discrepancy is the key . Intuitively , MSE allows for more accurate distillations since it punishes any numerical discrepancies between predicted scores and labels on the instance level . However , in practice , we found that using MSE for bi-tocross distillation aggravates the overfitting issue for cross-encoder : the cross-encoder , with its strong sentence-pair modelling capability , completely overfits to the mapping between concatenated sentence pairs and the pseudo scores . This elimination of discrepancy between model predictions and pseudo labels harms generalisation and the iterative learning cycle can not continue ( as the predicted scores by the student model will be the same as the teacher model ) . We thus use BCE loss for bi-to-cross since BCE is a more forgiving loss function compared with MSE . According to our gradient derivation ( App . §A.2 ) , BCE essentially is a temperature-sharpened version of MSE , which is more tolerant towards numerical discrepancies . This prevents cross-encoders from overfitting to the pseudo labels completely . Similar issue does not exist for the cross-to-bi distillation as for bi-encoders , two input sequences are separately encoded and the model does not easily overfit to the labels . We have a more thorough discussion in App . §A.2 and Tab . 8 to highlight the rationales of the current configuration .

Summary: # Summary  - This paper highlights the fact that bi-encoders (sentence encoders) and cross-encoders (sentence-pair encoders) have been considered somewhat separately in the literature of sentence-pair modeling, and that cyclic knowledge distillation from one to another (and vice versa) can be effective for improving the performance of both the encoder stereotypes. - The authors combine the two sentence modeling paradigms into an iterative joint framework, which is called TRANS-ENCODER, to simultaneously learn enhanced bi- and cross-encoders. - Specifically, given the original SimCSE or Mirror-BERT checkpoints provided by prior work, the proposed framework first (pseudo-)labels sentence pairs sampled from the task of interest by utilizing the off-the-shelf sentence encoders. And then, the labeled sentence pairs are utilized to improve the quality of cross-encoders. Further, the fine-tuned cross-encoders can also be exploited to label other sentence pairs which are again able to be used as a dataset for tuning the bi-encoders (sentence encoders). The framework repeats this process until both the bi- and cross-encoders become satisfactory in terms of their performance. - The proposed method is evaluated on semantic textual similarity (STS) tasks and binary classification datasets such as QQP, QNLI, and MRPC. The results demonstrate that the method improves both the bi- and cross-encoders to be more suitable for sentence-pair modeling tasks.