E mbed T ext N et: Dimension Reduction with Weighted Reconstruction and Correlation Losses for Efficient Text Embedding (2023)

Abstract: The size of embeddings generated by large language models can negatively affect system latency and model size in certain downstream practical applications (e.g. KNN search). In this work, we propose EmbedTextNet, a light add-on network that can be appended to an arbitrary language model to generate a compact embedding without requiring any changes in its architecture or training procedure. Specifically, we use a correlation penalty added to the weighted reconstruction loss that better captures the informative features in the text embeddings, which improves the efficiency of the language models. We evaluated EmbedTextNet on three different downstream tasks: text similarity, language modelling, and text retrieval. Empirical results on diverse benchmark datasets demonstrate the effectiveness and superiority of EmbedTextNet compared to state-of-art methodologies in recent works, especially in extremely low dimensional embedding sizes. The developed code for reproducibility is included in the supplementary material.

## introduction significant advances in computational resources along with the existence of the huge amount of data rendered large language models (llm) ubiquitous in various fields, such as natural language understanding (nlu), and vision-language multimodels @xcite . even with their relative success in diverse applications, the practical deployment of such models remains challenging with one of the main concerns being the text embedding size @xcite . for example, loading word embedding matrices with 2.5 million tokens and 300 dimensions consumes 6 gb memory on a 64-bit system @xcite . this may breach memory or latency budget in certain highly-constrained practical settings, such as nlu on mobile devices or knn search in a large index. text embeddings, in general, are generated from unlabeled text corpora and applied in several natural language processing (nlp) and information retrieval applications. embeddings are often precomputed and used in downstream applications. many post-processing algorithms have been considered over the years to reduce the dimensionality of the resulting vector space. one approach to reduce embedding dimensionality is re-training the language models with the desired dimensionality, this, however, can be costly and reduce model performance. another approach is to use unsupervised algorithms directly on the pre-computed text embeddings, without requiring any access to labels. several works exist in the literature including principal component analysis (pca) jolliffe and cadima (2016), algo @xcite , and uniform manifold approximation and projection (umap) @xcite . while most unsupervised approaches have made noticeable progress, they still suffer from many disadvantages according to our findings. first, they do not perform consistently in different applications. for example, algo works well, compared to pca, for dimensionality reduction (dr) in a word embedding. however, the roles are switched in the case of a multilingual sentence (m-sentence) embedding. second, the absolute performance of the baseline methods is significantly lower when extreme reduction is employed (e.g. 16). in this work, we propose a framework that efficiently reduces the dimensionality of text embedding. we show that the proposed embedtextnet, which uses a variational autoencoder (vae) with modified loss functions, can be used as an add-on module that improves the quality of the text embedding. we also demonstrate that embedtextnet can efficiently reduce the size of a text embedding considerably while preserving the quality of the embedding. our extensive experiments on three downstream applications depict that our method surpasses the original embeddings and state-of-theart (sota). ## related work model compression and dr of embedding are receiving more attention nowadays because of their ability to speed the inference process and reduce the model storage space. to compress a model, one can use pruning approaches that aim to repeatedly eliminate the redundant parameters in a model @xcite . quantization is another common way that works by limiting the number of bits used to represent the model weights @xcite . similarly, authors in @xcite compressed the input and output embedding layers with a parameter-sharing method. another approach for model compression is knowledge distillation, where the objective is to make a small model mimic a large pre-trained model by matching the ground-truth and soft labels from the large model @xcite . recently, researchers considered a block-wise low-rank approximation @xcite which employs statistical features of words to form matrix approximations for embedding and softmax layers. the extension of it was shown in @xcite which is based on word weighting and the k-means clustering methods. authors in @xcite considered parameterizing the embedding layers by tensor train decomposition to compress a model with a negligible drop in the result. all the aforementioned methods have the limitation of either requiring re-training the model to accommodate for the changes performed in the network or compromising the accuracy of the learning model. further, some of these techniques are not studied with pretrained llm which limits the scalability. unsupervised methods for dr of embedding have gained attention because of their label-free approach that transforms a high-dimensional embedding into a reduced one. the most common starting point is to utilize linear transformation approaches such as pca or independent component analysis (ica) @xcite ; hyvärinen (2013). these methods aim to project the original high-dimensional embedding into a lowerdimensional one by either maximizing the variance or the mutual information. another route is to employ nonlinear dr methods including kernel pca (kpca) @xcite , t-distributed stochastic neighbor embedding (t-sne) @xcite , and umap @xcite . the kpca employs a kernel function (e.g. gaussian) to map the original embeddings into a nonlinear space and performs pca on the projected embeddings. t-sne performs a nonlinear reduction method to capture the local structure of the high-dimensional embeddings while maintaining a global structure simultaneously. umap embeds data points in a nonlinear fuzzy topological representation by neighbor graphs and then, learns a low-dimensional representation while preserving maximum information of this space. the use of pre-trained llm is the backbone for several nlp applications such as text retrieval, question-answer, and semantic similarity between pairs. large pre-trained models such as bert @xcite , roberta @xcite and @xcite have shown to provide sota performance. however, when we necessitate a smaller embedding size for resource-constrained devices and prompt inference in downstream tasks like text retrieval, it becomes crucial to retrain pretrained llms and fine-tune them to attain the same level of performance as the original model. using dr, we can circumvent this process while also saving space that would otherwise be reserved for storing retrieval models (i.e. index). one simple approach to reduce the embedding size of these large models is to re-train the model while changing the last hidden layer size. this is inefficient since the re-training process will still, obviously, require enormous resources. a more natural and convenient mechanism is to work directly on the pre-trained large embedding which will eliminate the need for re-training @xcite . recently, the authors in @xcite combined pca with a post-processing @xcite to obtain effective dr on word embeddings. however, existing methods fail to work in extreme cases such as a 30-times reduction of the original size. therefore, in this work, we propose an algorithm that works directly on pre-computed embeddings to reduce their size even in extreme cases. this is validated on three different applications which are text similarity, language modelling and text retrieval to highlight the effectiveness of the proposed approach in different aspects such as similarity score, inference time and model size. ## dr for text embedding ensuring a fast inference (either between devices or user-device interaction) is critical, especially for applications related to nlu such as voice assistants. dr has practical applications in two distinct tasks: (1) direct application using pre-computed embeddings like glove and (2) retrieval task with a builtin index where we use pre-computed embeddings to generate an index for searching. these embeddings can be either contextual or non-contextual, depending on the specific application. to this end, we first introduce embedtextnet with modified reconstruction and additional correlation losses. we then show the three most common dr methods as the baselines. ## results and discussion we showcase the dr usefulness of embedtextnet on three different applications: text similarity for measuring the mimicry of reduced embedding with regards to the original one, language modelling for understanding the reproducibility of original data from dr and text retrieval for finding the effectiveness of dr in terms of frugality and performances. we also conduct an ablation study for the hyperparameters and structure of embedtextnet. finally, the computational cost for embedtextnet is detailed in appendix a. in general, during training, embedtextnet takes a longer time compared to pca and algo, and comparable time to umap. yet, at inference, it has a similar inference time with superior performance. this emphasizes the capabilities and potential of embedtextnet for dr. ## limitations there are three predictable limitations in the developed embedtextnet. first, while we have performed a thorough evaluation of embedtextnet on various downstream tasks, it is still a generalpurpose approach and its effectiveness on specific tasks or in specific domains may vary. thus, further research is needed to fully understand its capabilities and limitations in different contexts. second, as mentioned, embedtextnet is most suitable for scenarios where the embedding is saved during inference, such as text retrieval or similarity measurement when the fixed embedding is saved with a vocabulary (e.g. glove). however, it may not be as effective in scenarios where the embedding needs to be decoded back to its original form, such as text generation. third, the effectiveness of embedtextnet is evident on a large embedding dimension, and it may decrease when working with a small embedding dimension even if it was still better than other sota in our experiments (e.g. glove-50d → 10d in