Dimension Reduction for Efficient Dense Retrieval via Conditional Autoencoder (2022)

Abstract: Dense retrievers encode queries and documents and map them in an embedding space using pre-trained language models. These embeddings need to be high-dimensional to fit training signals and guarantee the retrieval effectiveness of dense retrievers. However, these high-dimensional embeddings lead to larger index storage and higher retrieval latency. To reduce the embedding dimensions of dense retrieval, this paper proposes a Conditional Autoencoder (ConAE) to compress the high-dimensional embeddings to maintain the same embedding distribution and better recover the ranking features. Our experiments show that ConAE is effective in compressing embeddings by achieving comparable ranking performance with its teacher model and making the retrieval system more efficient. Our further analyses show that ConAE can alleviate the redundancy of the embeddings of dense retrieval with only one linear layer. All codes of this work are available at https://github.com/NEUIR/ConAE.

## introduction as the first stage of numerous multi-stage ir and nlp tasks @xcite @xcite , dense retrievers @xcite have shown lots of advances in conducting semantic searching and avoiding the vocabulary mismatch problem @xcite . dense retrievers usually encode queries and documents as high-dimensional embeddings, which are necessary to guarantee retrieval effectiveness during training @xcite . nevertheless, high dimensional embeddings usually exhaust the memory to store the index and lead to longer retrieval latency @xcite . the research of building efficient dense retrieval systems has been stimulated recently @xcite . to reduce the dimensions of document embeddings, existing work reserves the principle dimensions or compresses query and document embeddings for building more efficient retrievers @xcite . there are two challenges in compressing embeddings of dense retrievers: the compressed embeddings should share a similar distribution with the original embeddings, making the low-dimensional embedding space uniform and the document embeddings distinguishable; all the compressed embeddings should have the ability to maintain the maximal information for matching related queries and documents during retrieval, which helps better align the related query-document pairs. this paper proposes a conditional autoencoder (conae), which aims to build efficient dense retrieval systems by reducing the embedding dimensions of queries and documents. conae first encodes high-dimensional embeddings into a lowdimensional embedding space and then generates embeddings that can be aligned to related queries or documents in the original embedding space. in addition, conae designs a conditional loss to regulate the low-dimensional embedding space to mimic the embedding distribution of highdimensional embeddings. our experiments show that conae is effective to compress the highdimensional embeddings and avoid redundant ranking features by achieving comparable retrieval performance with vanilla dense retrievers and better visualizing the embedding space with t-sne. ## related work dense retrievers use a bi-encoder architecture to encode queries and documents and map them in an embedding space for retrieval @xcite @xcite @xcite . to learn an effective embedding space, dense retrievers are forced to maintain high-dimensional embeddings to fit training signals. the most direct way to reduce the dimension of embeddings is that retaining parts of the dimensions of high-dimensional embeddings @xcite . some work uses the first 128 dimensions to encode both queries and documents @xcite or utilizes pca to retain the primary dimensions to recover most information from the raw embeddings @xcite . other work @xcite proposes a supervised method, which uses neural networks to compress the high-dimensional embeddings as lowerdimensional ones. these supervised models provide a better dimension reduction way than unsupervised models by avoiding missing too much information. to optimize the encoders, some work @xcite continuously trains dense retrievers with the contrastive training strategies @xcite . ## methodology this section introduces our conditional autoencoder (conae). we first introduce the preliminaries of dense retrieval (sec. 3.1), and then describe the architecture of conae (sec. 3.2). ## experimental methodology this section describes the datasets, evaluation metrics, baselines and implementation details of our experiments. dataset. four datasets are used to evaluate the retrieval effectiveness of different dimension reduction models, including ms marco (passage ranking) @xcite , nq @xcite , trec dl @xcite and trec-covid @xcite . in our experiments, we randomly sample 50,000 queries from the raw training set of ms marco as the development set and use ms marco (dev) as the testing set. the dimension reduction models that are trained on ms marco are also evaluated on two benchmarks, trec dl and trec-covid, aiming to evaluate their generalization ability. all data statistics are shown in evaluation metrics. ndcg@10 is used as the evaluation metric on three benchmarks, ms marco, trec dl and trec-covid. ms marco also uses mrr@10 as the primary evaluation metric @xcite . for the nq dataset, the hit accuracy on top20 and top100 is used as the evaluation metric, which is the same as previous work @xcite . baselines. in our experiments, we compare conae with two baselines from previous work @xcite , principle component analysis (pca) and ce. pca reduces the embedding dimension by retaining the principle dimensions that can keep most of the variance within the original representation. ce model uses two linear layers w q and w d without biases to transform dense representations of queries and documents into lower embeddings @xcite . we also start from ce models and continuously train the whole model to implement our ance models to generate query and document embeddings of different dimensions. implementation details. the rest describes our implementation details. all embedding dimension reduction models base on one of the best dense retrievers ance @xcite and build document index with exact matching (flat index), which is implemented by faiss @xcite . during training conae, we set the hyperparameter λ as 0.1 and search top100 documents using vanilla ance to construct the d top collection for each query. for our ce and ance models, we sample 7 negative documents for each query to contrastively train these models and sample 1 negative document to train conae. in our experiments, we set the batch size to 2 and accumulate step to 8 for ance. the batch size and accumulate step are 128 and 1 for other models. all models are implemented with pytorch and tuned with adam optimizer. the learning rates of ance and other models are set to 2e -6 and 0.001, respectively. ## evaluation result four experiments are conducted in this section to study the effectiveness of conae in reducing embedding dimensions for dense retrieval. ## conclusion this paper presents conae, which reduces the embedding dimension of dense retrievers. our experiments show that conae can achieve comparable retrieval performance with the teacher model, significantly reduce the index storage and accelerate the searching process. our further analyses show that the high-dimensional embeddings of dense retrievers are usually redundant and conae helps to alleviate such redundancy and visualize the embedding space more intuitively and effectively.