Understanding In-Context Learning Beyond Transformers: An Investigation of State Space and Hybrid Architectures

Título: Compreendendo o Aprendizado em Contexto (ICL) Além dos Transformers: Uma Investigação de Arquiteturas de Espaço de Estado e Híbridas

1. Problema e Motivação

O Aprendizado em Contexto (In-Context Learning - ICL) é a capacidade emergente de Grandes Modelos de Linguagem (LLMs) de aprenderem a realizar tarefas a partir de poucos exemplos (few-shot) apresentados no prompt, sem atualizar os parâmetros do modelo.

• Contexto Atual: Estudos de interpretabilidade mecânica já atribuíram o ICL em modelos baseados em Transformers a mecanismos específicos, como "Cabeças de Vetor de Função" (Function Vectors - FVs) e "Cabeças de Indução".
• A Lacuna: Com o surgimento de modelos mais eficientes, como Modelos de Espaço de Estado (SSMs) (ex: Mamba, Mamba2) e Arquiteturas Híbridas (que combinam Transformers e SSMs), pouco se sabe sobre os mecanismos internos que habilitam o ICL nessas novas arquiteturas.
• Objetivo: Investigar se os mecanismos de ICL descobertos em Transformers (especificamente FVs) são universais ou se variam entre SSMs puros e modelos híbridos, e como eles se comportam em diferentes tipos de tarefas de conhecimento.

2. Metodologia

Os autores empregaram uma abordagem combinada de análise comportamental e interpretabilidade mecânica sobre uma variedade de modelos e tarefas.

• Modelos Analisados:
  • Transformers: Qwen2.5-1.5B, Llama-3.2-1B, Gemma-3-1B-PT.
  • SSMs (Puros): Mamba-1.4B, Mamba2-1.3B.
  • Híbridos: Hymba-1.5B-Base (camadas paralelas de Atenção e Mamba), Zamba2-1.2B (camadas sequenciais: Mamba2 seguido de Atenção).
• Tipos de Tarefas (ICL Baseado em Conhecimento):
  1. Recuperação de Conhecimento Paramétrico: Tarefas que exigem recuperar fatos armazenados nos pesos do modelo (ex: "Qual a capital da França?").
  2. Compreensão de Conhecimento Contextual: Tarefas que exigem inferir regras ou padrões a partir do texto fornecido no prompt (ex: classificação de sentimentos, detecção de discurso de ódio).
• Técnicas de Análise:
  • Experimentos de Randomização de Rótulos: Para testar a robustez e a capacidade de aprendizado com exemplos corretos vs. incorretos.
  • Experimentos de "Label Flipping" (Inversão de Rótulos): Para verificar se o modelo aprende a nova associação contextual ou apenas memoriza o conhecimento prévio.
  • Identificação de Vetores de Função (FVs): Cálculo do Efeito Indireto Médio (AIE) para identificar quais cabeças (heads) contribuem mais para a mudança de logits em direção à resposta correta.
  • Intervenção Mecânica:
    • Steering (Direcionamento): Adicionar o vetor de função médio à saída da cabeça durante a inferência.
    • Ablação: Remover ou zerar a ativação das cabeças identificadas como FVs para observar a queda de desempenho.

3. Principais Resultados e Descobertas

A. Comportamento Geral e Robustez:

• Embora diferentes arquiteturas apresentem desempenho qualitativamente similar em tarefas gerais, seus mecanismos internos são distintos.
• Modelos baseados em Mamba2 e Zamba2 mostraram desempenho marginal em tarefas de compreensão contextual quando os exemplos eram perfeitos, mas sofreram degradação significativa quando os exemplos estavam errados, sugerindo uma dependência diferente de memorização versus aprendizado real.
• No entanto, sob o experimento de label flipping, todos os modelos (incluindo Mamba2) conseguiram aprender associações contrafactuais, indicando que todos possuem capacidade de ICL, mas com eficiências variadas.

B. Localização e Papel das Cabeças de Vetor de Função (FVs):

• Transformers e Mamba (v1): As FVs são os principais impulsionadores do ICL.
• Mamba2: Diferente dos outros, o Mamba2 não depende significativamente de FVs para realizar ICL. A intervenção (steering) em suas cabeças tem pouco impacto, sugerindo que ele utiliza um mecanismo diferente (possivelmente relacionado à sua arquitetura de estado seletivo com dimensões menores por cabeça).
• Modelos Híbridos (Hymba e Zamba2): O ICL é primariamente controlado pelas FVs localizadas nas camadas de Atenção (Self-Attention), independentemente de as camadas serem empilhadas em série ou em paralelo. As camadas de Mamba nos híbridos contribuem menos para a recuperação de conhecimento paramétrico via FVs.

C. Diferença entre Tipos de Tarefas:

• Recuperação Paramétrica: As FVs são altamente concentradas e cruciais. Em modelos híbridos, as FVs estão fortemente localizadas nas camadas de atenção.
• Compreensão Contextual: As FVs têm menos impacto e são menos concentradas.
  • Em tarefas contextuais, a intervenção nas camadas de Mamba de modelos híbridos frequentemente reduz o desempenho, sugerindo que essas camadas servem a outras funções além da recuperação de FVs.
  • Os conjuntos de FVs para tarefas paramétricas e contextuais não se sobrepõem significativamente na maioria dos modelos (exceto no Mamba-1.4B).

D. Localização por Camada:

• O direcionamento (steering) de FVs em camadas intermediárias e tardias tende a melhorar o desempenho de ICL, especialmente em modelos Transformers e híbridos.
• Em Zamba2, as camadas de atenção intermediárias desempenham um papel crítico.

4. Contribuições Chave

1. Extensão da Interpretabilidade: O trabalho expande o conceito de Vetores de Função (FVs), anteriormente limitado a Transformers, para SSMs e arquiteturas híbridas, revelando que o mecanismo não é universal.
2. Descoberta do Mecanismo do Mamba2: Evidencia que o Mamba2 opera com um mecanismo de ICL distinto, menos dependente de FVs tradicionais, possivelmente devido à sua estrutura de cabeças com dimensões ocultas menores.
3. Distinção de Tarefas: Demonstra que "Recuperação de Conhecimento" e "Compreensão Contextual" dependem de subconjuntos diferentes de cabeças e mecanismos dentro do mesmo modelo.
4. Validação de Arquiteturas Híbridas: Confirma que, em modelos híbridos, a capacidade de ICL é ancorada nas camadas de atenção, mesmo que a maior parte do modelo seja composta por SSMs.

5. Significado e Impacto

• Para a Pesquisa de LLMs: Este estudo desafia a visão de que o ICL é um fenômeno monolítico. Ele sugere que a eficiência computacional de modelos como Mamba2 pode vir com uma mudança fundamental na forma como o aprendizado em contexto é realizado internamente.
• Para o Design de Modelos: Para melhorar o ICL em modelos SSMs ou híbridos, pode ser necessário focar no aprimoramento das camadas de atenção ou investigar novos mecanismos além das FVs para modelos puramente baseados em SSMs (como o Mamba2).
• Metodológico: O artigo reforça a necessidade de combinar análises comportamentais (desempenho) com análises mecânicas (intervenção interna) para obter uma compreensão completa das capacidades dos modelos.

Em resumo, o trabalho revela que, embora o desempenho superficial possa ser similar, a "engrenagem" interna do aprendizado em contexto varia drasticamente entre Transformers, SSMs clássicos, SSMs de nova geração (Mamba2) e modelos híbridos, exigindo estratégias de análise e otimização específicas para cada arquitetura.