Hessian-Enhanced Token Attribution (HETA): Interpreting Autoregressive LLMs

1. El Problema

Los métodos actuales de atribución (explicabilidad) para modelos de lenguaje (LLMs) presentan limitaciones críticas cuando se aplican a arquitecturas autoregresivas de solo decodificador (decoder-only):

• Dependencia de aproximaciones lineales: La mayoría de las técnicas existentes (como Integrated Gradients, LIME o KernelSHAP) se basan en derivadas de primer orden o aproximaciones lineales. Estas fallan en capturar la complejidad no lineal y las interacciones semánticas profundas inherentes a la generación autoregresiva.
• Falta de causalidad en la atención: Los métodos basados en pesos de atención a menudo no reflejan la influencia causal real, ya que la atención indica "dónde mira" el modelo, no necesariamente "qué afecta" la salida. Además, ignoran las rutas de influencia indirectas a través de conexiones residuales y capas MLP.
• Inestabilidad: Las atribuciones actuales son inestables ante cambios en los hiperparámetros de decodificación (temperatura, top-p) y reestructuraciones sintácticas, lo que cuestiona su fiabilidad.
• Brecha arquitectónica: Los métodos diseñados para modelos codificadores (bidireccionales) no se transfieren bien a modelos generativos, donde la atribución debe ser condicional al token objetivo y respetar la máscara causal.

2. Metodología: HETA (Hessian-Enhanced Token Attribution)

El autores proponen HETA, un marco unificado diseñado específicamente para modelos de lenguaje de solo decodificador. HETA combina tres componentes complementarios para generar atribuciones a nivel de token que son conscientes del contexto, causalmente fieles y semánticamente fundamentadas:

A. Influencia de Transición Semántica (Semantic Transition Influence)

Para garantizar la causalidad dirigida al objetivo, HETA rastrea el flujo de valores ponderado por la atención que termina en la posición del token objetivo ($T$).

• Mecanismo: Calcula un "despliegue" (rollout) de atención condicionado al objetivo, agregando solo las rutas que llegan a $T$ a través de las capas y cabezas de atención.
• Fórmula: Se define una masa de transición semántica $M_T[i]$ que asigna importancia a los tokens que tienen un camino causal hacia el token objetivo, normalizando sobre el simplex. Esto actúa como una "puerta" causal que descarta tokens irrelevantes.

B. Sensibilidad Basada en Hessianos (Hessian-Based Sensitivity)

Para capturar efectos de segundo orden y no linealidades donde los gradientes pueden ser cero (regiones planas o saturadas), HETA utiliza la curvatura de la superficie de log-verosimilitud.

• Mecanismo: Estima la sensibilidad de segundo orden mediante productos Hessianos-Vectores (HVP) utilizando estimadores de Hutchinson (vectores Rademacher).
• Ventaja: A diferencia de los gradientes, el Hessian captura cómo cambia la función ante perturbaciones finitas incluso cuando la derivada primera es nula, modelando interacciones no lineales entre tokens.

C. Impacto de Información Teórica (KL Divergence)

Para cuantificar la contribución informativa de un token, se mide la pérdida de información cuando dicho token se enmascara.

• Mecanismo: Se calcula la divergencia de Kullback-Leibler (KL) entre la distribución original del token objetivo y la distribución cuando el token de entrada $x_i$ se enmascara.
• Fórmula: $I(x_i \to x_T) = D_{KL}(P_{orig} || P_{masked}^{(i)})$. Esto mide cuánto cambia la incertidumbre predictiva al eliminar la información de un token específico.

Puntuación Final de Atribución

La puntuación final combina estos tres elementos:
$$\text{Attr}(x_i \to x_T) = M_T[i] \cdot (\beta S_i^{(T)} + \gamma I(x_i \to x_T))$$
Donde:

• $M_T[i]$ es la puerta causal (transición semántica).
• $S_i^{(T)}$ es la sensibilidad basada en el Hessian.
• $I(x_i \to x_T)$ es la contribución de información (KL).
• $\beta$ y $\gamma$ son pesos que equilibran la sensibilidad de curvatura y la contribución informativa.

3. Contribuciones Clave

1. Nuevo Marco HETA: Un método de atribución novedoso para modelos de solo decodificador que integra flujo semántico causal, sensibilidad de segundo orden (Hessian) e impacto informativo (KL).
2. Nuevo Dataset Curado: Creación y liberación de un dataset de 2,000 instancias (combinando NARRATIVEQA y SCIQ) diseñado específicamente para evaluar la calidad de la atribución en tareas de generación. Incluye una métrica nueva llamada Dependent Sentence Attribution (DSA) para medir la alineación con evidencia predictiva real frente a distractores.
3. Evaluación Exhaustiva: Experimentos extensos en cuatro modelos de decodificador diversos (GPT-J, LLaMA-3.1, Phi-3, Qwen2.5) que demuestran la superioridad de HETA sobre líneas base fuertes (ContextCite, Integrated Gradients, ReAgent, etc.).

4. Resultados Experimentales

Los resultados muestran que HETA supera consistentemente a los métodos existentes en múltiples dimensiones:

• Fidelidad y Robustez: HETA logra las puntuaciones más altas en métricas de fidelidad (Soft-NC y Soft-NS), indicando que las atribuciones son estables bajo perturbaciones de entrada. En GPT-J 6B, HETA supera a la siguiente mejor metodología (ReAgent) por más de un factor de 2 en LongRA.
• Alineación Humana (DSA): En el dataset curado, HETA obtiene puntuaciones DSA significativamente más altas (ej. 4.80 en GPT-J vs. 3.60 para ReAgent), demostrando que asigna correctamente la importancia a los tokens de evidencia y no a distractores semánticamente ricos pero irrelevantes.
• Invarianza a Hiperparámetros: HETA es notablemente robusto a cambios en los hiperparámetros de decodificación (temperatura, top-p, top-k). Mientras las líneas base fluctúan entre un 2% y un 5%, HETA mantiene una variación menor al 1%.
• Análisis de Componentes: Las pruebas de ablación confirman que los tres componentes (flujo semántico, Hessian y KL) son complementarios y necesarios. Eliminar cualquiera de ellos degrada significativamente el rendimiento.
• Eficiencia: Aunque el cálculo del Hessian es costoso, los autores proponen variantes de aproximación (Hessian de bajo rango y ventanas deslizantes) que reducen el tiempo de ejecución en un 40-50% y el uso de memoria en un 30-40%, manteniendo la mayor parte de la fidelidad.

5. Significado e Impacto

El trabajo de HETA establece un nuevo estándar para la interpretabilidad en modelos de lenguaje generativos.

• Superación de la linealidad: Demuestra que las aproximaciones de primer orden son insuficientes para modelos modernos y que la incorporación de información de segundo orden (curvatura) es crucial para capturar la verdadera influencia causal.
• Causalidad Estricta: Al integrar una puerta de flujo semántico condicionada al objetivo, HETA resuelve el problema de asignar importancia a tokens que no tienen una conexión causal real con la salida generada.
• Aplicabilidad Práctica: La propuesta de aproximaciones eficientes (bajo rango y ventanas) hace viable el uso de métodos de atribución de alta fidelidad en contextos de larga longitud y modelos grandes, facilitando el depuramiento, la auditoría ética y la comprensión de modelos en dominios de alto riesgo.

En resumen, HETA proporciona una explicación más fiel, robusta y semánticamente coherente de cómo los tokens de entrada influyen en las predicciones de los LLMs, superando las limitaciones fundamentales de las técnicas de atribución tradicionales.