Causal Evidence of Stack Representations in Modeling Counter Languages Using Transformers

Resumo Técnico: Evidência Causal de Representações de Pilha na Modelagem de Linguagens de Contador usando Transformers

Enunciado do Problema
Embora a interpretabilidade mecanística tenha identificado com sucesso características latentes em modelos de linguagem (LMs) treinados em linguagens formais, uma lacuna crítica permanece no estabelecimento do papel causal dessas características. Trabalhos anteriores, como Tiwari et al. (2025), demonstraram que transformers treinados em linguagens de contador (ex: Dyck-1, Shuffle-k) aprendem representações consistentes com uma estrutura de pilha subjacente. No entanto, esses estudos não determinaram se essas representações são meramente correlacionadas com a tarefa ou se são causalmente necessárias para a computação do modelo. Este artigo aborda a questão: as representações de pilha aprendidas estão causalmente envolvidas na capacidade do transformer de prever o próximo token?

Metodologia
O estudo emprega um pipeline experimental de três estágios envolvendo treinamento, sondagem (probing) e intervenção causal:

1. Treinamento do Modelo: Um transformer apenas codificador (encoder-only) é treinado para a previsão do próximo token em linguagens Shuffle-k (onde $k \in \{2, 4, 6, 8\}$). O Shuffle-k é definido como o entrelaçamento de $k$ linguagens Dyck-1 independentes, exigindo que o modelo rastreie a profundidade de $k$ pilhas distintas simultaneamente. A arquitetura do modelo consiste em uma única camada com 4 cabeças de atenção, uma dimensão de embedding de 32 e uma dimensão oculta de 64. O modelo alcança precisão de validação perfeita até o 25º epoch, garantindo que todas as análises subsequentes sejam realizadas em um modelo totalmente convergido.
2. Sondagem da Profundidade da Pilha: Sondas lineares são treinadas para prever a profundidade de pilhas específicas ($c_i(t)$) a partir dos estados ocultos do modelo ($h^{(l)}_t$) na camada $l$. Para garantir que as sondas capturem informações estruturais genuínas em vez de correlações espúrias, uma sonda de controle é treinada em rótulos perturbados. A métrica Seletividade é definida como a diferença de precisão entre a sonda experimental e a sonda de controle. Uma alta seletividade indica que a profundidade da pilha está genuinamente codificada no espaço oculto.
3. Intervenção Causal (Ablação):
   • Extração de Direção: A matriz de pesos $W$ da sonda é decomposta via Decomposição de Valores Singulares (SVD). O vetor singular à direita superior, $w$, é identificado como a direção principal correspondente à representação da profundidade da pilha.
   • Ablação: Durante a inferência, o componente do estado oculto alinhado com $w$ é removido usando a operação $h' = h - \alpha(w \cdot h)w$, onde $\alpha \in [0, 1]$ representa a força da ablação.
   • Controle: Direções aleatórias são de forma semelhante abladas para servir como uma linha de base (baseline).
   • Avaliação: Os estados ablados são passados diretamente para o cabeçalho do decodificador (sem reexecutar as camadas internas do transformer) para medir o impacto na Precisão Posicional ($A_{pos}$) e na Precisão Sequencial ($A_{seq}$).

Resultos Principais

• Existência da Representação: Sondas lineares extraem com sucesso informações de profundidade de pilha com alta seletividade (ex: 85,8% para $k=6$, Pilha 2), confirmando que o modelo aprende representações consistentes com uma estrutura de pilha.
• Necessidade Causal: A ablação direcionada da direção principal da pilha ($w$) causa um colapso dramático no desempenho do modelo. À medida que a força da ablação $\alpha$ aumenta de 0 para 1, a Precisão Sequencial cai para próximo de 0%.
• Especificidade: Ablar direções aleatórias não tem efeito significativo na precisão, confirmando que o colapso de desempenho é específico à remoção da representação da pilha e não um artefato geral de manipulação de vetores.
• Sensibilidade da Métrica: O estudo observa que, embora a precisão posicional sofra uma queda, a precisão sequencial colapsa de forma mais acentuada. Isso é atribuído à natureza cumulativa dos erros na geração de sequências, onde um único token incorreto invalida toda a string.

Significância e Alegações
O artigo afirma fornecer forte evidência empírica de que as representações de pilha em transformers treinados em linguagens de contador não são apenas artefatos emergentes, mas são causalmente necessárias para a computação interna do modelo e para a previsão do próximo token.

Ao demonstrar que a remoção de uma direção de representação específica e aprendida leva a uma falha quase total, os autores argumentam que o transformer utiliza explicitamente essas estruturas do tipo pilha para resolver as tarefas de raciocínio lógico exigidas pelas linguagens formais. Isso move o campo além da análise correlacional (sondagem) em direção ao entendimento mecânico causal, validando que o comportamento de "caixa preta" do modelo depende de primitivas computacionais específicas e identificáveis.

Direções Futuras
Os autores sugerem que trabalhos futuros explorem:

• Intervir no meio da rede, passando estados ablados ou editados através da rede feed-forward para observar mudanças nos outputs subsequentes.
• Investigar o papel de diferentes arquiteturas na emergência, redundância e causalidade dessas representações.
• Compreender como os transformers lidam com a manipulação de outras primitivas computacionais emergentes, como contadores.