Context Attribution with Multi-Armed Bandit Optimization

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Pan, Deng, Murugesan, Keerthiram, Hua, Ting, Moniz, Nuno, Chawla, Nitesh
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866910155067621376
author Pan, Deng
Murugesan, Keerthiram
Hua, Ting
Moniz, Nuno
Chawla, Nitesh
author_facet Pan, Deng
Murugesan, Keerthiram
Hua, Ting
Moniz, Nuno
Chawla, Nitesh
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2506_19977
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Context Attribution with Multi-Armed Bandit Optimization
Pan, Deng
Murugesan, Keerthiram
Hua, Ting
Moniz, Nuno
Chawla, Nitesh
Artificial Intelligence
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.
title Context Attribution with Multi-Armed Bandit Optimization
topic Artificial Intelligence
url https://arxiv.org/abs/2506.19977