Resumo Técnico: Sondagem Causal em Tempo de Inferência em LLMs (HDMI)

1. Declaração do Problema

O artigo aborda as limitações dos métodos atuais de sondagem causal utilizados para compreender como os Modelos de Linguagem de Grande Escala (LLMs) utilizam representações internas para previsão.

• Limitação da Sondagem Correlacional: Os métodos tradicionais treinam classificadores auxiliares ("sondas") em estados ocultos para detectar propriedades linguísticas (por exemplo, número do sujeito). Embora esses métodos mostrem o que está codificado, eles não provam como o modelo utiliza essa informação para gerar o próximo token.
• Limitação da Sondagem Causal: As abordagens causais existentes tentam intervir em estados ocultos para alterar propriedades (por exemplo, mudar de singular para plural) e observar mudanças na saída. No entanto, esses métodos geralmente dependem do treinamento de uma sonda separada para determinar a direção da intervenção (por exemplo, usando o gradiente da sonda para inverter um rótulo).
  • Desvantagens: Isso introduz custos adicionais de treinamento, requer supervisão específica da propriedade e arrisca desalinhamento. A fronteira de classificação da sonda pode não coincidir com a "geometria preditiva" interna do modelo, levando a intervenções que não direcionam efetivamente a geração do modelo ou alteram inadvertidamente propriedades não relacionadas.

Os autores propõem a Sondagem Causal em Tempo de Inferência, um framework onde as intervenções são aplicadas inteiramente no momento da inferência, sem reentrenar o modelo ou treinar sondas auxiliares.

2. Metodologia: Intervenção de Margem Impulsionada por Estado Oculto (HDMI)

A contribuição central é o HDMI, uma técnica baseada em gradiente e livre de sondas que aproveita diretamente o cabeçote de saída nativo do modelo generativo para direcionar estados ocultos.

2.1 Mecanismo Central

O HDMI opera definindo um objetivo de margem entre um token "fonte" (a continuação preferida original do modelo) e um token "alvo" (a continuação contrafactual desejada).

• Objetivo: Maximizar a margem de logit $L(x) = \phi(x)_\tau - \phi(x)_\sigma$, onde $\phi(x)$ é o vetor de logits, $\tau$ é o índice do token alvo e $\sigma$ é o índice do token fonte.
• Intervenção: O método realiza ascensão de gradiente no estado oculto $h_\ell(x)$ em uma camada específica $\ell$ para aumentar a probabilidade do token alvo enquanto diminui a do token fonte.
• Eficiência Computacional: Para a camada final, o gradiente é um produto matriz-vetor de forma fechada: $\nabla_{h_L} L(x) = W_U^\top (e_\tau - e_\sigma)$. Isso evita a necessidade de treinamento iterativo de sondas.
• Extensão Multi-token: O método suporta conjuntos de tokens (por exemplo, múltiplos verbos plurais válidos) somando logits sobre conjuntos de alvos ($T^+$) e conjuntos de fontes ($T^-$), definidos como $L_{set}(x) = \sum \phi(x)_{\tau_i} - \sum \phi(x)_{\sigma_j}$.

2.2 HDMI com Antecipação (LA-HDMI) para Edição de Texto

Para abordar a edição de texto, onde um usuário fornece uma sequência modificada, os autores introduzem o LA-HDMI.

• Desafio: A geração autoregressiva padrão não permite que estados ocultos no passo $t$ estejam cientes de futuras edições em $t' > t$.
• Solução: O LA-HDMI introduz um mecanismo de "antecipação" que retropropaga gradientes através dos embeddings de tokens esperados.
  • Ele desacopla a passagem frontal (usando baixa temperatura $\beta_f$ para manter a fluência) do caminho de gradiente (usando alta temperatura $\beta_g$ para calcular embeddings esperados).
  • A função objetivo $J_t(h_t)$ acumula objetivos de margem para futuras posições de edição (até $S_{max}$ passos à frente), ajustando o estado oculto atual $h_t$ para direcionar gerações futuras para as edições do usuário, enquanto preserva a fluência em regiões não editadas.

3. Framework de Avaliação

O artigo avalia intervenções usando duas métricas primárias, seguindo o framework de Canby et al. [2024]:

1. Completude: Mede se a intervenção altera com sucesso a propriedade visada (por exemplo, mudar o número do sujeito) no estado oculto, verificado por uma sonda de validação separada.
2. Seletividade: Mede se propriedades não relacionadas (variáveis de incômodo, por exemplo, número de substantivos não sujeitos) permanecem estáveis após a intervenção.
3. Confiabilidade: A média harmônica de Completude e Seletividade, servindo como a medida geral da qualidade da intervenção.

4. Resultados Chave

Os autores avaliaram o HDMI no Meta-Llama-3-8B-Instruct e no Pythia-70M usando o corpus de concordância LGD e o benchmark CausalGym.

• Desempenho: O HDMI alcança consistentemente maior Confiabilidade do que as linhas de base impulsionadas por sondas (AlterRep, FGSM, PGD) na maioria das tarefas.
  • Em suites de concordância (por exemplo, concordância sujeito-verbo), o HDMI frequentemente alcança Completude perfeita e alta Seletividade.
  • Ao contrário das linhas de base, o desempenho do HDMI é independente da qualidade de uma sonda intervencionista, pois não utiliza uma. Linhas de base como AlterRep e PGD mostraram quedas significativas de desempenho em modelos menores (Pythia-70M) quando a precisão da sonda era menor.
• Robustez: O HDMI mantém confiabilidade consistente entre diferentes tamanhos de modelos e tarefas, enquanto métodos baseados em sondas mostraram resultados mistos sensíveis à propriedade específica e ao modelo.
• Edição de Texto: O LA-HDMI demonstrou com sucesso a capacidade de direcionar a geração para edições especificadas pelo usuário (por exemplo, mudar "girl" para "owl" ou "was" para "were"), mantendo fluência gramatical e ajustando o contexto necessário (por exemplo, mudar "a" para "an").

5. Significado e Alegações

O artigo alega que o HDMI representa uma mudança da sondagem causal dependente de sondas para nativa do modelo.

• Alinhamento: Ao usar o próprio cabeçote de saída do modelo para definir a direção da intervenção, o HDMI garante que a perturbação esteja alinhada com a geometria preditiva real do modelo, evitando os riscos de desalinhamento de sondas externas.
• Eficiência: O método é computacionalmente leve, exigindo apenas ascensão de gradiente em estados ocultos, sem a sobrecarga de treinar classificadores auxiliares.
• Generalidade: A abordagem estende-se naturalmente a propriedades multi-token e oferece um mecanismo para edição fluente de texto via gradientes de antecipação.

Os autores reconhecem limitações, incluindo sensibilidade a hiperparâmetros (tamanho do passo, temperatura, horizonte) e a potencial atenuação de gradientes em tarefas de edição de texto de longo alcance. Eles sugerem trabalhos futuros em seleção adaptativa de camadas e agendamento de tamanho de passo para melhorar ainda mais a seletividade sem sacrificar a completude.