🤖 AI Summary
Production-grade Retrieval-Augmented Generation (RAG) systems often struggle to balance varying query difficulties under a fixed retrieval budget, compromising either answer quality or efficiency while adhering to strict Service Level Objectives (SLOs). This work proposes SAGE—the first general-purpose, online fine-tuning-free, SLO-aware adaptive retrieval strategy—which dynamically predicts the required number of retrieved passages per query via imitation learning, using lightweight initial retrieval features such as score distributions and rank gaps. SAGE incurs negligible overhead at inference time with no additional large model calls. Experiments show that under a 5-second P95 latency SLO, SAGE boosts compliance from 30% to 95% on Natural Questions, reduces latency by 36%, cuts retrieval cost by 51%, and sacrifices only 2 points in Exact Match. A single SAGE policy generalizes across multiple datasets and four major LLM families, improving SLO compliance by 45–52 percentage points without any loss in answer quality.
📝 Abstract
Retrieval-Augmented Generation (RAG) systems in production operate under strict service level objectives (SLOs) on tail latency and infrastructure cost. However, standard retrieval pipelines rely on fixed retrieval budgets that ignore query difficulty, over-retrieving for easy queries and under-serving hard ones, forcing operators to trade answer quality against SLO compliance. This paper proposes SAGE, a learned SLO-aware adaptive retrieval policy that dynamically selects the number of passages k per query. SAGE uses lightweight features derived from initial retrieval (e.g., score distributions, rank gaps, lexical signals) and is trained offline via imitation learning from an oracle that approximates optimal latency-quality trade-offs. At inference, it adds no LLM calls and minimal overhead. On Natural Questions, under a 5s P95 latency SLO, SAGE achieves 95% SLO compliance versus 30% for the best static baseline (k=20), reduces P95 latency by 36% and retrieval cost by 51% with only 2 percentage points Exact Match (EM) loss. A single policy trained on Natural Questions generalizes across HotpotQA, UnSeenTimeQA, and four LLM families (Llama, Qwen, Mistral, Gemma), consistently yielding +45-52 point SLO improvements without quality degradation.