Resumo Técnico: MapFormers

Declaração do Problema

Os sistemas de IA baseados em Transformers atuais, apesar do desempenho impressionante em tarefas dentro da distribuição, exibem fragilidade significativa ao enfrentar cenários fora da distribuição (OOD). Frequentemente, eles falham em generalizar para situações novas que humanos e animais lidam com facilidade. Os autores atribuem essa lacuna à falta de vieses estruturais que permitem aos cérebros biológicos construir mapas cognitivos: modelos internos que desentrelaçam a estrutura abstrata e invariante de um ambiente de seu conteúdo específico e variável.

As tentativas existentes de modelar mapas cognitivos, como a Máquina Tolman-Eichenbaum (TEM) que utiliza matrizes dependentes de ação em RNNs, sofrem de duas limitações críticas:

1. Falta de Escalabilidade: Elas exigem multiplicação de matrizes sequencial para atualizar posições, impedindo o processamento paralelo.
2. Ações Pré-especificadas: Frequentemente exigem que o conjunto de ações (por exemplo, direções cardeais) seja codificado manualmente em vez de aprendido a partir dos dados.

Além disso, os codificadores posicionais padrão (como RoPE) são estáticos ou dependem de heurísticas que não suportam a composição dinâmica de ações necessária para a integração de caminho em ambientes desconhecidos.

Metodologia

O artigo introduz os MapFormers, uma classe de arquiteturas Transformer projetadas para aprender mapas cognitivos e realizar integração de caminho via previsão do próximo token sem supervisão. A inovação central reside no uso de matrizes dependentes de entrada para atualizar codificações posicionais, derivadas da teoria dos grupos de Lie.

Estrutura Teórica

• Formalismo de Grupo de Lie: Os autores modelam ações como elementos de um grupo de Lie (especificamente grupos compactos como $SO(2)$ para rotações). Em vez de aprender uma matriz separada para cada ação possível, o modelo aprende geradores de álgebra de Lie.
• Integração de Caminho via Mapas Exponenciais: As ações são representadas como geradores infinitesimais na álgebra de Lie. A integração de caminho (compor uma sequência de ações) é realizada somando esses geradores na álgebra (acumulação linear) e, em seguida, aplicando um único mapa exponencial para obter a transformação finita no grupo de Lie.
  • Para mapas comutativos (por exemplo, navegação em grade 2D), isso permite que a integração de caminho seja computada em paralelo via soma cumulativa, seguida por uma única exponencial de matriz.
  • Essa abordagem reduz o viés indutivo do número de ações para a complexidade estrutural da álgebra.

Variantes Arquiteturais

O artigo propõe duas variantes de MapFormers, unificando codificação posicional absoluta e relativa com modelos de memória distintos:

1. MapEM (Memória Episódica):

   • Usa embeddings posicionais absolutos para representar um mapa cognitivo.
   • Desentrelaça conteúdo (observações) e estrutura (posição) em populações neurais separadas.
   • Computa atenção na conjunção de conteúdo e posição ($Q_G = \text{vec}(Q^T \cdot P)$), atuando efetivamente como uma máscara de atenção onde as codificações posicionais atualizam-se independentemente do conteúdo.
   • Adequado para tarefas que exigem recordação robusta de estados específicos.

2. MapWM (Memória de Trabalho):

   • Usa codificação posicional relativa (análoga ao RoPE, mas dinâmica).
   • Representa a estrutura implicitamente girando a atividade neural (chaves e consultas) em subespaços distintos via matrizes de rotação dependentes de entrada ($R_{\theta PI}$).
   • O ângulo de rotação $\theta$ é derivado de uma projeção de baixo posto do token de entrada, permitindo que o modelo aprenda quais tokens atuam como ações (atualizando a posição) e quais são observações (deixando a posição inalterada).

Mecanismo Dependente de Entrada

Ambas as arquiteturas empregam uma projeção em dois estágios para determinar a "duração de integração" ($\Delta_t$) para cada token:

1. Projeção Interna: Mapeia o token de entrada para uma representação de ação de baixa dimensão.
2. Projeção Externa: Projeta essa representação para uma duração de integração por bloco.
   O ângulo de rotação final é $\theta_t = \omega \odot \Delta_t$, onde $\omega$ é um vetor de velocidade angular aprendível. A integração de caminho é computada como a soma cumulativa desses ângulos ($\theta_{0 \to t} = \text{cumsum}(\theta_t)$).

Contribuições Principais

1. Arquitetura MapFormers: Uma variante Transformer novel que aprende mapas cognitivos desentrelaçando estrutura de conteúdo usando codificações posicionais dependentes de entrada baseadas em geradores de álgebra de Lie.
2. Integração de Caminho Paralela: Ao aproveitar grupos de Lie comutativos (por exemplo, $SO(2)^k$), os modelos alcançam integração de caminho paralela, superando o gargalo sequencial de modelos anteriores de mapas cognitivos baseados em RNN.
3. Unificação de Modelos de Memória: O trabalho fornece um link teórico entre codificação posicional absoluta/relativa e memória episódica/de trabalho, mostrando como ambos podem ser derivados dos mesmos princípios teóricos de Lie.
4. Escalabilidade para Mapas Não-Comutativos: O framework é estendido para problemas não-comutativos (por exemplo, árvores genealógicas) usando múltiplos geradores, embora isso exija integração de caminho sequencial.

Resultados Experimentais

Os autores avaliaram os MapFormers em três tarefas formais que visam capacidades cognitivas distintas:

1. Cópia Seletiva (Gating):

   • Tarefa: Copiar uma sequência ignorando tokens de distração.
   • Resultado: Os MapFormers (tanto EM quanto WM) alcançaram generalização OOD perfeita, enquanto as linhas de base padrão (RoPE, TAPE) falharam. Isso demonstra a capacidade de aprender mecanismos de gating dinâmicos.

2. Navegação em Grade (1D e 2D):

   • Tarefa: Prever observações em uma sequência de ações alternadas e observações aleatórias, exigindo que o modelo infera a topologia latente e integre a posição.
   • Resultado: Os MapFormers alcançaram precisão quase perfeita em sequências OOD (variando comprimento e densidade), enquanto as linhas de base (incluindo CoPE e PathAtt) falharam em generalizar.
   • Análise Neural: Os modelos aprenderam a projetar tokens de observação para rotação zero (identidade) e tokens de ação para direções de rotação específicas, efetivamente desentrelaçando conteúdo de estrutura.

3. Dyck-2 (Hierarquias Aninhadas):

   • Tarefa: Prever continuições válidas de sequências de parênteses, exigindo o rastreamento de uma pilha.
   • Resultado: Os MapFormers generalizaram perfeitamente para sequências mais longas e profundas do que as vistas no treinamento usando apenas uma única camada. As linhas de base exigiram múltiplas camadas e ainda falharam em dados OOD.
   • Mecanismo: Os modelos embutiram a pilha diretamente na codificação posicional usando dimensões ortogonais para diferentes tipos de parênteses e direções opostas para parênteses de abertura/fechamento.

4. Escalonamento em Linguagem Natural:

   • Tarefa: Pré-treinamento no OpenWebText.
   • Resultado: Um modelo MapWM de 12 camadas mostrou melhorias consistentes na perplexidade sobre uma linha de base RoPE ($18,79$ vs $19,14$) e melhor extrapolação de comprimento. No entanto, no benchmark sintático BLIMP, o MapWM não mostrou ganhos significativos sobre o RoPE, sugerindo que a integração de caminho comutativa sozinha pode não ser suficiente para toda a modelagem sintática.

Significado e Afirmações

O artigo afirma que os MapFormers fornecem um viés estrutural crítico que permite generalização OOD sistemática ao capturar estruturas latentes via integração de caminho, em vez de depender de heurísticas.

• Robustez: Os modelos demonstram que desentrelaçar relações abstratas do conteúdo permite generalização robusta para comprimentos de sequência não vistos, profundidades e densidades ambientais.
• Escalabilidade: O uso de computação paralela em mapas comutativos permite que esses princípios escalem para domínios de grande escala e do mundo real, conforme evidenciado pelos resultados em linguagem natural.
• Plausibilidade Biológica: A abordagem preenche a lacuna entre mapas cognitivos biológicos (que fatorizam "o quê" de "onde") e a aprendizagem profunda moderna, oferecendo um caminho para sistemas de IA que podem raciocinar através de problemas novos entendendo a geometria subjacente de seu ambiente.

Os autores permanecem modestos quanto às limitações, observando que a integração de caminho não-comutativa não é paralelizável e que a restrição comutativa atual pode limitar a modelagem de certas estruturas sintáticas complexas na linguagem. Eles sugerem que trabalhos futuros poderiam explorar arquiteturas híbridas ou extensões não-comutativas para raciocínio relacional mais complexo.