Context Attribution with Multi-Armed Bandit Optimization

1. Problem Statement

In Retrieval-Augmented Generation (RAG), Large Language Models (LLMs) often generate answers based on retrieved external contexts. However, ensuring these answers are genuinely grounded in the provided context is difficult, as models frequently hallucinate or rely on ungrounded information.

The Core Challenge:

• Attribution: Identifying exactly which segments of the retrieved context (sentences or paragraphs) contributed to a specific generated token or answer.
• Computational Cost: Existing post-hoc attribution methods (e.g., SHAP, LIME, ContextCite) rely on perturbation-based approaches. They require systematically masking or removing context segments and re-querying the LLM to measure the impact on output probability.
• Scalability: For long contexts with $N$ segments, the number of possible subsets is exponential ($2^N$). Uniform or exhaustive sampling is computationally prohibitive, especially for black-box models accessed via APIs where query costs and latency are high.

2. Methodology: CAMAB Framework

The authors propose CAMAB (Context Attribution with Multi-Armed Bandit), a framework that reformulates context attribution as a Combinatorial Multi-Armed Bandit (CMAB) problem to optimize query efficiency.

A. Problem Formulation

• Arms: Each context segment $s_i$ is treated as an "arm."
• Actions (Super-arms): An action consists of selecting a subset of segments $S \subseteq C$.
• Reward: The reward $V(S)$ is defined as the average log-probability of the original response tokens $R$ given the subset $S$:
  $$V(S) = \frac{1}{T} \sum_{t=1}^{T} \log P_M(r_t | Q, S, r_{1:t-1})$$
  This measures how well the subset $S$ supports the generation of the original answer. Higher log-probability indicates better support.

B. Linear Thompson Sampling (LinTS)

To navigate the combinatorial space efficiently without exhaustive search, CAMAB employs Linear Thompson Sampling:

1. Linear Assumption: The reward is modeled as a linear function of the segments included in the subset:
   $$V(S) = \mathbf{w}^\top \mathbf{x} + \epsilon$$
   Where $\mathbf{x}$ is a binary vector indicating segment inclusion, and $\mathbf{w}$ represents the latent contribution weights of each segment.
2. Bayesian Inference: The algorithm maintains a posterior distribution over the weight vector $\mathbf{w}$. It uses a precision matrix $B_t$ to track correlations between segments.
3. Adaptive Exploration:
   • At each round $t$, the algorithm samples a weight vector $\tilde{\mathbf{w}}^{(t)}$ from the current posterior.
   • It constructs a subset $S_t$ by including all segments where $\tilde{w}_j > 0$.
   • It queries the LLM with $S_t$ to observe the reward and updates the posterior distribution.
4. Implicit Interaction Modeling: Although the model assumes linearity, the precision matrix $B_t$ implicitly captures interactions (synergies) between segments through off-diagonal correlations, allowing the model to learn that certain segments perform better together without explicit non-linear modeling.

C. Output

After a fixed query budget ($T_{max}$), the posterior mean weights $\hat{\mathbf{w}}$ serve as the final attribution scores. Segments are ranked by these scores to identify the most influential context.

3. Key Contributions

1. Novel Formulation: The first work to frame segment-level context attribution as a Combinatorial Multi-Armed Bandit problem.
2. Efficient Algorithm: Introduction of CAMAB, which utilizes Linear Thompson Sampling to adaptively explore the space of context subsets, significantly reducing the number of model queries required compared to uniform sampling.
3. Black-Box Compatibility: The method relies only on token-level log-probabilities (available via most commercial APIs), making it applicable to both open-source and proprietary black-box models without requiring internal model access (gradients/logits).
4. Implicit Interaction Capture: Demonstrates that the precision matrix in LinTS can implicitly capture segment dependencies and substitutability, overcoming limitations of simple linear assumptions.

4. Experimental Results

The authors evaluated CAMAB on three diverse benchmarks (HotpotQA, CNN/DailyMail, TyDi QA) using LLaMA-3.1-8B and Qwen2.5-7B.

• Query Efficiency: CAMAB achieves up to a 30% reduction in model queries compared to existing methods while matching or exceeding their performance.
  • In information-seeking tasks (HotpotQA, TyDi QA), CAMAB at a budget of $s=40$ often outperforms SHAP at $s=60$.
  • CAMAB converges to near-optimal performance much faster than SHAP or ContextCite.
• Attribution Quality:
  • Log-Probability Drop: CAMAB consistently yields higher drops when top-ranked segments are removed, indicating it correctly identifies critical evidence.
  • BERTScore: CAMAB achieves lower BERTScore (better semantic consistency) when removing attributed segments, confirming high fidelity.
  • Gold Fact Alignment: On HotpotQA, CAMAB outperforms SHAP in Precision@1 and F1 scores when compared against human-annotated gold supporting facts.
• Robustness: The method remains effective even under extremely tight budgets ($s=20$), where SHAP performance degrades significantly.

5. Significance and Impact

• Scalability for RAG: CAMAB provides a scalable solution for interpreting RAG systems in resource-constrained environments (e.g., high-cost API calls or real-time applications).
• Trustworthiness: By efficiently identifying the specific context segments responsible for an answer, it enhances the interpretability and trustworthiness of LLMs, helping to detect hallucinations.
• Practicality: Unlike methods requiring full logit access or massive computational resources, CAMAB is designed for the "black-box" reality of modern LLM deployment, making it immediately applicable to commercial systems.

Limitations: The method assumes a linear reward structure (though interactions are captured implicitly) and requires access to token-level log-probabilities, which is not universally available from all API providers. In scenarios with no budget constraints, traditional exhaustive methods might still yield slightly more plausible attributions.