Title: Context Attribution with Multi-Armed Bandit Optimization
Primary area: interpretability and explainable AI
Keywords: Large language model, attribution method, explainable machine learning, reinforcement learning

Abstract: Understanding which parts of the retrieved context contribute to a large language model’s generated answer is essential for building interpretable and trustworthy generative QA systems. We propose a novel framework that formulates context attribution as a combinatorial multi-armed bandit (CMAB) problem. Each context segment is treated as a bandit arm, and we employ Combinatorial Thompson Sampling (CTS) to efficiently explore the exponentially large space of context subsets under a limited query budget. Our method defines a reward function based on normalized token likelihoods, capturing how well a subset of segments supports the original model response. Unlike traditional perturbation-based attribution methods (e.g., SHAP), which sample subsets uniformly and incur high computational costs, our approach adaptively balances exploration and exploitation by leveraging posterior estimates of segment relevance. This leads to substantially improved query efficiency while maintaining high attribution fidelity. Extensive experiments on diverse datasets and LLMs demonstrate that our method achieves competitive attribution quality with fewer model queries.

TL;DR: The paper introduces CAMAB, a bandit-based framework that efficiently identifies which context segments support LLM outputs, achieving high attribution fidelity with far fewer queries than traditional perturbation methods