Context Attribution with Multi-Armed Bandit Optimization

Resumen Técnico: CAMAB

1. Planteamiento del Problema

En los sistemas de Generación Aumentada por Recuperación (RAG), es fundamental entender qué partes del contexto recuperado contribuyen realmente a la respuesta generada por un Modelo de Lenguaje Grande (LLM). Sin esta atribución, es difícil garantizar la veracidad y evitar alucinaciones.

Los métodos existentes enfrentan dos limitaciones principales:

• Métodos de entrenamiento: Requieren modificar el modelo para que cite fuentes, pero la fiabilidad de estas citas no es verificable.
• Métodos post-hoc (perturbación): Técnicas como SHAP o ContextCite evalúan la importancia de segmentos de texto eliminándolos o enmascarándolos y midiendo el impacto en la salida. Sin embargo, estos métodos requieren un muestreo exhaustivo o uniforme de subconjuntos de contexto, lo que genera costos computacionales prohibitivos (miles de llamadas al LLM), haciéndolos poco prácticos para contextos largos o modelos de caja negra con presupuestos de consultas limitados.

El objetivo es desarrollar un marco que identifique los segmentos de contexto más influyentes con el mínimo número de consultas al modelo.

2. Metodología: CAMAB

Los autores proponen CAMAB (Context Attribution with Multi-Armed Bandit), un marco que reformula el problema de atribución como un Problema de Bandido Multi-Brazo Combinatorio (CMAB).

Formulación del Problema:

• Brazos: Cada segmento del contexto (ej. una oración) se trata como un "brazo" base.
• Acción: Seleccionar un subconjunto de segmentos (un "super-brazo") para evaluar.
• Recompensa: La calidad de la respuesta del LLM dado ese subconjunto. Se mide mediante la probabilidad logarítmica promedio de los tokens de la respuesta original dada la subsecuencia de contexto seleccionada (Ecuación 1).
  • Si un subconjunto mantiene alta la probabilidad de la respuesta, indica que esos segmentos son informativos.

Algoritmo: Muestreo de Thompson Lineal (LinTS)
En lugar de probar subconjuntos al azar o exhaustivamente, CAMAB utiliza Muestreo de Thompson Lineal para explorar el espacio de subconjuntos de manera adaptativa:

1. Suposición Lineal: Se asume que la recompensa es una función lineal de las contribuciones marginales de los segmentos individuales más un término de sesgo.
2. Muestreo Bayesiano: El algoritmo mantiene una distribución de probabilidad posterior sobre los pesos latentes (importancia) de cada segmento.
3. Exploración vs. Explotación: En cada paso, muestrea un vector de pesos plausible. Selecciona los segmentos cuyos pesos muestreados son positivos para formar el siguiente subconjunto a evaluar.
4. Actualización: Tras consultar al LLM y obtener la recompensa, actualiza la matriz de precisión y el vector de información para refinar la estimación de importancia.
5. Interacciones Implícitas: Aunque el modelo es lineal, la matriz de precisión captura correlaciones entre segmentos (interacciones) sin necesidad de modelado no lineal explícito.

3. Contribuciones Clave

1. Reformulación del Problema: Transforman la atribución de contexto a nivel de segmento en un problema de CMAB, permitiendo una búsqueda eficiente en un espacio combinatorio exponencial.
2. Eficiencia de Consultas: Utilizan LinTS para priorizar adaptativamente los subconjuntos informativos basándose en estimaciones posteriores, reduciendo drásticamente el número de llamadas al modelo en comparación con métodos uniformes.
3. Aplicabilidad Universal: El método funciona tanto con modelos de código abierto como con APIs de caja negra (como OpenAI o Together), siempre que se tenga acceso a las probabilidades logarítmicas de los tokens.
4. Validación Empírica: Demuestran que el enfoque es robusto en escenarios realistas con presupuestos de recursos limitados.

4. Resultados Experimentales

Los experimentos se realizaron en tres conjuntos de datos diversos (HotpotQA, CNN/DailyMail, TyDi QA) utilizando modelos LLM de vanguardia (LLaMA-3.1-8B y Qwen2.5-7B).

• Eficiencia de Muestreo: CAMAB logra reducir las consultas al modelo en hasta un 30% en comparación con SHAP y ContextCite, manteniendo o superando la calidad de la atribución.
  • En tareas de búsqueda de información (HotpotQA, TyDi QA), CAMAB con un presupuesto de 40 consultas supera a SHAP con 60 consultas.
• Calidad de Atribución:
  • Log-Probability Drop: CAMAB identifica mejor los segmentos críticos, causando una mayor caída en la probabilidad de la respuesta al eliminarlos.
  • BERTScore: Mantiene una consistencia semántica superior al comparar respuestas generadas con contextos perturbados.
  • Alineación con Hechos: En HotpotQA, CAMAB supera a SHAP en métricas de precisión (P@1) y F1 al alinearse con los hechos de apoyo anotados por humanos.
• Robustez: El método converge más rápido a soluciones óptimas, siendo especialmente superior en presupuestos muy ajustados (ej. 20 consultas), donde otros métodos degradan su rendimiento significativamente.

5. Significado e Impacto

El trabajo de CAMAB es significativo porque aborda el cuello de botella principal de la interpretabilidad en LLMs: el costo computacional.

• Escalabilidad: Permite la atribución de contexto en escenarios de larga duración y aplicaciones comerciales donde el número de llamadas a la API es un factor limitante de costos y latencia.
• Confianza: Proporciona una herramienta para verificar la "groundedness" (sustento en el contexto) de las respuestas generadas, crucial para aplicaciones críticas en medicina, derecho o noticias.
• Nueva Dirección: Establece un paradigma donde la optimización bayesiana y los algoritmos de bandit pueden reemplazar a los métodos de perturbación exhaustiva para tareas de explicabilidad en IA generativa.

En conclusión, CAMAB demuestra que es posible obtener atribuciones de alta fidelidad en sistemas de generación de lenguaje con recursos limitados, equilibrando la exploración de la incertidumbre con la explotación de los hallazgos más prometedores.