Compressing Sentence Representation for Semantic Retrieval via Homomorphic Projective Distillation (2022)

Abstract: How to learn highly compact yet effective sentence representation? Pre-trained language models have been effective in many NLP tasks. However, these models are often huge and produce large sentence embeddings. Moreover, there is a big performance gap between large and small models. In this paper, we propose Homomorphic Projective Distillation (HPD) to learn compressed sentence embeddings. Our method augments a small Transformer encoder model with learnable projection layers to produce compact representations while mimicking a large pre-trained language model to retain the sentence representation quality. We evaluate our method with different model sizes on both semantic textual similarity (STS) and semantic retrieval (SR) tasks. Experiments show that our method achieves 2.7-4.5 points performance gain on STS tasks compared with previous best representations of the same size. In SR tasks, our method improves retrieval speed (8.2×) and memory usage (8.0×) compared with state-of-the-art large models. Our implementation is available at https://github.com/XuandongZhao/HPD .

## introduction it is a fundamental problem to learn compact yet effective sentence representations. good representations have wide applications in nlp, including web search @xcite , question answering @xcite , knowledge inference @xcite , and machine translation @xcite . sentence embedding models take a sentence as the input and output a fixed-length continuous vector representation. based on bert @xcite and roberta @xcite , recent sentence embedding models such as sentence-bert (sbert) @xcite and simcse @xcite , are fine-tuned on sentence pair scoring tasks to learn better sentence representations, which show much improvement @xcite . in contrast to the teacher model, which is a large pre-trained language model with a fixed pca dimension reduction module, the student model is a smaller language model with a learnable projection layer. the mean square error of both dimension reduction results is used to train the student model. in downstream tasks. however, these models are big in two aspects. 1) they contain hundreds of millions to billions of parameters, which requires large memory and powerful machine to serve in production; 2) their resulting embeddings are high dimensional (e.g. 1024), requiring huge database to store and index, which cause high search latency. therefore, it is challenging to directly use these large models in real-world applications with strict throughput/latency requirement and bounded hardware resources. our work focuses on reducing both the model size and the representation size. there has been several works to reduce model size and retain the superior model performance. recent studies @xcite @xcite have used knowledge distillation (kd) on large language models to derive compressed compact models with decent performance. tinybert @xcite performs layer-to-layer transformer distillation at pre-training and task-specific learning stage utilizing the teacher's hidden states and self-attention distributions. minilm @xcite proposes task-agnostic transformer distillation, which uses self-attention distributions and value relations to help the student deeply mimic the teacher's self-attention modules. nevertheless, directly fine-tuning small transformer models for sentence embedding shows less desirable results than large ones @xcite . can we learn a compact yet highly performant sentence representation? in this work, we propose hpd: a dimension reduced sentence embedding model via projected knowledge distillation. the key idea is to start from a pre-trained large model and distill its knowledge into a small one. the large model is fine-tuned on natural language inference (nli) datasets first. then the small and large ones are augmented with linear projection layer and principal component analysis (pca) @xcite respectively to reduce final representation dimension. in this way, the small model is expected to produce semantic meaningful representations (semantically similar sentences will have close embeddings), where it mimics the power of large models through homomorphic mappings. we evaluate our model on semantic textual similarity (sts) and semantic retrieval (sr) tasks. empirical results show that our model can attain 2.7-4.5 points of performance gain on sts tasks compared to other dimension reduction approaches and achieve competitive retrieval performance against large sentence embedding models while significantly improving retrieval speed (8.2×) and memory usage (8.0×) in sr tasks. ## related work embedding techniques are used to represent complex data mathematically @xcite @xcite . sentence embedding is a well-researched topic with a plethora of proposed approaches. early works @xcite build upon the distributional hypothesis and train the models to predict the surrounding sentences. sent2vec @xcite generate sentence embeddings using word vectors along with n-gram embeddings. @xcite propose to fine-tune a siamese network on nli datasets, which is then further extended to pre-trained models in sentence-bert @xcite . simcse @xcite proposes a con-trastive learning method and achieves state-of-theart performance on sts tasks. recently, @xcite combines pca based dimensionality reduction with a postprocessing algorithm to address the latency and capacity issues of large dimensionality models. @xcite proposed extended robust pca (ex-rpca) to do dimension reduction. but both of them only focus on word embedding. @xcite find that the whitening operation can enhance the isotropy of sentence distribution and reduce the dimensionality of the sentence representation, which optimizes the memory storage and accelerates the retrieval speed. we use this approach as one of our baselines. ## method the overview of our approach is illustrated in figure 1. given a set of sentences @xmath0, our goal is to obtain efficient sentence embedding models f : @xmath1, where d is the embedding dimension. the teacher model f t is trained on the same nli dataset as @xcite ; @xcite ; @xcite , where there are three types of sentence pairs (entailment/neutral/contradiction). we follow the supervised contrastive training framework in simcse @xcite and take a cross-entropy objective with in-batch negatives and hard negatives. the idea is based on the assumption that a good semantic representation should be able to bring similar sentences together while pushing dissimilar ones apart. let (e i , e + to the teacher model after its average pooling layer and we add a projection layer, h s @xmath2, to the student model, where h t i , h s @xmath3. w @xmath4 algorithm 1 shows how to conduct the pca over a set of initial sentence embeddings for the teacher model. we sample m sentences from the training dataset and get the embeddings {e t i } m @xmath5. we then construct a centered matrix e of d ′ t × m size, where d ′ t is the initial embedding dimension. thus the i-th column of e corresponds to e t i -ēt , i.e. row means have been subtracted. then we perform the singular value decomposition (svd) of e (line 2 in algorithm 1). because only the first d principle components are needed, we reserve the first d columns of v (i.e. the first d eigenvectors), which we denote as weight matrix w t . the transformed embeddings h t i in the new pc space are given by the i-th column of w t ⊤ e (line 5 in algorithm 1). note that pca only requires m sample sentences and calculating their initial embeddings. so during the distillation process, the teacher's transformer parameters θ t and pca weight matrix w t are fixed, while the student's transformer parameters θ s , projection weight w s , and projection bias b s can be updated. we minimize the distance between final embeddings h t i and h s i by taking the mean squared loss: ## experiment we conduct our experiments on standard semantic textual similarity (sts) tasks using the senteval toolkit @xcite for evaluation. we also test mean reciprocal rank (mrr), memory usage, and retrieval speed on semantic retrieval (sr) tasks. all of our experiments are tested on a server with intel i7-5930k cpu @ 3.50ghz, nvidia titan xp gpu, cuda 11.3 and cudnn. ## conclusion and discussion in this paper, we propose an effective method to compress sentence representation using homomorphic projective distillation. we demonstrated that this approach successfully enables small language models to achieve competitive high-quality sentence representations compared with large ones while keeping a small embedding size to optimize the memory storage and retrieval latency in downstream tasks. our results show that knowledge distillation with augmented data improves the student model's capability to cover and understand more complex sentence variances. the learned projection layer with contrastive loss for sentence embedding can outperform other dimension reduction methods. we also try adding whitening transformation on hpd's output and the performance is slightly dropped (appendix b). since we find that smaller dimensions can have slightly better results than larger ones, we will check over the optimal projected layer size to enhance the isotropy of sentence representation distribution for semantic similarity tasks in our future work.