SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shashank Rajput*
University of Wisconsin-Madison	Nikhil Mehta*
Google DeepMind	Anima Singh
Google DeepMind
Raghunandan Keshavan
Google	Trung Vu
Google	Lukasz Heldt
Google	Lichan Hong
Google DeepMind
Yi Tay
Google DeepMind	Vinh Q. Tran
Google	Jonah Samost
Google	Maciej Kula
Google DeepMind
Ed H. Chi
Google DeepMind	Maheswaran Sathiamoorthy
Google DeepMind
Modern recommender systems perform large-scale retrieval by embedding queries and item candidates in the same unified space, followed by approximate nearest neighbor search to select top candidates given a query embedding. In this paper, we propose a novel generative retrieval approach, where the retrieval model autoregressively decodes the identifiers of the target candidates. To that end, we create semantically meaningful tuple of codewords to serve as a Semantic ID for each item. Given Semantic IDs for items in a user session, a Transformer-based sequence-to-sequence model is trained to predict the Semantic ID of the next item that the user will interact with. We show that recommender systems trained with the proposed paradigm significantly outperform the current SOTA models on various datasets. In addition, we show that incorporating Semantic IDs into the sequence-to-sequence model enhances its ability to generalize, as evidenced by the improved retrieval performance observed for items with no prior interaction history.
Recommender systems help users discover content of interest and are ubiquitous in various recommendation domains such as videos [4, 43, 9], apps [3], products [6, 8], and music [18, 19]. Modern recommender systems adopt a retrieve-and-rank strategy, where a set of viable candidates are selected in the retrieval stage, which are then ranked using a ranker model. Since the ranker model works only on the candidates it receives, it is desired that the retrieval stage emits highly relevant candidates.
Figure 1: Overview of the Transformer Index for Generative Recommenders (TIGER) framework. With TIGER, sequential recommendation is expressed as a generative retrieval task by representing each item as a tuple of discrete semantic tokens.
* Equal contribution. Work done when SR was at Google. Correspondence to rajput.shashank11@gmail.com, nikhilmehta@google.com, nlogn@google.com.
There are standard and well-established methods for building retrieval models. Matrix factorization [19] learns query and candidate embeddings in the same space. In order to better capture the non-linearities in the data, dual-encoder architectures [39] (i.e., one tower for the query and another for the candidate) employing inner-product to embed the query and candidate embeddings in the same space have become popular in recent years. To use these models during inference, an index that stores the embeddings for all items is created using the candidate tower. For a given query, its embedding is obtained using the query tower, and an Approximate Nearest Neighbors (ANN) algorithm is used for retrieval. In recent years, the dual encoders architectures have also been extended for sequential recommendations [11, 24, 41, 17, 32, 6, 44] that explicitly take into account the order of user-item interactions.
We propose a new paradigm of building generative retrieval models for sequential recommendation. Instead of traditional query-candidate matching approaches, our method uses an end-to-end generative model that predicts the candidate IDs directly. We propose to leverage the Transformer [36] memory (parameters) as an end-to-end index for retrieval in recommendation systems, reminiscent of Tay et al. [34] that used Transformer memory for document retrieval. We refer to our method as Transformer Index for Generative Recommenders (TIGER). A high-level overview of TIGER is shown in Figure 1. TIGER is uniquely characterized by a novel semantic representation of items called "Semantic ID" - a sequence of tokens derived from each item's content information. Concretely, given an item's text features, we use a pre-trained text encoder (e.g., SentenceT5 [27]) to generate dense content embeddings. A quantization scheme is then applied on the embedding of an item to form a set of ordered tokens/codewords, which we refer to as the Semantic ID of the item. Ultimately, these Semantic IDs are used to train the Transformer model on the sequential recommendation task.
Representing items as a sequence of semantic tokens has many advantages. Training the transformer memory on semantically meaningful data allows knowledge sharing across similar items. This allows us to dispense away with the atomic and random item Ids that have been previously used [33, 42, 11, 8] as item features in recommendation models. With semantic token representations for items, the model is less prone to the inherent feedback loop [1, 26, 39] in recommendation systems, allowing the model to generalize to newly added items to the corpus. Furthermore, using a sequence of tokens for item representation helps alleviate the challenges associated with the scale of the item corpus; the number of items that can be represented using tokens is the product of the cardinality of each token in the sequence. Typically, the item corpus size can be in the order of billions and learning a unique embedding for each item can be memory-intensive. While random hashing-based techniques [16] can be adopted to reduce the item representation space, in this work, we show that using semantically meaningful tokens for item representation is an appealing alternative. The main contributions of this work are summarized below:
1. We propose TIGER, a novel generative retrieval-based recommendation framework that assigns Semantic IDs to each item, and trains a retrieval model to predict the Semantic ID of an item that a given user may engage with.
2. We show that TIGER outperforms existing SOTA recommender systems on multiple datasets as measured by recall and NDCG metrics.
3. We find that this new paradigm of generative retrieval leads to two additional capabilities in sequential recommender systems: 1. Ability to recommend new and infrequent items, thus improving cold-start recommendations, and 2. Ability to generate diverse recommendations using a tunable parameter.
Paper Overview. In Section 2, we provide a brief literature survey of recommender systems, generative retrieval, and the Semantic ID generation techniques we use in this paper. In Section 3, we explain our proposed framework, and outline the various techniques we use for Semantic ID generation. We present the result of our experiments in Section 4, and conclude the paper in Section 5.
Sequential Recommenders. Using deep sequential models in recommender systems has developed into a rich literature. GRU4REC [11] was the first to use GRU based RNNs for sequential recommendations. Li et al. [24] proposed Neural Attentive Session-based Recommendation (NARM), where an attention mechanism along with a GRU layer is used to track long term intent of the user.
Figure 2: An overview of the modeling approach used in TIGER.
AttRec [41] proposed by Zhang et al. used self-attention mechanism to model the user's intent in the current session, and personalization is ensured by modeling user-item affinity with metric learning. Concurrently, Kang et al. also proposed SASRec [17], which used self-attention similar to decoder-only transformer models. Inspired by the success of masked language modeling in language tasks, BERT4Rec [32] and Transformers4Rec [6] utilize transformer models with masking strategies for sequential recommendation tasks.  $S^3$ -Rec [44] goes beyond just masking by pre-training on four self-supervised tasks to improve data representation. The models described above learn a high-dimensional embedding for each item and perform an ANN in a Maximum Inner Product Search (MIPS) space to predict the next item. In contrast, our proposed technique, TIGER, uses Generative Retrieval to directly predict the Semantic ID of the next item.
P5 [8] fine-tunes a pre-trained large language models for multi-task recommender systems. The P5 model relies on the LLM tokenizer (SentencePiece tokenizer [29]) to generate tokens from randomly assigned item IDs. Whereas, we use Semantic ID representation of items they are learned based on the content information of the items. In our experiments (Table 2), we demonstrate that recommendation systems based on Semantic ID representation of items yield much better results than using random codes.
Semantic IDs. Hou et al. proposed VQ-Rec [12] to generate "codes" (analogous to Semantic IDs) using content information for item representation. However, their focus is on building transferable recommender systems, and do not use the codes in a generative manner for retrieval. While they also use product quantization [15] to generate the codes, we use RQ-VAE to generate Semantic IDs, which leads to hierarchical representation of items (Section 4.2). In a concurrent work to us, Singh et al. [31] show that hierarchical Semantic IDs can be used to replace item IDs for ranking models in large scale recommender systems improves model generalization.
Generative Retrieval. While techniques for learning search indices have been proposed in the past [20], generative retrieval is a recently developed approach for document retrieval, where the task is to return a set of relevant documents from a database. Some examples include GENRE [5], DSI [34], NCI [37], and CGR [22]. A more detailed coverage of the related work is in Appendix A. To the best of our knowledge, we are the first to propose generative retrieval for recommendation systems using Semantic ID representation of items.
Our proposed framework consists of two stages:
1. Semantic ID generation using content features. This involves encoding the item content features to embedding vectors and quantizing the embedding into a tuple of semantic codewords. The resulting tuple of codewords is referred to as the item's Semantic ID.
2. Training a generative recommender system on Semantic IDs. A Transformer model is trained on the sequential recommendation task using sequences of Semantic IDs.
Figure 3: RQ-VAE: In the figure, the vector output by the DNN Encoder, say  $\boldsymbol{r}_0$  (represented by the blue bar), is fed to the quantizer, which works iteratively. First, the closest vector to  $\boldsymbol{r}_0$  is found in the first level codebook. Let this closest vector be  $\boldsymbol{e}_{c_0}$  (represented by the red bar). Then, the residual error is computed as  $\boldsymbol{r}_1 := \boldsymbol{r}_0 - \boldsymbol{e}_{c_0}$ . This is fed into the second level of the quantizer, and the process is repeated: The closest vector to  $\boldsymbol{r}_1$  is found in the second level, say  $\boldsymbol{e}_{c_1}$  (represented by the green bar), and then the second level residual error is computed as  $\boldsymbol{r}_2 = \boldsymbol{r}_1 - \boldsymbol{e}_{c_1}'$ . Then, the process is repeated for a third time on  $\boldsymbol{r}_2$ . The semantic codes are computed as the indices of  $\boldsymbol{e}_{c_0}, \boldsymbol{e}_{c_1}$ , and  $\boldsymbol{e}_{c_2}$  in their respective codebooks. In the example shown in the figure, this results in the code (7, 1, 4).
3.1 Semantic ID Generation
In this section, we describe the Semantic ID generation process for the items in the recommendation corpus. We assume that each item has associated content features that capture useful semantic information (e.g. titles or descriptions or images). Moreover, we assume that we have access to a pre-trained content encoder to generate a semantic embedding  $\pmb{x} \in \mathbb{R}^d$ . For example, general-purpose pre-trained text encoders such as Sentence-T5 [27] and BERT [7] can be used to convert an item's text features to obtain a semantic embedding. The semantic embeddings are then quantized to generate a Semantic ID for each item. Figure 2a gives a high-level overview of the process.
We define a Semantic ID to be a tuple of codewords of length  $m$ . Each codeword in the tuple comes from a different codebook. The number of items that the Semantic IDs can represent uniquely is thus equal to the product of the codebook sizes. While different techniques to generate Semantic IDs result in the IDs having different semantic properties, we want them to at least have the following property: Similar items (items with similar content features or whose semantic embeddings are close) should have overlapping Semantic IDs. For example, an item with Semantic ID (10, 21, 35) should be more similar to one with Semantic ID (10, 21, 40), than an item with ID (10, 23, 32). Next, we discuss the quantization schemes which we use for Semantic ID generation.
RQ-VAE for Semantic IDs. Residual-Quantized Variational AutoEncoder (RQ-VAE) [40] is a multi-level vector quantizer that applies quantization on residuals to generate a tuple of codewords (aka Semantic IDs). The Autoencoder is jointly trained by updating the quantization codebook and the DNN encoder-decoder parameters. Fig. 3 illustrates the process of generating Semantic IDs through residual quantization.
RQ-VAE first encodes the input  $\mathbf{x}$  via an encoder  $\mathcal{E}$  to learn a latent representation  $\mathbf{z} \coloneqq \mathcal{E}(\mathbf{x})$ . At the zero-th level  $(d = 0)$ , the initial residual is simply defined as  $\mathbf{r}_0 \coloneqq \mathbf{z}$ . At each level  $d$ , we have a codebook  $\mathcal{C}_d \coloneqq \{\pmb{e}_k\}_{k=1}^K$ , where  $K$  is the codebook size. Then,  $\mathbf{r}_0$  is quantized by mapping it to the nearest embedding from that level's codebook. The index of the closest embedding  $e_{c_d}$  at  $d = 0$ , i.e.,  $c_0 = \arg \min_i \| \pmb{r}_0 - \pmb{e}_k \|$ , represents the zero-th codeword. For the next level  $d = 1$ , the residual is defined as  $\mathbf{r}_1 \coloneqq \mathbf{r}_0 - e_{c_0}$ . Then, similar to the zero-th level, the code for the first level is computed by finding the embedding in the codebook for the first level which is nearest to  $\mathbf{r}_1$ . This process is repeated recursively  $m$  times to get a tuple of  $m$  codewords that represent the Semantic ID. This recursive approach approximates the input from a coarse-to-fine granularity. Note that we chose to use a separate codebook of size  $K$  for each of the  $m$  levels, instead of using a single,  $mK$ -sized codebook. This was done because the norm of residuals tends to decrease with increasing levels, hence allowing for different granularities for different levels.
Once we have the Semantic ID  $(c_0, \ldots, c_{m-1})$ , a quantized representation of  $z$  is computed as  $\widehat{z} := \sum_{d=0}^{m-1} e_{c_i}$ . Then  $\widehat{z}$  is passed to the decoder, which tries to recreate the input  $x$  using  $\widehat{z}$ . The RQ-VAE loss is defined as  $\mathcal{L}(x) := \mathcal{L}_{\mathrm{recon}} + \mathcal{L}_{\mathrm{rqvae}}$ , where  $\mathcal{L}_{\mathrm{recon}} := \|x - \widehat{x}\|^2$ , and  $\mathcal{L}_{\mathrm{rqvae}} := \sum_{d=0}^{m-1} \|s g[r_i] - e_{c_i}\|^2 + \beta \|r_i - s g[e_{c_i}]\|^2$ . Here  $\widehat{x}$  is the output of the decoder, and  $s g$  is the stop-gradient operation [35]. This loss jointly trains the encoder, decoder, and the codebook.
As proposed in [40], to prevent RQ-VAE from a codebook collapse, where most of the input gets mapped to only a few codebook vectors, we use k-means clustering-based initialization for the codebook. Specifically, we apply the k-means algorithm on the first training batch and use the centroids as initialization.
Other alternatives for quantization. A simple alternative to generating Semantic IDs is to use Locality Sensitive Hashing (LSH). We perform an ablation study in Subsection 4.2 where we find that RQ-VAE indeed works better than LSH. Another option is to use k-means clustering hierarchically [34], but it loses semantic meaning between different clusters [37]. We also tried VQ-VAE, and while it performs similarly to RQ-VAE for generating the candidates during retrieval, it loses the hierarchical nature of the IDs which confers many new capabilities that are discussed in Section 4.3.
Handling Collisions. Depending on the distribution of semantic embeddings, the choice of codebook size, and the length of codewords, semantic collisions can occur (i.e., multiple items can map to the same Semantic ID). To remove the collisions, we append an extra token at the end of the ordered semantic codes to make them unique. For example, if two items share the Semantic ID (12, 24, 52), we append additional tokens to differentiate them, representing the two items as (12, 24, 52, 0) and (12, 24, 52, 1). To detect collisions, we maintain a lookup table that maps Semantic IDs to corresponding items. Note that collision detection and fixing is done only once after the RQ-VAE model is trained. Furthermore, since Semantic IDs are integer tuples, the lookup table is efficient in terms of storage in comparison to high dimensional embeddings.
3.2 Generative Retrieval with Semantic IDs
We construct item sequences for every user by sorting chronologically the items they have interacted with. Then, given a sequence of the form  $(\mathsf{item}_1,\ldots ,\mathsf{item}_n)$ , the recommender system's task is to predict the next item  $\mathsf{item}_{n + 1}$ . We propose a generative approach that directly predicts the Semantic ID of the next item. Formally, let  $(c_{i,0},\dots ,c_{i,m - 1})$  be the  $m$ -length Semantic ID for  $\mathsf{item}_i$ . Then, we convert the item sequence to the sequence  $(c_{1,0},\dots ,c_{1,m - 1},c_{2,0},\dots ,c_{2,m - 1},\dots ,c_{n,0},\dots ,c_{n,m - 1})$ . The sequence-to-sequence model is then trained to predict the Semantic ID of  $\mathsf{item}_{n + 1}$ , which is  $(c_{n + 1,0},\dots ,c_{n + 1,m - 1})$ . Given the generative nature of our framework, it is possible that a generated Semantic ID from the decoder does not match an item in the recommendation corpus. However, as we show in appendix (Fig. 6) the probability of such an event occurring is low. We further discuss how such events can be handled in appendix E.
Datasets. We evaluate the proposed framework on three public real-world benchmarks from the Amazon Product Reviews dataset [10], containing user reviews and item metadata from May 1996 to July 2014. In particular, we use three categories of the Amazon Product Reviews dataset for the sequential recommendation task: "Beauty", "Sports and Outdoors", and "Toys and Games". We discuss the dataset statistics and pre-processing in Appendix C.
Evaluation Metrics. We use top-k Recall (Recall@K) and Normalized Discounted Cumulative Gain (NDCG@K) with  $K = 5, 10$  to evaluate the recommendation performance.
RQ-VAE Implementation Details. As discussed in section 3.1, RQ-VAE is used to quantize the semantic embedding of an item. We use the pre-trained Sentence-T5 [27] model to obtain the semantic embedding of each item in the dataset. In particular, we use item's content features such as title, price, brand, and category to construct a sentence, which is then passed to the pre-trained Sentence-T5 model to obtain the item's semantic embedding of 768 dimension.
The RQ-VAE model consists of three components: a DNN encoder that encodes the input semantic embedding into a latent representation, residual quantizer which outputs a quantized representation, and a DNN decoder that decodes the quantized representation back to the semantic input embedding
(a) The ground-truth category distribution for all the items in the dataset colored by the value of the first codeword  $c_{1}$ .
(b) The category distributions for items having the Semantic ID as  $(c_{1},*,*)$ , where  $c_{1} \in \{1,2,3,4\}$ . The categories are color-coded based on the second semantic token  $c_{2}$ .
Figure 4: Qualitative study of RQ-VAE Semantic IDs  $(c_{1}, c_{2}, c_{3}, c_{4})$  on the Amazon Beauty dataset. We show that the ground-truth categories are distributed across different Semantic tokens. Moreover, the RQVAE semantic IDs form a hierarchy of items, where the first semantic token  $(c_{1})$  corresponds to coarse-level category, while second/third semantic token  $(c_{2} / c_{3})$  correspond to fine-grained categories.
space. The encoder has three intermediate layers of size 512, 256 and 128 with ReLU activation, with a final latent representation dimension of 32. To quantize this representation, three levels of residual quantization is done. For each level, a codebook of cardinality 256 is maintained, where each vector in the codebook has a dimension of 32. When computing the total loss, we use  $\beta = 0.25$ . The RQ-VAE model is trained for 20k epochs to ensure high codebook usage ( $\geq 80\%$ ). We use Adagrad optimizer with a learning rate of 0.4 and a batch size of 1024. Upon training, we use the learned encoder and the quantization component to generate a 3-tuple Semantic ID for each item. To avoid multiple items being mapped to the same Semantic ID, we add a unique  $4^{th}$  code for items that share the same first three codewords, i.e. two items associated with a tuple (7, 1, 4) are assigned (7, 1, 4, 0) and (7, 1, 4, 1) respectively (if there are no collisions, we still assign 0 as the fourth codeword). This results in a unique Semantic ID of length 4 for each item in the recommendation corpus.
Sequence-to-Sequence Model Implementation Details. We use the open-sourced T5X framework [28] to implement our transformer based encoder-decoder architecture. To allow the model to process the input for the sequential recommendation task, the vocabulary of the sequence-to-sequence model contains the tokens for each semantic codeword. In particular, the vocabulary contains 1024  $(256 \times 4)$  tokens to represent items in the corpus. In addition to the semantic codewords for items, we add user-specific tokens to the vocabulary. To keep the vocabulary size limited, we only add 2000 tokens for user IDs. We use the Hashing Trick [38] to map the raw user ID to one of the 2000 user ID tokens. We construct the input sequence as the user Id token followed by the sequence of Semantic ID tokens corresponding to a given user's item interaction history. We found that adding user ID to the input, allows the model to personalize the items retrieved.
We use 4 layers each for the transformer-based encoder and decoder models with 6 self-attention heads of dimension 64 in each layer. We used the ReLU activation function for all the layers. The MLP and the input dimension was set as 1024 and 128, respectively. We used a dropout of 0.1. Overall, the model has around 13 million parameters. We train this model for 200k steps for the "Beauty" and "Sports and Outdoors" dataset. Due to the smaller size of the "Toys and Games" dataset, it is trained only for 100k steps. We use a batch size of 256. The learning rate is 0.01 for the first 10k steps and then follows an inverse square root decay schedule.
4.1 Performance on Sequential Recommendation
In this section, we compare our proposed framework for generative retrieval with the following sequential recommendation methods (which are described briefly in Appendix B): GRU4Rec [11], Caser [33], HGN [25], SASRec [17], BERT4Rec [32], FDSA [42],  $S^3$ -Rec [44], and P5 [8]. Notably