Making Transformer Decoders Better Differentiable Indexers
Authors: Wuchao Li
Kai Zheng
Defu Lian
Qi Liu
Wentian Bao
Enyun Yu
Yang Song
Han Li
Kun Gai
Venue: ICLR 2025 Poster (2025)
Keywords: Generative Retrieval
Generative Index
End-to-end Recommender System
Information Retrieval
Transformer
Primary area: optimization

Abstract: Retrieval aims to find the top-k items most relevant to a query/user from a large dataset. Traditional retrieval models represent queries/users and items as embedding vectors and use Approximate Nearest Neighbor (ANN) search for retrieval. Recently, researchers have proposed a generative-based retrieval method that represents items as token sequences and uses a decoder model for autoregressive training. Compared to traditional methods, this approach uses more complex models and integrates index structure during training, leading to better performance. However, these methods remain two-stage processes, where index construction is separate from the retrieval model, limiting the model's overall capacity. Additionally, existing methods construct indices by clustering pre-trained item representations in Euclidean space. However, real-world scenarios are more complex, making this approach less accurate. To address these issues, we propose a \underline{U}nified framework for \underline{R}etrieval and \underline{I}ndexing, termed \textbf{URI}. URI ensures strong consistency between index construction and the retrieval model, typically a Transformer decoder. URI simultaneously builds the index and trains the decoder, constructing the index through the decoder itself. It no longer relies on one-sided item representations in Euclidean space but constructs the index within the interactive space between queries and items. Experimental comparisons on three real-world datasets show that URI significantly outperforms existing methods.

Published as a conference paper at ICLR 2025
MAKING TRANSFORMER DECODERS BETTER DIFFER- ENTIABLE INDEXERS
Wuchao Li', Kai Zheng", Defu Lian', Liu', Wentian Bao Yun En Yu?, Yang Song,, Han Li,, Kun Gai
University of Science and Technology of China, Kuaishou, 3 Indepeedeent
(1 { mliwchhac, giliu67]@nan @mail. www.t..uutte.e........ee...eeeet wustc.edu... edu. en, (kaizheng, 1ihan08 ( @kuai shou ccm 1 Liandofu@usted edu. en, wb2328@columbia.cda yuenyune126 (onnynee22..cmm com, ys@sonyls mo mo gai. kwwkqq.ee.......... com
ABSTRACT
Retrieval aims to find the top-k items most relevant to query/user from large dataset. Traditional retrieval models represent queries/users and items as embed- ding vectors and use Approximate Nearest Neighbor (ANN) search for retrieval. Recently, researchers have proposed a generative-based retrieval method that rep- resents items as token sequences and uses decoder model for autoregressive training. Compared to traditional methods, this approach uses more complex mod- els and integrates index structure during training, leading to better performance. However, these methods remain two-stage processes, where index construction is separate from the retrieval model, limiting the model's overall capacity. Addi- tionally, existing methods construct indices by clustering pre- pre-trained item repre- sentations in Euclidean space. However, real-world scenarios are more complex, making this approach less accurate. To address these issues, we propose a Unified framework for Retrieval and Indexing, termed URI URI ensures strong consis- tency between index construction and the retrieval model, typically Transformer decoder. URI simultaneously builds the index and trains the decoder, construct- ing the index through the decoder itself. It no longer relies one-sided item representations Euclidean space but constructs the index within the interactive space between queries and items. Experimental comparisons three real-world datasets show that URI significantly outperforms existing methods.
INTRODUCTION
Retrieval is critical upstream stage in various tasks, such as information retrieval and recommender systems,
the upper bound for downstream ranking tasks. Given a query, the retrieval model identifies the top-k most relevant items from a large-scale dataset. Throughout the paper, "query" will be used to represent both query and user.
1.1 PRIOR METHODS AND LIMITATIONS
Traditional retrieval models generate representations for queries and items, using simple similarity measures like the inner product. After index construction, the Approximate Nearest Neighbor search retrieves the relevant items. However, this approach has notable performance limitations due to the linearity of the similarity metric, which constrains model complexity, and the decoupling of model training from index construction.
To overcome these limitations, researchers have proposed new method, Generative Retrieval (Tay et al., 2022; Wang al. 2022) Rajput et al., 2024 Feng et al. 2022). introducing Transformer decoders /Vaswani 2017 into the retrieval process. These methods involve multiple stages: first constructing a hierarchical clustering index for items, representing them sequences of cluster IDs, and then training decoder to generate these sequences. Compared to traditional methods, these approaches more complex model architectures and establish connections between model training and the constructed index, resulting in improved performance. However, index construction still
"Corresoondigg author.

---

Published as a conference paper at ICLR 2025
disconnected from the retrieval model, relying only item information without considering queries, and heavily dependent on simplistic similarity measures and clustering algorithms in Euclidean space, greatly limiting index effectiveness.
1.2 OUR SOLUTION
In this work, we propose a Unified framework for the Retriever and Indexer, termed URL In URI, the Retriever and Indexer share the same Decoder model, so the index is not constructed using simple similarity measures but is directly generated by the Retriever. In previous methods, the Decoder only memorizes the pre-constructed index, whereas in URI, the Decoder acts as differentiable index constructor.
We introduce a new metric for generative retrieval indices, called Average Partition Entropy (APE), which is model-inen mrodel-independent and can be evaluated immediately after construction. Using this metric, we formally define the index construction task and propose theoretically effective EM (Dempster et al., 1977) algorithm. We then employ a machine learning approach to simulate the conditions and outcomes of this algorithm by designing loss functions. In previous works, after index construction, an encoder model is only required on the query side to extract information, while the item side is represented by sequence of index ID tokens. However, in URI, index construction and model training occur simultaneously, extracted from the query and item by their respective encoders is input to the Decoder, resulting in probability distribution over & tokens for each. Our loss functions trained based these distributions, enabling the query to identify positive items and the item to learn its assigned token.
we encoders both the query and item sides. Information
1.3
OUR CONTRIBUTIONS
We propose novel unified framework for generative retrieval, URI, which constructs the index simultaneously with training the Transformer Decoder Retriever.
We introduce new evaluation metric for generative indices, Average Partition Entropy (APE), which allows direct assessment of index quality immediately after construction, without model training.
Using APE, formally define the generative index construction task and provide theoretical explanation for the rationale behind URI's loss function.
We conduct extensive experiments on three real-world datasets, showing that URI achieves state- of the-art performance as a generative retrieval method. Moreover, we design various experiments explain its superior performance from multiple perspectives.
RELATED WORK
Traditional Retrieval Traditional retrieval represents both queries and items embedding vec- tors, typically measuring their similarity using the inner product, Euclidean distance, or cosine sim- ilarity. The method for learning these embeddings varies by task. In document retrieval, pre- trained NLP models like BERT (Kenton & Toutanova 2019) and ALBERT (Lan) 2019) are commonly used. In recommender systems, dual-ttowr models mostly used, including MF (He et al.) (016). Deep- ICF (Xue et al. 2019). and NCF (He et al.) 2017) for general recommendations, as well as SASRec (Kang & McAuley 2018), YouTubeDNN (Covington et al. 2016), and GRU4Rec (Hidasi 2015) for sequential recommendations. Once trained, various indices are constructed for the item set, including graph- based (Makov & Yashunin 2 Jayaram Subramanya et al. (019)) quantization- based (Jegou et al.] 2010), LSH-based (Datar et al. 2004). and tree- rreebased (Bentley 1975) indices. Approximate Nearest Neighbor Search (ANNS) is then used to complete retrieval.
Generative Retrieval Existing generative retrieval methods use multi-stage training approach. They first construct a hierarchical clustering index based on pre-trained representations, use cluster- ID sequences to represent items, and then train the model to generate these sequences. DSI (Tay) et al. 2022) and NCI (Wang et al., 2022) both construct indices by performing hierarchical K- text representations. Additionally, NCI uses Prefix-Awar Weight-Adapiig Decoder to adapt to different codebook levels. RecForest (Feng et al.) 2022) pre-trains a recom- mendation model and constructs the index using item ID) representations, enhancing performance
means clustering
2

---

Published as a conference paper at ICLR 2025
by integrating multiple trees. TIGER (Rajput et al. 2024 adopts RQ-VAE (Zegidoorr et al. 2021 to build the index from item text representations and uses the item's 's semantic ID sequence to model user historical behavior sequences. DR (Gao et al. 2021) uses an MLP for inference and iteratively updates the index with the EM (Dempster et al. 1977) algorithm.
A MEASURE OF INDEX: APE
Inspired by the Mutual Information (Dhillnn al. 2003), we propose a metric to measure the index in generative methods. We consider the case where the index has single layer, meaning the semantic length is In later sections, we extend this to an index with L layers. Given query corresponding to positive items / {y' ..... y" the n items are distributed into k buckets after indexing. The query then obtains a probability distribution Q = Iq....., qel over the & buckets through the generative model's decoder. To the best of our knowledge, all existing work uses cross- entropy loss to ensure that the input query a retrieves buckets containing samples from y Let Idx(-) denote the constructed index that maps item specific bucket, with Idx(y:) b,. The loss function is defined as
Lee (2, )12t =
Σ
log qb. =
10b, Σ ny log Uj, where ny = E1(b, = j).
(I)
2=1
1=1
It can be easily proven that the loss attains its minimum value when qj = nsyn using the method of Lagrange multipliers. Substituting qj. we can get
LoolE,
Σ ny log
log
n
7
The lower bound of the loss function is established once the index is constructed. It depends only on the concentration of positive items corresponding to a query within the index. Let (Idxx) = DE r(Idx())==)) .... Σ. [((x((u)==))
1 ...., mt] represent the frequency distribu- 12 items corresponding to the query x. The entropy of Q.. tion of the buckets containing the positive ((xxx(dx))
ny significantly affects the retrieval performance of the query x af- ter cross-entropy training. Inspired by this, we propose a new evaluation metric for the index in generative retrieval models, called Average Partition Entropy (APE).
ΣS nh log n
APE(Idx) E (Q..
APE measures the concentration of relevant items in the index throughout the dataset. It depends solely the dataset and the constructed index, independent of the neural network model, allowing APE to preliminarily evaluate index quality. A smaller APE value indicates better index construc- tion. Balance crucial in index construction, and all methods account for when building indices. The proposed APE metric is meaningful only when the index is relatively balanced. For extremely unbalanced indices, such as when all items are assigned to one bucket, APE reaches its minimum value, but making it meaningless. In generative retrieval, most models use a hierarchical, multi-layer index structure. For index with L layers and width k, obtain Q. R²* at the leaf level.
4
METHODOLOGY
4.1 PROBLEM FORMULATION AND MOTIVATION
Before introducing the proposed framework URI, we first formulate the task of index construction in generative retrieval and discuss how solve it. Given dataset D consisting of a query set X and an item set y of size N, where each query x E X' has its corresponding set of positive samples 2ot V. the task is allocate all items evenly into k buckets, with each bucket containing N/k items. The mapping from items to buckets is denoted as Idx(-). The optimal index is defined as:
arg min AP APE(Idx)
(4)
Even allocation is introduced to prevent all items from being assigned to a single bucket, which would result in APF = and render it meaningless. Assuming N is divisible by k, and
3

---

Published as a conference paper at ICLR 2025
N/k, there are NI/(k!- 2!) possible allocation schemes (indices). This is an NP-hard allocation problem, making direct optimization infeasible due to high computational cost. Existing methods attempt to approximate the optimal solution by clustering in the item representation space; however, a significant gap remains. These methods assume that items close in Fuclidean distance yield similar model outputs, based the premise of a small Lipschitz constant. However, as models increase in complexity, the cumulative effect of the Lipschitz constant becomes more pronounced (Fazlyab al. 2019) challenging this assumption.
To address this issue, we propose constructing the index directly from the final retrieval model, rather than in the shallow representation space. Some previous methods like DR (Gao et al. (02T) and JTTM (Zhu et al.] (219)) iteratively update the index based on the model after training, using the EM (Dmpsterr et al. 1977) algorithm to improve index quality. This because the model memorizes the index while capturing the rich collaborative filtering information within the dataset the same time. In Section B] we define 2n under the ideal assumption of the model's 's memorization of the index. However, in practice, 2a is influenced by complex collaborative filtering relationships between queries and items, preventing convergence to the minimum point of (x, (xt) The actual Qx contains all the information absorbed by the model during training. Reconstructing the index based on the distribution 2₂ obtained after training allows the new index to outperform the original. We propose simple greedy algorithm for reconstruction.
Greedy Algorithm: For each item y with a relevant query set x = {=I...., vn (for each x E X y E )),, we calculate the expected distribution E E2 [2.] of these queries. The item is assigned to the token with the highest probability in its expected distribution. Under certain conditions, this greedy algorithm can be proven to correctly assign the majority of items.
Theorem 1. Given dataset DD consisting of m queries and N items, where the N items are evenly allocated into k buckets, and the probability of each query being related to an item is p. Given the probability distribution { Juex of each query across the k buckets, using the aforementioned greedy algorithm can recover the allocation scheme of each item. Given
④ (1 8)12 (1+ 2N +(1++))
p(1 P)
the probability of successfully allocating each item Peorrett 2 1 - 5, where t- is the inverse CDF of the standard normal distribution.
For detailed proof, see Appendix D In real-world datasets, statistical values often meet the theo- rem's 's conditions, indicating that this greedy algorithm is reliable. Therefore, the optimal index prob- lem can be optimized using the EM algorithm: (1) At the E-step, the model optimizes Lee (2., )+ based on pre-constructed index and obtains converged Q.. (2) At the M-step, the index is recon- structed using the greedy algorithm described above, based on Q.-
We simulate the M-step using a machine-learning approach. As discussed Section training with cross-entropy aligns the query distribution with the expected distribution of related items. By swapping the roles of items and queries, the item distribution can be similarly trained using cross- entropy. Considering an item y and its relevant queries Xx+
}, for z' E X let 2₂ iqi...., g..l Given the probability distribution Py = [p......PP. of item y over & tokens from the decoder, the cross-entropy is:
n
k
(,, Σ
o3 log
Σ j=1
log
(5)
It is straightforward to prove that the loss function reaches its minimum value when Pj (Er //) This method completes the first step of the greedy algorithm. Next, we assign the item to the token with the highest probability in the distribution. We introduce entropy as the loss function to approximate the item distribution to a one-hot distribution:
Crommit (y) Σ Py log Pj-
Note that index construction requires a constraint that items must be evenly allocated. To meet this constraint, we also use machine-learning approach, introducing the expected distribution to
4

---

Published as a conference paper at ICLR 2025
Lisonrry (()) ((C).p))
pendy
gondl
proda
proby
prib
prob
sup p (()
Rf
Layert
(PP).)) H(P) (R, FIPDD
Query'''s probability distribution
Py: Ite''s prubability distrbutiion
Lineir Sofnns
Croaa Alcition
ID Emb
Text
Masked Seir Attertion
share
Misked Sei Amcotion
Emb
B05
Laye
Embedking Layer
Text Ficooder
Text Encoder
Emhedding Layer
QuaryUUar]
Ban(DDoc)
Step1: Training for Layer]
Stap2: Assignment for Layeri
Step3: Trzining for Layer2
Step3: Assignment for Layer2
Figure 1: Overview Framework of URI's Retriever and Indexer.
approximate a uniform distribution. The specific form is given by:
Lbatance (2) - H(U*, Eney(Pa)) Σ log Eyey
Uk is a uniform distribution over tokens. Leommit aims to approximate the distribution one-hot distribution, while Lbatance aims to approximate the distribution to a uniform distribution. However, these objectives do not conflict, as Ccommit acts on individual items, while Lbalance acts on the entire item set y.
4.2 DECODER AS INDEXER AND RETRIEVER
Building on the previous discussion of index construction, we propose unified framework for the retriever and indexer, called URL. URI simulates the above EM algorithm using differentiable machine-learning approach, enabling end-to-end training. URI uses the same decoder model to generate distributions for both the query and item over tokens. We denote the distributions the /th token obtained from the decoder for the query and item as P/₃ and P/, respectively. H(:,)) represents cross-entropy, and H (-) represents entropy. The loss functions are summarized as follows:
C'rsstep Co. H(P,, Pb)
LM_step (v, Ct Clommit + Chatance HP')) Pt) H(P)) EPP)), Note that the two losses are trained simultaneously: the terms "E-step" and "M-step" are used for clarity. From gradient optimization perspective, the two L.n terms ensure that Po and P, of corresponding queries and items are as close as possible. Lrommii prevents PQ and P; from con- verging to uniform distribution, while Lhatance prevents them from converging to the same one-hot distribution, avoiding index collapse.
We use positive sample pairs (query, item) to train the two loss functions. However, relying solely on positive samples introduces bias, items not in the positive sample set cannot be properly learned. To address this, we incorporate randomly sampled items to train Lcommit and Loatance- These random items do not participate in optimizing Lee. they lack corresponding matching queries. In each training batch, we sample number of random items equal to the positive samples. Furthermore, Lbatance lacks adaptability across different tokens. To mitigate this, we introduce weighting mechanisms to adjust the distribution across tokens:
Chatance (U* E(PH)
Σ sgd( (.)) log Epeylpe) wy -log Eoggrey 11""1 1 +1 (10)
5

---

Published as a conference paper at ICLR 2025
The weight's magnitude relates to the difference between E[P/] and the uniform distribution. A larger difference indicates higher imbalance, requiring a larger weight. A constant factor of ensures weight stability, preventing it from approaching zero.
Assignment The Transformer Decoder uses autoregressive training, meaning that when training the next layer's token, previous tokens must be input. In our approach, since each layer's tokens unknown before training, train the model layer by layer. Once the model converges, items are explicitly assigned to k tokens, preparing for training the next layer. To assign items, we first pass all items through the Decoder obtain each item's probability distribution. Based on this probability distribution, we use a balanced assignment algorithm (Li et al.] (023). Specifically, we set a maximum load for each token, typically the average allocation size Ník. The probability distribution of each item is sorted in descending order, and tokens are selected sequentially. If the number of items selected token reaches the maximum load, the next token is chosen.
During training, we first train the model on the first layer, optimizing P_step + C step Once the model converges, we assign first-layer tokens the items. Next, we train the second-layer tokens, using the first-layer tokens assigned to the items as inputs to the decoder. To prevent the model from forgetting information learned the first layer, we introduce memory reinforcement loss, which is the same as the loss used in prior works for training the decoder after index construction:
C memory Σ (H(C*, P PD):; Decnder L'Esste L'M
C Ememory (11)
where CK denotes the one-hot label of the token assigned to the item layer i. Let C memory 0.
We train Decoder layer by layer from 1 to = L, where a hyperparameter.
Figure shows the training process at Layer 2. The figure illustrates the URI framework, with its core being the decoder component, serving as both the Retriever and the Indexer. Different encoders can be used for different types of queries and items. For example, in sequential recommendation scenario, the query part can use a sequential model as its encoder. Additionally, features extracted by the encoder input to the decoder sequentially, allowing any desired representations to be freely added, enriching the index construction with more information.
Although URI requires layer-by- training. the number of nodes increases exponentially. There- fore, in practice, training two to three layers is generally sufficient. Moreover, tokens learned using this method are more appropriately named, and their sequences exhibit greater consistency across different layers. This occurs because when learning tokens of later layers, all assigned tokens from preceding layers are used as inputs. Consequently, the information represented by the same token across layers becomes more consistent, enhancing model performance. We validate this through experiments (See 5.3)
4.3 RANKER
After the training of the Retriever, a ranker will be trained to identify the topk items for the final retrieval. In previous work, most methods (Tay et al. 2022; Wang et al., 2022) Rajput et al., 20044] Feng et al., 2022) in- volved randomly adding a new token different items in the final layer to serve as the ranker, while the DR(Gao et al. 202T) model utilizes soft- max loss over the entire item set. We follow the approach of DR, taking the inner product of query and item as the score, using sampled softmax loss (Covington et al. 2016) for training, where 10 negative samples are drawn for each positive example. Furthermore, we incorporate hard neg- ative mining. Specifically, we first sample 9 random negative examples, then select one hard negative sample from the sibling nodes of the positive sample's ss node. This is feasible due to the high quality of the index we construct. We summarize the ranker structure in Figure
Sangle Soltess
Suzi Pookes
Sus Peling
ID Seat Text Trae
bibb
Text Bisb
Query (Ulve)
Ilc (lem)
Figure 2: Ranker.

---

Published as a conference paper at ICLR 2025
4.4 INFERENCE
In the index, each item is represented by a sequence of token IDs, transforming it into a sequence generation task. The Transformer Decoder employs beam search to complete this task, ultimately generating the top-b sequences. These top-b sequences represent the top-b buckets in the final index layer. The union of all items in these buckets forms the candidate set. The final top-k retrieval results are obtained by computing the inner product between the query's representation and item representations in the candidate set.
4.5 COMPLEXITY ANALYSIS
Let the time for single pass through the Decoder be O(t), the dataset size be (D), and the number of index layers be L. Then, the training complexity of URI is ((2L Compared to other methods, URI introduces
additional constant factor of 2L due to layer-by-layer training. with each training step involving separate forward passes for both the query and item. This may slightly increase training time; however, this overhead is acceptable, does not require index construction before training, eliminating that overhead. During inference, URI's complexity is the same as existing methods, dominated by the Decoder with a complexity of O(BL ). where B is the beam size and L is the complexity required for the Decoder to generate sequence length J..
L is typically small in practice, and URI
EXPERIMENT
We conduct extensive experiments to answer the following questions:
RQ1 What is the overall performance of URI compared to other Generative methods?
RQ2
Why does URI perform so well?
RQ3 How to set the hyper- syper-parameters for URI index?
RQ4
How is the index progressively built during URI training?
5.1
EXPERIMENTAL SETTINGS
Datasets We select three widely-used real-world datasets, each corresponding trieval tasks: document retrieval, general recommendation, and sequential recommendation. The KuaiSAR (Sun et al. 2023) dataset is short video search dataset, where the query is text and the item is a short video with a text description, corresponding to the document retrieval task. The Beauty and Toys and Games datasets are both from the Amazon platform (He & McAuley] 2016)) We use the Beauty dataset for the general recommendation task and the Toys and Games dataset for the sequential recommendation task, where the sequence length is set to 20.
of three
Baselines and Settings In this paper, we focus on demonstrating that our method is the optimal generative retrieval approach. To that end, we selected five representative works in this field: DSI (Tay al. 2022). NCI (Wang al. 2022). RecForest (Feng et al. 2022). TIGER (Rajput et al. 2024), and DR et al.) 2021). Among these, DSI and NCT are designed for document retrieval tasks, while the other three address recommendation tasks. All five methods follow multi-stage approach, where the index first constructed, followed by model training. To ensure faimess in comparison, we maintain consistency with the original papers in terms of the retrieval model archi- tecture and index construction methods. For text extraction, we utilize the 'bert-base-cased (Kenton) & Toutanova 2019) model from Huggingface (Jain 2022) across all methods. We set the represen- tation dimension to 96, the number of layers in the Decoder model to 2, the beam size to 20, and the learning rate to 0.001. For NCI and DSI, we set the number of tokens in the final layer 100, referred to as in the original papers. For KuaiSAR dataset, we set the index width k = 64 and depth 2. For Beauty and Toys datasets, we set 32 and 2.
Evaluation Metrics Following most prior works, we evaluate the performance of each method using Recall@K K and NDCG@K (where K = 10 or 20). Additionally, we use our newly proposed APE metric assess the effectiveness of the indices constructed by different methods.

---

Published as a conference paper at ICLR 2025
5.2 OVERALL PERFORMANCE
Table 1: Recall(R)@1 @ K and NDCG(N) @k K Results of URI compared with baseline methods
Dutaset
KuaiSAR
Amazon Beauty
Toys and Games
Metric
R@10
N@
R@20
N@20
R@ 10
N@
2@
N@20
R@
N@
R@
N@20
RecForest
0.0384
0.0194
0.0626
0.0254
0.0102
0,0046
0.0179
0,0065
0.0344
0.0254
0.0399
0.0267
w/ URI Index
0.0546
0.0270
0.0785
0,0330
0.0129
0,0062
0,0214
0,0084
0.0518
0.0371
0.0586
0.0381
DSI
0.0976
0,0519
0,1459
0.0641
0.0105
0,0050
0.0176
0.0068
0.0486
0.0345
0.0536
0.0358
URE Index
0.1027
0.0560
0.1591
0.0700
0.0125
0.0058
0.0216
0.0081
0.0518
0.0373
0.0579
0.0391
NCI
0.0568
0.0276
0,1053
0.0397
0.0061
0.0028
0,0110
0.0040
0.0325
0.0206
0.0398
0.0225
w/ URI Index
0.0650
0.0312
0.1174
0.0444
0.0079
0.0037
0.014]
0.0052
0.0362
0.0220
0.0448
0.0242
TIGER
0.0376
0.0187
0,0624
0.0250
0.0092
0.0041
0.0167
0.0058
0.0322
0.0201
0.0383
0.0231
w/ URI Index
0,0546
0.0270
0.0785
0,0330
0,0129
0.0062
0.0214
0,0084
0.0509
0.0360
0.0582
0.0373
DR
0.0604
0.0358
0,0829
0.0415
0.0105
0.0049
0.0190
0.0075
0.0429
0.0271
0.0478
0.0295
w/ URI index
0.1124
0,0589
0.1533
0.0692
0.0125
0.0058
0.0203
0.0080
0.0523
0.0312
0.0607
0.0384
URI
0.1793
0.0925
0.2594
0.1127
0.0157
0.0067
0.0294
0.0097
0.0573
0.0373
0.0685
0.0401
In this section, address RQ1 by comparing the overall performance of URI against baseline methods using Recall@ 10, Recall@21 20, NDCG@ 10, and NDCG@ 20. Additionally, we replace all baselines' original indices with the index constructed by URI while keeping their model structures unchanged. Table D summarizes all results, from which we make the following observations.
URI outperforms all generative retrieval methods significantly across all datasets. For Recall @ 10, URI improves upon the best baseline by 83.7%, 49.5%, and 33.6% on the KuaiSAR, Beauty, and Toys and Games datasets, respectively. On the KuaiSAR dataset, URI achieves over 100% im- provement compared to most baselines. This is because the text in the dataset is token- andnminnonymzzed, preventing the direct use of pre-trained large language models for representation extraction. Con- sequently, we train a small language model from scratch on the limited data, restricting the effec- tiveness of indices built with these representations and resulting in poorer baseline performance. However, URI significantly mitigates this issue by learning query-item correlations, demonstrating the potential of constructing indices directly from the interaction space of queries and items. Ad- ditionally, when all baselines the index constructed by URL, performance improves to varying degrees, indicating URI's index is highly effective across different model structures.
5.3 EFFECTIVENESS ANALYSIS
We analyze the reasons behind URI's performance from multiple perspectives to answer RQ2.
APE Comparison We evaluated the APE of indices constructed by different methods, including Random, K-means, VQ-VAE (Van Den Oord et al. 2017), RQ-VAE (Zeghidour al.] [02T), and URI. For each dataset, we select different values of A under a typical two-layer setting, allowing us compare indices with varying partition granularities. Since the KuaiSAR dataset has larger number of items, we chose relatively larger values of k. All results are summarized in Table As shown in the table, the APE values of the indices constructed by URI are significantly lower than those of other methods, indicating that related items for the same query are more tightly clustered. This leads to more precise retrieval performance after training.
Token Consistency In previous index construction methods, clusters are named randomly. For multi-layer indices, this approach has little impact on the first layer. However, for subsequent lay- ers, relationships between clusters at different levels are not considered during naming, and tokens are randomly assigned, limiting the index's 's effectiveness. In URI, layer-by-layer training allows information from all previously constructed tokens to be incorporated when constructing the current layer. Consequently, during naming, similar clusters are assigned the same token. To demonstrate this, we design two variants of the URI index, in which token naming is degraded.
Variant I: A uniform token transformation is applied to all tokens the final layer, meaning transformations are the same for nodes under different parent nodes.
Variant 2: A transformation is applied to tokens in the final layer, but transformations differ for nodes under different parent nodes.
8

---

Published as a conference paper at ICLR 2025
Table 2: Comparison of the APE values for indices constructed by different methods.
Datasct
k
Random
Kmeans
VQ-VAE
RQ-VAE
URI
32
1.222
0.915
0.922
0,909
0.654
KuaiSAR
64
1.224
1.037
1.042
1.032
0,954
128
1.224
1.141
1.144
1.143
1.135
16
1.127
1.080
1.088
1.075
0.988
Beauty
32
1.134
1.106
1.112
1.110
1.040
64
1.135
1.125
1.128
1.126
1.107
16
1.756
1.662
1.672
1.662
1.594
Toys and Games
1.770
1.729
1.733
1.728
1.660
1.771
1.763
1.765
1.761
1.728
Figure 5 illustrates these two variants. Note that both variants retain the same APE as the original URI index. Table 3 presents the @ecall 10 results. Compared to the original index, the poorer results of Variant suggest that the original index maintains consistency actoss clusters with the same token across different layers. Furthermore, the worse performance of Variant compared to Variant indicates that, within the same layer, token naming for nodes under different parent nodes also exhibits consistency.
Variart
Vanam12
Table 3: Recall @ 10 of URI and variants
Dataset
KuaiSAR
Beauty
Toys and Games
Variant I
0.1633
0.0138
0.0561
Variant 2
0.1581
0.0124
0,0519
Original URI
0.1793
0.0157
0.0573
Figure 3: Variant Illustration.
propose adaptive weight loss Liatance- Here, we demonstrate its necessity through experiments. Table&] that removing the adaptive weight significantly declines URI's performance. When training with the original Lbalance- many tokens are ignored, resulting in an "empty bucket" state. Introducing adaptive weight ensures the optimizer focuses on these buckets.
Adaptive Weight in Lialanee_w In Equation 10
Table 4: Ablation results of Recall @ 10 for Adaptive Weight and Hard Negative Mining
Dataset
KuaiSAR
Beauty
Toys and Games
w/o Adaptive Weight
0.1301
0.0122
0.0517
w/o Hard Negative Mining Original URI
0.1528
0.0143
0.0558
0.1793
0.0157
0.0573
Hard Negative Mining In Section 4.3 we propose incorporating hard negative samples during ranker training by sampling from the sibling nodes of positive sample's node. The Recall 10 results in Table 4 indicate that removing hard negative samples decreases URI's overall perfor- mance, demonstrating the elfectivvness of hard negative sampling. This also provides evidence of the exploitable similarity among sibling nodes within the URI index.
5.4 SYPEE-PARAMMEN SENSITIVITY ANALYSIS
In this section, we compare results under different index hyperparameter settings to address RQ- For generative retrieval models, larger implies more parameters and leaf nodes at the same depth L, resulting in finer-grained item partitioning and more possible item combinations during retrieval, thus increasing the model's upper performance bound. However, finer partitioning reduces the number of items within the same node, increasing the risk of overfitting, Therefore, it is crucial to select & and L to balance accuracy and robustness in the index. Figure shows the Recall 10 curves as the candidate size varies under different k and L settings for the three datasets. The

---

Published as a conference paper at ICLR 2025
results indicate that k 64 and L optimal for KuaiSAR, while 32 and L are most suitable for Beauty and Toys and Games. Under these settings, each leaf node contains approximately 10 to 100 items. In practice, this conclusion can guide the selection of & and L.
Beauty
Toys and Games
0014
Recalig
0.40
0301
36.1
26-
32
16.1-
Wo
#CCandidate Item
#Candidate
Figure 4: The Recall @ curve with respect the number of candidate items on different indices.
5.5 CASE STUDY
In this section, we address RQ4 by observing how the index is constructed, monitoring the token training process of KuaiSAR in the first layer with lt We plot the output probability distributions of specific query (with search session id: 2) and its relevant positive items as pie charts after model processing. Jisure3]sshows that probability distributions are initially uniform, indicating that the model has no prior knowledge. As training progresses, the item distributions gradually converge to one-hot. Finally, Iteml, Item2, Item4, and Item5 are assigned to Token '4', while Item3 is assigned to Token '2'. The query's distribution closely aligns with the expected value of these five samples' distributions, consistent with the discussion in Section 3. In the original dataset, Iteml, Item2, Item4, and Item5 belong to the category dimo-dimensional,"" while Item3 belongs Short Series." This information not introduced during training, indicating that the model learns similar knowledge solely through collaborative filtering signals between the query and items.
"Film and
Token '1'
Token '2'
Token '3'
Tokca '4'
0 epoch
epechs
5 epochs
Query
Iteml [Tossincniional]
Ioem2 [Toodimnnioonal
Itemž
Item4
Itemš
[Film and aimt scries]
[Two-dincsional] [Two miv-ddimnnionall
Figure 5: Case Study.
6 CONCLUSION AND FUTURE WORK
In this paper, we propose Averaged Partition Entropy (APE), new metric for evaluating indices in generative models, and introduce a novel generative retrieval framework, URI, where model train- ing and index construction are completed simultaneously within the same Decoder. Compared to previous methods, where the Decoder only memorizes an already constructed index. URI treats the Decoder as differentiable Indexer. Comparative experiments with baselines show that URI achieves state-of-the-an (SOTA) performance. Analytical experiments explain why URI outper- forms from various perspectives, including lower APE for the index and token consistency across layers. In the future, we will focus on developing methods to enable URI to generate multi-level tokens simultaneously to accelerate training.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The work was supported by grants from the National Natural Science Foundation of China (No. U24A20253).
REFERENCES
Jon Louis Bentley. Multidimensional binary search trees used for associative searching. Communi- cations of the ACM, 18(9):509-55) -17, 1975.
Paul Covington, Jay Adams, and Emre Sargin. Deep neural networks for youtube recommendations. In Proceedings of the 10th ACM conference on recommender systems, pp. 191-198, 2016.
Mayur Datar, Nicole Immorlica, Piotr Indyk, and Vahab S Mirrokni. Locality-sensitie sensitive hashing scheme based on p-stable distributions. In Proceedings of the twentieth annual symposiun on Computational geomelry, pp. 253-262, 2004.
Arthur P Dempster, Nan M Laird, and Donald Rubin. Maximum likelihood from incomplete data via the em algorithm. Journal of the royal statisticul society; series B (methodologicul), 39(1): 1-22, 1977.
Inderjit S Dhillon, Subramanyam Mallela, and Rahul Kumar. A divisive information-theoretie fea- ture clustering algorithm for text classification. Journal of machine learning research, 3(Mar): 1265-1287, 2003.
Mahyar Fazlyab, Alexander Robey, Hamed Hassani, Manfred Morari, and George Pappas. Efficient and accurate estimation of lipschitz constants for deep neural networks. Advances in neural information processing systems, 32, 2019.
Chao Feng, Wuchao Li, Defu Lian, Zheng Liu, and Enhong Chen. Recommender forest for efficient retrieval. Advances in Neural Information Processing Systems, 35:38912-38999 2022.
Weihao Gao, Xiangjun Fan, Chong Wang, Jiankai Sun, Kai Jia, Wenzi Xiao, Ruofan Ding, Xingyan Bin, Hui Yang, and Xiaobing Liu. Learning an end-to-end structure for retrieval in large-scale recommendations. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pp. 524-533, 2021.
Ruining He and Julian McAuley. Ups and downs: Modeling the visual evolution of fashion trends with one-class collaborative filtering. In proceedings of the 25th international conference on world wide web, pp. 507-517, 2016.
Xiangnan He, Hanwang Zhang, Min-Ye Yen Kan, and Tat-Seng Chua. Fast matrix factorization for online recommendation with implicit feedback. In Proceedings of the 39th International ACM SIGIR conference on Research and Development in Information Retrieval, pp. 549-558, 2016.
Xiangnan He, Lizi Liao, Hanwang Zhang, Liqiang Nie, Xia Hu, and Tat-Seng Chua. Neural col- laborative filtering. In Proceedings of the 26th international conference world wide web, pp- 173-182, 2017.
B Hidasi. Session-based recommendations with recurrent neural networks. arXiv preprint arXiv:1511.06939, 2015.
Shashank Mohan Jain. Hugging face. In Introduction to transformers for NLP: With the hugging face library and models to solve problems, pp. 51-67. Springer, 2022.
Suhas Jayaram Subramanya, Fnu Devvrit, Harsha Vardhan Simhadri. Ravishankar Krishnawamy, and Rohan Kadekodi. Diskann: Fast accurate billion point nearest neighbor search on a single node. Advances in Neural Information Processing Systems, 32, 2019.
Herve Jegou. Matthijs Douze, and Cordelia Schmid. Product quantization for nearest neighbor
search. IEEE transactions on pattern analysis and machine intelligence, 33(1):117-128 2010.

---

Published as a conference paper at ICLR 2025
W'aggChhnn Kang and Julian McAuley. Self-attentive sequential recommendation. In 2018 IEEE international conference on data mining (ICDM), pp. 197-206. IEEE, 2018.
Jacob Devlin Ming-War Wei Chang Kenton and Lee Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of naacL-IILT, volume 1, pp. 2. Minneapolis, Minnesota, 2019.
Z Lan. Albert: A lite bert for self-supervised learning of language representations. arXiv preprint arXiv: 1942, 2019.
Wuchao Li, Chao Feng, Defu Lian, Yuxin Xic, Haifeng Liu, Yong Ge, and Enhong Chen. Learning balanced tree indexes for large-scale vector retrieval. In Proceedings the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 3353-1362, 2023.
Yu A Malkov and Dmitry A Yashumin. Efficient and robust approximate nearest neighbor search using hicrarchical navigable small world graphs. IEEE transactions on pattern analysis and machine intelligence, 42(4))824-836, 2018.
Shashank Rajput, Nikhil Mehta, Anima Singh, Raghunandan Hulikal Keshavan, Trung Vu, Lukasz Heldt, Lichan Hong, Tay, Vinh Tran, Jonah Samost, al. Recommender systems with gener- ative retrieval. Advances in Neural Information Processing Systems, 36, 2024.
Zhongxiang Sun. Zihua Si, Xiaoxue Zang, Dewei Leng, Yanan Niu, Yang Song, Xiao Zhang, and Jun Xu. Kuaisar: unified search and recommendation dataset. In Proceedings of the 32nd ACM International Conference Information and Knowledge Management, pp. 5407-5411, 2023.
Yi Tay, Vinh Tran, Mostafa Dehghani, Jianmo Ni, Dara Bahri, Harsh Mehta, Zhen Qin, Kai Hui, Zhe Zhao, Jai Gupta, et al. Transformer memory differentiable search index. Advances in Neural Information Processing Systems, 35:21831-2188h 2022.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Yujing Wang, Yingyan Hou, Haonan Wang. Ziming Miao, Shibin Wu, Qi Chen, Yuqing Xia, Cheng- min Chi, Guoshuai Zhao, Zheng Liu, et al. A neural corpus indexer for document retrieval. Ad- vances Neural Information Processing Systems, 35:25600-2566 2022.
Feng Xue, Xiangnan He, Xiang Wang, Jiandong Xu, Kai Liu, and Richang Hong. Deep item-based collaborative filtering for top-n recommendation. ACM Transactions Information Systems (TOIS), 37(3):1-25, 2019.
Neil Zeghidour, Alejandro Luebs, Ahmed Omran, Jan Skoglund, and Marco Tagliasacchi. Sound- stream: An end-to-end neural audio codec. IEEF/ACM Transactions Audio, Speech, and Language Processing, 30:495-507, 2021.
Han Zhu, Daqing Chang, Ziru Xu, Pengye Zhang, Xiang Li. Jie He. Han Li, Jian Xu, and Kun Gai. Joint optimization of tree-based index and deep model for recommender systems. Advances in Neural Information Processing Systems, 32, 2019.
12

---

Published as a conference paper at ICLR 2025
APPENDIX
A NOTATIONS
All notations in the main text and their descriptions are summarized in Table
Notation
Description
r, x
Query, Query set Item, Item Set
v.
x+
Queries related to item y Items related to query æ
V
D Q
Dataset consisting of X' and
The distribution of items related to in the index
p
Item-query relevance probability
m
The number of queries
The number of items
5
A tolerance level The CDF of standard normal distribution The distribution at the i-th token of the query The distribution at the i-th token of the item The token assigned to the item at the i-th Layer
2(-)
P
Pr
C
Ha,,b) b)
Cross-entropy of a and b
H(6)
Entropy of A uniform distribution over A: classes
U
k
The width of the index (The number of tokens) The depth of the index (The number of Layers)
Table 5: Notations
B DATASET STATISTICS
Dataset
#Query
#Item
# Interaction
KuaiSAR
191330
112388
1093920
Beauty
41895
15817
162713
Toys and Games
13271
25357
90557
Table 6: Dataset Statistics
We summarize the statistics of the three datasets in Table Setting the tolerance level to 0.95 and substituting the dataset density for p. all three datasets satisfy the conditions Theorem 1. This indicates that our proposed greedy algorithm is applicable most datasets and demonstrates the rationality of URI's approach in simulating the greedy algorithm through machine learning.
C
SUPPLEMENT
D PROOF OF THEOREM
Before deriving Theorem 1, can transform the problem of recovering the allocation scheme into a matrix equation-solving problem. Let W be the interaction matrix between queries and items, with size of x N. The value given position is if the query is related the item, and
13

---

Published as a conference paper at ICLR 2025
otherwise. In recommender system, W represents the purchase matrix. Let H be the item allocation matrix, with a size of N k, where each row contains exactly one 1, and the sum each column is N/k. Define M W x H. where M has a size of x k. representing the frequency distribution each query over the & buckets. Thus, the problem is transformed into solving for H given M and W. The conclusions of the theorem and the greedy algorithm are also correspondingly transformed, as described below.
D.1 THEOREM
Theorem 2. (Correct Reconstruction Probability)
Given the following conditions:
Tolerance Level: tolerance level & (0,
Parameters: Positive integers N and k, and a probability p E (0,
Random Matrix W: An N binary matrix where each element W,, is independently and identically distributed (i.i.d)), satisfying:
P(Wr, p. P(W,,
Structured Matrix H: An N x k hinary matrix satisfying:
Row Constraint: Each row contains exactly one I, and all other elements are 0.
h
H. 1, =i= 1.2..... NN
1=1
Coolmm Constraint: Each column contains exactly N₂ divisible by k).
N E
ones (assuming N is
Σ H N2, vý= = 1,2,....
Product Matrix M: Defined as M W. H.
If the number of rows m in matrix W satisfies:
(1 0)12 (1+ 2N p(1 + p)) & p(1 - pp)
m 2
Then, using the following greedy algorithm, the probability that each row of matrix H is correctly reconstructed is:
Poorrect 2 1 1-
D.2 GREEDY ALGORITHM DESCRIPTION
For each row index = =,,,, 1,,,...,
1. Identify Set Sy: Find all rows in matrix W where the i-th column is 1.
S = (r Wr, = 1}
2. Compute Vector 8;: For each column j ...,.... calculate the sum of corresponding entries matrix M:
Σ Mrj
rES;
3. Select Maximum: Determine the column j" with the highest sum:
arg max
4. Update Matrix H: Assign the the i-th row of to column j":
Hi
14

---

Published as a conference paper at ICLR 2025
D.3 PROOF
For given row index and any incorrect column j + j (where j" is the correct column), define the difference:
As (
8/ (
Using the structure of the problem:
As C2 EH..
Hig)es
where: Ct2 the number rows in matrix W where column has a value of Cr is the number of rows in matrix W where both columns i and have value of 1.
(a) For Ctp
Since C₁ counts the number of times the i-th column in W contains 1, it can be written as:
C
W
where W,, is an independent Bernoulli random variable with probability Therefore:
Expectation:
E|c,,l = m
Variance:
Varic = m p(1 p)
(b) For e, (when * i):
Similarly, q, counts the number of times both columns and in matrix W contain 1. Since W and W., are independent Bernoulli variables:
Expectation:
Elca,11 m p²
Variance:
Var(c,,1 m m² p²)
We now calculate the expectation and variance of (J)
(a) Expectation of As
Using the linearity of expectation and the fact that the expectation of Hi.-- - Ht,j zero (as both binary variables):
E[As(1) (1) = Elc,al = m m.p(1-1 -p)
(b) Variance of As!!):
Since G, and og are independent, the variance As is the sum of the variances of the terms:
Var(As (6) Varlc,,l + > [Hi H² Var[cru]
Since (Hr,5- H2,5) takes values of (with occurring when one column contains 1 and
the other does not), let n be the number of times (Hi H1,;)² We have:
N
N
2N
+
k
k
Thus, the variance becomes:
Var (ms (1)
Var[ci,] n. Var(cy)
15

---

Published as a conference paper at ICLR 2025
Substitute the values of Varles,d and Var[c,,]]:
Var[As ( p(1 p) + n -221 r²)
By the Central Limit Theorem, since m is a fairly large number, we approximate As, as normally distributed variable:
As N (ElAs!",, VarlAs!
Thus, the probability of correct assignment is:
1|
P A () 0) VarlAo!
= +
where & is the CDF of the standard normal distribution.
Substitute the expressions for E +[As,"1 and Var[As (1) into the probability formula:
mp(1 p)
P A 0) = + mmpp1 p) + nmpp (1 - p2)
Simplify the Expression:
v /mmp(1 - p)
P As ()) > >0) =
1 + mp(1
$ (1-5) (1+ p(1+t))
Substitute n S K m
p(1-p)
((s) (1) > >0) + (1+ 11pppp(1 + p)
V mmp(1 p)
21-8
16


Reviews:
[{"review_id": "KQ6G53MC2m", "reviewer": "Reviewer_Qbkc", "summary": "This paper proposes URI, a variant of generative retrieval methods which learns the retrieval model and index construction jointly. Besides, this paper also proposes averaged partition entropy (APE) as a measure for index construction quality. In URI, the same decoder is used to generate tokens for both query and items and is trained in a using an EM algorithm. In E-step, the model is optimized through minimizing the cross entropy loss.  In M-step, To avoid model collapse while minimizing APE, two loss functions apart from the cross entropy loss are introduced to encourage that the token distribution w.r.t. each item follows a peak distribution while the expected token distribution w.r.t. the whole item space follows a uniform distribution. Finally, an additional dual encoder model is trained to identify the topk items as the final retrieved result. Experiments on KuaiSAR and Amazon Beauty/Toys/Games demonstrates the effectiveness of URI compared to other generative retrieval methods.", "questions": "My major concerns are listed in the weakness section.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "1. Learning to construct index jointly with retrieval models is an important problem for generative retrieval. The proposed method is well motivated.\n2. Performance is impressive, especially on the KuaiSAR dataset.", "weaknesses": "1. Writing can be further improved, especially the experiment section. The difference between URI and GR methods w/ URI index is unclear in Table 1. The analysis of token consistency is hard to understand and it will be better to demonstrate their differences to original URI through notations.\n2. Insufficient analysis and comparison on other generative methods (e.g., ASI [1] and GenRet [2] ) that learns both retrieval models and indexing jointly. \n3. The assumption of Theorem 1 seems very strong and there lacks empirical analysis on real-world dataset to verify the rationality. \n4. The comparison in experiments is not unfair. The candidate size of URI is larger than other GR methods which requires the mapping between the token representation and the item is bijective while URI allows each leaf node contains more than 1 items. Besides, URI is equipped with an additional ranker.\n\n[1] Yang T, Song M, Zhang Z, et al. Auto Search Indexer for End-to-End Document Retrieval[C]//Findings of the Association for Computational Linguistics: EMNLP 2023. 2023: 6955-6970.\n[2] Sun W, Yan L, Chen Z, et al. Learning to tokenize for generative retrieval[J]. Advances in Neural Information Processing Systems, 2024, 36.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "g3n6JA9DPh", "reviewer": "Reviewer_esfx", "summary": "This paper introduces URI (Unified Retrieval and Indexing), a framework for generative retrieval that integrates index construction with Transformer Decoder training. URI enhances consistency between the index and retrieval model. The authors also introduce Average Partition Entropy (APE), a model-independent metric for evaluating generative indices after index construction. Targeting this metric, they also propose an optimization algorithm with a theoretical explanation. Experiments confirm the effectiveness of the proposed method.", "questions": "Please see the above weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The motivation is clear. The authors argue that recent methods separate index construction from retrieval  in traditional and generative models, which limits the overall performance. Thus, this paper attempts to unify index construction and retrieval.\n- This paper provides a theoretical guarantee of the proposed greedy algorithm.\n- The experiments on multiple benchmark datasets validate the effectiveness of the proposed method compared with the baseline models.", "weaknesses": "- The authors overlook recent retrieval work such as [1,2]. The authors are encouraged to discuss them in the related work.\n- Since the final goal of the proposed method has many hyper-parameters in (11), more sensitivity analysis also is crucial. \n- Theorem 1 lacks the corresponding empirical results, so more evidence is needed.\n- Including both index construction and retrieval in the unified framework may increase computation and storage overhead, affecting efficiency. Therefore, an efficacy analysis (e.g., memory usage, training/inference time, and FLOPs) is necessary.\n\n[1] Improved Diversity-Promoting Collaborative Metric Learning for Recommendation. IEEE TPAMI, 2024.\n\n[2]Hierarchical latent relation modeling for collaborative metric learning,”inRecSys,2021", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EUGUJsmLgU", "reviewer": "Reviewer_qpuA", "summary": "This paper proposes URI (Unified framework for Retrieval and Indexing), a novel approach that unifies index construction and retrieval model training for generative retrieval systems. The key innovation is using the same Transformer decoder both as the retriever and the indexer, constructing the index simultaneously with model training rather than as a separate pre-processing step. The paper introduces Average Partition Entropy (APE) as a new metric for evaluating generative indices and provides theoretical analysis using EM algorithm to justify their approach. The authors demonstrate URI's effectiveness through experiments on three real-world datasets.", "questions": "- How does the computational complexity of URI compare to traditional two-stage approaches? Is there a significant overhead in training time or memory usage?\n- How sensitive is the method to the choice of initial conditions? Does the unified training require special initialization strategies?\n- Could the authors provide more details about how URI handles cold-start scenarios where new items or queries are added to the system?\n- How does URI perform in scenarios with highly imbalanced data distributions? Are there any special considerations or modifications needed?\n- The theoretical analysis assumes certain statistical conditions - how often are these conditions met in practice? What happens when they are not met?\n- Could the authors elaborate on potential extensions of URI to handle dynamic indices that need to be updated over time?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper identifies clear limitations in existing two-stage approaches where index construction is separated from retrieval model training. \n- The approach of using the decoder itself as an indexer is creative and original and introduction of APE as an evaluation metric is well-justified and useful\n- Addresses fundamental limitations in generative retrieval systems leads to potentially applicable across various retrieval tasks\n- Could influence future work in information retrieval and recommender systems", "weaknesses": "Limited Empirical Analysis\n- While experiments are conducted on three datasets, more details about these datasets would be helpful\n- Ablation studies could better isolate the impact of different components\n- Comparison with more baseline methods would strengthen the evaluation\n\nScalability Concerns\n- The paper doesn't thoroughly discuss computational complexity\n- Practical considerations for large-scale deployment are not addressed\n- Memory requirements for the unified approach could be significant", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ibuo8URdJ6", "reviewer": "Reviewer_1tQz", "summary": "The paper tackles item (document) retrieval problem, i.e., the algorithm should return the indices (and scores) of most relevant items given a query. In particular, the paper focuses on generative (decoder) based subset of algorithms. Current decoder-based methods follow a two-step procedure where first the index is learned over pre-trained features (e.g., DSI) and then the model is trained to generate in this space. The paper argues that the aforementioned procedure can be sub-optimal, and it proposes URI (a unified framework for generative retrieval) that simultaneously learns the index as well as the decoder. \n\nThe model is trained in a layer-by-layer manner (in contrast to existing ones). The first layer is trained by optimizing $\\mathcal{L_{E step}} + \\mathcal{L_{M step}} $ - the similarity between corresponding queries and items is optimized while avoiding index collapse. The buckets are then assigned for each item via balanced assignment. First layer tokens are then passed onto second layer which is trained via an additional memory reinforcement loss term to avoid forgetting. Finally, a ranker is trained that assigns score to query item pairs on the basis of inner product between pooled embeddings.  \n\nThe results are reported on KuaiSAR, Amazon Beauty and NQ-320K datasets with focus on demonstraing the benefit of URI index as compared to the baselines (DSI, NCI, TIGER etc.,).", "questions": "1. Please look into the points above concerning end-to-end baselines and training/inference time. \n2. The final scores are computed via a re-ranker that uses the inner product (on pooled query and doc embeddings). It may be worthwhile to try out a cross-encoder based re-ranker. \n3. Can you please comment on scalability of URI? How does the performance change with increase in number of docs / items?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper is decently written and easy to follow.\n2. The paper reports numbers on KuaiSAR, Amazon Beauty and NQ-320K datasets and it compares against leading methods including DSI and NCI. \n3. The proposed method yields good benefits when the index of baselines is replaced with the proposed URI based index as well as in terms of overall numbers. For e.g., the performance of DSI and NCI improved by ~3% on the NQ-320K datasets. The benefits are even starker on the KuaiSAR dataset where the overall performance improved by upto 14% as compared to the baselines.  \n4. The paper analyses the impact of URI index as compared to other popular indices including kMeans and VQ-VAE. The APE values are reported in Table 2. The benefits of layer wise training (token consistency), Adaptive weights ($\\mathcal{L}_{balance\\ w}$) and hard negative mining (during ranking). \n5. The paper provides hyper-parameter details in 5.4.", "weaknesses": "1. The results section compares against various methods including DR, DSI and NCI. It shows the impact of changing the original index with URI index alongside the overall comparison. However, the overall numbers don't seem to be for the best version of respective algorithms. It is perhaps done for consistency across algorithms? For e.g., the R@20 for NQ-320K can be in range of 0.56 - 0.89 for methods including ANCE, DSI, and NCI whereas URI is 0.057 (?). Please see https://arxiv.org/pdf/2206.02743. Please clarify if I am misinterpreting something. \n2. The paper is missing the training and inference times of the proposed method. These stats will provide a fair comparison against the baselines (in addition to already included theoretical complexity).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TWNSpzCK4V", "reviewer": "Reviewer_YJMb", "summary": "This paper proposes an end-to-end retriever and index, It simultaneously clusters query and relevant items and ensure that they both follow similar distribution over cluster nodes at each level. The URI (proposed approach) ensures that clusters are balanced and accurate.", "questions": "Please see the weakness section and answer the following:\n- How does method compares with EHI, what are the similarity and differences between EHI and URI\n- APE is standard practice to measure cluster quality why is it a contribution in your paper. (See ECLARE)\n- Add comparison on BIER benchmark datasets\n\nECLARE: Extreme Classification with Label Graph Correlations, WWW 2021.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This paper proposes an end-to-end retriever and indexer.\nResults are state of the art", "weaknesses": "- This paper looks a lot like EHI: End-to-end Learning of Hierarchical Index for Efficient Dense Retrieval with similar loss functions, TMLR 2024\n- Results on standard bier benchmark datasets are missing\n- In real world scenario new items are always popping in and out of existence, how does URI deals with zero shot items", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)