Resumo Técnico: Rastreamento de Circuitos Markovianos para Dinâmica de Estados em Transformers

Declaração do Problema

A interpretabilidade mecânica tem focado principalmente na análise de circuitos locais, características e cabeças de atenção para explicar comportamentos isolados em transformers. No entanto, tarefas de sequência parcialmente observadas exigem um paradigma computacional diferente: o modelo deve manter e atualizar um resumo interno (estado) ao longo do tempo para informar previsões futuras. Embora os Modelos Ocultos de Markov (HMMs) forneçam um framework matematicamente rigoroso para tais tarefas — onde a solução ótima é um vetor de crença bayesiana sobre estados latentes —, há uma falta de benchmarks controlados para testar se as ativações de transformers codificam, de fato, essas dinâmicas de transição de estado. Trabalhos existentes frequentemente estudam componentes (sondagem, correção, extração de estado) isoladamente, sem acesso simultâneo à verdade fundamental exata para estados latentes, vetores de crença e alvos contrafactuais.

Metodologia: Rastreamento de Circuitos Markovianos (MCT)

O artigo introduz o Rastreamento de Circuitos Markovianos (MCT), um pipeline de diagnóstico projetado para testar se as ativações de transformers contêm estrutura de transição de estado grosseira. A metodologia opera dentro de um benchmark sintético controlado usando seis famílias distintas de HMMs com complexidades variadas de transição e emissão.

O Benchmark

O benchmark treina transformers causais minúsculos (2 camadas, largura 128, ~275k parâmetros) em emissões visíveis de HMMs onde os estados latentes, matrizes de transição e probabilidades de emissão são conhecidos exatamente. Esta configuração fornece:

1. Verdade Fundamental Exata: Sequências de estados latentes, vetores de crença bayesiana e matrizes de transição.
2. Alvos Exatos: Previsões de próximo token ótimas de Bayes e alvos contrafactuais de estado forçado (a distribuição do próximo token se o estado latente fosse forçado a um valor específico).

O Pipeline MCT

Dado um transformer treinado em emissões de HMM, o MCT executa as seguintes etapas:

1. Extração de Estado: Um extrator de estado $g$ mapeia vetores de ativação $a_{\ell,t}$ para estados candidatos discretos $z_t$. O extrator de referência utiliza agrupamento K-means em ativações residuais.
2. Recuperação de Crença: Uma sonda de ridge linear mapeia ativações para o vetor de crença bayesiana exato para testar se informações de crença estão presentes.
3. Recuperação de Transição: A sequência de estados extraída é usada para estimar uma matriz de transição, que é comparada à matriz de transição verdadeira do HMM usando divergência KL por linha.
4. Teste de Ordem Markoviana: Modelos empíricos de ordem 0, 1 e 2 são ajustados aos estados extraídos para avaliar a força do sinal markoviano de primeira ordem.
5. Forçamento de Estado (Intervenção): Este é o teste causal central. O centróide de um cluster de estado recuperado é corrigido (patched) no modelo em uma posição específica de token. A saída do modelo é então comparada ao alvo contrafactual exato do HMM (a distribuição resultante do forçamento do estado latente). O desempenho é medido pela divergência KL entre a saída do modelo corrigido e o alvo exato.

Resultados Principais

Aprendizado da Tarefa

Em 18 execuções (6 famílias de HMMs $\times$ 3 sementes), os transformers aprenderam a tarefa de previsão efetivamente, alcançando uma perda excedente média sobre o preditor ótimo de Bayes de apenas 0,0138.

Recuperação de Estado e Transição

• Informação de Crença: Sondas lineares confirmaram que informações de crença são recuperáveis das ativações em todas as famílias, embora a qualidade varie significativamente por família de HMM.
• Estrutura de Transição: Abstrações de estado baseadas em resíduos e PCA recuperaram estrutura de transição parcial. O sinal foi mais forte nos regimes persistentes (transições adesivas) e fáceis de separar (baixa entropia de emissão).
• Limitações: A recuperação foi mais fraca nos regimes de emissão ambígua e seis estados. O agrupamento exato de crenças permaneceu como um limite superior mais forte do que o agrupamento derivado de ativações, sugerindo que o modelo representa um objeto contínuo semelhante a uma crença, em vez de uma cópia direta do rótulo latente discreto.

Forçamento de Estado (Evidência Causal)

A intervenção de forçamento de estado forneceu a evidência mais clara de utilidade causal:

• Desempenho: Corrigir um centróide de estado recuperado reduziu a divergência KL para o alvo contrafactual exato de 0,1957 (sem correção) para 0,0532 em média.
• Controles: Esta melhoria superou significativamente os controles, incluindo centróides de estado errado, ativações médias, ativações aleatórias, centróides de rótulos embaralhados e até mesmo centróides de estado verdadeiro em algumas métricas agregadas.
• Implicação: A abstração de estado recuperada é causalmente significativa e mais próxima do estado preditivo interno do modelo do que o rótulo latente amostrado em certos contextos.

Contribuições

1. Benchmark Controlado: Um framework unificado fornecendo rótulos latentes exatos, crenças bayesianas, matrizes de transição e alvos contrafactuais para testar métodos de interpretabilidade de transformers.
2. Pipeline de Diagnóstico (MCT): Um pipeline de referência para mapear ativações em abstrações de estado, testar ordem markoviana e validar relevância causal via forçamento de estado.
3. Descobertas Empíricas: Evidência de que pequenos transformers aprendem preditores próximos aos ótimos de Bayes e que suas ativações residuais contêm estrutura de transição parcial e grosseira, particularmente em regimes com sinais de transição fortes.

Significado e Alegações

O artigo posiciona o MCT não como um algoritmo completo para recuperar estados ocultos, mas como um framework de diagnóstico para dinâmicas de estado em transformers. Os autores afirmam que:

• O MCT detecta com sucesso estrutura de transição preditiva parcial em configurações sintéticas.
• Identifica onde os sinais de interpretabilidade são fortes (regimes persistentes/fáceis) e onde são fracos (regimes ambíguos/alta entropia).
• Testes de forçamento de estado usando contrafactuais exatos de HMM fornecem um método rigoroso para avaliar a relevância causal de abstrações de estado extraídas.

O trabalho é modesto em escopo, reconhecendo que as tarefas são sintéticas, os modelos são minúsculos e o extrator de referência é simples. Não afirma que o raciocínio em linguagem natural possui estruturas de estado idênticas, mas estabelece uma linha de base para testar se abstrações derivadas de ativações podem expor dinâmicas preditivas grosseiras contra uma verdade fundamental exata.