Context Attribution with Multi-Armed Bandit Optimization

1. Il Problema: Attribuzione del Contesto nella RAG

Nelle moderne applicazioni di Retrieval-Augmented Generation (RAG), i Modelli Linguistici di Grande Dimensione (LLM) utilizzano contesti esterni per rispondere a domande. Tuttavia, è fondamentale capire quali parti del contesto recuperato hanno effettivamente contribuito alla risposta generata per garantire interpretabilità e fiducia.

Le sfide principali identificate sono:

• Allucinazioni: Gli LLM possono generare informazioni non supportate dal contesto.
• Costo Computazionale: I metodi esistenti per l'attribuzione (come SHAP o ContextCite) si basano su perturbazioni sistematiche o campionamento uniforme di sottoinsiemi di contesto. Questo richiede un numero elevato di query al modello (inference), rendendoli proibitivi per contesti lunghi o modelli "black-box" con budget limitati.
• Inefficienza: Metodi come SHAP non adattano la strategia di campionamento, esplorando lo spazio delle combinazioni in modo inefficiente.

2. Metodologia: CAMAB e Thompson Sampling Lineare

Gli autori propongono CAMAB (Context Attribution with Multi-Armed Bandit), un nuovo framework che riformula il problema dell'attribuzione come un Problema di Braccio Multi-Armed Combinatorio (CMAB).

Formulazione del Problema

• Bracci (Arms): Ogni segmento del contesto (es. una frase) è considerato un "braccio" base.
• Azione (Action): Selezionare un sottoinsieme di segmenti (un "super-braccio").
• Ricompensa (Reward): La qualità della risposta generata dal modello quando viene fornito solo quel sottoinsieme di contesto. La ricompensa $V(S)$ è definita come la media dei log-probabilità dei token della risposta originale data la risposta generata con il sottoinsieme $S$:
  $$V(S) = \frac{1}{T} \sum_{t=1}^{T} \log P_M(r_t | Q, S, r_{1:t-1})$$
  Questa metrica è applicabile sia a modelli open-source che a API black-box (come OpenAI o Together), purché forniscano le log-probabilità dei token.

Algoritmo: Linear Thompson Sampling (LinTS)

Per navigare efficientemente nello spazio combinatorio esponenziale dei sottoinsiemi, CAMAB utilizza il Linear Thompson Sampling:

1. Assunzione Lineare: Si assume che la ricompensa sia una funzione lineare dei contributi marginali dei segmenti inclusi, più un termine di bias (contesto vuoto).
   $$V(S) = w^T x + \epsilon$$
   Dove $x$ è un vettore binario che indica quali segmenti sono inclusi e $w$ sono i pesi latenti da stimare (che rappresentano l'importanza attribuita).
2. Aggiornamento Bayesiano: Il modello mantiene una distribuzione a posteriori sui pesi $w$ (distribuzione Gaussiana multivariata).
3. Esplorazione vs. Sfruttamento:
   • Ad ogni iterazione, si campiona un vettore di pesi $\tilde{w}$ dalla distribuzione a posteriori.
   • Si costruisce il sottoinsieme di contesto $S_t$ includendo solo i segmenti con pesi campionati positivi ($\tilde{w}_j > 0$).
   • Si interroga l'LLM con $S_t$, si osserva la ricompensa e si aggiorna la matrice di precisione $B_t$ e il vettore di informazione.
4. Vantaggio Chiave: La matrice di precisione $B_t$ cattura implicitamente le interazioni tra i segmenti (effetti di sostituzione o sinergia) senza richiedere modellazioni non lineari esplicite.

3. Contributi Chiave

1. Riformulazione del Problema: Prima formulazione dell'attribuzione del contesto come problema CMAB, permettendo un'ottimizzazione bayesiana adattiva.
2. Efficienza delle Query: L'uso di LinTS riduce drasticamente il numero di query necessarie rispetto ai metodi di perturbazione uniforme, adattando dinamicamente la selezione dei sottoinsiemi basandosi sulle stime posteriori.
3. Applicabilità Black-Box: Il metodo funziona utilizzando solo le log-probabilità dei token, rendendolo compatibile con le API commerciali dove non è possibile accedere ai gradienti interni o alle distribuzioni complete dei logit.
4. Validazione Empirica: Dimostrazione che l'approccio supera o eguaglia gli stati dell'arte con un budget di query significativamente inferiore.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su tre benchmark diversi (HotpotQA, CNN/DailyMail, TyDi QA) utilizzando due modelli LLM (LLaMA-3.1-8B e Qwen2.5-7B) accessibili via API.

• Efficienza delle Query: CAMAB raggiunge prestazioni di attribuzione pari o superiori a SHAP e ContextCite con fino al 30% in meno di query. Ad esempio, su HotpotQA, CAMAB con 40 query supera SHAP con 60 query.
• Metriche di Valutazione:
  • Log-Probability Drop: CAMAB mostra un calo maggiore nella probabilità quando i segmenti critici vengono rimossi, indicando una migliore identificazione delle prove chiave.
  • BERTScore: I risultati generati rimuovendo i segmenti attribuiti da CAMAB sono semanticamente più distanti dall'originale rispetto ai metodi basati, confermando che i segmenti rimossi erano più influenti.
• Convergenza: CAMAB converge molto più rapidamente. Mentre SHAP necessita di un budget elevato (es. 100 query) per stabilizzarsi, CAMAB raggiunge prestazioni quasi ottimali già con 40-60 query.
• Allineamento con Fatti Gold: Su HotpotQA, CAMAB ha mostrato una precisione superiore (P@1 = 0.780 vs 0.680 di SHAP) nell'identificare i fatti di supporto annotati manualmente.

5. Significato e Implicazioni

Il lavoro di CAMAB è significativo per diversi motivi:

• Scalabilità: Rende l'attribuzione del contesto praticabile per contesti lunghi e scenari reali dove il costo delle API LLM è un vincolo critico.
• Interpretabilità Adattiva: Sposta il paradigma da un campionamento "cieco" e uniforme a un'esplorazione intelligente e guidata dai dati, ottimizzando il trade-off tra accuratezza e costo computazionale.
• Robustezza: Funziona efficacemente anche in ambienti "black-box" reali, senza richiedere accesso ai pesi del modello o ai gradienti, a differenza di molti metodi di interpretabilità tradizionali.

Limitazioni: Il metodo dipende dalla disponibilità delle log-probabilità dei token (non sempre esposte da tutti i provider) e potrebbe convergere su soluzioni sub-ottimali in ambienti estremamente rumorosi se il bilanciamento esplorazione-sfruttamento non è calibrato correttamente. Tuttavia, per la maggior parte delle applicazioni RAG con budget limitato, rappresenta una soluzione superiore.