Context Attribution with Multi-Armed Bandit Optimization

Título: Atribuição de Contexto com Otimização de Bandit Multi-Armed (CAMAB)

1. Problema e Motivação

A Geração Aumentada por Recuperação (RAG) tornou-se fundamental para melhorar a precisão factual de Grandes Modelos de Linguagem (LLMs) ao fornecer contexto externo. No entanto, um desafio crítico é garantir que as respostas geradas sejam genuinamente fundamentadas nesse contexto, evitando alucinações.

• O Desafio: É essencial identificar quais segmentos específicos do contexto recuperado contribuíram para a resposta final do modelo.
• Limitações das Abordagens Atuais:
  • Métodos de Pós-Hoc (ex: SHAP, LIME, ContextCite): Envolvem perturbar ou mascarar segmentos de contexto e avaliar o impacto na saída. Embora precisos, eles exigem um número massivo de consultas ao modelo (inference calls), tornando-os computacionalmente proibitivos para contextos longos e caros, especialmente em cenários de API "caixa-preta".
  • Métodos de Treinamento: Modelos treinados para citar fontes não garantem que as citações reflitam o raciocínio real durante a inferência.

O objetivo do trabalho é criar um método de atribuição que seja altamente eficiente em termos de consultas (query-efficient), mantendo alta fidelidade, mesmo com orçamentos de consulta limitados.

2. Metodologia: CAMAB

Os autores propõem o CAMAB (Context Attribution with Multi-Armed Bandit), um novo framework que formula a atribuição de contexto como um problema de Bandit Multi-Armed Combinatório (CMAB).

• Formulação do Problema:

  • Cada segmento de contexto (ex: uma frase) é tratado como um "braço" (arm) do bandit.
  • Uma "ação" consiste em selecionar um subconjunto (super-arm) desses segmentos.
  • O objetivo é identificar os segmentos mais influentes dentro de um orçamento limitado de consultas ao LLM.

• Função de Recompensa:

  • A importância de um segmento é medida pela probabilidade de log dos tokens da resposta original quando o segmento está incluído ou excluído.
  • A recompensa $V(S)$ para um subconjunto $S$ é a média do log-probabilidade dos tokens de resposta gerados com aquele contexto parcial:
    $$V(S) = \frac{1}{T} \sum_{t=1}^{T} \log P_M(r_t | Q, S, r_{1:t-1})$$
  • Isso permite que o método funcione tanto em modelos open-source quanto em APIs comerciais (caixa-preta), desde que forneçam log-probabilidades.

• Algoritmo: Linear Thompson Sampling (LinTS):

  • Em vez de amostrar subconjuntos uniformemente (como no SHAP), o CAMAB utiliza o Linear Thompson Sampling para explorar o espaço de combinações de forma adaptativa.
  • Suposição Linear: Assume-se que a recompensa é uma soma linear das contribuições marginais dos segmentos, mais um termo de viés.
  • Mecanismo: O algoritmo mantém uma distribuição de probabilidade (Bayesiana) sobre os pesos latentes de cada segmento. A cada rodada:
    1. Amostra-se um vetor de pesos plausível.
    2. Seleciona-se o subconjunto de segmentos com pesos amostrados positivos.
    3. Consulta-se o LLM e observa-se a recompensa.
    4. Atualiza-se a distribuição posterior (matriz de precisão e vetor de informação).
  • Vantagem: A matriz de precisão captura implicitamente interações e dependências entre segmentos sem modelagem não-linear explícita, equilibrando exploração (testar segmentos incertos) e exploração (focar nos melhores).

3. Contribuições Principais

1. Formulação CMAB: Transformar a atribuição de nível de segmento em um problema de Bandit Multi-Armed Combinatório para otimização de orçamento.
2. Eficiência com LinTS: Uso de Linear Thompson Sampling para explorar adaptativamente o espaço de subconjuntos de contexto, reduzindo drasticamente o número de consultas necessárias.
3. Validação em Cenários Realistas: Experimentos extensivos em três benchmarks diversos (HotpotQA, CNN/DM, TyDi QA) usando dois LLMs de ponta via APIs comerciais, demonstrando eficácia em configurações de caixa-preta.

4. Resultados Experimentais

Os experimentos compararam o CAMAB com SHAP, ContextCite e uma linha de base aleatória, utilizando modelos como LLaMA-3.1-8B e Qwen2.5-7B.

• Eficiência de Amostragem (Orçamento de Consultas):

  • O CAMAB alcançou uma redução de até 30% no número de consultas ao modelo em comparação com métodos existentes, mantendo ou superando a qualidade da atribuição.
  • Em tarefas de busca de informação (HotpotQA, TyDi QA), o CAMAB com orçamento de 40 consultas frequentemente superou o SHAP com orçamento de 60 consultas.
  • O CAMAB convergiu para desempenho quase ótimo muito mais rápido que os métodos de perturbação uniforme.

• Qualidade da Atribuição:

  • HotpotQA e TyDi QA: O CAMAB superou consistentemente os baselines em métricas de Top-k Log-Probability Drop e BERTScore, indicando uma melhor identificação das evidências críticas.
  • CNN/DailyMail (Resumo): Desempenho competitivo, quase idêntico ao SHAP, devido à estrutura de viés de liderança em resumos de notícias que facilita a convergência.
  • Alinhamento com Fatos Ouro (Gold Facts): No HotpotQA, o CAMAB teve maior precisão (P@1) e F1 em comparação com o SHAP ao alinhar os segmentos atribuídos com os fatos de suporte anotados por humanos.

• Custo Computacional:

  • A atualização da matriz posterior do CAMAB ($O(N^3)$) adicionou um atraso de parede (wall-clock) insignificante em comparação com o custo de inferência do LLM. Na verdade, o CAMAB foi ligeiramente mais rápido no tempo total por amostra devido à seleção adaptativa de subconjuntos menores.

5. Significado e Conclusão

O trabalho CAMAB estabelece um novo padrão para a interpretabilidade de RAG em cenários de recursos limitados.

• Impacto Prático: Permite a atribuição de contexto em modelos de caixa-preta (APIs comerciais) de forma economicamente viável, reduzindo custos de inferência e latência.
• Inovação Teórica: Demonstra que a formulação de problemas de atribuição como otimização Bayesiana sequencial (Bandits) é superior à amostragem aleatória ou exaustiva para espaços combinatórios grandes.
• Limitações: O método depende do acesso a log-probabilidades de tokens (não disponível em todas as APIs) e pode convergir para soluções locais subótimas em ambientes extremamente ruidosos se o equilíbrio exploração-exploração não for bem calibrado.

Em resumo, o CAMAB oferece uma solução escalável e robusta para entender "por que" um LLM gerou uma resposta específica, tornando os sistemas de IA mais transparentes e confiáveis sem o custo proibitivo de métodos tradicionais.