BRo-JEPA: Learning Modular Arithmetic in Latent Space

Resumo Técnico: BRo-JEPA – Aprendendo Aritmética Modular no Espaço Latente

Definição do Problema
O artigo investiga uma questão fundamental no aprendizado profundo: se as redes neurais realmente aprendem regras algébricas abstratas ou se apenas memorizam padrões de treinamento. Embora os modelos atuais se destaquem em tarefas matemáticas dentro de intervalos observados, eles frequentemente falham ao extrapolar para intervalos numéricos ou operações não vistas. Os autores postulam que modelos supervisionados padrão e modelos JEPA (Arquitetura de Predição de Embedding Conjunto) existentes dependem de correlações estatísticas em vez de uma lógica aritmética subjacente. Especificamente, ao serem incumbidos da aritmética modular (ex: $y = (x + k) \mod 10$) usando imagens de dígitos MNIST como estados, os modelos padrão ajustam-se às operações observadas, mas falham em generalizar para operações não observadas (generalização zero-shot estrita).

Metodologia
Os autores utilizam o framework JEPA, onde um modelo prevê a representação latente de um dígito alvo dado um dígito de contexto (imagem) e uma operação aritmética (ação), em vez de prever saídas de pixels brutos. O estudo compara várias variações arquiteturais:

1. Baselines Supervisionados: Modelos supervisionados (MLP e ResNet-18) utilizando perda de entropia cruzada e embeddings de ação composicionais aditivos. Esses embeddings representam operações como múltiplos escalares de um vetor base (ex: $op(k) = k \cdot v$) para impor uma estrutura aritmética no espaço de embedding.
2. JEPA com Embeddings Aditivos: Um modelo JEPA que utiliza os mesmos embeddings aditivos, aumentado com uma "Perda de Consistência de Composição" (Composition Consistency Loss). Esta perda auxiliar impõe que aplicar uma operação composta (ex: $+2$) em um passo deve resultar no mesmo resultado latente que aplicar operações primitivas sequencialmente (ex: $+1$ e depois $+1$).
3. BRo-JEPA (Block-Rotation JEPA): A principal contribuição é um preditor inovador que substitui o MLP padrão não restrito por um preditor de rotação de bloco condicionado à operação.
   • Mecanismo: Inspirado nos Rotary Position Embeddings (RoPE), este preditor trata a dinâmica do espaço latente como uma série de rotações planares 2D. Para um vetor latente $z$ e uma operação $k$, a previsão é $z_p = R_{block}(k \cdot \theta) \cdot z$.
   • Viés Indutivo: Esta arquitetura impõe explicitamente a estrutura cíclica da aritmética módulo-10 sobre o espaço latente. Os ângulos de rotação $\theta$ podem ser fixos (definidos como $2\pi/10$) ou aprendíveis.
   • Variantes: Os autores testam a Rotação de Frequência Única (SFR), onde todos os blocos latentes compartilham um único ângulo, e a Rotação de Múltiplas Frequências (MFR), onde cada bloco possui um ângulo independente.
   • Regularização: Para evitar o "colapso latente" (onde diferentes entradas mapeiam para representações indistinguíveis), o modelo utiliza Regularização de Variância-Invariância-Covariância (VICReg) e um codificador de alvo por Média Móvel Exponencial (EMA).

Principais Contribuições

1. Definição de Testbed: A introdução da aritmética modular de MNIST como um ambiente controlado para estudar o aprendizado de regras em modelos de mundo latentes.
2. Análise de Falha: Demonstração de que os baselines supervisionados e os modelos JEPA com embeddings aditivos, apesar de ajustarem as operações observadas, falham na extrapolação zero-shot estrita. Eles capturam os sinais das operações, mas falham em codificar a magnitude ou a geometria composicional.
3. Proposta do BRo-JEPA: O desenvolvimento de um modelo de dinâmica latente do tipo JEPA que apresenta um preditor de rotação de bloco condicionado à operação, que se alinha com a estrutura cíclica da aritmética modular.
4. Análise Geométrica: Análise mostrando que o BRo-JEPA organiza as representações latentes em um manifesto geométrico estruturado (um manifesto circular), permitindo que o modelo internalize a topologia da aritmética modular.

Resultos
Os experimentos foram conduzidos com treinamento nas operações $\{-1, +1\}$ e avaliação nas operações não vistas $\{\pm 2, \dots, \pm 9\}$.

• Desempenho dos Baselines: Os baselines supervisionados alcançaram alta precisão nas operações observadas (~97-98%), mas precisão próxima de zero nas operações não vistas (2,42% para ResNet-18). Visualizações de PCA revelaram que esses modelos comprimiram arbitrariamente as operações não vistas no espaço de representação das operações observadas.
• JEPA com Embeddings Aditivos: Embora esses modelos tenham mostrado alta precisão de rollout (aplicando $+1$ repetidamente para simular $+k$), falharam na previsão direta zero-shot de operações não vistas (precisão ~10%). A perda de consistência de composição melhorou isso para ~99%, mas é considerada um "zero-shot fraco" porque o modelo foi indiretamente exposto a operações compostas durante o treinamento.
• Desempenho do BRo-JEPA:
  • Precisão Zero-Shot: O BRo-JEPA baseado em ResNet com Rotação de Múltiplas Frequências (MFR) alcançou 99,46% de precisão zero-shot nas operações não vistas.
  • Precisão de Rollout: O modelo também alcançou 99,46% de precisão em rollouts de múltiplos passos, confirmando a estabilidade das dinâmicas latentes aprendidas.
  • Estrutura Latente: Visualizações de PCA confirmaram que o BRo-JEPA conseguiu coagir o espaço latente em um manifesto elíptico/circular contínuo onde as classes de dígitos estão dispostas sequencialmente. A variante MFR produziu protótipos de classe melhor separados (menor similaridade de cosseno fora da diagonal) em comparação com a SFR, levando a uma generalização superior.
  • Ângulos Aprendíveis: Quando os ângulos de rotação eram aprendíveis, o MFR convergiu rapidamente para uma generalização eficaz, enquanto o SFR exibiu atrasos significativos e instabilidade, frequentemente convergindo para soluções harmônicas subótimas.

Significância e Alegações
O artigo afirma que modelos de mundo latentes podem aprender regras de transformação simbólica, mas apenas quando o viés arquitetural coincide com a estrutura do problema. Os resultados sugerem que:

• Preditores neurais padrão carecem das restrições necessárias para preservar a geometria composicional para operações não vistas.
• Modelar explicitamente a natureza cíclica da aritmética modular via rotações de bloco no espaço latente permite uma forte generalização zero-shot sem exigir a exposição às operações específicas durante o treinamento.
• Esta abordagem oferece uma direção promissora para o aprendizado de representações além da aritmética, sugerindo que incorporar estruturas de transformação apropriadas em modelos de transição latente é crucial para o aprendizado de regras abstratas.

Os autores mantêm a modéstia, observando que seus resultados são específicos para o cenário altamente controlado da aritmética modular de MNIST com uma estrutura de grupo cíclico conhecida. Eles não alegam que a arquitetura irá generalizar automaticamente para tarefas de raciocínio simbólico arbitrárias ou grupos de transformação desconhecidos sem vieses estruturais semelhantes.