Compact Token Representations with Contextual Quantization for Efficient Document Re-ranking (2022)

Abstract: Transformer based re-ranking models can achieve high search relevance through context- aware soft matching of query tokens with document tokens. To alleviate runtime complexity of such inference, previous work has adopted a late interaction architecture with pre-computed contextual token representations at the cost of a large online storage. This paper proposes contextual quantization of token embeddings by decoupling document-specific and document-independent ranking contributions during codebook-based compression. This allows effective online decompression and embedding composition for better search relevance. This paper presents an evaluation of the above compact token representation model in terms of relevance and space efficiency.

## introduction modern search engines for text documents typically employ multi-stage ranking. the first retrieval stage extracts top candidate documents matching a query from a large search index with a simple ranking method. the second stage or a later stage uses a more complex machine learning algorithm to re-rank top results thoroughly. recently neural reranking techniques from transformer-based architectures have achieved impressive relevance scores for top k document re-ranking, such as @xcite . however, using a transformer-based model to rank or re-rank is extremely expensive during the online inference @xcite . various efforts have been made to reduce its computational complexity (e.g. @xcite ). a noticeable success in time efficiency improvement is accomplished in colbert @xcite which conducts late interaction of query terms and document terms during runtime inference so that token embeddings for documents can be pre-computed. using colbert re-ranking after a sparse retrieval model called @xcite can further enhance relevance. similarly becr @xcite , cedr-knrm @xcite , and prettr @xcite have also adopted the late interaction architecture in their efficient transformer based re-ranking schemes. while the above work delivers good search relevance with late interaction, their improvement in time efficiency has come at the cost of a large storage space in hosting token-based precomputed document embeddings. for example, for the ms marco document corpus, the footprint of embedding vectors in colbert takes up to 1.6tb and hosting them in a disk incurs substantial time cost when many embeddings are fetched for re-ranking. it is highly desirable to reduce embedding footprints and host them in memory as much as possible for fast and high-throughput access and for i/o latency and contention avoidance, especially when an online re-ranking server is required to efficiently process many queries simultaneously. the contribution of this paper is to propose a compact representation for contextual token embeddings of documents called contextual quantization (cq). specifically, we adopt codebook-based quantization to compress embeddings while explicitly decoupling the ranking contributions of document specific and document-independent information in contextual embeddings. these ranking contributions are recovered with weighted composition after quantization decoding during online inference. our cq scheme includes a neural network model that jointly learns context-aware decomposition and quantization with an objective to preserve correct ranking scores and order margins. our evaluation shows that cq can effectively reduce the storage space of contextual representation by about 14 times for the tested datasets with insignificant online embedding recovery overhead and a small relevance degradation for re-ranking passages or documents. ## problem definition and related work the problem of neural text document re-ranking is defined as follows. given a query with multiple terms and a set of candidate documents, rank these documents mainly based on their embeddings and query-document similarity. with a bert-based reranking algorithm, typically a term is represented by a token, and thus in this paper, word "term" is used interchangeably with "token". this paper is focused on minimizing the space cost of token embeddings for fast online re-ranking inference. deep contextual re-ranking models. neural re-ranking has pursued representation-based or interaction-based algorithms @xcite @xcite . embedding interaction based on query and document terms shows an advantage in these studies. the transformer architecture based on bert @xcite has been adopted to re-ranking tasks by using bert's [cls] token representation to summarize query and document interactions @xcite @xcite @xcite . recently bert is integrated in late term interaction @xcite @xcite which delivers strong relevance scores for re-ranking. efficiency optimization for transformer-based re-ranking. several approaches have been proposed to reduce the time complexity of transformerbased ranking. for example, architecture simplification @xcite , late interaction with precomputed token embeddings @xcite , early exiting @xcite , and model distillation @xcite @xcite . we will focus on the compression of token representation following the late-interaction work of colbert @xcite and becr @xcite as they deliver fairly competitive relevance scores for several well-known ad-hoc trec datasets. these late-interaction approaches follow a dual-encoder design that separately encodes the two sets of texts, studied in various nlp tasks @xcite @xcite @xcite . several previous re-ranking model attempted to reduce the space need for contextual token embeddings. colbert has considered an option of using a smaller dimension per vector and limit-ing 2 bytes per number as a scalar quantization. becr @xcite uses lsh for hashingbased contextual embedding compression @xcite . prettr (macavaney et al., 2020) uses a single layer encoder model to reduce the dimensionality of each token embedding. following prettr, a contemporaneous work called sdr in @xcite considers an autoencoder to reduce the dimension of representations, followed by an off-the-shelf scalar quantizer. for the autoencoder, it combines static bert embeddings with contextual embeddings. inspired by this study, our work decomposes contextual embeddings to decouple ranking contributions during vector quantization. unlike sdr, cq jointly learns the codebooks and decomposition for the document-independent and dependent components guided by a ranking loss. vector quantization. vector quantization with codebooks was developed for data compression to assist approximate nearest neighbor search, for example, product quantizer (pq) from @xcite , optimized product quantizer (opq) from @xcite ; residual additive quantizer(rq) from @xcite and local search additive quantizer (lsq) from @xcite . recently such a technique has been used for compressing static word embeddings @xcite and document representation vectors in a dense retrieval scheme called jpq @xcite . none of the previous work has worked on quantization of contextual token vectors for the re-ranking task, and that is the focus of this paper. ## contextual quantization applying vector quantization naively to token embedding compression does not ensure the ranking effectiveness because a quantizer-based compression is not lossless, and critical ranking signals could be lost during data transformation. to achieve a high compression ratio while maintaining the competitiveness in relevance, we consider the ranking contribution of a contextual token embedding for soft matching containing two components: 1) document specific component derived from the self attention among context in a document, 2) document-independent and corpus-specific component generated by the transformer model. since for a reasonable sized document set, the second component is invariant to documents, its storage space is negligible compared to the first component. thus the second part does not need compres-figure 1 : offline processing and online ranking with contextual quantization sion. we focus on compressing the first component using codebooks. this decomposition strategy can reduce the relevance loss due to compression approximation, which allows a more aggressive compression ratio. our integrated vector quantizer with contextual decomposition contains a rankingoriented scheme with an encoder and decoder network for jointly learning codebooks and composition weights. thus, the online composition of decompressed document-dependent information with document-independent information can retain a good relevance. ## concluding remarks our evaluation shows the effectiveness of cq used for colbert in compressing the space of token embeddings with about 14:1 ratio while incurring a small relevance degradation in ms marco passage and document re-ranking tasks. the quantized token-based document embeddings for the tested cases can be hosted in memory for fast and high-throughput access. this is accomplished by a neural network that decomposes ranking contributions of contextual embeddings, and jointly trains context-aware decomposition and quantization with a loss function preserving ranking accuracy. the online time cost to decompress and recover embeddings is insignificant with 1ms for the tested cases. the cq implementation is available at https://github.com/yingruiyang/contextualquantizer. our cq framework is also applicable to the contemporaneous work @xcite . using unicoil scores for the first-stage sparse retrieval and colbertv2+cq (@xmath0, @xmath1,000 passage reranking, we achieve 0.387 mrr@10 on the msmarco passage dev set, 0.746 ndcg@10 on trec dl19, and 0.726 ndcg@10 on dl20 with about 10.2gb embedding space footprint. notice that colbertv2 achieves a higher mrr@10 number 0.397 for the passage dev set when used as a standalone retriever @xcite and dense retrieval with such a multi-vector representation is likely to be much more expensive than retrieval with a sparse representation on a large dataset. the previous work in dense retrieval has often employed faster but approximate search, but that comes with a visible loss of relevance @xcite . thus the above relevance number using colbertv2+cq for re-ranking with uni-coil sparse retrieval is fairly strong, achievable with a reasonable latency and limited computing resource. its embedding space size is 2.8x smaller than the 29gb space cost in the standalone @xcite for ms marco passages. our future work is to investigate the above issue further and study the use of cq in the other late-interaction re-ranking methods.