What Are You Token About? Dense Retrieval as Distributions Over the Vocabulary (2023)

Abstract: Dual encoders are now the dominant architecture for dense retrieval. Yet, we have little understanding of how they represent text, and why this leads to good performance. In this work, we shed light on this question via distributions over the vocabulary. We propose to interpret the vector representations produced by dual encoders by projecting them into the model’s vocabulary space. We show that the resulting projections contain rich semantic information, and draw connection between them and sparse retrieval. We find that this view can offer an explanation for some of the failure cases of dense retrievers. For example, we observe that the inability of models to handle tail entities is correlated with a tendency of the token distributions to forget some of the tokens of those entities. We leverage this insight and propose a simple way to enrich query and passage representations with lexical information at inference time, and show that this significantly improves performance compared to the original model in zero-shot settings, and specifically on the BEIR benchmark.

## introduction dense retrieval models based on neural text representations have proven very effective @xcite @xcite , improving upon strong traditional sparse models like bm25 @xcite . however, when applied off-theshelf (i.e., in out-of-domain settings) they often experience a severe drop in performance @xcite @xcite . moreover, the reasons for such failures are poorly understood, as the information captured in their representations remains under-investigated. figure 1 : an example of our framework. we run the question "where was michael jack born?" through the question encoder of dpr @xcite , and project the question representation e q to the vocabulary space using bert's masked language modeling head @xcite . the result is a distribution over the vocabulary, q. we apply the same procedure for passages as well. these projections enable reasoning about and improving retrieval representations. in this work, we present a new approach for interpreting and reasoning about dense retrievers, through distributions induced by their query foot_0 and passage representations when projected to the vocabulary space, namely distributions over their vocabulary space (figure 1 ). such distributions enable a better understanding of the representational nature of dense models and their failures, which paves the way to simple solutions that improve their performance. ## background in this work, we suggest a simple framework for interpreting dense retrieves, via projecting their representations to the vocabulary space. this is done using the (masked) language modeling head of their corresponding pretrained model. we begin by providing the relevant background. ## vocabulary projections we now describe our framework for projecting query and passage representations of dense retrievers to the vocabulary space. given a dense retrieval model, we utilize the mlm head of the model it was initialized from to map from encoder output representations to distributions over the vocabulary (eq. 1). for example, for dpr @xcite we take bert's mlm head, as dpr was initialized from bert. given a query q, we use the query encoder enc q to obtain its representation e q as in eq. 2. similarly, for a passage p we apply the passage encoder enc p to get e p . we then apply the mlm head as in eq. ( 1 ) to obtain the vocabulary projection: note that it is not clear a-priori that q and p will be meaningful in any way, as the encoder model has been changed since pretraining, while the mlmhead function remains fixed. moreover, the mlm function has not been trained to decode "pooled" sequence-level representations (i.e., the results of cls or mean pooling) during pretraining. despite this intuition, in this work we argue that p and q are actually highly intuitive and can facilitate a better understanding of dense retrievers. ## experiment setup to evaluate our framework and method quantitatively, we consider several dense retrieval models and datasets. ## analyzing dense retrievers via vocabulary projections in section 3, we introduce a new framework for interpreting representations produced by dense retrievers. next, we describe empirical findings that shed new light on what is encoded in these representations. via vocabulary projections, we draw connections between dense retrieval and well-known concepts from sparse retrieval like lexical overlap ( §5.1), query prediction ( §5.2) and query expansion ( §5.3). ## token amnesia the analysis in section 5 shows that vocabulary projections of passages (i.e., p ) predict which of the input tokens are likely to appear in relevant questions. however, in some cases these predictions utterly fail. for example, in figure 2 the token "michael" is missing from the top-k of the passage projection p . we refer to such cases as token amnesia. here we ask, do these failure in query prediction hurt retrieval? next, we demonstrate that token amnesia indeed correlates with well-known failures of dense retrievers ( §6.1). to overcome this issue, we suggest a lexical enrichment procedure for dense representations ( §6.2) and demonstrate its effectiveness on downstream retrieval performance ( §6.3). ## related work projecting representations and model parameters to the vocabulary space has been studied previously mainly in the context of language models. the approach was initially explored by nostalgebraist (2020). @xcite showed that feedforward layers in transformers can be regarded as key-value memories, where the value vectors induce distributions over the vocabulary. @xcite view the token representations themselves as inducing such distributions, with feed-forward layers "updating" them. @xcite suggest to project all transformer parameters to the vocabulary space. dense retrieval models, however, do not have any language modeling objective during fine-tuning, yet we show that their representations can still be projected to the vocabulary. despite the wide success of dense retrievers recently, interpreting their representations remains under-explored. @xcite analyze neural retrieval models (not only dense retrievers) via diagnostic probes, testing characteristics like sensitivity to paraphrases, styles and factuality. @xcite decode the query representations of neural retrievers using a t5 decoder, and show how to "move" in representation space to decode better queries for retrieval. language models (and specifically mlms) have been used for sparse retrieval in the context of termweighting and lexical expansion. for example, @xcite and @xcite learn such functions over bert's vocabulary space. we differ by showing that dense retrievers implicitly operate in that space as well. thus, these approaches may prove effective for dense models as well. while we focus in this work on dense retrievers based on encoder-only models, our framework is easily extendable for retrievers based on autoregressive decoder-only (i.e., left-to-right) models like gpt @xcite , e.g., @xcite and @xcite . ## conclusion in this work, we explore projecting query and passage representations obtained by dense retrieval to the vocabulary space. we show that these projections facilitate a better understanding of the mechanisms underlying dense retrieval, as well as their failures. we also demonstrate how projections can help improve these models. this understanding is likely to help in improving retrievers, as our lexical enrichment approach demonstrates.