NAIL : Lexical Retrieval Indices with Efficient Non-Autoregressive Decoders (2023)

Abstract: Neural document rerankers are extremely effective in terms of accuracy. However, the best models require dedicated hardware for serving, which is costly and often not feasible. To avoid this servingtime requirement, we present a method of capturing up to 86% of the gains of a Transformer cross-attention model with a lexicalized scoring function that only requires 10 ^{-6} % of the Transformer’s FLOPs per document and can be served using commodity CPUs. When combined with a BM25 retriever, this approach matches the quality of a state-of-the art dual encoder retriever, that still requires an accelerator for query encoding. We introduce nail (Non-Autoregressive Indexing with Language models) as a model architecture that is compatible with recent encoder-decoder and decoder-only large language models, such as T5, GPT-3 and PaLM. This model architecture can leverage existing pre-trained checkpoints and can be fine-tuned for efficiently constructing document representations that do not require neural processing of queries.

## introduction we attempt to answer the following question: to what extent can the computationally-intensive inference in modern neural retrieval systems be pushed entirely to indexing time? neural networks have revolutionized information retrieval, both with powerful reranking models that cross-attend to query and document, and with dual-encoder models that map queries and documents to a shared vector space, leveraging approximate nearest neighbor search for top-k retrieval. the strongest systems typically use a dual-encoder for retrieval followed by a cross-attention reranker to improve the ordering. however, both these components tend to be built on increasingly large transformers @xcite @xcite and thus rely on dedicated accelerators to process queries quickly at serving time. in many application settings, this may be impractical or costly, and as we will show, potentially unnecessary. in particular, we explore a retrieval paradigm where documents are indexed by predicted query token scores. as a result, scoring a query-document pair (q, d) simply involves looking up the scores for the tokens in q associated with d in the index. while the scores are predicted by a neural network, the lookup itself involves no neural network inference so can be faster than other approaches. however, this also means that there can be no crossattention between a specific query and document or even a globally learned semantic vector space. given these shortcomings, it is unclear that such a model, which offloads all neural network computation to indexing time, can be a practical alternative to its more expensive neural counterparts. in addition, while large pre-trained language models have been shown to generalize well over a number of language and retrieval tasks @xcite @xcite @xcite , a key challenge is that they have universally adopted a sequence-to-sequence architecture which is not obviously compatible with precomputing query scores. naive approaches are either computationally infeasible (scoring all possible queries), or rely on sampling a small, incomplete set of samples (such as in @xcite . to overcome this challenge, we introduce a novel use of non-autoregressive decoder architecture that is compatible with existing transfomerbased language models (whether encoder-decoder or decoder-only, @xcite . it allows the model, in a single decode step, to score all vocabulary items in parallel. this makes document indexing with our model approximately as expensive as indexing with document encoders used in recent dual-encoder retrieval systems @xcite @xcite . we call the retrieval system based on this proposed system cross-attention enc. finally, by using bm25 as a retriever and nail as a reranker, we can match state-of-the-art dualencoders @xcite with 10 -4 % of the inference-time flops. 3. we propose our model as a preferred solution when significant compute is available at indexing time, but not on-demand at serving time, and we provide a cost analysis that illustrates when our approach could be preferred to previous work that harnesses llms. ## related work there has been much work in information retrieval leveraging neural networks, which we cannot adequately cover in this paper. for a comprehensive overview, we refer the reader to the survey by hambarde and proenca 2023. here, we focus on methods that minimize the use of expensive neural methods at query inference time (typically methods of sparse retrieval) and on those that leverage llms. lm-based term weighting bag-of-words models, such as tf-idf and bm25 @xcite , use term weighting based on corpus statistics to determine relevance of document terms to query terms. our work can be seen as a way to construct document term weights that are both (1) unconditional with respect to the query, and (2) indexed using lexicalized features (specifically, we use a vector of token scores). as a result, this type of document representation can be precomputed (at indexing time) and does not require expensive computation at query-time. prior work on leveraging language models to produce such lexicalized term weighting can be roughly divided into two groups: those with just document-side encoders, and those with query-side and documentside encoders. examples of the first group include deepct @xcite , deeptr @xcite , and deepimpact @xcite , tilde v2 @xcite , and spladedoc @xcite . these systems are examples of the model paradigm we are exploring, in which all neural network computation happens at indexing time. our work can be seen as an attempt to update these systems (which use word2vec embeddings or encoder-only language models) to modern encoder-decoder architectures. splade-doc is the most recent (and performant) of these, so is in many cases the most useful point of comparison for our work. we include results for the best version of splade-doc @xcite . examples of the second group include sparta @xcite , colbert @xcite , colbert v2 @xcite , coil @xcite , splade @xcite , and splade v2 @xcite . these sparse dual-encoders have proven themselves competitive with dense dual-encoders, and have some advantages like improved interpretability. we demonstrate comparable performance without the need for any query-side encoder. ## nail model a major goal of this work is to investigate retrieval methods that forego neural computation and the need for specialized accelerator hardware at query time. as such, we focus on a method that uses a large neural model to precompute the required representations of the retrieval items (documents) ahead of time. then, at retrieval time, the method performs only basic featurization (e.g., tokenization) of the queries. specifically, we investigate query-document scoring functions that score the compatibility of a query-document pair with the inner-product of separate featurizations of the query φ q (q) and document this form is familiar from both traditional lexicalized retrieval and from more recent work on dense retrieval. in lexicalized retrieval, (e.g., tf-idf and bm25) @xcite , φ q and φ d assign non-zero scores to sub-strings of q and d. on the other hand, in dense retrieval @xcite @xcite , φ q and φ d are neural networks that map q and d to dense vectors. note that this formulation does not allow for deeper interactions between d and q, such as cross-encoder scorers, as these cannot be computed efficiently and without an accelerator at query time. we investigate an alternative formulation of equation 1 than either traditional lexicalized retrieval or dense retrieval. in this formulation, φ d can be an arbitrarily complex neural network, but φ q must be a sparse featurization that can be quickly computed on commodity cpus. this way, it is possible to push all costly neural network inference to indexing time, and avoid the need for accelerators at serving-time. for this paper, we choose φ q to be a simple tokenizer, but we believe that our results could also extend to more complex sparse featurizations. ## model training and experiments to train the nail model, we have empirically found it beneficial to perform two stages of training (1) a pre-training stage the uses self-supervised tasks over a large, unlabeled text corpus, and (2) a finetuning stage that relies on question-answering data via explicit hard negatives. we present the details of each of the training steps in sections 4.1 and 4.2. our model is implemented within the t5x framework @xcite and we initialize model weights with published t5.1.1 checkpoints @xcite . unless otherwise noted, the nail model size used in the experiments is xl, with roughly 3 billion parameters. we saw no further gains from increasing parameters further. to be compatible with t5 checkpoints, we also adopt the t5 vocabulary and attendant senten-cepiece tokenizer @xcite . the vocabulary consists of 32,000 tokens extracted from a english-focused split of common crawl. ## experimental evaluation we compare nail to other systems that have published results on beir. to compare with some sparse systems that have not been evaluated on beir datasets, we also make of use the ms-marco passage ranking task. we focus on answering the following questions: ## concluding remarks we introduce a new model for sparse, lexicalized retrieval, called nail that adapts expensive pretrained sequence-to-sequence language models for document indexing. the main elements of nail are (1) the use of a non-autoregressive decoder, (2) the use of vocabulary based representation for documents and queries, (3) a self-supervised training approach that is critical for good performance. with nail, we focus on offloading all neural computation to indexing time, allowing serving to operate cheaply and without the use of accelerators. evaluating retrieval on beir, we show that the nail approach is as effective as recent dualencoder systems and captures up to 86% of the performance gains of a cross-attention model on ms-marco while being able to serve requests on commodity cpus.