CAMAB: Context Attribution with Multi-Armed Bandit Optimization

Official code for "Context Attribution with Multi-Armed Bandit Optimization", accepted to the Findings of the Association for Computational Linguistics: ACL 2026 (arXiv:2506.19977).

CAMAB formulates segment-level context attribution in retrieval-augmented generation as a combinatorial multi-armed bandit problem and solves it with Linear Thompson Sampling (LinTS). Each context segment is an arm, a subset of segments is a super-arm, and the reward is the average log-probability of the original response under that subset. LinTS adaptively explores the space of subsets, yielding attribution scores that match or exceed SHAP / ContextCite with up to 30% fewer model queries.

Installation

git clone https://github.com/pd90506/camab.git
cd camab
pip install -r requirements.txt
pip install -e .

Copy the example config and fill in your Together API key:

cp config.example.json config.json
# edit config.json and paste your TOGETHER_API_KEY

config.json is gitignored. Attribution and evaluation for CAMAB / SHAP / LOO / Random run against the Together API (black-box, token-level log-probabilities only); ContextCite additionally needs a local GPU copy of the model because it regresses on full logits.

Quickstart

Run CAMAB attribution + evaluation on 10 HotpotQA samples, budget 40:

python llmexp/experiments/run_evaluation.py \
  --explainer mab --dataset_name hotpot --model llama \
  --n_sample 10 --n_perturb 40 --mode both

Outputs:

• output/attr_cache/llama_hotpot_mab_iter_40_n10_max40.jsonl — per-sample attribution scores.
• output/results/llama_hotpot_mab_iter_40_n10_max40.csv — Top-k log-probability drop and BERTScore at k ∈ {1, 3, 5}.

Reproducing the paper

All results are regenerated end-to-end from the Together API — no attribution caches are shipped with this repository. The paper's main experiments use --n_sample 10000. Expect the full sweep to take hours per setting; scale --n_sample down for a faster sanity run.

Tables 1–2 (main results, LLaMA-3.1-8B, s ∈ {20, 40, 60})

for ds in hotpot cnn tydi; do
  for s in 20 40 60; do
    for method in mab shap contextcite loo random; do
      python llmexp/experiments/run_evaluation.py \
        --explainer $method --dataset_name $ds --model llama \
        --n_sample 10000 --n_perturb $s --mode both
    done
  done
done

Table 4 (Qwen2.5-7B, s=40)

for ds in hotpot cnn tydi; do
  for method in mab shap contextcite loo random; do
    python llmexp/experiments/run_evaluation.py \
      --explainer $method --dataset_name $ds --model qwen \
      --n_sample 10000 --n_perturb 40 --mode both
  done
done

Table 6 (larger query budget, s=100, HotpotQA)

for method in mab shap; do
  python llmexp/experiments/run_evaluation.py \
    --explainer $method --dataset_name hotpot --model llama \
    --n_sample 10000 --n_perturb 100 --mode both
done

Table 3 and Figure 1 (alignment with HotpotQA gold supporting facts)

# Build the 200-sample HotpotQA distractor/validation set
python scripts/analysis/build_hotpot_val_samples.py

# Run MAB / SHAP / Random on that set
python scripts/analysis/run_attribution.py --explainer mab    --n_perturb 40
python scripts/analysis/run_attribution.py --explainer shap   --n_perturb 40
python scripts/analysis/run_attribution.py --explainer random

# Compute P@k / F1 / AUROC / AP against gold supporting facts
python scripts/analysis/evaluate_supporting_facts.py
python scripts/analysis/aggregate_report.py

# Figure 1: marginal attribution bars + posterior correlation heatmap
python scripts/analysis/heatmap_interaction.py

A pre-computed Figure 1 (output/results/heatmap_interaction.png) and its diagnostics (output/results/heatmap_interaction_diagnostics.npz) are included so the plot can be re-rendered with heatmap_interaction.py --dry_run without any API calls.

Repository layout

camab/
├── llmexp/                       # Core package
│   ├── data/                     # hotpot.py, cnn.py, tydi.py + EvalDataset
│   ├── llm/                      # Together-API and local HF LLM wrappers
│   ├── explainer/                # CAMAB (MAB) + SHAP + LOO + ContextCite
│   ├── experiments/              # run_evaluation.py unified entrypoint
│   └── utils/                    # attribution / evaluation loops, caching
├── scripts/
│   └── analysis/                 # Gold-supporting-facts eval + Figure 1
├── output/
│   ├── test_data/                # HotpotQA distractor/val sample for §3.3
│   ├── attr_cache/               # Per-run attribution caches (created on first run)
│   └── results/                  # Metric CSVs + Figure 1 PNG / diagnostics NPZ
├── requirements.txt
├── setup.py
├── config.example.json           # Template for TOGETHER_API_KEY
└── LICENSE                       # MIT

Command reference

llmexp/experiments/run_evaluation.py flags:

Citation

@inproceedings{pan2026camab,
  title         = {Context Attribution with Multi-Armed Bandit Optimization},
  author        = {Pan, Deng and Murugesan, Keerthiram and Hua, Ting and
                   Moniz, Nuno and Chawla, Nitesh V.},
  booktitle     = {Findings of the Association for Computational Linguistics: ACL 2026},
  year          = {2026},
  eprint        = {2506.19977},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CL},
  url           = {https://arxiv.org/abs/2506.19977}
}

License

MIT — see LICENSE.