Paper: IRGen: Generative Modeling for Image Retrieval (ICLR-2024)

Reviewer comment (Presentation):
The paper is written in a way that is hard to follow and understand, both the method and the contributions. Figure 1 is ambiguous or wrong.

Author response:
Figure 1 illustrates our AR training methodology that given an input image, we feed it to our transformer (encoder and decoder) and the output is supervised by the image identifier that is obtained from the query's nearest neighbor through semantic image tokenizer (which is trained before AR model). To address any misunderstandings, we kindly request more specific details regarding your concerns. 
We itemize our contribution as follows for clarity:
(1) To the best of our knowledge, our work represents the first effort in employing a sequence-to-sequence generative model to address image retrieval in a fully end-to-end manner that directly generates the nearest neighbor's identifier givent the query.
(2) We introduce a semantic image tokenizer explicitly designed for image retrieval to obtain image identifiers, facilitating autoregressive training. It is essential to note that this tokenizer differs significantly from previous quantization models, as our identifier solely represents the image and is not employed for distance computation, as is the case in quantization.
(3) Extensive experiments show that our approach exhibits substantial improvements in results, achieving new state-of-the-art results. We have demonstrated its scalability to million-scale datasets and provided throughput number, showcasing its near-real-time performance.

Reviewer comment (Novelty):
A lot of relevant work is missing and therefore so is proper discussion on the method's relation to other works, especially in the area of deep quantization and hashing. Some examples are listed, but many more can be found through the works listed.

Author response:
We provide a more in-depth discussion comparing our approach with prior joint learning methods. In the existing literature, early methods predominantly focused on three distinct aspects: embedding learning for enhancing search accuracy, embedding compression (quantization/hashing) for reducing memory cost, and clustering (often involving quantization) or neighborhood graphs for improving search efficiency. However, as each of these components is essential for a large-scale retrieval system, integrating individual efforts may result in suboptimal solutions. Consequently, there has been a growing interest in studying joint learning. However, all these prior works (please refer to our general response for references to prior works) focus on end-to-end training of encoding and compression together while overlooking the indexing part (meaning non-exhaustive search), which is challenging to formulate as differentiable. They leverage quantization (compression index) for ANN search (with distance approximation) but still rely on linear scan in their experiments. In contrast, our approach reframes retrieval as a completely novel end-to-end task, eliminating explicit distinctions between embedding, compression, and indexing. Our method involves offline storage of image identifiers acquired from the proposed tokenizer. Notably, this differs conceptually from quantization, as we do not use the tokenized identifier for distance comparison as in quantization. Retrieving a query's nearest neighbor in our method only involves passing the query through the autoregressive (AR) model. We will incorporate this discussion into the final paper.

Reviewer comment (Novelty):
The large AR model on top acts as a differentiable index. There is no proper discussion on how this compares to other differentiable indexes from a technical perspective. The only difference mentioned to other recent differentiable indexing methods at the end of section 4 is that the input features (document identifiers) are obtained in a different way, but this has nothing to do with the actual differentiable index modeling part.

Author response:
Differentiable indexing in our approach means non-exhaustive search, whereas deep quantization methods use differentiable index for embedding compression. Our focus on image retrieval marks the first effort. We introduce a novel semantic tokenizer specifically designed for image retrieval. We add a direct comparison using NCI's hierarchical k-means tokenizer versus ours, emphasizing the effectiveness of our proposed semantic tokenizer.

| Method \ Inshop   | Prec@1   | Prec@10  | Prec@20  | 
|:-------- |:--------:| --------:| --------:| 
| IRGen     |  92.4    |  87.0    |  86.6    | 
| NCI (HKM)  | 89.0 |81.6 |79.8 |

Reviewer comment (Novelty):
The training objective is not much different from pairwise objectives on top of quantized codes used in quantization/hashing works. Although one has to pass through the index, only image pair codes are used for computing the loss. This is why beam search is needed at test time, i.e., there is a clear train-test discrepancy here, same as training deep hashing. It is unclear how this is more "end2end", as the objective of the loss is not a full database ranking objective. Also, beam search is a complex data structure that keeps multiple paths and includes "a prefix tree containing valid codes". One could argue that beam search is itself as complex as any other indexing structure.

Author response:
There is a serious misunderstanding. Our model does not use image pair codes for computing the loss, as indicated in your comment. Instead, we aim to generate the codes for the nearest neighbor corresponding to the given input. The distinctions between our model and previous deep quantization methods have been elaborated in the response above. Our more end-to-end property stems from that we directly find the nearest neighbor through the learned generative model, compared to deep quantization that needs to calculate distance between query and each database's codes. When only the top-1 sample is required, beam search is unnecessary, and generation can be performed by straightforwardly selecting the identifier with the highest probability, ensuring no train-test discrepancy. However, for top-K retrieval, we leverage beam search with a configured beam size of K. In this scenario, a train-test discrepancy emerges. One potential solution is the further integration of beam search into the training process, a direction that is worth pursuing.

Reviewer comment (Experiments):
Experimental validation is generally lacking. As mentioned above, the contribution of the two modules separately should be studied. Also, reporting only the precision and recall metrics in isolation is not ideal for image retrieval. Results with mAP (or at least mAP100) for all datasets should be presented. Moreover, results on classic retrieval datasets like Oxford and Paris [ir1] are missing.

Author response:
It appears that certain experiments presented in the paper may have been overlooked. The contributions of the two modules have been examined, and please refer to point (2) in the 3rd response. We have conducted a mAP comparison on two million-scale datasets, ImageNet and Places365, as in Table 3 of the paper. Regarding the remaining three datasets in Table 1 and 2, a direct mAP or mAP@100 comparison is not reasonable due to the limited number of ground truth samples, all less than 100 for each query. It is important to note that previous works, including the compared baselines, only report recall@K metric in their experiments. To provide a comprehensive comparison, we reported precision@K metric in our paper. Furthermore, we have evaluated on three widely used benchmarks and two million-scale datasets. The classical retrieval datasets, Oxford and Paris, are relatively small scale, with 5062 images for Oxford and 6412 images for Paris. These datasets are commonly employed for zero-shot retrieval scenario, particularly for pre-training models. Notwithstanding, we add comparison to Paris, which consistently validates the effectiveness of our proposed method.


| Method \ RParis |   mAP for M  | mAP for H |
| :------| ------: | -------: |
|   IRGen   | 75.7 |40.4 |
| FT-CLIP| 71.4 |27.1 |

Reviewer comment (Experiments):
Inference speed experiments are lacking and misleading. The proposed method has a large computational overhead, i.e., a large transformer decoder instead of a simple index. A comparison versus FT-CLIP plus some indexing is needed, but the proposed approach could be even orders of magnitude slower. Excluding feature extraction with CLIP (the proposed method uses CLIP as the encoder for the tokenizer), the proposed method at test time needs to 1) tokenize, 2) pass through a large ViT decoder (12 transformer decoder blocks with dimension 768 for the small datasets, 24 layers for ImageNet), and 3) perform beam search. In contrast, FT-CLIP with a basic IVFPQ index only needs to search the index, which is extremely fast.

Author response:
There is a serious misunderstanding that needs clarification. At test time, our model does not use the tokenizer, as this step is exclusive to database images and can be executed offline. The query only undergoes a transformer encoder and decoder process to search for its nearest neighbor. Similarly, other baselines like FT-CLIP and IRT also involve passing the query through a transformer encoder to obtain the query embedding, resulting in same time cost for the encoder at test time. Secondly, despite our decoder being either 12-layer or 24-layer, its efficiency stands out due to the notably short sequence length. Lastly, beam search is conducted concurrently with the decoding process. We have provided throughput comparison on Inshop dataset in Figure 5 in the paper. Here we provide a comprehensive comparison focusing on the ImageNet dataset. It is crucial to emphasize that a retrieval system's storage considerations extend beyond just the model; it must also store features for all database vectors, constituting a substantial portion of storage requirements. In terms of inference cost, our CPU inference cost is 0.30s per image, which is not orders of magnitude larger compared to FT-CLIP with linear scan. Although FT-CLIP with SPANN lowers the inference cost, it comes at the expense of a substantial drop in performance. Our GPU inference cost, tested on a Tesla P100 PCIe 16GB GPU, is 0.037s/image, which is near-real-time, considering practical retrieval systems at the order of 10 milliseconds. Moreover, our approach demonstrates competitive performance and significantly reduces the memory cost. We anticipate that the prevalent reliance on GPUs for generative models will swiftly lead to newer GPU generations with increased processing power, more efficient memory management, and improved hardware support, thereby enhancing the inference speed of our autoregressive retrieval. Additionally, our model, being GPT-based, can leverage rapidly developed acceleration techniques designed for GPT.

|  Method | Model (#Params) |  Database storage  |
| -------:| ----------------- |  --------------:|
|      IRGen | ViT-base encoder (86M)+ 24-layer decoder (307M)  | 4.9MB： 4 bytes (id)/image with 1281167 samples    |
| FT-CLIP | ViT-base (86M)    | 3.67GB：  768x4 bytes (feature)/image with 1281167 samples  |


|  Method | Search         | Inference cost | mAP@100    |
| -------:| -------------- | --------------:| --- |
|      IRGen | Autoregressive |   0.037s/image |   76.0  |
| FT-CLIP | Linear scan (SPANN)   |   0.14s/image (0.007s/image) |   77.0 (65.5)  |

Reviewer comment (Experiments):
The model is highly task-specific: a different trained AR model/index is needed for every database. It is also unclear how such an index can handle changes in the database.

Author response:
Certain experiments presented in our study seem to have been overlooked. We have conducted experiments to demonstrate the capability to handle distribution change by adding new, unseen data within the same class, as shown in Table 4 in the paper. In scenarios where more new data come in and the distribution of gallery data has changed drastically, all existing data-dependent retrieval methods, including embedding model in image retrieval as well as dense retrieval methods in document retrieval, have to re-train the model and re-construct the data index. Essentially, dynamic updates only support a small number of new data that do not substantially alter the distribution.  Nevertheless, it's important to highlight that data-dependent models continue to be preferred due to their superior performance. In future work we may extend IRGen to an adaptive structure capable of capturing the evolving distribution of new data during training.

Reviewer comment (Theory):
While emphasizing global feature extraction from class tokens might improve efficiency, there's a risk of overlooking crucial spatial information present in other parts of the image, possibly leading to incomplete or less accurate retrieval results.

Author response:
Image retrieval is a high-level vision task, akin to extreme classification problem, reliant on a global understanding of the image, where global features are pivotal. Our method proves effective not only in general retrieval, such as ImageNet, but also excels in fine-grained tasks like Inshop clothes retrieval, CUB retrieval, and car retrieval, showcasing its ability to capture necessary details for fine-grained search using global features. Recognizing scenarios where similarity is defined over small image areas requiring spatial information, we suggest a solution involving a pre-processing step that detects these important regions. For instance, person re-identification is such a case utilizing detection before search. Alternatively, incorporating both global and local features ensures a comprehensive image representation but may entail complex models with increased computational demands. Striking a balance between efficiency and efficacy is a longstanding challenge in retrieval. This paper has successfully present a novel approach that is conceptually different from previous joint learning methods, and demonstrate state-of-the-art performance on widely used benchmarks with near-real-time inference speed.

Reviewer comment (Experiments):
The paper proposes the use of beam search for efficient top-K retrieval, but this method can be computationally intensive, especially for large image databases. Additionally, validating each generated image identifier can be a time-consuming process, even with their proposed prefix tree optimization.

Author response:
First, the proposed prefix tree only contains valid codes and thus eliminates the need of validating each generated image identifier, largely enhancing the inference speed. We then provide a comprehensive comparison of model capacity and inference cost in the table below, focusing on the million-scale ImageNet dataset. It is crucial to emphasize that a retrieval system's storage considerations extend beyond just the model; it must also store features for all database vectors, constituting a substantial portion of storage requirements, especially in the context of a large database. While compression can alleviate this concern, the performance is severely compromised due to quantization. In terms of inference cost, our CPU inference cost is 0.30s per image, which is not orders of magnitude larger compared to FT-CLIP with linear scan. Although FT-CLIP with SPANN lowers the inference cost, it comes at the expense of a substantial drop in performance. Our GPU inference cost, tested on a Tesla P100 PCIe 16GB GPU, is 0.037s/image, which is near-real-time, considering practical retrieval times at the order of 10 milliseconds. Moreover, our approach demonstrates competitive performance compared to FT-CLIP with linear scan, and significantly reduces the memory cost. For inference speed, we anticipate that the prevalent reliance on GPUs for generative models will swiftly lead to newer GPU generations with increased processing power, more efficient memory management, and improved hardware support, thereby enhancing the inference speed of our autoregressive retrieval. Additionally, our model, being GPT-based, can leverage rapidly developed acceleration techniques designed for GPT.

Reviewer comment (Experiments):
The approach, especially the beam search with constraints, may face scalability issues when dealing with vast and diverse image datasets. As databases grow, the efficiency and accuracy of the method may be challenged.

Author response:
Our current objective is to showcase the effectiveness of the proposed generative based model. We have conducted experiments on three widely used retrieval benchmarks, and show the scalability on other two wellknown million-scale datasets. Further scaling to handle billion-scale datasets poses a considerable challenge due to the substantially increased demand for resources, including storage and computation, which currently exceeds our available capacity.  In future work, we may address billion-scale scenario with a hybrid solution where billions of images can be clustered into millions of semantic clusters, and our proposed model IRGen can then be utilized to retrieve the relevant cluster identifiers. This approach would enable us to leverage current success in million-scale datasets and handle billion-scale datasets in a more efficient manner.

Reviewer comment (Experiments):
While the method aims to capture high-level semantics, it might not generalize well across diverse datasets with varying characteristics. The paper does not address how the model would adapt to datasets with vastly different semantic structures or image content.

Author response:
In scenarios where more new data come in with vastly different semantic structure, leading to the distribution of gallery data change drastically, all existing data-dependent retrieval methods, including embedding model in image retrieval as well as dense retrieval methods in document retrieval, have to re-train the model and re-construct the data index. Essentially, generalization (or dynamic updates) only supports a small number of new data that do not substantially alter the distribution. Our experiments have demonstrated this case in Table 4 in the paper, showcasing the model's adaptability to minor distribution changes. Nevertheless, it's important to highlight that data-dependent models continue to be preferred due to their superior performance. In future we may extend IRGen to an adaptive structure capable of capturing the evolving distribution of new data during training.

Reviewer comment (Experiments):
There seems to be a limited discussion on how this method performs compared to other state-of-the-art techniques. Without comprehensive comparative studies, it's challenging to ascertain the model's effectiveness and superiority in real-world scenarios.

Author response:
In our experiments, we extensively compared our approach with state-of-the-art techniques. The dataset leaderboards, accessible at the paperwithcode website, highlight CGD as the leader. Our paper has provided a detailed comparison of our approach with it, demonstrating a significant improvement margin. We will make it more clear in the final revision.

Reviewer comment (Experiments):
How do you ensure that the emphasis on global feature extraction from class tokens doesn't compromise the finer spatial details of the image, which could be vital for certain retrieval scenarios?

Author response:
Already presented in the weakness part. Moreover, as we fine-tune the CLIP image encoder concurrently with our model, it is plausible that the acquired features will dynamically adapt to capture intricate details essential for effective similarity search. This joint optimization process enhances the model's ability to discern subtle nuances, as verified by the fine-grained retrieval.

Reviewer comment (Reproducibility):
Given the complexities of beam search, especially for large image databases, what specific optimizations have you implemented to ensure real-time or near-real-time retrieval performances?

Author response:
We introduce a prefix tree that eliminates nodes lacking children, ensuring that only valid image IDs can be reached through the prefix tree. This design significantly accelerates beam search, as there is no need to check the validity of each generated image identifier. A detailed comparison of model capacity and inference cost is provided in the weaknesses part, demonstrating that our approach achieves near-real-time retrieval performance.

Reviewer comment (Experiments):
How does the proposed model handle noisy or imperfect image datasets? Are there specific preprocessing steps or augmentations recommended to enhance retrieval accuracy?

Author response:
We would like to point out that to demonstrate the effectiveness of the proposed framework, our approach, same to all previous methods, evaluate the retrieval performance on the widely used benchmarks tailored to minimize the presence of noisy images. In real-world scenarios, however, noise and imperfections are inevitable. Two strategies can be employed. 1) Data Filtering: utilize a designed mechanism for outlier removal or employ a simple classifier trained on labeled data (good quality vs. low quality) to filter extreme noise from the dataset. 2) Model regularization such as data augmentation (rotation, flip, etc) to enhance the model robustness, and dropout regularization to prevent overfitting.

Reviewer comment (Experiments):
How adaptable is the autoregressive encoder-decoder architecture to other types of multimedia content, such as videos or audio? Are there potential modifications or extensions to the proposed method for such content?

Author response:
We would like to break the extension into two scenarios, one is extension to retrieval in another datatype such as videos or audio, and the other is extension to multi-modal retrieval in multi-modal data. For the first case, the autoregressive encoder-decoder architecture can be modified for alternative data types. Both the AR model and the semantic image tokenizer would need adjustments to accommodate video or audio inputs, presenting potential challenges in the process. In the second, more intriguing scenario, an additional operation is necessary to align the feature space from multiple modalities. We plan to delve into this aspect in future research.

Reviewer comment (Theory):
Given the model's dependency on class tokens, how does it handle images that may not fit neatly into predefined classes or those that belong to multiple overlapping classes?

Author response:
For the first question, we've highlighted in the weaknesses part that all existing data-dependent retrieval methods require re-training and re-construction of the data index to handle new images with significantly different semantic structures. To address this, we envision potential modifications to IRGen, transforming it into an adaptive structure capable of capturing the evolving distribution of new data during training in the future. Regarding the second question, when dealing with images belonging to multiple overlapping classes, the primary challenge is defining similarity between two images. Typically, a score reflecting the degree of overlap between two images can be assigned. This score indicates that if the two images share more class labels, their similarity score is higher. Once similarity scores are determined, we can adapt the training in the AR model to utilize triplet ranking loss.

Reviewer comment (Experiments):
You have mentioned using classification loss for training. Were other loss functions explored during the development? If so, how did they impact the model's performance?

Author response:
The classification loss is used for our image identifier training. In addition to that, we propose to utilize M levels of partial reconstruction loss, which is shown in Equation (3) in the paper. We add ablation experiments to show how it impact the model's performance.

| Model \ Inshop   | Prec@1   | Prec@10  | Prec@20  | Prec@30  |
|:-------- |:--------:| --------:| --------:| --------:|
| Full loss   |  92.4    |  87.0    |  86.6    |  86.5    |
| Only classification loss     |  92.1    | 85.1    |  83.4   |  82.5    |

Reviewer comment (Theory):
retrieval with autoregressive means the retrieval may not enjoy the benefit of maximum inner-product search

Author response:
We find the noted weakness a bit perplexing, particularly in the absence of specific details regarding the particular benefit that may not be enjoyed by retrieval with an autoregressive approach. Our interpretation is based on our best understanding of the context. Maximum inner-product search (MIPS) pertains to a similarity search problem where the objective is to efficiently locate the data point in a collection that maximizes the inner product with a given query vector. Nonetheless, if the ultimate goal is to retrieve similar data, our methodology allows us to formulate and address this challenge optimally in an end-to-end manner, rather than a two-step process involving initial feature embedding followed by maximum inner-product search. We acknowledge the potential scenarios where data embedding becomes imperative, such as compact feature learning for transmission purposes. In such cases, one could also treat the embedded data as input and train an autoregressive model to identify the query's nearest neighbor, bypassing the tedious tuning over  the joint integration of compression and indexing.

Reviewer comment (Reproducibility):
as the current retrieval is done with GPU, if using CPU, the decoder process will be significantly slowed down and may affect the retrieval speed

Author response:
We provide a comprehensive comparison of model capacity and inference cost in the table below, focusing on the ImageNet dataset. It is crucial to emphasize that a retrieval system's storage considerations extend beyond just the model; it must also store features for all database vectors, constituting a substantial portion of storage requirements, especially in the context of a large database. In terms of inference cost, our CPU inference cost is 0.30s per image, which is not orders of magnitude larger compared to FT-CLIP with linear scan. Although FT-CLIP with SPANN lowers the inference cost, it comes at the expense of a substantial drop in performance. Our GPU inference cost, tested on a Tesla P100 PCIe 16GB GPU, is 0.037s/image, which is near-real-time, considering practical retrieval systems at the order of 10 milliseconds. Moreover, our approach demonstrates competitive performance compared to FT-CLIP with linear scan, and significantly reduces the memory cost. We anticipate that the prevalent reliance on GPUs for generative models will lead to newer GPU generations with increased processing power, more efficient memory management, and improved hardware support, thereby enhancing the inference speed of our autoregressive retrieval. Additionally, our model, being GPT-based, can leverage rapidly developed acceleration techniques designed for GPT. Our model also opens avenues for the integration of GPT in the development of more intelligent and context-aware information retrieval and generation systems.

|  Method | Model (#Params) |  Database storage  |
| -------:| ----------------- |  --------------:|
|      IRGen | ViT-base encoder (86M)+ 24-layer decoder (307M)  | 4.9MB： 4 bytes (id)/image with 1281167 samples    |
| FT-CLIP | ViT-base (86M)    | 3.67GB：  768x4 bytes (feature)/image with 1281167 samples  |


|  Method | Search         | Inference cost | mAP@100    |
| -------:| -------------- | --------------:| --- |
|      IRGen | Autoregressive |   0.037s/image |   76.0  |
| FT-CLIP | Linear scan (SPANN)   |   0.14s/image (0.007s/image) |   77.0 (65.5)  |

Reviewer comment (Reproducibility):
does the semantic tokenizer share the weights with the encoder?

Author response:
No. They are separate models with different weights.

Reviewer comment (Experiments):
seems like we can directly use the code from semantic tokenizer for image retrieval -- similar to how product quantization performs retrieval, what is the necessity of employing the decoder? What is the disadvantange/how much performance degrades if we just use the semantic tokenizer's code for retrieval?

Author response:
We add the experimental results that directly use the codes from our semantic tokenizer for image retrieval. Specifically, we use asymmetric distance computation (ADC) to approximate the distance and use linear scan for exhaustive search, similar to product quantization. The observed significant degradation in performance underscores the imperative need for our autoregressive decoder.

| Method \ Inshop   | Prec@1   | Prec@10  | Prec@20  | Prec@30  |
|:-------- |:--------:| --------:| --------:| --------:|
| IRGen     |  92.4    |  87.0    |  86.6    |  86.5    |
| RQ-ADC|     62.5    |  48.0  |  44.8   |  43.3  |

Reviewer comment (Experiments):
what if you train a normal classification model, then perform PQ to obtain the "tokens" instead of the proposed semantic tokenizer?

Author response:
We add the experimental comparison on Inshop dataset with the counterpart that utilizes the PQ tokenizer (performing PQ after training a normal classification model) instead of our proposed semantic tokenizer. The inferior performance of the counterpart underscores the effectiveness of our semantic tokenizer.

| Model \ Inshop   | Prec@1   | Prec@10  | Prec@20  | Prec@30  |
|:-------- |:--------:| --------:| --------:| --------:|
| IRGen (proposed tokenizer)     |  92.4    |  87.0    |  86.6    |  86.5    |
| IRGen (PQ tokenizer)      |  73.8    |  66.1    |  61.8    |  59.2    |

Reviewer comment (Experiments):
is a two-stage training process? can they be trained end-to-end?

Author response:
Indeed, our current process involves training a tokenizer to obtain the image identifier and subsequently training an AR model for end-to-end search. However, it's crucial to note that our model is specifically optimized with an end-to-end retrieval target, eliminating the train-inference mismatch presented in previous retrieval models. Looking forward, we envision the potential for a joint learning of both stages, further enhancing the image identifier's awareness of the retrieval process. This is an avenue we plan to explore in future work.

Reviewer comment (Novelty):
The concept of joint training the embedding model and index structure is not entirely novel.

Author response:
We offer a detailed comparison of our approach with previous joint learning methods. Early methods in the literature primarily concentrated on three key aspects: embedding learning to enhance search accuracy, embedding compression (quantization/hashing) to reduce memory cost, and clustering (often involving quantization) or neighborhood graphs to improve search efficiency. However, integrating these essential components individually may lead to suboptimal solutions for large-scale retrieval systems. As a result, there is a growing interest in the study of joint learning. However, all prior works (please refer to our general response for references to prior works) focus on end-to-end training of encoding and compression together while overlooking the indexing part (meaning non-exhaustive search), which is challenging to formulate as differentiable. They leverage quantization (compression index) for ANN search (with distance approximation) but still rely on linear scan in their experiments. In contrast, our approach reframes retrieval as a completely novel end-to-end task, eliminating explicit distinctions between embedding, compression, and indexing. Our method involves offline storage of image identifiers acquired from the proposed tokenizer. Notably, this differs conceptually from quantization, as we do not use the tokenized identifier for distance comparison as in quantization. Retrieving a query's nearest neighbor in our method only involves passing the query through the autoregressive (AR) model.

Reviewer comment (Experiments):
The approach scales well for million scale datasets, but not billions.

Author response:
Thank you for acknowledging our results with million-scale datasets. Handling billion-scale datasets presents a significant challenge, given the heightened resource demands, including storage and computation, that currently surpass our available capacity. Our primary goal is to highlight the effectiveness of the proposed generative model. To demonstrate this, we conducted experiments on three retrieval benchmarks and two additional million-scale datasets. In future endeavors, we aim to address the billion-scale scenario through a hybrid solution. This involves clustering billions of images into millions of semantic clusters, allowing our model IRGen to efficiently retrieve relevant cluster identifiers. This approach would enable us to handle billion-scale datasets in a more efficient manner.

Reviewer comment (Novelty):
The concept of jointly training the retrieval / embedding model and ANN index structures (or search models) is not novel. It would be great to have a more in depth review of these approaches. A quick example: "Joint Learning of Deep Retrieval Model and Product Quantization based Embedding Index", SIGIR'21. https://dl.acm.org/doi/10.1145/3404835.3462988

Author response:
Refer to the weaknesses section for an in-depth explanation. The mentioned example also involves joint learning of encoding and compression (referred to as embedding index in the paper), relying on linear scan search in their experiments. In contrast, our approach defines retrieval as a novel end-to-end task, directly retrieving a query's nearest neighbor by passing through the AR model, without explicit modeling for embedding, compression, and indexing.

Reviewer comment (Evaluation):
Figure 4: The precision of the proposed approach increases as top-k increases. The expectation is that there should be a trade-off.

Author response:
Traditional methods tend to witness a decline in precision with an increasing retrieval set size (K). While they initially achieve high precision by prioritizing the most relevant images in the top ranks, the precision drops rapidly as K grows due to the retrieval of irrelevant and confusing images. This indicates a struggle to retrieve harder positive images, often necessitating a reranking stage after the initial retrieval. In contrast, our AR decoder consistently retrieves relevant samples as the retrieval list lengthens. Our top-K performance even exhibits a slight increase with a larger K, showcasing its proficiency in handling challenging examples. This capability can be attributed to our approach of randomly sampling codes from images in the same class as the query image during training, enhancing the model's capacity to handle difficult samples. This advantage potentially mitigates the need for a post-reranking stage in our model to some extent.

Reviewer comment (Experiments):
Table 4: The generalization capability of the approach is demonstrated by taking out 5% of the ImageNet dataset during training and using them for inference/testing. However, those 5% examples are coming from the same ImageNet dataset distribution. The model already learned the semantic classes and their visual variations. Therefore, this experiment does not fully demonstrate the out-of-domain capability of the proposed method. Also, it is not clear if the same treatment was applied for the other baseline model compared. Need more clarity and justification for this experiment.

Author response:
In scenarios where more new data come in and the distribution of gallery data has changed drastically, all existing data-dependent retrieval methods, including embedding model in image retrieval as well as dense retrieval methods in document retrieval, have to re-train the model and re-construct the data index. Essentially, dynamic updates only support a small number of new data that do not substantially alter the distribution. Our experiments have demonstrated the capability of handling this case through adding new, unseen data within the same class, showcasing the model's adaptability to minor distribution changes. Nevertheless, it's important to highlight that data-dependent models continue to be preferred due to their superior performance. In future work, there is a potential extension of IRGen towards an adaptive structure capable of capturing evolving distribution during training.  For our experiment, the other baseline models employed the same treatment for fair comparison, and we will provide further clarification in the revised version.

Reviewer comment (Experiments):
There is a need for a table comparing the capacity and inference cost of the baseline models to the proposed solution in the Appendix (for Apples-to-Apples comparisons). A similar comparison is also needed for the retrieval stage. At million scale, even linear search is quite practical. It would be also great to have ablations where the model capacity is varied and the overall performance is evaluated.

Author response:
We provide a comprehensive comparison of model capacity and inference cost in the table below, focusing on the ImageNet dataset. It is crucial to emphasize that a retrieval system's storage considerations extend beyond just the model; it must also store features for all database vectors, constituting a substantial portion of storage requirements. While compression can alleviate this concern, the performance is severely compromised due to quantization. In terms of inference cost, our CPU inference cost is 0.30s per image, which is not orders of magnitude larger compared to FT-CLIP with linear scan. Although FT-CLIP with SPANN lowers the inference cost, it comes at the expense of a substantial drop in performance. Our GPU inference cost, tested on a Tesla P100 PCIe 16GB GPU, is 0.037s/image, which is near-real-time, considering practical retrieval systems at the order of 10 milliseconds. Moreover, our approach demonstrates competitive performance compared to FT-CLIP with linear scan, and significantly reduces the memory cost. For inference speed, we anticipate that the prevalent reliance on GPUs for generative models will lead to newer GPU generations with increased processing power, more efficient memory management, and improved hardware support, thereby enhancing the inference speed of our autoregressive retrieval. Additionally, our model, being GPT-based, can leverage rapidly developed acceleration techniques designed for GPT. 
To explore the scaling law concerning model capacity, we conducted experiments using an AR model with a 12-layer decoder on the ImageNet dataset, given current time limit and computational constraints. The results in the table below indicate that our approach achieves improved performance as the model capacity increases. We intend to further investigate the performance of our method in scenarios involving larger model sizes and expanded database sizes in future work. We will add this discussion in the Appendix.

|  Method | Model (#Params) |  Database storage  |
| -------:| ----------------- |  --------------:|
|      IRGen | ViT-base encoder (86M)+ 24-layer decoder (307M)  | 4.9MB： 4 bytes (id)/image with 1281167 samples    |
| FT-CLIP | ViT-base (86M)    | 3.67GB：  768x4 bytes (feature)/image with 1281167 samples  |


|  Method | Search         | Inference cost | mAP@100    |
| -------:| -------------- | --------------:| --- |
|      IRGen | Autoregressive |  0.037s/image|   76.0  |
| FT-CLIP | Linear scan (SPANN)   |   0.14s/image (0.007s/image) |   77.0 (65.5)  |



|  Method | Model (#Params) |  mAP@100  |
| -------:| ----------------- |  --------------:|
|      IRGen | ViT-base encoder (86M)+ 24-layer decoder (307M)  |  76.0  |
| IRGen | ViT-base encoder (86M) + 12-layer decoder (86M)   |   74.1 |

Reviewer comment (Experiments):
Residual Vector Quantization also uses beam search and RVQ codes may be used for efficiently knn search using prefix trees (or better with Aggregating Tree). Therefore the inference time operations are quite similar to the IRGen paper except that the previous paper does not train an encoding for the search task. This could also serve as a natural baseline which would demonstrate the need for joint training of encoder and indexing/search structures. OR, after training the encoder, could we combine and use the below paper for a more scalable (billions) search? A discussion should suffice. 
Liu et al, "Generalized Residual Vector Quantization and Aggregating Tree for Large Scale Search", IEEE T on Multimedia, 2017

Author response:
We applied beam search to the prefix tree derived from RQ codes learned by the tokenizer. However, the performance of this approach was found to be significantly inferior to our proposed method, highlighting the importance of our end-to-end training approach. Notably, our way of accelerating beam search using a prefix tree containing only valid codes aligns with the concept of aggregating tree in the referenced work. This involves merging nodes along a path when there are no branches, effectively eliminating nodes lacking children (invalid codes in our case). Such an approach markedly reduces computational overhead. This technical design is pivotal for enhancing beam search efficiency, particularly advantageous in billion-scale search scenarios. We will incorporate this reference in the final paper.


| Method \ ImageNet  | mAP@100  |
|:-------- |:--------:|
| IRGen|  76.0   |
| RQ prefix-tree|  56.7  |