A retrieval conditioned rebinding circuit for dynamic entity tracking in large language models

Resumo Técnico: Um Circuito de Revinculação Condicionado à Recuperação para Rastreamento Dinâmico de Entidades em Grandes Modelos de Linguagem

Problema
Grandes Modelos de Linguagem (LLMs) devem interpretar corretamente o contexto e recuperar informações relevantes ao vincular entidades aos seus atributos e atualizar esses vínculos conforme os estados mudam. Embora trabalhos anteriores tenham demonstrado que os LLMs podem realizar o rastreamento de entidades, seu comportamento permanece inconsistente em diferentes configurações, e os mecanismos computacionais subjacentes para lidar com transições de estado explícitas em contextos naturalísticos não são totalmente compreendidos. Especificamente, não está claro se os modelos mantêm um "modelo de situação" globalmente atualizado após cada mudança de estado (por exemplo, uma operação de troca/swap) ou se empregam um mecanismo mais direcionado para atualizar os vínculos apenas quando a recuperação é necessária.

Metodologia
Os autores investigam este problema usando uma tarefa de rastreamento de estado dinâmico envolvendo caixas e objetos, adaptada de Kim e Schuster (2023). O estudo utiliza intervenções causais nas ativações de quatro modelos ajustados por instrução: Gemma2-9B-it, Gemma3-12B-it, Llama3.2-3B-it e Llama3.1-8B-it. A análise procede em quatro estágios:

1. Análise de Mediação Causal: Usando intervenções de intercâmbio, os autores rastreiam o fluxo de informação ao nível do token. Eles constroem entradas contrafatuais para determinar como a informação associada às caixas de referência, caixas de alvo da troca e objetos alvo influencia a predição final.
2. Path Patching (Remendagem de Caminho): Para localizar o mecanismo no nível da cabeça de atenção, os autores aplicam path patching (V-patching e Q-patching) para identificar um circuito compacto. Eles selecionam iterativamente caminhos que mais impactam negativamente o desempenho quando corrompidos, isolando um subconjunto de cabeças de atenção responsáveis pela tarefa.
3. Intervenções de Intercâmbio Específicas de Papel: Os autores testam os papéis funcionais dos grupos de cabeças identificados testando o patching de padrões de atenção entre execuções originais e contrafatuais. Isso distingue entre cabeças que medeiam a seleção de endereço (roteamento) e aquelas que transferem o conteúdo do objeto.
4. Intervenções de ID de Vinculação: Adaptando o framework de Feng e Steinhardt (2024), os autores realizam intervenções causais nos subespaços Query (Q) e Key (K) de grupos de cabeças específicos. Eles estimam direções de ID de vinculação e vetores de deslocamento para testar se o roteamento de recuperação depende do correspondência Q/K de identificadores de vinculação.

Principais Contribuições e Resultados

• Evidência de Revinculação Condicionada à Recuperação: A análise de mediação causal revela que a informação sobre as caixas de referência e de alvo da troca não flui diretamente de suas menções iniciais, mas emerge após a operação de troca e se propaga em direção à predição. Por outro lado, a informação do objeto é roteada de forma mais direta. Essa dissociação apoia a hipótese de "Revinculação Condicionada à Recuperação": o modelo não recodifica globalmente todo o estado pós-troca. Em vez disso, ele preserva os vínculos originais e aplica um remapeamento sobre os IDs de vinculação apenas durante a recuperação da resposta.
• Identificação de um Circuito Compacto: O path patching identifica um circuito de cabeça de atenção específico composto por 3–10% dos componentes totais do modelo (por exemplo, 38 cabeças em Gemma-9B) que recupera o desempenho quase total do modelo completo (0,89 de acurácia vs. 1,00 para o modelo completo). Este circuito consiste em cinco grupos funcionais:
  • Grupo E (Âncora de Vinculação): Escreve a informação de identidade da caixa.
  • Grupo D (Transmissor de Posição de Troca): Propaga a informação da troca.
  • Grupo C (Atualizador de Posição): Molda a atenção em direção às posições de vinculação atualizadas.
  • Grupo B (Desreferenciador): Atua como um ponteiro, mediando a seleção de endereço.
  • Grupo A (Recuperador de Resposta): Recupera o token do objeto alvo.
• Mecanismo de Roteamento de Endereço: Intervenções de intercâmbio confirmam que o Grupo B funciona como um ponteiro, sensível à seleção de endereço. O patching de seus padrões de atenção desloca o foco do modelo para a posição de origem correta sem transferir o conteúdo contrafatual do objeto.
• Assinaturas de Subespaço Q/K: Intervenções de ID de vinculação revelam que o mecanismo de ponteiro depende da correspondência Q/K, mas a implementação varia conforme a família do modelo:
  • Modelos Gemma: Exibem uma assinatura clara de correspondência Q/K onde a informação de vinculação é expressa tanto nos subespaços de query quanto de key. Deslocar o vetor de query redireciona efetivamente a atenção.
  • Modelos Llama: Dependem primordialmente de vetores de key para a informação de vinculação. Intervenções no lado da Key produzem os maiores deslocamentos de atenção, enquanto os efeitos no lado da query são mais fracos.
• Generalização: O mecanismo de revinculação condicionada à recuperação é observado nas famílias Gemma e Llama, sugerindo que é uma estratégia computacional compartilhada para o rastreamento de estado dinâmico, apesar das diferenças nas assinaturas específicas do circuito.

Significância e Alegações
O artigo afirma fornecer uma explicação interpretável, ao nível de circuito, de como os LLMs realizam o rastreamento de estado dependente de contexto. As descobertas sugerem que os LLMs não mantêm necessariamente uma simulação mental completa e globalmente atualizada do mundo após cada mudança de estado. Em vez disso, eles utilizam uma estratégia de "revinculação" onde a informação relacionada à troca é preservada e aplicada localmente no momento da recuperação para redirecionar a atenção ao objeto correto via IDs de vinculação atualizados.

Os autores observam que, embora a estratégia computacional (revinculação condicionada à recuperação) pareça consistente entre as famílias de modelos, a implementação ao nível do circuito (especificamente a dependência de subespaços Q vs. K) difere. Isso implica que o mecanismo pode ser uma solução compartilhada para o problema do rastreamento dinâmico, mas sua realização depende de arquiteturas de modelos ou distribuições de pré-treinamento específicas.

Limitações
Os autores reconhecem que sua configuração de tarefa é altamente controlada (troca única, objetos de um único token, caixas nomeadas alfabeticamente) e pode não capturar a gama total de mecanismos no discurso natural. Além disso, as direções de ID de vinculação usadas nas intervenções são estimadas sob uma suposição de representação linear aproximada, o que pode subestimar mecanismos se a informação de vinculação for representada de forma não linear ou distribuída por múltiplos subespaços.