ASI++: Towards Distributionally Balanced End-to-End Generative Retrieval

TL;DR

ASI++ advances fully end-to-end generative retrieval by introducing distribution balancing, representation bottleneck, and information consistency criteria that jointly optimize discrete ID assignments and dense representations. It extends the semantic indexing module with neural quantization, differentiable PQ, and RQ variants, and demonstrates SOTA-like performance on MS MARCO and strong practical results on ADS, including robust handling of new documents. The approach yields interpretable, hierarchical indexing space and improved ID utilization, addressing previous inefficiencies from long-tailed ID distributions. Overall, ASI++ offers a scalable, end-to-end framework that enhances retrieval accuracy while maintaining balanced, efficient indexing suitable for large-scale industrial deployments.

Abstract

Generative retrieval, a promising new paradigm in information retrieval, employs a seq2seq model to encode document features into parameters and decode relevant document identifiers (IDs) based on search queries. Existing generative retrieval solutions typically rely on a preprocessing stage to pre-define document IDs, which can suffer from a semantic gap between these IDs and the retrieval task. However, end-to-end training for both ID assignments and retrieval tasks is challenging due to the long-tailed distribution characteristics of real-world data, resulting in inefficient and unbalanced ID space utilization. To address these issues, we propose ASI++, a novel fully end-to-end generative retrieval method that aims to simultaneously learn balanced ID assignments and improve retrieval performance. ASI++ builds on the fully end-to-end training framework of vanilla ASI and introduces several key innovations. First, a distributionally balanced criterion addresses the imbalance in ID assignments, promoting more efficient utilization of the ID space. Next, a representation bottleneck criterion enhances dense representations to alleviate bottlenecks in learning ID assignments. Finally, an information consistency criterion integrates these processes into a joint optimization framework grounded in information theory. We further explore various module structures for learning ID assignments, including neural quantization, differentiable product quantization, and residual quantization. Extensive experiments on both public and industrial datasets demonstrate the effectiveness of ASI++ in improving retrieval performance and achieving balanced ID assignments.

Figures (6)

• Figure 1: Schematic diagram of training and inferring in our proposed ASI++. Given a search query $q$ or document $d$, ASI++ first encodes them into dense latent representations via encoder: $e_q$, $e_d$. Then ASI++ obtains discrete numeric index $id_q$, $id_d$ via semantic indexing module or decoder, respectively.
• Figure 2: Summary of main contributions of ASI++. (Upper) Previous end-to-end learning of GenIR face challenges in efficiently utilizing the dense latent and sparse indexing space. (Down) ASI++ optimizes the distribution of ID distributions, by (i) improving the quality and space utilization of dense representations to distinguish similar documents easier with decoder ($\mathcal{L}_{bot}$), (ii) fostering the utilization of indexing space to better match relevant query-documents ($\mathcal{L}_{di}$), and (iii) compressing most crucial information into dense bottleneck representations with information-theoretic $\mathcal{L}_{ib}$.
• Figure 3: Ablation of ASI++ on MS MARCO. We report the expectation of R@5/10.
• Figure 4: Performance of different variants of ASI++ on MS MARCO. We explore multiple implementation of the semantic indexing module: product quantization (PQ), residual quantization (RQ), and MLP-based neural quantization.
• Figure 5: Qualitative case study on the interpretability of IDs generated from ASI++ on MS MARCO.