Making Transformer Decoders Better Differentiable Indexers

Many thanks for your constructive and detailed comments. We believe that your comments will make our work more robust and convincing. The detailed responses to your comments are listed as follows:

Weaknesses1: This paper looks a lot like EHI. What are the similarities and differences?

Response: The similarity is that both EHI and URI utilize a unified retriever and indexer and are trained end-to-end. However, there are several key differences:

1. The two methods use different paradigms for retrieval. URI employs a generative approach to predict the next token, while EHI uses a discriminative approach to predict the relevant nodes at the current layer of the tree index.
2. URI employs Transformer Decoder model with causal masks which makes the hierarchical retrieval more effective, whereas EHI uses a simple MLP.
3. The design of the loss functions differs; URI's loss function simulates the EM algorithm, merging the two stages into one and does not involve negative sampling, while EHI's loss function primarily relies on contrastive loss between positive and negative samples.

The original EHI paper does not seem to provide code. We attempted to reproduce the EHI algorithm, but found that the loss function was difficult to converge. Due to the lack of explicit balancing constraints, the index easily collapses. We will cite this paper in the camera-ready version and provide a detailed comparison of the two algorithms.

Weaknesses2: Results on standard beir benchmark datasets are missing

Response: In our paper, we selected three datasets to demonstrate the applicability of URI in the fields of Information Retrieval, General Recommendation, and Sequential Recommendation. We will now include the NQ320K dataset from the BEIR benchmark. The results are summarized in the table below. From the table, we can see that URI outperforms other methods on NQ320K dataset.

Weaknesses3: How does URI deal with zero shot items?

Response: We propose two feasible methods. The first method is ID-independent: for new items, we only use their textual embedding and directly input them into the trained decoder to generate their code in the index. The second method is ID-dependent: we use the textual embedding to find the most similar item to the new item and use the ID embedding of the item to initialize the ID embedding of the new item. We use 10% of the items in the dataset as cold-start items, while the other items are used for normal training. The results of Recall@10 obtained are as follows. "Random" means that we randomly assign codes to cold-start items. "No Cold-Start" means that all the items are used for normal training.

From the table, we can see that using ID-dependent methods for cold-start items yields the best results, but this methods require searching, which is time-consuming. For information retrieval tasks, using only textual features can achieve good results. However, for recommendation tasks, the information gain from IDs is relatively higher. The choice of method can be tailored to the specific scenario. This is a very meaningful topic, and we will conduct more detailed experiments and analysis in the camera-ready version.

Q2: APE is standard practice to measure cluster quality why is it a contribution in your paper. (See ECLARE)

Response: We are sorry for the confusion. In ECLARE, the metric most similar to APE is LMI, but APE and LMI are not the same. The calculation of LMI is related to the prediction model and measures the mutual information between the true clusters and the predicted clusters. In contrast, the calculation of APE is independent of the prediction model and is only related to the index. It measures the concentration of items related to the same query within the index.

Thank you for your helpful and detailed comments. We believe that your comments will make our work more robust and convincing. The detailed responses to your comments are listed as follows:

Weaknesses1: The results on NQ-320K are not the best version of respective algorithms.

Response: We apologize for the confusion caused by the typo error. NQ320K was mentioned only in Table 1, but the actual dataset used was Toys and Games. The cause of this error is that we initially chose NQ320K dataset, but since both it and the KuaiSAR dataset are IR tasks, we decided to replace it with the Toys and Games dataset to observe URI's performance across more tasks. All the results related are based on the Toys and Games dataset. We apologize again for our oversight in not correcting this error in Table 1. We will correct this error in the camera-ready version. To further address your concerns, we now provide the experimental results for NQ320K. Due to resource constraints, we use the Base version from the NCI paper. From the table, we can see that URI outperforms other methods on NQ320K dataset.

Weaknesses2: The paper is missing the training and inference time of the proposed method.

Response: Thank you for your suggestion. We will include the efficacy analysis in the camera-ready version. The table below presents the analysis on the KuaiSAR dataset. From the table, it can be seen that the training time for URI is relatively long due to the hierarchical training process. However, this is acceptable considering its superior performance. A higher QPS (Queries Per Second) indicates shorter inference time, and both URI and DR have high QPS because they use ranker sorting, whereas other methods require an additional final layer for beam search, which is more time-consuming. In terms of space usage, both DR and URI have higher model parameters due to the presence of an item encoder.

Q2: The final scores are computed via a re-ranker that uses the inner product. It may be worthwhile to try out a cross-encoder based re-ranker.

Response: The design of the re-ranker is not the main contribution of this paper. We chose to add an additional ranker model because our experiments showed that this approach performed better than adding extra tokens in the final layer. However, since you mentioned using the cross-attention method, we tried this approach on the sequential recommendation task for the Toys and Games dataset, and indeed, the results are better than using the dot product.

Q3: Can you please comment on scalability of URI? How does the performance change with increase in number of docs / items?

Response: The training of URI relies on positive query-item pairs. The more such samples there are, the better URI can capture the collaborative filtering information in the data distribution, leading to improved performance. Therefore, as the number of items increases, the effectiveness of URI improves. This is evident from the improvement observed across the three datasets mentioned in the paper, where the kuaiSAR dataset, which has the highest number of items, shows the greatest improvement with URI.

Q2: How sensitive is the method to the choice of initial conditions? Does the unified training require special initialization strategies?

Response: We have attempted to warm-start the model parameters. Specifically, we trained a two-stage model and then used the parameters of this model to warm-start URI. We found that this approach can accelerate the convergence of the model, but the final performance after convergence is not as good as training from scratch. The table below shows a comparison of the results on KuaiSAR between cold-start and warm-start.

Q3: Could the authors provide more details about how URI handles cold-start scenarios where new items or queries are added to the system?

Response: We propose two feasible methods for cold start items. The first method is ID-independent: for new items, we only use their textual embedding and directly input them into the trained decoder to generate their code in the index. The second method is ID-dependent: we use the textual embedding to find the most similar item to the new item and use the ID embedding of the item to initialize the ID embedding of the new item. We use 10% of the items in the dataset as cold-start items, while the other items are used for normal training. The results of Recall@10 obtained are as follows. "Random" means that we randomly assign codes to cold-start items. "No Cold-Start" means that all the items are used for normal training.

From the table, we can see that using ID-dependent methods for cold-start items yields the best results, but this method require searching, which is time-consuming. For information retrieval tasks, using only textual features can achieve good results. However, for recommendation tasks, the information gain from IDs is relatively higher. We obtain conclusions similar to those in our response to Weakness2. The choice of method can be tailored to the specific scenario.

Q4: How does URI perform in scenarios with highly imbalanced data distributions? Are there any special considerations or modifications needed?

Response: The biggest impact of imbalanced data distribution is that it can lead to an imbalanced index, which ultimately affects performance. However, during training, we applied both soft and hard balancing constraints to the index. The soft balancing constraint refers to a balance loss term included in our loss function, while the hard balancing constraint refers to limiting the load of each token node during item assignment. Both of these measures effectively prevent the issue of index imbalance.

Q5: How often are the conditions in Theorem met in practice? What happens when they are not met?

Response: In fact, the conditions of Theorem 1 can be easily satisfied, as we have demonstrated in Appendix B. For a tolerance level of 0.95, all three datasets used in the paper meet the conditions of Theorem 1. This indicates that the greedy algorithm described in Theorem 1 can correctly allocate at least 95% of the items for these datasets. When the conditions are not met, the greedy algorithm described in Theorem 1 may lead to inaccurate item allocation. However, this is not the final method in our paper. We use an end-to-end training approach. Although it simulates the greedy algorithm, fundamentally it captures the collaborative filtering signals in the dataset, which allows us to avoid bad allocation.

Q6: Could the authors elaborate on potential extensions of URI to handle dynamic indices that need to be updated over time?

Response: In our response to Q3, we provided methods for handling cold-start items, which addresses the issue of dynamic index insertion. For the deletion of outdated items, entries for those items can be directly removed from their respective buckets. If there is a large influx of new data, such as user-item interaction data, the index needs to be retrained. Warm-starting the training with existing parameters can lead to rapid convergence of the index, as it essentially only involves adjusting the allocation of a part of items.

Many thanks for your helpful and detailed comments. The detailed responses to your questions are listed as follows:

Weaknesses1: More details about these datasets would be helpful.

Response: We have provided detailed information about the dataset in Appendix B.

Weaknesses2: Ablation studies could better isolate the impact of different components.

Response: We will include more ablation studies. In the loss function of URI, each term is crucial, and directly ablating any term would cause the model to fail to converge, making it hard to evaluate. Here, we will first supplement ablation study on the importance of ID features and text features. We evaluated the performance of URI using only one type of feature on the KuaiSAR and Toys datasets. The table below shows the results for Recall@10.

From the table, it is evident that both ID features and text features play a crucial role in the training of URI. The importance of these features varies across different tasks. For IR tasks, text features are more important, whereas for recommendation tasks, ID features are more significant. This is a very meaningful topic, and we will conduct more detailed experiments and analysis in the camera-ready version.

Weaknesses3: Comparison with more baseline methods would strengthen the evaluation.

Response: We will add ASI and GenRet as two additional baselines. Both ASI[1] and GenRet[2] are designed for the Information Retrieval task. We have added a comparative experiment between these two methods and URI on the KuaiSAR dataset. The results are as follows:

ASI enhances the accuracy of the codebook through reconstruction, while GenRet designs a loss function to align the codes of the Query and Document, which makes them better than other baselines. However, both methods still rely on constructing indices in Euclidean space. The indexer(codebook) and the final retrieval model (Decoder) still exhibit some incompatibility, which limits their performance.

[1] Yang T, Song M, Zhang Z, et al. Auto Search Indexer for End-to-End Document Retrieval

[2] Sun W, Yan L, Chen Z, et al. Learning to tokenize for generative retrieval

Weaknesses4-6: Please see the responses to Q1.

Q1: How does the computational complexity of URI compare to two-stage approaches? Is there a significant overhead in training time or memory usage?

Response: In Section 4.5 of the paper, we provided a complexity analysis of URI and the two-stage model. From a complexity perspective, the training time complexity of URI is 2$L$ times that of the two-stage model. Since $L$ is usually small ($L=2$ in our paper), the additional time overhead is acceptable. Moreover, URI does not require pre-training, which saves that portion of the time overhead. As for memory usage, URI adds an encoder for the item part compared to the two-stage model, which is a completely acceptable overhead.

The table below presents the analysis on the KuaiSAR dataset. From the table, it can be seen that the training time for URI is relatively long due to the hierarchical training process. However, this is acceptable considering its superior performance. A higher QPS (Queries Per Second) indicates shorter inference time, and both URI and DR have high QPS because they use ranker sorting, whereas other methods require an additional final layer for beam search, which is more time-consuming. In terms of memory usage, both DR and URI have more model parameters due to the presence of an item encoder.

Thank you for your insightful suggestions! We appreciate your encouragement and provide detailed responses to your comments below:

Weaknesses1: The authors overlook recent retrieval work such as HLR and DPCML. The authors are encouraged to discuss them in the related work.

Response: Thank you for your suggestion! Both of these works are based on Collaborative Metric Learning. HLR proposes a hierarchical CML model that jointly captures latent user-item and item-item relationships from implicit data. DPCML proposes a method called Diversity-Promoting Collaborative Metric Learning, which considers the commonly ignored minority interests of users. We will include a discussion of these two works in the related work section of the camera-ready version.

Weaknesses2: Since the final goal of the proposed method has many hyper-parameters in (11), more sensitivity analysis of hyper-parameters is crucial.

Response: In Loss Function 11, there are only two parameters, $k$ and $l$. Here, $k$ represents the width of the index, and $l$ represents the current layer being trained, with $l$ having a maximum value of $L$, which is the depth of the index. We have already provided a sensitivity analysis of these two parameters in Section 5.4 of the paper. The other symbols are not hyperparameters, and their meanings are explained in the sections where they are introduced. Additionally, we have further summarized the meaning of each symbol in the Appendix.

Weaknesses3: Theorem 1 lacks the corresponding empirical results, so more evidence is needed.

Response: The conditions of Theorem 1 can be easily satisfied, as we have demonstrated in Appendix B. For a tolerance level of 0.95, all three datasets used in the paper meet the conditions of Theorem 1. This indicates that the greedy algorithm described in Theorem 1 can correctly allocate at least 95% of the items in these datasets.

Weaknesses4: An efficacy analysis (e.g., memory usage, training/inference time, and FLOPs) is necessary.

Response: Thank you for your suggestion! We will include the efficacy analysis in the camera-ready version. The table below presents the analysis on the KuaiSAR dataset. From the table, it can be seen that the training time for URI is relatively long due to the hierarchical training process. However, this is acceptable considering its superior performance. A higher QPS (Queries Per Second) indicates shorter inference time, and both URI and DR have high QPS because they use ranker sorting, whereas other methods require an additional final layer for beam search, which is more time-consuming. In terms of memory usage, both DR and URI have more model parameters due to the presence of an item encoder.

Many thanks for your helpful and detailed comments. We believe that your comments will make our work more robust and convincing. The detailed responses to your comments are listed as follows:

Weaknesses1: Writing can be further improved: 1.difference between URI and GR methods w/ URI index; 2. The analysis of token consistency

Response: For the two points you mentioned, we will add the following explanations to the paper:

1. “GR methods w/ URI index” indicates that we retain the model structure and training methods of the GR methods, but only change the index. For example, for NCI, we use the PAWA network as described in the original paper. However, the index is no longer obtained through hierarchical clustering; instead, it is replaced with the URI-trained index.
2. “Token consistency” refers to the consistency of the same code across different levels. We will add symbolic representations to clarify this. For instance, in an index with 2 layers, the codes for items might be $[3^1, 2^2]$, $[4^1, 2^2]$, and $[2^1, 3^2]$, where the superscripts indicate the layer number. “Token consistency” has two meanings: The first meaning is that the same code across different levels has similarity, such as $2^2$ and $2^1$ in $[3^1, 2^2]$ and $[2^1, 3^2]$. The second meaning is that the same code under different parent nodes also has similarity, such as $2^2$ in $[3^1, 2^2]$ and $[4^1, 2^2]$ .

For other parts of the experiment, we will also use symbols or diagrams to make the presentation clearer. If you have any specific suggestions for modifications, please let us know.

Weaknesses2: Insufficient analysis and comparison on other generative methods (e.g., ASI and GenRet)

Response: Both ASI and GenRet are designed for the Information Retrieval task. We have conducted a comparative experiment between these two methods and URI on the KuaiSAR and NQ320K dataset. The results of Recall@10 are as follows:

ASI enhances the accuracy of the codebook through reconstruction, while GenRet designs a loss function to align the codes of the Query and Document, making them better than other baselines. However, both methods still rely on constructing indices in Euclidean space. The indexer (codebook) and the final retrieval model (Decoder) still exhibit some incompatibility, limiting their performance. We will include these two baselines in the camera-ready version and provide experiments on additional datasets.

Weaknesses3: The assumption of Theorem 1 seems very strong and there lacks empirical analysis on real-world dataset.

Response: The conditions of Theorem 1 can be easily satisfied, as demonstrated in Appendix B. For a tolerance level of 0.95, all three datasets used in the paper meet the conditions of Theorem 1. This indicates that the greedy algorithm described in Theorem 1 can correctly allocate at least 95% of the items in these datasets.

Weaknesses4: The comparison in experiments is not unfair. The candidate size of URI is large and URI is equipped with an additional ranker.

Response: We apologize for the misunderstanding. In fact, the candidate size is the same for all methods, as the beam size is uniformly set to 20. For URI and DR, this means selecting the top 20 buckets, while for other methods, it means selecting the top 20 parent nodes. Regardless of whether they are buckets or parent nodes, the candidate items attached are the same. The difference lies in the ranking process: URI and DR use an additional ranker, while other models still use the probabilities output by the Decoder for sorting. We use an additional ranker because it yields better results. Additionally, in Table 1, we compared the performance of the model structures of all methods on both the original index and the URI index. The results show a significant improvement in performance on the URI index. This phenomenon is unrelated to the ranker and demonstrates the inherent superiority of the URI index.

Dear Reviewer,

We would like to express our sincere gratitude for all the help you have provided in reviewing our work. Your feedback has been invaluable in helping us improve our submission. Since the deadline for the author-reviewer discussion is approaching and we have yet to receive your further feedback, we hope to kindly request your help once again. We would greatly appreciate it if we could have further discussions with you to ensure that our submission meets the high standards of ICLR.

Best regards,

Authors