Resumo Técnico: Decomposição de Mecanismos de Predição para Recall em Contexto

Definição do Problema

Este artigo introduz uma nova família de problemas lúdicos (toy problems) projetados para preencher a lacuna entre tarefas de aprendizado em contexto contínuo (como regressão linear) e tarefas de recall associativo discreto (como o Multi-Query Associative Recall, ou MQAR). A tarefa central envolve prever o estado de um sistema de séries temporais discretas e intercaladas.

Os dados subjacentes consistem em trajetórias sem ruído geradas por sistemas dinâmicos lineares ortogonais aleatórios ($x_{i+1} = U x_i$, onde $U$ é uma matriz ortogonal aleatória). Essas trajetórias são segmentadas e intercaladas dentro de uma única janela de contexto, demarcadas por "rótulos de pontuação simbólica" discretos (tokens de abertura e fechamento). O modelo é treinado para realizar a predição do próximo token em sequências intercaladas. O desafio exige que o modelo:

1. Identifique qual estado de sistema específico deve ser recuperado com base em um rótulo simbólico.
2. Aplique a dinâmica linear subjacente correta a esse estado para prever observações subsequentes.
3. Continue a predição para a sequência retomada.

Os autores postulam que o sucesso exige a distinção entre duas capacidades distintas: aprender a dinâmica de um sistema conforme ele é visto pela primeira vez (ICL) e recuperar a dinâmica de um sistema visto anteriormente no contexto (recall associativo).

Metodologia

Geração de Dados

Os autores geram uma biblioteca de treinamento de 40.000 matrizes ortogonais ($U \in \mathbb{R}^{5 \times 5}$) e estados iniciais correspondentes. As trajetórias de treinamento são construídas intercalando segmentos dessas séries temporais.

• Intercalação: Segmentos são cortados e remontados com rótulos simbólicos aleatórios de "abertura" e "fechamento" atribuídos a cada instância do sistema. Crucialmente, o mapeamento entre um sistema e seu rótulo simbólico é randomizado para cada exemplo de treinamento, impedindo que o modelo memorize associações estáticas de rótulo-sistema.
• Distribuição: O número de sistemas por trajetória segue uma distribuição de Zipf e o número de cortes segue uma distribuição de Poisson, criando um cenário de "agulha no palheiro" (needle-in-a-haystack) onde o modelo deve recuperar um sistema específico em um contexto contendo vários outros.

Arquitetura do Modelo e Treinamento

Os autores treinam modelos do tipo GPT-2 (variando de 212K a 10,7M de parâmetros) usando predição de próximo token via erro quadrático médio (MSE). A dimensão do embedding de entrada é 57 (50 para rótulos simbólicos, 1 para início, 1 para flag de carga útil, 5 para vetores de estado). A camada de saída prevê o vetor de estado de 5 dimensões.

Protocolos de Avaliação

1. Não Intercalado (Uninterleaved): Avalia a capacidade do modelo de aprender a dinâmica do primeiro sistema visto no contexto.
2. Teste Agulha no Palheiro (Needle-in-a-Haystack - NiH): Avalia o recall associativo. Um "palheiro" de $N$ sistemas é apresentado, seguido por um símbolo de consulta indicando qual sistema deve ser retomado ("a agulha"). O desempenho é medido no primeiro token após a consulta (1-após) e em tokens subsequentes (2-após, etc.).
3. Experimentos Fora de Distribuição (OOD): Quatro experimentos específicos são desenhados para testar o mecanismo de recall:
   • Direcionamento Equivocado para Sequência Incorreta: Trocar o rótulo de consulta para apontar para um sistema diferente no palheiro.
   • Rotações Sincronizadas: Forçar todos os sistemas no palheiro a terem o mesmo estado no primeiro índice do segmento de teste, removendo a desambiguação baseada no estado.
   • Direcionamento Equivocado para Sequência Não Vista: Usar um rótulo de consulta que corresponde a um sistema não presente no palheiro.
   • Direcionamento Equivocado para Sequência Vista (Reinício):* Usar um rótulo que aponta para um sistema antigo enquanto o segmento de teste real é um novo sistema não visto.
4. Análise de Circuito: A poda de arestas (edge pruning) é aplicada aos pesos treinados para isolar os circuitos computacionais específicos usados para prever o token "1-após" versus o token "2-após".

Principais Contribuições e Resultados

1. Dinâmicas de Aprendizado Distintas para ICL e Recall

A dinâmica de treinamento revela uma separação clara sobre quando diferentes capacidades emergem:

• Aprendizado do Primeiro Sistema: A capacidade de aprender a dinâmica do primeiro sistema visto no contexto desenvolve-se gradualmente ao longo do treinamento, convergindo para o preditor de pseudoinversa ótimo.
• Reinício de ICL: A capacidade de reiniciar o ICL para um novo sistema (após ver outros sistemas) emerge relativamente cedo, mas ainda precede o recall associativo total.
• Recall Associativo: A capacidade de recuperar um sistema visto anteriormente baseando-se apenas em seu rótulo simbólico emerge abruptamente (uma transição de fase) muito mais tarde no treinamento, especificamente para o primeiro token após a consulta.

2. A Hipótese de Múltiplos Mecanismos (C3)

A descoberta central é que o modelo não utiliza um único mecanismo para toda a tarefa. Em vez disso, emprega dois mecanismos distintos que emergem em tempos diferentes e utilizam informações diferentes:

• Mecanismo 1 (Baseado em Rótulo): Usado exclusivamente para prever o primeiro token após um símbolo de consulta. Este mecanismo depende do rótulo simbólico discreto para recuperar o estado correto do sistema. Evidência: Quando o rótulo é direcionado incorretamente para um sistema errado, o erro para o primeiro token aumenta imediatamente.
• Mecanismo 2 (Baseado em Observação/Bayesiano): Usado para prever o segundo token e além. Este mecanismo ignora amplamente o rótulo simbólico e, em vez disso, usa o estado observado do primeiro token para inferir a dinâmica subjacente do sistema (uma predição de estilo "Bayesiano"). Evidência: Quando o rótulo é direcionado incorretamente, o erro para os tokens subsequentes permanece baixo porque o modelo infere o sistema correto a partir do estado observado.

3. Desentranhamento de Circuitos

A poda de arestas confirma que esses dois mecanismos são fisicamente distintos nos pesos do modelo. Os circuitos responsáveis por prever o token "1-após" e o token "2-após" possuem 0% de sobreposição de arestas, indicando que o modelo aprende sub-redes separadas para iniciar um episódio de tarefa versus continuá-lo.

4. Generalização para Linguagem Natural

Os autores validam este fenômeno em uma tarefa de linguagem natural (tradução Inglês-Espanhol) usando checkpoints do OLMo. Eles observam um fenômeno semelhante de transição de fase: a capacidade do modelo de iniciar uma tradução baseada em um rótulo de tarefa simbólico (recall associativo em contexto) emerge mais tarde e de forma mais abrupta do que sua capacidade de continuar a tradução uma vez iniciada. Quando rótulos semânticos (ex: "Spanish:") são usados em vez de simbólicos (ex: "X:"), a lacra desaparece, confirmando que o atraso é específico para o aprendizado em contexto da definição da tarefa.

Significância e Alegações

O artigo afirma fornecer uma explicação mecanística de como os Transformers realizam tarefas complexas em contexto, desafiando a visão do ICL como uma capacidade monolítica.

• Decomposição do ICL: Os autores argumentam que o "aprendizado em contexto" não é uma habilidade única, mas um composto de habilidades distintas (reconhecimento de tarefa vs. continuação de tarefa) que emergem sequencialmente.
• Subespecificação Local: O artigo sugere que, como o treinamento autorregressivo minimiza a perda média sobre os tokens, o modelo pode ter sucesso usando um mecanismo "bom o suficiente" (inferência baseada em observação) para a maioria dos tokens (2º, 3º, etc.), o que reduz a pressão do gradiente para aprender o mecanismo mais complexo (recuperação baseada em rótulo) exigido para o primeiro token.
• Transições de Fase: Os resultados apoiam a existência de distintas transições de fase no aprendizado, onde a emergência de uma capacidade (continuação) não garante a emergência imediata de outra (iniciação), mesmo quando a informação teoricamente necessária é similar.
• Interpretabilidade Mecanística: O trabalho demonstra que a poda de arestas pode isolar com sucesso esses circuitos funcionais distintos, fornecendo uma ferramenta para entender como diferentes partes de uma tarefa são tratadas por diferentes partes da rede.

Os autores concluem que suas descobertas oferecem uma nova dimensão de nuance à compreensão do ICL, sugerindo que, para tarefas de múltiplos tokens, os modelos podem operar simultaneamente em modos diferentes (baseado em rótulo vs. baseado em observação) dependendo da posição na sequência.