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

Título: Hessian-Enhanced Token Attribution (HETA): Interpretando LLMs Autoregressivos

1. O Problema

As técnicas atuais de atribuição (explicabilidade) para modelos de linguagem (LLMs) enfrentam limitações críticas quando aplicadas a arquiteturas decoder-only (apenas decodificador) em tarefas generativas:

• Dependência de Aproximações Lineares: A maioria dos métodos existentes (como Gradientes Integrados, LIME, SHAP) baseia-se em aproximações de primeira ordem (gradientes) ou lineares. Eles falham em capturar a complexidade não linear e as interações semânticas profundas inerentes à geração autoregressiva.
• Falhas em Regiões Planas: Em regiões onde o gradiente é zero (ou próximo de zero) devido a ativações saturadas ou funções de ativação como ReLU, os métodos baseados em gradiente atribuem importância nula a tokens que, na verdade, têm influência causal significativa no output.
• Inadequação da Atenção: Métricas baseadas puramente em pesos de atenção não são indicadores confiáveis de influência causal, pois refletem "onde o modelo olha" e não necessariamente "o que afeta a saída". Além disso, a agregação pós-hoc da atenção frequentemente ignora caminhos de influência indiretos através de conexões residuais e camadas MLP.
• Inconsistência: Métodos existentes produzem atribuições instáveis e inconsistentes sob mudanças de hiperparâmetros de decodificação (temperatura, top-p) ou reescrita sintática.

2. Metodologia: HETA

Os autores propõem o Hessian-Enhanced Token Attribution (HETA), um framework unificado projetado especificamente para modelos decoder-only. O HETA combina três componentes complementares para gerar uma pontuação de atribuição que é causal, sensível à curvatura e informada pela teoria da informação:

A. Influência de Transição Semântica (Causal Gating)

Para garantir a causalidade direcionada ao token alvo ($x_T$), o método rastreia o fluxo de valores ponderado pela atenção que termina na posição alvo, respeitando a máscara causal do decoder.

• Calcula-se uma "rotação" (rollout) de atenção-condicionada ao alvo, agregando apenas os caminhos que terminam em $T$.
• Isso gera um vetor de influência semântica ($M_T$) que atua como um "portão" (gate), restringindo a atribuição apenas a tokens com caminhos causais reais para o token alvo.

B. Sensibilidade Baseada em Hessiana (Efeitos de Segunda Ordem)

Para capturar interações não lineares e efeitos de curvatura que os gradientes de primeira ordem perdem:

• O método utiliza o Hessiano da log-verossimilhança em relação aos embeddings de entrada ($H_T = \nabla^2_X \log P_\theta(x_T | x_{<T})$).
• Como calcular o Hessiano completo é computacionalmente proibitivo, o HETA estima a sensibilidade usando Produtos Hessiano-Vetor (HVP) via estimadores de Hutchinson (vetores de Rademacher).
• Isso permite detectar influência mesmo quando o gradiente é zero, capturando a curvatura local da superfície de decisão do modelo.

C. Impacto Teórico da Informação (KL Divergence)

Para quantificar a contribuição informacional de um token:

• O método mede a mudança na distribuição preditiva do token alvo quando um token de contexto específico ($x_i$) é mascarado.
• Utiliza a Divergência de Kullback-Leibler (KL) entre a distribuição original e a distribuição mascarada: $I(x_i \to x_T) = D_{KL}(P_{orig} || P_{masked}^{(i)})$.
• Isso fornece uma interpretação probabilística robusta da importância do token.

Fórmula Final de Atribuição

A pontuação final de atribuição para um token $x_i$ é uma combinação ponderada:
$$\text{Attr}(x_i \to x_T) = M_T[i] \cdot (\beta \cdot S_i^{(T)} + \gamma \cdot I(x_i \to x_T))$$
Onde:

• $M_T[i]$: O gate causal (transição semântica).
• $S_i^{(T)}$: A sensibilidade baseada no Hessiano (curvatura).
• $I(x_i \to x_T)$: A contribuição de informação (KL).
• $\beta, \gamma$: Pesos hiperparâmetros que equilibram os termos.

3. Contribuições Principais

1. Framework HETA: Um novo método de atribuição para LLMs decoder-only que integra fluxo semântico causal, sensibilidade de segunda ordem (Hessiana) e impacto informacional (KL), superando as limitações das abordagens lineares.
2. Novo Dataset Curado: Os autores criaram e lançaram um dataset de 2.000 instâncias mistas (NARRATIVEQA + SCIQ) projetado especificamente para avaliar a qualidade da atribuição em geração. O dataset contrasta contextos narrativos irrelevantes com segmentos de evidência diagnóstica, permitindo uma avaliação precisa de alinhamento causal.
3. Métrica DSA (Dependent Sentence Attribution): Uma nova métrica para quantificar o quão bem o método concentra a massa de atribuição na evidência relevante versus contexto distrator.
4. Análise de Robustez e Escalabilidade: Demonstração de que o HETA é robusto a mudanças de hiperparâmetros de decodificação e sintaxe, com técnicas de aproximação (Hessiana de baixo rank e janelamento) para escalabilidade em contextos longos.

4. Resultados Experimentais

Os experimentos foram conduzidos em quatro modelos decoder-only diversos (GPT-J 6B, Phi-3 14B, LLaMA-3.1 70B, Qwen2.5 3B) e múltiplos datasets (LongRA, TellMeWhy, WikiBio, e o dataset curado).

• Fidelidade e Robustez: O HETA superou consistentemente todos os baselines (incluindo ContextCite, Integrated Gradients, Peering, TDD, ReAgent) nas métricas de fidelidade (Soft-NC, Soft-NS). O HETA mostrou-se estável com variações menores que 1% sob mudanças de hiperparâmetros de decodificação, enquanto os baselines variaram entre 2-5%.
• Alinhamento Humano (DSA): No dataset curado, o HETA alcançou pontuações DSA significativamente superiores (ex: 4.80 a 5.10 vs. valores negativos ou baixos para outros métodos), indicando que ele identifica corretamente os tokens causais e ignora distratores.
• Estudos de Ablação: A remoção de qualquer um dos três componentes (Fluxo Semântico, Hessiano ou KL) resultou em degradação significativa em todas as métricas, confirmando que os três pilares são complementares e essenciais.
• Validação Causal: Testes de intervenção (patching) mostraram que o HETA tem a maior correlação com efeitos causais reais, superando métodos baseados apenas em atenção ou gradientes.

5. Significado e Limitações

• Significado: O HETA estabelece um novo padrão para interpretabilidade em modelos generativos, provando que a incorporação de informações de segunda ordem (curvatura) e restrições causais rigorosas é necessária para atribuições fiáveis em LLMs. Ele resolve o problema de "alucinação de explicação" onde métodos anteriores atribuem importância a tokens irrelevantes.
• Limitações: O principal custo é o overhead computacional e de memória devido ao cálculo de produtos Hessiano-Vetor. Embora o papel proponha aproximações de baixo rank e janelamento que reduzem o custo em 40-50% mantendo a qualidade, o método ainda é mais lento que métodos baseados apenas em gradientes ou atenção.
• Futuro: O trabalho sugere otimizações futuras em aproximações de baixo rank e amostragem de camadas para melhorar a escalabilidade em contextos extremamente longos.

Em resumo, o HETA oferece uma abordagem matematicamente fundamentada e empiricamente superior para entender por que e como os tokens de entrada influenciam a geração em modelos de linguagem modernos, superando as limitações das abordagens lineares tradicionais.