Resumo Técnico: Mecanismos de Mistura na Recuperação de Entidades em Contexto

Declaração do Problema

Os modelos de linguagem (LMs) destacam-se no raciocínio em contexto, uma capacidade fundamentalmente dependente da habilidade de vincular entidades relacionadas dentro de um texto (por exemplo, vincular "Ann" a "pie" em "Ann loves pie") e subsequentemente recuperá-las com base em uma consulta (por exemplo, responder "Who loves pie?"). Embora pesquisas anteriores tenham estabelecido que os LMs utilizam um mecanismo posicional — recuperando entidades com base em seu índice na sequência de contexto —, essa visão foi amplamente derivada de estudos envolvendo listas curtas de grupos de entidades (por exemplo, $n=2$ ou $3$).

O problema central abordado neste trabalho é a confiabilidade desse mecanismo posicional em configurações mais complexas e realistas. À medida que o número de grupos de entidades vinculadas aumenta, os LMs têm dificuldade em raciocinar sobre contextos longos, exibindo fenômenos como o efeito "lost-in-the-middle" (perdido no meio). Os autores investigam se o mecanismo posicional se generaliza para sequências mais longas ou se os LMs empregam estratégias alternativas para manter a precisão na recuperação.

Metodologia

Os autores empregam técnicas de interpretabilidade mecanicista, especificamente intervenções de intercâmbio, para sondar a estrutura causal interna dos LMs.

1. Configuração da Tarefa: O estudo utiliza tarefas de raciocínio em contexto baseadas em modelos, envolvendo "Grupos de Entidades" (tuplas de entidades que preenchem papéis específicos, por exemplo, Nome + Comida). O modelo é consultado para recuperar uma entidade-alvo com base em uma entidade de consulta dentro de um grupo específico.
2. Design Contrafactual: Para distinguir entre hipóteses concorrentes, os autores constroem conjuntos de dados emparelhados de entradas originais e contrafactuais. Esses pares são projetados de modo que três mecanismos hipotéticos — Posicional, Lexical e Reflexivo — prevejam saídas distintas se uma intervenção de intercâmbio for realizada em suas respectivas variáveis intermediárias.
   • Mecanismo Posicional ($P$): Recupera a entidade no índice do grupo da consulta contrafactual.
   • Mecanismo Lexical ($L$): Recupera a entidade vinculada à entidade de consulta contrafactual (por exemplo, se a consulta mudar de "Tim" para "Ann", recupera o que "Ann" ama).
   • Mecanismo Reflexivo ($R$): Recupera a entidade-alvo por meio de um ponteiro direto originado na própria entidade-alvo, desreferenciado através da consulta.
3. Protocolo de Intervenção: Os autores realizam intervenções de intercâmbio no fluxo residual do último token em camadas específicas ($\ell$) identificadas como o ponto imediatamente antes da ocorrência da recuperação. Eles aplicam (patch) o vetor do fluxo residual da execução contrafactual na execução original e observam a distribuição do próximo token resultante.
4. Modelagem Causal: Um modelo causal ($M$) é desenvolvido para simular o comportamento do LM. Este modelo combina os três mecanismos como uma mistura ponderada:
   • O termo posicional é modelado como uma distribuição Gaussiana centrada no índice da consulta, com variância dependente do índice.
   • Os termos lexical e reflexivo são modelados como distribuições one-hot.
   • O modelo é treinado para prever a distribuição de logits do LM usando a Divergência de Jensen–Shannon (JSD) como métrica de perda.
5. Escopo: Os experimentos cobrem nove modelos de três famílias (Llama, Gemma, Qwen), variando de 2B a 72B parâmetros, avaliados em dez tarefas distintas de vinculação.

Principais Contribuições

1. Identificação de Três Mecanismos Distintos

O artigo desafia a visão predominante de que a recuperação é exclusivamente posicional. Ele demonstra que os LMs empregam uma mistura de três mecanismos:

• Posicional: Eficaz para entidades no início e no final do contexto, mas torna-se ruidoso e difuso para posições intermediárias à medida que o número de grupos ($n$) aumenta.
• Lexical: Recupera o parceiro vinculado da entidade de consulta (por exemplo, "pie" $\to$ "Ann"). Este mecanismo fornece um sinal agudo e one-hot.
• Reflexivo: Utiliza um ponteiro direto da entidade-alvo para si mesma, que é recuperado através da consulta. Isso é particularmente crucial quando o alvo aparece antes da consulta em uma cláusula (por exemplo, "Tim loves tea" $\to$ "Who loves tea?"), um cenário onde a cópia lexical é arquitetonicamente difícil para modelos autoregressivos.

2. O Padrão de Confiabilidade em "U"

Através de extensos experimentos de intervenção, os autores descobrem um padrão consistente:

• Primeiras e Últimas Posições: A recuperação é dominada pelo mecanismo posicional, que é altamente concentrado e confiável.
• Posições Intermediárias: O sinal posicional degrada-se significativamente. Para compensar, os LMs dependem fortemente dos mecanismos lexical e reflexivo, que fornecem sinais mais agudos para desambiguar a entidade correta.
• Interação: Os mecanismos exibem "sinergia competitiva". Quando os índices se alinham, os sinais se amplificam; quando divergem, suprimem-se mutuamente.

3. Modelo Causal de Previsão do Próximo Token

Os autores constroem um modelo causal que combina esses três mecanismos. Este modelo alcança 95% de concordância (medida pela Similaridade de Jensen–Shannon) com as distribuições reais de próximo token dos LMs em várias tarefas e tamanhos de modelo. Em contraste, modelos que dependem exclusivamente do mecanismo posicional (a visão predominante) alcançam pontuações significativamente mais baixas (~45%), performando pior do que uma linha de base uniforme em algumas configurações.

4. Generalização para Entradas Naturalistas

O estudo valida que essas descobertas se mantêm em configurações mais naturalistas, onde grupos de entidades são intercalados com texto "preenchedor" (frases abertas). Embora a adição de tokens de preenchimento aumente o ruído e enfraqueça ligeiramente o mecanismo lexical, a dependência fundamental na mistura de três mecanismos persiste, oferecendo uma explicação mecanicista para o efeito "lost-in-the-middle".

Resultados

• Prevalência dos Mecanismos: Em posições intermediárias de contextos longos, o mecanismo posicional accounts por apenas ~20% do comportamento do modelo, enquanto os mecanismos lexical e reflexivo tornam-se os principais impulsionadores da recuperação correta.
• Robustez do Modelo: O padrão de mistura é consistente entre as famílias de modelos (Llama, Gemma, Qwen) e escalas (de 2B a 72B parâmetros).
• Validação Reflexiva: Através de designs contrafactuais específicos onde a resposta contrafactual não existe no contexto original, os autores provam a existência de um mecanismo de ponteiro reflexivo distinto do próprio token de resposta.
• Estudos de Ablação: Remover qualquer mecanismo único do modelo causal degrada significativamente o desempenho, confirmando que todos os três são necessários para uma abstração precisa do comportamento do LM.

Significado e Alegações

O artigo afirma estabelecer uma imagem mais completa de como os LMs vinculam e recuperam entidades em contexto. Ao ir além da visão posicional simplista, os autores fornecem uma explicação mecanicista tanto para as forças (recuperação robusta via múltiplos sinais) quanto para as fragilidades (ruído em posições intermediárias levando a erros "lost-in-the-middle") dos LLMs em configurações de contexto longo.

O trabalho sugere que o fenômeno "lost-in-the-middle" não é meramente uma falha de atenção, mas uma degradação específica do sinal posicional, que é apenas parcialmente compensada pelos mecanismos lexical e reflexivo. Os autores disponibilizam seu código e dados para facilitar pesquisas futuras sobre os mecanismos causais do raciocínio em contexto.