Faster Retrieval for Better, More Efficient RAG

Maximum inner product search (MIPS) finds the database key most relevant to a query, a task central to recommendation, search, and retrieval-augmented generation. Exhaustive comparison becomes costly at scale, while conventional clustering methods route queries by distance to cluster centroids and can miss the best key. This is especially problematic when queries, such as short user questions, differ substantially from the longer documents represented by database keys. The proposed approach learns offline from the expected query distribution, paying an upfront training cost to improve retrieval speed and accuracy later. It introduces SupportNet, which estimates how well the best key in each cluster matches a query, and KeyNet, which predicts the best-matching key directly. Both are trained on query examples labeled using exhaustive search.

Tests of SupportNet used Natural Questions, with about 2.6 million passages divided into 128 clusters. The learned routers outperformed centroid routing: a small model reached 95% routing recall in 1.74 milliseconds, compared with 3.36 milliseconds for centroids, while a large model reached 97% recall in 1.74 milliseconds. When scanning only 5% of the database, learned routers identified the cluster containing the true best match about 90-93% of the time; the mid-size model achieved 92.6%, versus 87.5% for the baseline. KeyNet offers another way to improve search: it rewrites a query toward its predicted best match before passing it to an existing inverted-file index, allowing that optimized index to search more effectively.