Resumo Técnico: Sub-JEPA

Declaração do Problema

As Arquiteturas Preditivas de Incorporação Conjunta (JEPAs) oferecem uma estrutura para aprender modelos de mundo prevendo representações latentes futuras sem reconstruir pixels brutos. No entanto, treinar esses modelos de ponta a ponta enfrenta um compromisso viés-variância crítico. Sem restrições suficientes, o codificador colapsa para soluções triviais (mapeando todas as entradas para incorporações idênticas); inversamente, restrições excessivas suprimem a riqueza das representações aprendidas.

Trabalho recente, LeWorldModel (LeWM), abordou isso aplicando um prior Gaussiano isotrópico às incorporações latentes, permitindo treinamento estável de ponta a ponta. No entanto, os autores identificam uma limitação fundamental nessa abordagem: tarefas de controle natural tipicamente residem em variedades de baixa dimensão dentro de espaços ambientes de alta dimensão. Impor um prior Gaussiano isotrópico global diretamente no espaço ambiente de alta dimensão introduz um viés excessivamente forte que não corresponde à geometria intrínseca das dinâmicas subjacentes, potencialmente prejudicando o desempenho em tarefas com baixa dimensionalidade intrínseca.

Metodologia: Sub-JEPA

Os autores propõem Sub-JEPA (Regularização Gaussiana em Subespaço), um método que busca um ponto de operação favorável na fronteira viés-variância, relocando restrições Gaussianas do espaço ambiente para múltiplos subespaços de baixa dimensão.

Mecanismo Central

1. Projeção em Subespaço Ortogonal: Em vez de regularizar o vetor latente completo $D$-dimensional $z$, o método projeta $z$ em $K$ matrizes de projeção aleatória congeladas e de linhas ortonormais $\{P_k\}$. Cada projeção mapeia o vetor latente para um subespaço de dimensão inferior de dimensão $d_s$ (onde $d_s \approx D/K$).
   $$z^{(k)} = P_k z$$
2. Regularização Gaussiana Multi-Subespaço: A restrição Gaussiana é aplicada independentemente dentro de cada um dos $K$ subespaços.
   • Dentro de cada subespaço $k$, as incorporações projetadas são ainda projetadas em $M$ vetores unitários aleatórios.
   • A estatística de normalidade Epps–Pulley é computada para essas projeções 1D para impor a gaussianidade.
   • A perda total de regularização $L_{reg}$ é a média dessas estatísticas através de todos os subespaços e direções aleatórias.
3. Objetivo de Treinamento: O modelo é treinado de ponta a ponta usando uma perda combinada:
   $$L_{total} = L_{pred}(\hat{z}_{t+1}, z_{t+1}) + \lambda L_{reg}(Z)$$
   onde $L_{pred}$ é o erro de previsão do estado latente (MSE) e $\lambda$ é o peso da regularização.

Escolhas de Design

• Projeções Congeladas: As matrizes de projeção são inicializadas ortogonalmente e mantidas congeladas. Isso impede que o regularizador se adapte à distribuição latente em evolução, garantindo que a restrição permaneça estável e independente.
• Ortogonalidade: O uso de projeções ortogonais garante que os subespaços forneçam visões não redundantes e equilibradas da representação latente, preservando a estrutura geométrica.

Contribuições Principais

1. Identificação do Viés: O artigo identifica que o prior Gaussiano isotrópico global no LeWM introduz viés excessivo para tarefas com baixa dimensionalidade intrínseca, não correspondendo à geometria subjacente da tarefa.
2. Estrutura Sub-JEPA: Os autores propõem um método que relaxa o prior global aplicando regularização Gaussiana sobre múltiplas visões projetadas aleatórias, de baixa dimensão e de linhas ortonormais. Isso preserva o efeito anti-colapso enquanto permite que o espaço latente se conformede melhor à dimensionalidade intrínseca das dinâmicas.
3. Validação Empírica: Extensos experimentos em quatro ambientes de controle contínuo demonstram que o Sub-JEPA supera consistentemente o LeWM. Os ganhos de desempenho mostraram-se correlacionados com reduções no rank efetivo das incorporações aprendidas, confirmando que o método respeita melhor a dimensionalidade intrínseca.

Resultados Experimentais

O método foi avaliado em quatro benchmarks de controle contínuo: Two-Room (baixa dimensionalidade intrínseca), Reacher, PushT e OGB-Cube (3D visualmente rico).

• Desempenho de Planejamento: O Sub-JEPA superou consistentemente o LeWM em todos os ambientes. A melhoria foi mais pronunciada no Two-Room, onde o prior Gaussiano de espaço completo do LeWM estava mais desalinhado com a baixa dimensionalidade da tarefa. O Sub-JEPA também superou o PLDM (uma linha de base multi-objetivo) e o DINO-WM (uma linha de base com codificador congelado) em várias tarefas sem depender de backbones visuais pré-treinados.
• Análise de Rank Efetivo: Uma correlação direta foi observada entre a redução no rank efetivo (do LeWM para o Sub-JEPA) e o aumento na taxa de sucesso do planejamento. Tarefas com compressões de rank maiores (Two-Room, OGB-Cube) mostraram os maiores ganhos de desempenho.
• Decodificabilidade Física: Experimentos de sondagem no ambiente PushT mostraram que as representações do Sub-JEPA preservaram informações de estado físico (localização do agente, localização do bloco) tão bem ou melhor que o LeWM, com decodificabilidade não linear aprimorada.
• Qualidade da Trajetória Latente:
  • Coerência Temporal: Visualizações mostraram que o Sub-JEPA produziu trajetórias latentes mais coerentes temporalmente em comparação com a evolução menos organizada no LeWM.
  • Endireitamento de Caminhos: O Sub-JEPA alcançou maior "retidão" nos vetores de velocidade latente, indicando dinâmicas mais suaves e que evoluem de forma mais linear, o que é crucial para o planejamento.
  • Estabilidade de Longo Horizonte: Em execuções em malha aberta, o Sub-JEPA manteve a geometria da cena e exibiu menos deriva espacial ao longo de horizontes longos em comparação com o LeWM.

Significado e Alegações

O artigo alega que o Sub-JEPA oferece uma solução simples, porém eficaz, para o compromisso viés-variância em modelos de mundo baseados em JEPA. Ao mover a regularização Gaussiana para subespaços, o método:

• Relaxa o prior isotrópico global sem perder os benefícios anti-colapso.
• Corresponde melhor à dimensionalidade intrínseca das dinâmicas subjacentes, particularmente para tarefas onde o espaço de estados é de baixa dimensão.
• Serve como uma forte linha de base para futuras pesquisas em modelos de mundo baseados em JEPA, demonstrando que restrições estruturais podem ser simplificadas e tornadas mais flexíveis sem sacrificar a estabilidade.

Os autores enfatizam que sua abordagem mantém a simplicidade da receita de treinamento do LeWM (usando apenas perda de previsão e um regularizador) enquanto introduz um prior estrutural mais flexível que produz desempenho de planejamento superior e qualidade de representação.