Apple researchers train small networks to steer RAG searches

Apple researchers trained two small neural networks to steer queries toward likely maximum-inner-product matches, a core operation in retrieval-augmented generation. In the tested settings, the approach can reduce search work at inference time, but only after substantial offline training for a fixed vector database.

Maximum inner product search, or MIPS, finds the stored vector with the largest dot product for a query. The researchers frame the highest score across a database as a convex support function, whose gradient identifies the maximizing key. SupportNet learns that score function and ranks clusters to search. KeyNet takes another route: it predicts an approximation of the best key vector, which can then be passed to an existing approximate-search index.

That training requirement is the main tradeoff. The authors typically used 1 billion to 10 billion training samples in their experiments, generating each target through exhaustive MIPS. They also said augmented queries require new targets. Because the method depends on a representative query distribution and is learned for a fixed database, the research does not establish how often retraining would be needed as a corpus changes—or what that work would cost in time, energy or money.

In a Natural Questions experiment with 128 clusters, the authors reported that an extra-small SupportNet trained on 1 billion samples placed the exact target in the top-ranked cluster about 72% of the time, versus about 56% for centroid routing. Searching four clusters lifted the reported rate to about 95%. The centroid baseline needed 16 clusters to reach a similar level, using roughly four times the floating-point operations in that setup.

These findings come from retrieval benchmarks, not a deployed Apple product. The main evaluation covered four BEIR datasets with roughly 52,000 to 5.2 million keys; an appendix tested BioASQ at about 15 million keys. The paper did not demonstrate billion-key scaling. The authors also reported weaker performance when test queries drifted from the training distribution, although the learned method kept an advantage at a middle compute budget under moderate simulated drift.

The paper appears in the proceedings of the 43rd International Conference on Machine Learning. Apple also released a public implementation of the SupportNet and KeyNet approach.