# Context Attribution with Multi-Armed Bandit Optimization

arXiv: https://arxiv.org/abs/2506.19977 · HTML: https://mlsanity.com/paper/2506.19977/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in RAG robustness and adaptive retrieval, one of 15 directions in Retrieval-augmented generation.
- Builds on 23 of its 31 references that are in the map.
- First posted on arXiv 24 Jun 2025.

## Placed in

- Retrieval-augmented generation / [RAG robustness and adaptive retrieval](https://mlsanity.com/rag/rag-robustness-and-adaptive-retrieval/)

## Builds on (most cited first, 23 of 23)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [A Unified Approach to Interpreting Model Predictions](https://mlsanity.com/paper/1705.07874/) (2017)
- [BERTScore: Evaluating Text Generation with BERT](https://mlsanity.com/paper/1904.09675/) (2019)
- ["Why Should I Trust You?": Explaining the Predictions of Any Classifier](https://mlsanity.com/paper/1602.04938/) (2016)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [Retrieval-Augmented Generation for Large Language Models: A Survey](https://mlsanity.com/paper/2312.10997/) (2023)
- [WebGPT: Browser-assisted question-answering with human feedback](https://mlsanity.com/paper/2112.09332/) (2021)
- [Get To The Point: Summarization with Pointer-Generator Networks](https://mlsanity.com/paper/1704.04368/) (2017)
- [BARTScore: Evaluating Generated Text as Text Generation](https://mlsanity.com/paper/2106.11520/) (2021)
- [Enabling Large Language Models to Generate Text with Citations](https://mlsanity.com/paper/2305.14627/) (2023)
- [TyDi QA: A Benchmark for Information-Seeking Question Answering in Typologically Diverse Languages](https://mlsanity.com/paper/2003.05002/) (2020)
- [Teaching language models to support answers with verified quotes](https://mlsanity.com/paper/2203.11147/) (2022)
- [FastSHAP: Real-Time Shapley Value Estimation](https://mlsanity.com/paper/2107.07436/) (2021)
- [ContextCite: Attributing Model Generation to Context](https://mlsanity.com/paper/2409.00729/) (2024)
- [LongCite: Enabling LLMs to Generate Fine-grained Citations in Long-context QA](https://mlsanity.com/paper/2409.02897/) (2024)
- [Thompson Sampling for Combinatorial Semi-Bandits](https://mlsanity.com/paper/1803.04623/) (2018)
- [Generating Hierarchical Explanations on Text Classification via Feature Interaction Detection](https://mlsanity.com/paper/2004.02015/) (2020)
- [BanditSum: Extractive Summarization as a Contextual Bandit](https://mlsanity.com/paper/1809.09672/) (2018)
- [Training Language Models to Generate Text with Citations via Fine-grained Rewards](https://mlsanity.com/paper/2402.04315/) (2024)
- [Ada-SISE: Adaptive Semantic Input Sampling for Efficient Explanation of Convolutional Neural Networks](https://mlsanity.com/paper/2102.07799/) (2021)
- [Multi-Level Explanations for Generative Language Models](https://mlsanity.com/paper/2403.14459/) (2024)
- [Automation and Feature Selection Enhancement with Reinforcement Learning (RL)](https://mlsanity.com/paper/2503.11991/) (2025)
- [Fast Explanations via Policy Gradient-Optimized Explainer](https://mlsanity.com/paper/2405.18664/) (2024)

## Built on by (most cited first, 2 of 2)

- [In-Context Credit Assignment via the Core](https://mlsanity.com/paper/2605.06920/) (2026)
- [A Component-Based Survey of Interactions between Large Language Models and Multi-Armed Bandits](https://mlsanity.com/paper/2601.12945/) (2026)

## Abstract (arXiv)

Understanding which parts of the retrieved context contribute to a large language model's generated answer is essential for building interpretable and trustworthy retrieval-augmented generation. We propose a novel framework that formulates context attribution as a combinatorial multi-armed bandit problem. We utilize Linear Thompson Sampling to efficiently identify the most influential context segments while minimizing the number of model queries. Our reward function leverages token log-probabilities to measure how well a subset of segments supports the original response, making it applicable to both open-source and black-box API-based models. Unlike SHAP and other perturbation-based methods that sample subsets uniformly, our approach adaptively prioritizes informative subsets based on posterior estimates of segment relevance, reducing computational costs. Experiments on multiple QA benchmarks demonstrate that our method achieves up to 30% reduction in model queries while matching or exceeding the attribution quality of existing approaches. Our code is publicly available at https://github.com/pd90506/camab.