A Lightweight Library for Energy-Based Joint-Embedding Predictive Architectures

Título: EB-JEPA: Uma Biblioteca Leve para Arquiteturas Preditivas de Incorporação Conjunta Baseadas em Energia

1. O Problema

O campo de modelos de mundo (world models) e aprendizado de representações enfrenta desafios fundamentais:

• Ineficiência de Modelos Generativos: Métodos baseados em reconstrução de pixels (como MAE e VideoMAE) exigem recursos computacionais massivos para modelar detalhes irrelevantes para a tarefa, além de frequentemente aprenderem características não informativas para percepção e planejamento.
• Barreira de Entrada: Embora as Arquiteturas Preditivas de Incorporação Conjunta (JEPAs) tenham mostrado resultados promissores em aprendizado auto-supervisionado e modelagem de mundo, as implementações existentes são frequentemente complexas, projetadas para treinamento em grande escala e difíceis de navegar para fins educacionais ou experimentação rápida.
• Colapso de Representação: Um desafio crítico no treinamento de JEPAs é evitar o "colapso", onde o encoder aprende representações triviais e constantes. Estratégias existentes (como stop-gradient) podem ser instáveis ou difíceis de ajustar.

2. Metodologia: O Framework EB-JEPA

O artigo apresenta o EB-JEPA, uma biblioteca de código aberto que implementa JEPAs baseadas em Modelos Baseados em Energia (EBMs). A ideia central é aprender a prever no espaço de representações latentes em vez do espaço de pixels.

Princípios Fundamentais:

• Função de Energia: O modelo define uma função de energia $E(x, y)$ que mede a compatibilidade entre entrada e saída. O objetivo é minimizar a energia para pares corretos (previsão precisa) e maximizá-la para incorretos.
• Regularização para Prevenir Colapso: Diferente de métodos que usam stop-gradient ou amostras negativas, o EB-JEPA utiliza perdas de regularização explícitas para manter a diversidade das representações.
• Objetivo Unificado: A função de perda geral combina erro de previsão e regularização:
  $$L = L_{pred}(g_\phi(z, u), z') + \lambda R(z)$$
  Onde $z$ é a representação, $u$ é a condição (ex: ações), e $R$ é o regularizador.

Três Níveis de Complexidade Implementados:

1. Image-JEPA (Representação de Imagem): Aprende invariância a diferentes visualizações (cortes, jitter de cor) de uma mesma imagem.
2. Video-JEPA (Predição Temporal): Prevê a próxima representação latente baseada em um contexto de frames anteriores, capturando dinâmicas temporais sem acesso a frames futuros.
3. AC-Video-JEPA (Modelo de Mundo Condicionado a Ação): Estende a predição para incluir ações de controle. O modelo aprende a prever o próximo estado latente dado o estado atual e uma ação, permitindo planejamento.

Componentes Chave:

• Regularizadores: Implementação de VICReg (variância e covariância) e SIGReg (baseado em testes de Gaussianidade isotrópica). O SIGReg é destacado por sua estabilidade e simplicidade (apenas um hiperparâmetro).
• Treinamento com Rollout Multietapa: Para reduzir o viés de exposição (exposure bias) e melhorar a previsão de longo prazo, o modelo é treinado com previsões recursivas de $k$ passos, alinhando o treinamento com a inferência autossorregressiva.
• Planejamento: Utiliza o algoritmo MPPI (Model Predictive Path Integral) para otimizar sequências de ações que minimizam a energia acumulada até atingir um objetivo, em vez de apenas o estado final.

3. Principais Contribuições

1. Implementações Acessíveis: Três exemplos progressivos (imagem, vídeo, planejamento) treináveis em uma única GPU em poucas horas, democratizando o acesso a JEPAs.
2. Arquitetura Modular: Componentes reutilizáveis (encoders, preditores, regularizadores, planejadores) que facilitam a recombinação para novas aplicações.
3. Recursos Educacionais: Documentação clara e estrutura de código projetada para ajudar pesquisadores a entender os princípios de aprendizado auto-supervisionado baseado em energia.
4. Análise Abrangente: Estudos de ablação sistemáticos que revelam a importância crítica de cada componente de regularização.

4. Resultados Experimentais

Os experimentos foram realizados em três benchmarks: CIFAR-10, Moving MNIST e o ambiente "Two Rooms".

• Aprendizado de Representação (CIFAR-10):

  • O modelo alcançou 91% de precisão em linear probing (usando ResNet-18), competindo com métodos auto-supervisionados de ponta.
  • O uso de um projetor aprendido melhorou a precisão em ~3 pontos.
  • SIGReg demonstrou maior estabilidade em relação à sintonização de hiperparâmetros comparado ao VICReg.

• Predição de Vídeo (Moving MNIST):

  • O treinamento com rollouts de múltiplos passos ($k=4$) superou significativamente a predição de passo único em tarefas de detecção downstream, reduzindo o viés de exposição.
  • O modelo manteve previsões coerentes de movimento de dígitos em horizontes estendidos.

• Planejamento Condicionado a Ação (Two Rooms):

  • O modelo alcançou uma taxa de sucesso de 97% na tarefa de navegação em um ambiente com paredes aleatórias.
  • Ablação de Regularização (Crítico):
    • Remover a perda de Dinâmica Inversa (IDM) causou o colapso total do modelo (sucesso de 1%), devido a correlações espúrias.
    • Remover termos de variância ou covariância reduziu o sucesso em ~50%.
    • Remover similaridade temporal reduziu o sucesso em ~35%.
  • O uso de custo cumulativo ao longo do horizonte de planejamento foi superior ao custo apenas no estado final.

5. Significado e Impacto

O EB-JEPA preenche uma lacuna crítica entre a teoria de JEPAs e a prática acessível.

• Validação Teórica: Confirma que a abordagem baseada em energia com regularização explícita é eficaz para evitar colapso e aprender características úteis para planejamento, sem a necessidade de reconstrução de pixels.
• Ferramenta de Pesquisa: Permite a iteração rápida de inovações algorítmicas em pequena escala antes de escalar para ambientes complexos e distribuídos.
• Futuro: O trabalho aponta para direções como modelos de mundo hierárquicos (multi-escala temporal) e o aprendizado de funções de custo e valor específicas para tarefas, integrando modelagem de mundo com aprendizado por reforço baseado em recompensa.

Em resumo, o EB-JEPA não é apenas uma biblioteca de código, mas um manifesto pedagógico e prático que demonstra como construir modelos de mundo eficientes, estáveis e escaláveis focando em representações semânticas em vez de pixels.