ERATE : Efficient Retrieval Augmented Text Embeddings (2023)

Abstract: Embedding representations of text are useful for downstream natural language processing tasks. Several universal sentence representation methods have been proposed with a particular focus on self-supervised pre-training approaches to leverage the vast quantities of unlabelled data. However, there are two challenges for generating rich embedding representations for a new document. 1) The latest rich embedding generators are based on very large costly transformer-based architectures. 2) The rich embedding representation of a new document is limited to only the information provided without access to any explicit contextual and temporal information that could potentially further enrich the representation. We propose efficient retrieval-augmented text embeddings (ERATE) that tackles the first issue and offers a method to tackle the second issue. To the best of our knowledge, we are the first to incorporate retrieval to general purpose embeddings as a new paradigm, which we apply to the semantic similarity tasks of SentEval. Despite not reaching state-of-the-art performance, ERATE offers key insights that encourages future work into investigating the potential of retrieval-based embeddings.

## introduction state-of-the-art sentence embedding models @xcite have competed against one another to approach human-like performance in several nlp tasks. despite the gains observed in performance of sentence embeddings on public benchmarks such as senteval @xcite , the progress has come at a large computational expense. for example, the largest model amongst the sentence-t5 series consists of up to billions of parameters while gpt-3 * work done during internship. we propose a new paradigm that aims to maintain the benefits of high-complexity rich embedding models at reduced computational requirements. our novel paradigm investigates whether retrieval can be used to bypass the compute intensive embedding model in a similar manner to the application of retrieval for generation @xcite tasks for real world large scale use cases with latency and compute constraints. we propose to use a lightweight retrieval model combined with rich pre-computed representations, in order to approximate the richer representations of a large embedding model. we find retrieval-based embeddings struggle against standard text embedding models but their performance can be improved by aggregating neighbours from different light embedding representations and increasing the size of the datastore of precomputed embeddings. to our knowledge, this is the first attempt to use retrieval approaches for developing general purpose sentence embeddings. our main contributions can be summarised as follows: we hope other researchers will engage in this novel setup to develop more efficient sentence embeddings that will allow high-performing representations to be accessible to a broader community. our work focuses on developing lightweight embeddings that out-compete existing lightweight embeddings but we believe erate can be used for a wider range of applications. specifically, input documents often lack the full contextual information or temporal relevance to generate the necessary high-quality text embedding. erate offers the opportunity for the embedding of a given document to encapsulate information from other similar documents to increase the information content whilst also being more up-to-date with more recent documents added to a datastore. several works have looked at approaches to make less expensive sentence embedding representations. for example, embedding recycling @xcite for language models is proposed as a reduced compute approach for downstream tasks. this involves caching activations from intermediate layers in large pre-trained models such that when similar inputs are seen during inference time, the cached output can be used in order to skip a part of the model structure. embedding recycling has been demonstrated to out-compete distilled models, such as distilbert @xcite . in contrast, we investigate whether fixed em-bedding representations can be generated more efficiently using retrieval without any additional training, relying only on pre-computed embeddings. other works have investigated efficient methods for retrieval from a large set of documents such as colbert @xcite and plaid @xcite interaction models that use offline encoding of documents. rather than making the retrieval step more efficient, our work focuses on using retrieval as a tool for enhancing the development of general purpose embeddings. text generation and language modelling has seen several works involving performance boost with retrieval. @xcite investigates extending a pre-trained language model by including the k-nearest neighbours, which @xcite applies to questionanswering. similarly, @xcite introduced retrieval-augmented generation (rag) models where a pre-trained retriever and a pre-trained sequence-to-sequence model are fine-tuned end-toend. @xcite released retro as a successor of realm @xcite where an autoregressive language work is retrieval-enhanced by making the training documents explicitly available at inference time. finally, @xcite present atlas for retrieval-enhanced language modelling where the sequence-to-sequence model takes the retrieved documents and the query to generate the output text for knowledge-intensive tasks. we probe whether retrieval-incorporated approaches can bring similar benefits for the development of fixed embedding representations, not end-to-end sequence-to-sequence models. ## retrieval for text embeddings this section explains how efficient retrieval augmented text embeddings are developed. the main idea is that a query document only needs to be embedded using a light embedder and by outlining the nearest neighbours in the light space, the corresponding pre-computed embeddings can be combined to generate the rich query embedding. let d denote a new document, for which we want to determine the rich embedding representation, x. let f light (•) and f rich (•) be embedding generators that map a given document to the light and rich embedding spaces respectively: note, we assume that the operation f rich (d) is prohibitively compute intensive while ĥ = f light ( d) is feasible. instead, there exists a set of documents {d} k @xmath0, {x} d @xmath1, have been pre-computed. let g ranker (•) denote a retrieval system that ranks all embeddings (with pairwise cosine distance) in a set based upon a query embedding. hence, the ranks are: the final rich embedding can then be calculated as a combination of the rich embedding representations of the top r documents: the process is depicted in the pipeline of figure 1 . alternative approaches can be considered for the combination process of equation 3 1 . ## ablations this section presents three ablations: (1) using an alternative light embedder; (2) an attempt to align the light and rich embedding spaces; (3) distillation of a rich embedder onto a light embedder. erate relies on combining the rich embeddings of the neighbours identified from a light embedding space. the projection model consists of an input layer followed by a relu followed by a single hidden layer that predicts an embedding in the target embedding space with a cosine embedding loss. the vanilla datastore embeddings are used as the training data with 10% of the data cut-out for validation. a distillation inspired approach is considered where a light embedding model aims to mimic the embeddings of the rich sentence-t5 model as an alternative strategy to erate. distilbert is selected as the light model foot_5 . for every datastore embedding, the light model is finetuned (all parameters) to predict the output embedding from the rich model. the distilled model achieves an average score on the sts tasks of 45.6% which is lower than the light model from ## conclusions retrieval-based embeddings are proposed as er-ate that bypass inference through an expensive embedding generation model but hope to leverage its richness. however, the current set-up for er-ate achieves subpar performance on text similarity tasks with some gains observed from combining neighbours of a unique dropout mask approach and extending the datastore size of pre-computed light and rich embeddings for retrieval. we highlight multiple areas of future work. future work should investigate erate-based approaches in a hybrid setting: erate embeddings are used for sentences where they are likely to work effectively (neighbours are in a dense space allowing accurate approximations) while the default expensive embedder can be used when er-ate is unlikely to be effective. erate can be increasingly effective when only partial information is available in a query for which an embedding is desired as combining the embeddings of neighbouring documents can enrich the information content. however, sentence-level embeddings offer little opportunity to explore the gains by additional information and hence future work should investigate the scope of erate at the document-level; mteb @xcite potentially offers suitable tasks. we should also investigate alternative approaches for aligning the light and rich spaces and better combining neighbours' embeddings e.g. self-attention.