SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Weiwei Sun $^{1}$ , Lingyong Yan $^{2}$ , Zheng Chen $^{1}$ , Shuaiqiang Wang $^{2}$ , Haichao Zhu $^{2}$ , Pengjie Ren $^{1}$ , Zhumin Chen $^{1}$ , Dawei Yin $^{2}$ , Maarten de Rijke $^{3}$ , Zhaochun Ren $^{4*}$ $^{1}$ Shandong University, China  $^{2}$ Baidu Inc., China
<sup>3</sup>University of Amsterdam, The Netherlands <sup>4</sup>Leiden University, The Netherlands {sunnweiwei,lingyongy}@gmail.com yindawei@acm.org m.derijke@uva.nl z. ren@liacs.leidenuniv.nl
As a new paradigm in information retrieval, generative retrieval directly generates a ranked list of document identifiers (docids) for a given query using generative language models (LMs). How to assign each document a unique docid (denoted as document tokenization) is a critical problem, because it determines whether the generative retrieval model can precisely retrieve any document by simply decoding its docid. Most existing methods adopt rule-based tokenization, which is ad-hoc and does not generalize well. In contrast, in this paper we propose a novel document tokenization learning method, GENRET, which learns to encode the complete document semantics into docids. GENRET learns to tokenize documents into short discrete representations (i.e., docids) via a discrete auto-encoding approach. We develop a progressive training scheme to capture the autoregressive nature of docids and diverse clustering techniques to stabilize the training process. Based on the semantic-embedded docids of any set of documents, the generative retrieval model can learn to generate the most relevant docid only according to the docids' semantic relevance to the queries. We conduct experiments on the NQ320K, MS MARCO, and BEIR datasets. GENRET establishes the new state-of-the-art on the NQ320K dataset. Compared to generative retrieval baselines, GENRET can achieve significant improvements on unseen documents. Moreover, GENRET can also outperform comparable baselines on MS MARCO and BEIR, demonstrating the method's generalizability.
Document retrieval plays an essential role in web search applications and various downstream knowledge-intensive tasks by identifying relevant documents to satisfy users' queries. Recently, generative retrieval has emerged as a new paradigm for document retrieval [1, 5, 37, 41, 46, 47] that directly generates a ranked list of document identifiers (docids) for a given query using generative language models (LMs). Unlike dense retrieval [9, 13, 23, 42], generative retrieval presents an end-to-end solution for document retrieval tasks [37]. It also offers a promising approach to better exploit the capabilities of recent large LMs [1, 41].
As shown in Figure 1, document tokenization aims to tokenize each document in the corpus as a sequence of discrete characters, i.e., docids. Document tokenization plays a crucial role in generative retrieval, as it defines how the document is distributed in the semantic space [37]. And it is still an open problem how to define docids. Most previous generative methods tend to employ rule-based document tokenizers, such as generating titles or URLs [5, 46], or clustering results from off-the-shelf document embeddings [37, 41]. Such rule-based methods are usually ad-hoc and do not generalize well. In particular, the tokenization results potentially perform well on retrieving documents that have been seen during training, but generalize poorly to unlabeled documents [17, 20].
Figure 1: An overview of our proposed method. The proposed method utilizes a document tokenization model to convert a given document into a sequence of discrete tokens, referred to as a docid. This tokenization process allows for the reconstruction of the original document through a reconstruction model. Subsequently, an autoregressive generation model is employed to retrieve documents through the generation of their respective docids.
To address the above problem, we propose GENRET, a document tokenization learning framework that learns to tokenize a document into semantic docids in a discrete auto-encoding scheme. GEN-RET consists of a shared sequence-to-sequence-based document tokenization model, a generative retrieval model, and a reconstruction model. In the proposed auto-encoding learning scheme, the tokenization model learns to convert documents to discrete docids, which are subsequently utilized by the reconstruction model to reconstruct the original document. The generative retrieval model is trained to generate docids in an autoregressive manner for a given query. The above three models are optimized in an end-to-end fashion to achieve seamless integration.
There are usually two challenges when using auto-encoding to optimize a generative retrieval model: (i) docids with an autoregressive nature, and (ii) docids with diversity. To address the first challenge and also to stabilize the training of GENRET, we devise a progressive training scheme. This training scheme allows for a stable training of the model by fixing optimized prefix docids  $z_{<t}$ . To optimize the docids at each step, three proposed losses are utilized: (i) a reconstruction loss for predicting the document using the generated docid, (ii) a commitment loss for committing the docid and avoiding forgetting, and (iii) a retrieval loss for optimizing the retrieval performance end-to-end. To address the second challenge, we propose a parameter initialization strategy and a re-assignment of the docid based on a diverse clustering technique to increase the diversity of the generated docids.
We conduct extensive experiments on three well-known document retrieval benchmark datasets, NQ320K [15, 37], MS MARCO [4, 46], and BEIR [38]. GENRET attains superior retrieval performance against state-of-the-art generative retrieval models on NQ320K. GENRET achieves  $+14\%$  relative improvements on the unseen test set of NQ320K over the best generative retrieval baseline. Experiments on MS MARCO and six BEIR datasets also show that GENRET outperforms existing generative methods and achieves competitive results compared to popular dense retrieval models. Experiments on retrieving new documents, analytical experiments, and an efficiency analysis confirm the effectiveness of the proposed model.
We summarize our contributions as follows: (i) We propose GENRET, a generative retrieval model that represents documents as discrete semantic docids. To the best of our knowledge, this is the first tokenization learning method for document retrieval. (ii) We propose an auto-encoding approach, where the docids generated by our tokenization model are reconstructed by a reconstruction model to ensure the docids capture the semantic information of the document. (iii) We devise a progressive training scheme to model the autoregressive nature of docids and stabilize the training process. (iv) Experimental results demonstrate that GENRET achieves significant improvements, especially on unseen documents, over generative retrieval baselines.2
The document retrieval task can be formalized as the process of retrieving a relevant document  $d$  for a search query  $q$  from a collection of documents  $\mathcal{D}$ . Each document  $d\in \mathcal{D}$  is assumed to be a plain text consisting of a sequence of tokens, denoted as  $d = \{d_{1},\dots ,d_{|d|}\}$ , where  $|d|$  is the total number of tokens in the document. For generative retrieval models, it is usually challenging and computationally
2The code of this work is available at www.github.com/sunnweiwei/GenRet.
inefficient to directly generate original documents of typically long length. Therefore, most existing approaches rely on a technique named document tokenization, which represents a document  $d = \{d_1,\ldots ,d_{|d|}\}$  as a shorter sequence of discrete tokens (docid)  $z = \{z_{1},\dots,z_{t},\dots,z_{M}\}$ , where each token  $z_{t}$  is as a  $K$ -way categorical variable, with  $z_{t}\in [1,2,\ldots ,K]$ , and  $M$  is the length of the docid. See Figure 1 for an example of document tokenization with  $M = 3$  and  $K = 64$ .
As an alternative sequence of the original document, the tokenized docid  $z$  should satisfy the following two properties: (i) different documents have short but different docids; and (ii) docids capture the semantics of their associated documents as much as possible [37]. Because  $z$  is a sequence of a fixed length and usually shorter than the original document  $d$ , the model's training and inference can be simplified and more efficient. This paper employs a tokenization model  $Q \colon d \to z$  to map  $d$  to docid  $z$ . More details about  $Q$  are provided in Section 3.1. After tokenizing each document to docid  $z$ , a generative retrieval model  $P \colon q \to z$  learns to retrieve relevant documents by generating a query  $q$  to a docid  $z$  autoregressively [37].
Conventionally, document tokenization is done by a fixed pre-processing step, such as using the title of a document or the results of hierarchical clustering obtained from BERT [5, 37]. However, it has been observed that such ad-hoc document tokenization methods often fail to capture the complete semantics of a document. For example, the title of a web page often does not exist or has low relevance to the content of the web page, and the use of clustering-based docids arbitrarily defines the document in discrete space.
In this paper, we propose GENRET, a novel tokenization learning method based on discrete auto-encoding, to learn semantic docid in a fully end-to-end manner. Figure 1 gives an overview of the proposed method. The proposed GENRET comprises three main components: (i) a sequence-to-sequence based retrieval model  $P(z \mid q)$ , (ii) a document tokenization model  $Q(z \mid d)$ , and (iii) a reconstruction model  $R(d \mid z)$ . The document tokenization model tokenizes a document  $d$  into unique discrete variables  $z$ , and the retrieval model is trained to generate the latent variables  $z$  for a given query  $q$ . In addition, the reconstruction model is used to re-generate the original document from  $z$  to ensure  $z$  captures the semantics of the original document as much as possible.
Following DSI [37], we employ an encoder-decoder Transformer to implement the generative retrieval model. Specifically, given an input text  $d^3$ , the T5-based tokenization model encodes  $d$  and a prefix of docid  $z_{<t}$  and continuously produces latent representation  $\mathbf{d}_t$  of  $d$  at time step  $t$ :
$$
\mathbf {d} _ {t} = \operatorname {D e c o d e r} (\operatorname {E n c o d e r} (d), z _ {<   t}) \in \mathbb {R} ^ {D}, \tag {1}
$$
where  $D$  denotes the hidden size of the model.  $\mathrm{Encoder}(d)$  denotes the output of the Encoder.
Then, the tokenization model generates a token for each document based on  $\mathbf{d}_t$ . At each timestep  $t$ , we define an external embedding matrix named codebook  $\mathbf{E}_t \in \mathbb{R}^{K \times D}$ , where  $K$  is the size of the discrete latent space. There are  $K$  embedding vectors  $\mathbf{e}_{t,j} \in \mathbb{R}^D, j \in [K]$ , and each vector  $\mathbf{e}_{t,j}$  can be regarded as the centroid of a segmentation.
Based on the codebook  $\mathbf{E}_t$ , the discrete latent variable  $z_{t}$  at timestep  $t$  is calculated by a dot-product look-up using the codebook  $\mathbf{E}_t$ :
$$
Q \left(z _ {t} = j \mid z _ {<   t}, d\right) = \operatorname {S o f t m a x} _ {j} \left(\mathbf {d} _ {t} \cdot \mathbf {E} _ {t} ^ {\top}\right), \tag {2}
$$
where  $Q(z_{t} = j \mid z_{<t}, d)$  denotes the probability of tokenizing  $d$  to a particular value  $j \in [K]$  at timestep  $t$ , Softmax $_{j}$  is a softmax function to output the probability of axis  $j$ .
Document reconstruction model. The docid generated by the tokenization model  $Q$  is required to capture the semantic information of the document. To this end, we propose an auto-encoding training scheme, where a reconstruction model  $R \colon z \to d$  that predicts  $d$  using  $z$  is designed to force the tokenization model  $Q \colon d \to z$  to reproduce a docid  $z$  that can be reconstructed back-to-the original document.
3We use document  $d$  here for the denotation, noting that the computation is the same when  $q$  is input.
The input of the reconstruction model is docid  $z$ , and the output is its associated document  $d$ . We first embed  $z$  into representation matrix  $\mathbf{z} = \{\mathbf{z}_1, \dots, \mathbf{z}_M\} \in \mathbb{R}^{M \times D}$  using the codebook of the tokenization model:
$$
\mathbf {z} = \left\{\mathbf {e} _ {1, z _ {1}}, \mathbf {e} _ {2, z _ {2}}, \dots , \mathbf {e} _ {M, z _ {M}} \right\} \in \mathbb {R} ^ {M \times D}, \tag {3}
$$
where each  $t \in [M]$ ,  $\mathbf{z}_t = \mathbf{e}_{t,z_t} \in \mathbb{R}^D$  is the embedding vector of  $z_t$  in the  $t$ -step codebook  $\mathbf{E}_t$ .
We then devise a retrieval-based reconstruction model that predicts the target document  $d$  by retrieving it from document collection  $\mathcal{D}$ , based on the inputs  $\mathbf{z}$ . The relevance score between the input docid  $z$  and the target document  $d$  is defined as follows:
$$
R (d \mid \mathbf {z}) = \prod_ {t = 1} ^ {M} \frac {\exp \left(\mathbf {z} _ {t} \cdot \operatorname {s g} \left(\mathbf {d} _ {t} ^ {\top}\right)\right)}{\sum_ {d ^ {*} \in S \left(z _ {<   t}\right)} \exp \left(\mathbf {z} _ {t} \cdot \operatorname {s g} \left(\mathbf {d} _ {t} ^ {* \top}\right)\right)}, \tag {4}
$$
where  $S(z_{<t})$  is a sub-collection of  $\mathcal{D}$  consisting of documents that have a docid prefix that is the same as  $z_{<t}$ .  $d^{*} \in S(z_{<t})$  represents a document from the sub-collection  $S(z_{<t})$ .  $\mathbf{d}_t$  and  $\mathbf{d}_{t}^{*}$  are continuous representations of documents  $d$  and  $d^{*}$ , respectively, as defined in Eq. 1. The operator  $\mathrm{sg}(\cdot)$  is the stop gradient operator to prevent gradient back propagation. Intuitively,  $R(d \mid \mathbf{z})$  is designed to retrieve a specific document  $d$  from a set of documents  $S(z_{<t})$  at each timestep  $t$ . The set  $S(z_{<t})$  only includes those documents that are assigned the same docid prefix  $z_{<t}$  as the target document  $d$ . By utilizing this loss function, at each step  $t$ , the model is facilitated to learn the residual semantics of the documents not captured by the previous docid  $z_{<t}$ .
3.2 Model optimization
For the document tokenization model  $Q(z \mid d)$ , generative retrieval model  $P(z \mid q)$ , and reconstruction model  $R(d \mid z)$ , jointly optimizing these three models using auto-encoding is challenging due to the following reasons: (i) Learning docids in an autoregressive fashion. On one hand, the prediction of the  $z_{t}$  at time  $t$  relies on previously predicted docids  $z_{<t}$ , which is often under-optimized at the beginning and rapidly changes during training, making it difficult to reach convergence. On the other hand, simultaneously optimizing  $z$  makes it challenging to guarantee a unique docid assignment. Hence, to stabilize the training of GENRET, we devise a progressive training scheme (see Section 3.2.1). (ii) Generating docids with diversity. Optimizing the model using auto-encoding often leads to unbalanced docid assignment: a few major docids are assigned to a large number of documents and most other docids are rarely assigned. Such a sub-optimal distribution of docids affects the model distinguishability, which in turns triggers length increments of docids in order to distinguish conflicting documents. We introduce two diverse clustering techniques to ensure docid diversity (see Section 3.2.2).
3.2.1 Progressive training scheme
To optimize each of the three models listed above in an autoregressive manner, we propose a progressive autoencoding learning scheme, as illustrated in Figure 2. The whole learning scheme contains  $M$  learning steps with respect to the final docid in  $M$ -token. And the docid  $z_{T}$  at step  $T \in [M]$  is learned and optimized at the corresponding learning step. Besides, at each step  $T \in [M]$ , the docid  $z_{T}$  and the model parameters associated with  $z_{T}$  generation are updated, while previously produced docids  $z_{<T}$  and other parameters are kept fixed. By progressively performing the above process, we can finally optimize and learn our models.
Figure 2: Progressive training scheme.  $z_{t}$  (docid at timestep  $t$ ) is optimized at the  $t$ -th training step, while  $z_{<t}$  (docids before timestep  $t$ ) are kept fixed.
At each optimization step, say the  $T$ -step, we devise the learning objective for document tokenization consisting of three loss functions detailed below.
Reconstruction loss. We utilize the reconstruction model  $R(d \mid z)$  as an auxiliary model to learn to optimize the docid generation, whose main goal is capturing as much semantics in the docid as
possible. Therefore, we define a reconstruction loss function of step  $T$  as follows:
$$
\begin{array}{l} \mathcal {L} _ {\operatorname {R e c}} = - \log R (d \mid \hat {\mathbf {z}} _ {\leq T}), \quad \text {w h e r e} \hat {\mathbf {z}} _ {\leq T} = \left\{\mathrm {s g} (\mathbf {z} _ {1}), \dots , \mathrm {s g} (\mathbf {z} _ {T - 1}), \mathbf {z} _ {T} \right\} \in \mathbb {R} ^ {T \times D} \\ \forall t \in [ T ]: \mathbf {z} _ {t} = \mathbf {e} _ {t, j ^ {*}} \in \mathbb {R} ^ {D}, \quad \text {w h e r e} j ^ {*} = \underset {j} {\arg \max } Q \left(z _ {t} = j \mid z _ {<   t}, d\right), \tag {5} \\ \end{array}
$$
where  $\hat{\mathbf{z}}_{\leq T}$  is the first  $T$  representations of the  $z$ , and only the variable  $\mathbf{z}_T$  is optimized in step  $T$ .  $Q(z_{t} = j\mid z_{< t},d)$  is defined in Eq. 2. And the document tokenization model  $Q$  can therefore be optimized when minimizing  $\mathcal{L}_{\mathrm{Rec}}$ .
Of note, since the computation involves a non-differentiable operation (arg  $\max (\cdot)$ ), we apply straight-through gradient estimation to back-propagate the gradient from reconstruction loss to  $\mathbf{d}_T$  following [39, 44], which copies the gradient of  $\mathbf{z}_T$  directly to  $\mathbf{d}_T$ . Specifically, the gradients to document representation  $\mathbf{d}_T$  are defined as  $\frac{\partial\mathcal{L}_{\mathrm{Rec}}}{\partial\mathbf{d}_T} \coloneqq \frac{\partial\mathcal{L}_{\mathrm{Rec}}}{\partial\mathbf{z}_T}$ . And the gradients to the codebook embedding  $\mathbf{e}_{T,j}$  are defined as  $\frac{\partial\mathcal{L}_{\mathrm{Rec}}}{\partial\mathbf{e}_{T,j}} \coloneqq 1_{z_T = j}\frac{\partial\mathcal{L}_{\mathrm{Rec}}}{\partial\mathbf{z}_T}$ .
Commitment loss. In addition, to make sure the predicted docid commits to an embedding and to avoid models forgetting previous docid  $z_{<t}$ , we add a commitment loss as follows:
$$
\mathcal {L} _ {\mathrm {C o m}} = - \sum_ {t = 1} ^ {T} \log Q \left(z _ {t} \mid z _ {<   t}, d\right). \tag {6}
$$
Retrieval loss. For the generative retrieval model  $P$ , we jointly learn it together with the document tokenization model  $Q$ , where  $P$  learns to generate the docids of relevant documents  $d$  given a query  $q$ . Specifically, suppose  $(q, d)$  are a query and relevant document pair; we define the learning objective of retrieval model  $P$  as:
$$
\mathcal {L} _ {\text {R e t}} = - \log \frac {\exp \left(\mathbf {q} _ {T} \cdot \mathbf {d} _ {T}\right)}{\sum_ {d ^ {*} \sim B} \exp \left(\mathbf {q} _ {T} \cdot \mathbf {d} ^ {*} _ {T}\right)} - \sum_ {t = 1} ^ {T} \log P \left(z _ {t} \mid z _ {<   t}, q\right), \tag {7}
$$
where the first term is a ranking-oriented loss enhancing the model using  $(q,d)$  pair;  $d^{*}$  is an in-batch document sampled from the same training mini-batch  $B$ ;  $\mathbf{q}_T$  and  $\mathbf{d}_T$  denote the representation of  $q$  and  $d$  at timestep  $T$ . The second term is the cross-entropy loss for generating docid  $z$  based on  $q$ .
The final loss we use at step-  $T$  is the sum of reconstruction loss, commitment loss, and retrieval loss:
$$
\mathcal {L} = \mathcal {L} _ {\text {R e c}} + \mathcal {L} _ {\text {C o m}} + \mathcal {L} _ {\text {R e t}}. \tag {8}
$$
3.2.2 Diverse clustering techniques
To ensure diversity of generated docids, we adopt two diverse clustering techniques—codebook initialization and docid re-assignment at each progressive training step, where codebook initialization mainly aims to increase the balance of semantic space segmentation, and the docid re-assignment mainly aims to increase the balance of docid assignments.
Codebook initialization. In order to initialize the codebook for our model, we first warm-up the model by passing the continuous representation  $\mathbf{d}_T$  to the reconstruction model instead of the docid representation  $\mathbf{z}_T$  as defined in Eq. 3. During this warm-up phase, we optimize the model using the reconstruction loss  $\mathcal{L}_{\mathrm{Rec}}$  and commitment loss  $\mathcal{L}_{\mathrm{Com}}$ . Next, we collect the continuous representations  $\mathbf{d}_T$  of all documents in  $\mathcal{D}$ , and cluster them into  $K$  groups. The centroids of these clusters are then used as the initialized codebook  $\mathbf{E}_T$ . To balance the initialized docid distribution, we utilize a diverse constrained clustering algorithm, Constrained K-Means, which first normalizes the embeddings of each prefix group, and modifies the cluster assignment step (E in EM) by formulating it as a minimum cost flow (MCF) linear network optimization problem [2].
Docid re-assignment. In order to assign docids to a batch of documents, we modify the dot-product look-up results in Eq. 2 by ensuring that the docid for different documents in the batch are distinct following the method described in [6, 44]. Specifically, let  $\mathbf{D}_t = \{\mathbf{d}_t^{(1)},\dots,\mathbf{d}_t^{(B)}\} \in \mathbb{R}^{B\times D}$  denote the continuous representation of a batch of documents with batch size of  $B$ . The dot-product results are represented by  $\mathbf{H} = \mathbf{D}_t\cdot \mathbf{E}_t^\top \in \mathbb{R}^{B\times K}$ . To obtain distinct docids, we calculate an alternative  $\mathbf{H}^* = \mathrm{Diag}(\mathbf{u})\exp (\frac{\mathbf{H}}{\epsilon})\mathrm{Diag}(\mathbf{v})$ , where  $\mathbf{u}$  and  $\mathbf{v}$  are re-normalization vectors in  $\mathbb{R}^K$  and  $\mathbb{R}^B$ , respectively. The re-normalization vectors are computed via the iterative Sinkhorn-Knopp algorithm [8]. Finally,  $\mathbf{H}^*$  is used instead of  $\mathbf{H}$  in the Softmax and arg max (Eq. 2) operations to obtain the docid  $z_{t}$ .
4.1 Datasets and evalutaion metrics
We conduct experiments on three well-known document retrieval datasets: NQ320K [15], MS MARCO [4], and BEIR [38]. We divide the test set of NQ320K into seen test and unseen test, based on whether the target documents of the query have annotated queries in the training data, to test the generalization ability of the model on unlabeled documents. More details about data pre-processing and data statistics are listed in the Appendix A.
On NQ320K, we use Recall@{1,10,100} and Mean Reciprocal Rank (MRR)@100 as evaluation metrics, following [41]. On MS MARCO, we use Recall@{1, 10, 100} and MRR@10 as evaluation metrics, following [46]. On BEIR, we use nDCG@10 as the main metrics and calculate the average nDCG@10 values across multiple downstream sub-datasets as overall metrics.
We consider three types of baselines: sparse retrieval methods, dense retrieval methods, and generative retrieval methods. The sparse retrieval baselines are: BM25 [32] and DocT5Query [24]. The dense retrieval baselines are: DPR [13], ANCE [42], Sentence-T5 [22], GTR [23], and Contriever [11]. The generative retrieval baselines are: GENRE [5], DSI [37], SEAL [1], CGR-Contra [17], DSI-QG [47], NCI [41], and Ultron [46]. The following three baselines use the same pre-trained LM T5 as GENRET: (i) Sentence-T5 outputs continuous vectors, (ii) GENRE outputs document titles, and (iii) DSI-QG outputs clustering IDs, while GENRET outputs docids learned using the proposed tokenization method. See Appendix B for more details on the other baselines.
4.3 Implementation details