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: This paper proposes TRANS-ENCODER, an unsupervised approach to training bi-encoders and cross-encoders for sentence similarity tasks such as information retrieval, natural language inference, semantic textual similarity and clustering. The core idea of TRANS-ENCODER is self-distillation that alternatively trains a bi-encoder and a cross-encoder with pseudo-labels created from the other. The authors also propose a mutual-distillation extension to mutually bootstrap two self-distillation models trained in parallel. The effectiveness of TRANS-ENCODER is verified by empirical evidence.