AR-RAG: Autoregressive Retrieval Augmentation for Image Generation

We sincerely thank the reviewer for this insightful suggestion regarding zero-shot image generation and the importance of evaluating performance in such scenarios. In response, we conducted a targeted experiment to assess our method’s capabilities in this setting.

Specifically, we constructed an artificial retrieval database containing 30 images of Taylor Swift, collected from the internet, while ensuring no other human faces were included. The artificial retrieval database also contains 10,000 images from CC12M. (We filtered out images that contain human faces based on their captions) We then used the prompt “a woman is playing guitar on the stage” and generated 10 images each with both DAiD and FAiD. We observed that, when Taylor Swift’s images were present in the retrieval database, the generated faces in most images closely resembled her appearance. Conversely, after removing Taylor Swift’s images from the database and repeating the generation, the generated faces no longer displayed her features.

These results demonstrate that our retrieval-augmented approach can effectively leverage reference images for specific object or person generation, and is sensitive to the presence of such references in the retrieval set. We agree that further systematic exploration of zero-shot and few-shot settings is a promising direction and will add significant value to future versions of our work. We will make sure to include this new result in the revised version.

Question: The inference speed is notably slow, hindering practical application. Unlike Image RAG, which requires only one retrieval, AR-RAG performs O(N²) retrievals, where N is the number of generated patches. This quadratic complexity significantly increases computational overhead.

Answer: We would like to clarify that the retrieval complexity of our method is O(N), not O(N²) as stated in the review. Specifically, for N patches, our approach performs N retrievals (i.e., one per image patch) rather than a quadratic number of retrievals. For each patch, we retain only the top 2,048 retrieved candidates. Furthermore, we employ the Faiss library to construct our retrieval database, which is highly optimized for large-scale and hierarchical nearest neighbor search. As demonstrated in Table 4 of our submission, the average inference time for Janus-Pro + DAiD to generate 100 images increases by only 0.22% compared to the Janus-Pro baseline and most of the addtional inference time is attributed to the retrieval process during inference. This negligible increase underscores the efficiency and practical feasibility of our retrieval method. We hope this clarification addresses the concerns regarding computational overhead.

Question: The experiments are conducted against relatively weak baselines. For example, the "cat playing basketball" issue in the teaser might not exist in the latest methods like FLUX, indicating that the comparative evaluation could be more rigorous.

Answer: As suggested by the reviewer, we compare our method Janus-Pro + FAiD against two SOTA diffusion models including Flux-dev [1] and SD3-medium [2] on GenEval and MJHQ 30K. As reported in the Table below, these models consistently underperform relative to our proposed method, Janus-Pro + FAiD, despite their larger parameter sizes and pre-training on larger datasets.

Notably, the GenEval benchmark specifically assesses the ability to generate objects with rare attributes and to compose images containing multiple, less commonly co-occurring objects. The comparatively lower performance of Flux-dev and SD3-medium on these benchmarks suggests that these models have limitations in handling complex generation tasks—such as modeling spatial relations, attribute binding, and multi-object composition—even in their latest versions. We believe this comparative analysis is both fair and rigorous, and it demonstrates the unique strengths of our approach.

[1] Black Forest Labs. Flux, 2024.

[2] Scaling rectified flow transformers for high-resolution image synthesis.

Thanks for the authors' answers. My concern is resoloved. I tend to weakly accept this paper now.

We appreciate the reviewer's thoughtful feedback regarding computational efficiency and the source of performance improvements in our method.

To demonstrate that the performance improvements stem from the retrieval mechanism rather than merely from distributional similarities between training and evaluation databases, we conducted the suggested experiments that fine-tune the baseline model on JourneyDB and our mixture dataset. The results are presented below:

As expected, fine-tuning on in-domain datasets (JourneyDB and the Midjourney training set) improves performance, with JourneyDB fine-tuning outperforming the mixture dataset approach. However, our FAiD method significantly outperforms both fine-tuning approaches across all metrics, demonstrating clear benefits from our retrieval mechanism beyond what can be achieved through domain-specific fine-tuning alone.

Regarding computational overhead, we would like to first clarify that: (1) Our retrieval process is optimized using the FAISS library with GPU acceleration, resulting in minimal overhead. As shown in Table 4, the DAiD method adds only 0.22% additional time compared to the vanilla Janus-Pro model. (2) The primary computational cost in FAiD comes from the additional convolutional operations in the decoder. This overhead is expected to decrease proportionally when running on faster GPU devices.

Furthermore, we agree with the reviewer's suggestion to provide more explicit demonstrations of benefits versus costs. As suggested, we are conducting additional experiments with parameter-augmented baseline models. Specifically, we have added adapters to the vanilla Janus-Pro model and are fine-tuning it with the identical training setting. Due to time constraints, these experiments are still in progress. We commit to sharing these results during the discussion period and incorporating them in the next version of our paper.

Question: While the paper already evaluates the method on several standard benchmarks, including additional ones could strengthen the empirical support. Potential candidates include: a) T2I-Benchmark [A], to assess improvements in image compositionality. b) RareBench [B], to evaluate the benefits of patch-based retrieval in handling rare concepts.

Answer: We thank the reviewer for these constructive suggestions. In response, we have conducted additional experiments on both RareBench and T2I-Bench, using the official prompts and evaluation protocols. The results are presented below. RareBench Results:

T2I-Bench Results:

All models in above two tabels are the same models used in Tabel 1 in our paper including DAiD and FAiD. We adopt the official prompts and evaluation code for both benchmarks.

As observed on RareBench, our methods demonstrates substantial improvements over baselines, particularly for multi-object composition, indicating its effectiveness in generating rare or novel combinations not commonly found in the training data.

On T2I-Bench, our methods also outperform the baseline Janus model by a considerable margin. While our model does not surpass the diffusion-based RDM model in all evaluation sub-categories, we respectfully note that this may be partly attributable to the relatively lower performance of the original Janus architecture. Notably, our method outperforms all baselines by a substantial margin in the "color" and "numeracy" categories. We did not report results for 2d-spatial and 3d-spatial categories, as all evaluated models achieved a score of 0 on these sub-tasks.

Question: A qualitative visualization of retrieved patches during generation would help assess retrieval quality and provide insight into how patch-based retrieval aids compositional generation.

Answer: We thank the reviewer for the excellent suggestion to visualize the retrieved patches during generation. In our visualization experiment, we take an image, randomly select a patch position, and retrieve the most relevant patch from the retrieval database conditioned on the surrounding patches. We then replace the selected patch with the retrieved one and observe that the retrieved patches integrate smoothly with their surrounding context.

Due to conference policies, we are unable to submit additional PDF materials during the rebuttal phase. However, we will ensure that these visualizations are included in the revised version of the paper. Thank you again for this valuable suggestion.

Question: In addition to inference time, reporting GPU memory usage would offer a more comprehensive view of the computational overhead introduced by patch-based retrieval.

Answer: Here is the breakdown of inference memory usage for our models.

Loading Janus Pro: 4717MB

Performing a single image generation with Janus Pro: 5511MB

Loading Janus Pro+FAiD: 5153MB

Performing a single image generation with Janus Pro+FAiD: 5949MB

Hosting Retrieval database: 77346MB

As shown, incorporating FAiD modules results in only a modest increase in memory usage during both model loading and inference (approximately 400 MB). While the retrieval database itself requires additional memory, it is stored separately and can be efficiently managed. Overall, our retrieval approach introduces limited overhead relative to the improved performance it provides.

Question: Investigating the effect of the hyperparameter λ in DAiD (Eq. 4) would be valuable. An ablation study on this parameter could clarify the trade-off between model-predicted generation and retrieval-based guidance.

Answer: We thank the reviewer for pointing out the importance of ablation on the hyperparameter λ in DAiD (Eq. 4). We apologize for not including this analysis in the main paper. As noted in Figure 7 in Appendix C.2 (Hyperparameter Optimization), we performed a hyperparameter search on λ to study its effect. In the revised version, we will ensure that this ablation study and the associated trade-offs are highlighted and discussed more prominently in the main text.