Research map / Large language models / Retrieval and knowledge
Adapting to Non-Stationary Environments: Multi-Armed Bandit Enhanced Retrieval-Augmented Generation on Knowledge Graphs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a multi-objective, deep contextual multi-armed bandit router that picks among KG retrieval methods (SPARQL, LLM-agent, dense) for RAG, using online user feedback and a Generalized Gini Index aggregation to handle non-stationarity and accuracy-latency trade-offs.
Problem
KG-based RAG systems rely on a single retrieval method or static neural routers that need fully labeled data and periodic retraining, so they degrade under non-stationary query distributions and backend upgrades while needing to balance accuracy and response time.
Approach
Treat each KG retrieval method as a bandit arm and select it with a DistilBERT-based deep contextual bandit trained on partial (chosen-arm-only) feedback, updated online per query, with a Generalized Gini Index loss combining hit, recall, and retrieval-latency objectives.
Key results
- In the stationary setting, GGI-MAB reaches 86.64 hit and 75.60 recall on WebQSP and 79.35 hit and 72.02 recall on CWQ, the best among listed methods.
- Under domain shift (WebQSP to CWQ), the method achieves 76.35 hit vs 67.93 for retrieval ensemble, with 69.47 recall.
- Under retriever upgrade, it reaches 84.80 hit and 72.24 recall at 5.88 s/query, versus 83.74 hit and 67.52 recall at 15.00 s/query for ensemble.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use GGI-MAB instead of the strongest reported static router | KG-based RAG · WebQSP | Hit: 86.20 -> 86.64 | Helped |
| Use GGI-MAB instead of the strongest reported baseline | KG-based RAG · CWQ | Hit: 78.53 -> 79.35 | Helped |
| Use online MAB retrieval selection instead of retrieval ensemble | KG-based RAG under retriever update · Retriever update | Test Recall: 67.52 ± 0.77 -> 72.24 ± 0.71 | Helped |
| Use online MAB retrieval selection instead of the LLM Router | KG-based RAG under retriever update · Retriever update | Test Retrieval Delay: 10.36 ± 3.98 -> 5.88 ± 0.99 | Helped |
| Use online MAB retrieval selection instead of NN Router | KG-based RAG under domain shift · Domain shift | Test Hit: 69.57 ± 4.52 -> 76.35 ± 0.69 | Helped |
| Use online MAB retrieval selection instead of NN Router | KG-based RAG under domain shift · Domain shift | Test Retrieval Delay: 7.65 ± 1.81 -> 11.63 ± 0.28 | Hurt |
| Add GGI aggregation instead of learnable weight aggregation (ablation) | MO-MAB · WebQSP stationary | Test Hit: 85.31 ± 0.55 -> 86.64 ± 0.29 | Helped |
| Add GGI aggregation instead of learnable weight aggregation (ablation) | MO-MAB · WebQSP stationary | Test Recall: 74.38 ± 0.48 -> 75.60 ± 0.38 | Helped |
| Add GGI aggregation instead of learnable weight aggregation (ablation) | MO-MAB · WebQSP stationary | Test Retrieval Delay: 5.20 ± 0.58 -> 4.84 ± 0.81 | Helped |
| Enable online adaptation rather than keeping the model offline (ablation) | GGI-MO-MAB · System degradation | Test Hit: 61.70 ± 2.08 -> 77.55 ± 2.40 | Helped |
| Enable online adaptation rather than keeping the model offline (ablation) | GGI-MO-MAB · System degradation | Test Recall: 42.50 ± 2.62 -> 60.07 ± 3.29 | Helped |
| Enable online adaptation rather than keeping the model offline (ablation) | GGI-MO-MAB · System degradation | Test Retrieval Delay: 4.84 ± 0.81 -> 11.23 ± 3.21 | Hurt |
Surprising findings
- Retrieval ensemble of multiple retrievers does not beat the best single-router approach and is the slowest (15 s per query), as denoising heterogeneous retrieval results is hard.
Limitations
- Non-stationary evaluation is simulated by switching retrievers and datasets rather than live user feedback; hit is used as proxy for binary user feedback.
- Action space is limited to a few retrievers (DECAF, ChatKBQA, Reason-on-Graph in the implementation detail), and some baselines are reported from prior work with mismatched details.
Experimental setup
Datasets: WebQSP ComplexWebQuestions (CWQ)
Benchmarks: WebQSP ComplexWebQuestions (CWQ)
Metrics: hit recall retrieval delay (seconds per query)
Baselines: BGE DECAF StructGPT ChatKBQA Think-on-Graph Reason-on-Graph LLM Router NN Router UCB Thompson Sampling LinUCB SO-Deep-MAB
Builds on: Retrieval-Augmented Generation ChatKBQA Reason-on-Graph Think-on-Graph DistilBERT
Closest papers by meaning
- DynamicRAG: Leveraging Outputs of Large Language Model as Feedback for Dynamic Reranking in Retrieval-Augmented Generation (NeurIPS 2025 · similarity 0.90)
- Simple is Effective: The Roles of Graphs and Large Language Models in Knowledge-Graph-Based Retrieval-Augmented Generation (ICLR 2025 · similarity 0.89)
- Long-Context LLMs Meet RAG: Overcoming Challenges for Long Inputs in RAG (ICLR 2025 · similarity 0.89)
- Can Knowledge-Graph-based Retrieval Augmented Generation Really Retrieve What You Need? (NeurIPS 2025 · similarity 0.89)
- Scent of Knowledge: Optimizing Search-Enhanced Reasoning with Information Foraging (NeurIPS 2025 · similarity 0.89)
- RAGGED: Towards Informed Design of Scalable and Stable RAG Systems (ICML 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.