Resumo Técnico: UWM-JEPA

Declaração do Problema

As Arquiteturas de Predição de Incorporação Conjunta (JEPAs) preveem estados futuros no espaço latente em vez de reconstruir observações brutas, separando efetivamente a aprendizagem de representação da decodificação generativa. No entanto, em ambientes parcialmente observáveis, uma representação latente vetorial padrão carece da estrutura interna necessária para transportar "crença" sobre continuação ocultas. Quando o histórico de observações atual não determina unicamente o estado oculto, o futuro não é um único ponto no espaço de representação, mas uma distribuição sobre estados ocultos compatíveis.

Os preditores vetoriais padrão treinados contra alvos de Média Móvel Exponencial (EMA) podem minimizar a perda, mas a função de perda por si só não especifica quais incertezas devem ser preservadas, quais devem ser suavizadas ou como ações contrafactuais devem deslocar o estado de crença. Consequentemente, JEPAs com latente vetorial frequentemente falham em manter dinâmicas sensíveis à ação ou preservar informações relevantes para a tarefa durante "rolagem cega" (imaginar estados futuros sem novas observações).

Metodologia

Os autores introduzem o UWM-JEPA (Modelo de Mundo Unitário JEPA), uma modificação do framework JEPA que trata o latente como um estado de crença explícito em vez de um ponto determinístico.

1. Latente de Estado de Crença

Em vez de um vetor, o codificador online $E_\theta$ mapeia o contexto observado para uma matriz de densidade $\rho_t$ em um espaço latente bipartido $\mathcal{H}_T = \mathcal{H}_S \otimes \mathcal{H}_E$ (Sistema $\otimes$ Ambiente).

• Dimensões: O fator do sistema ($d_{sys}=8$) é lido pela perda e por sondas; o fator do ambiente ($d_{env}=2$) atua como um registrador para carregar modos ocultos e correlações sistema-ambiente.
• Restrições: A matriz de densidade é restrita a ser hermitiana, semi-definida positiva e com traço unitário. Essas condições são impostas por construção: o estado inicial é maximamente misturado, as atualizações de observação aplicam operadores de medição positivos e os passos de predição usam conjugação unitária.

2. Preditor Unitário

O preditor avança o estado latente via conjugação unitária:
$$\hat{\rho}_{t+k|t} = U_k \rho_t (U_k)^\dagger$$
onde $U = \exp(-iH\Delta t)$ e $H$ é um gerador hermitiano.

• Condicionamento à Ação: O gerador é linear na ação $a$: $H(a) = H_0 + a H_1$. Isso permite que a trajetória unitária seja direcionada pelas ações.
• Objetivo de Treinamento: O modelo utiliza a estrutura padrão JEPA: um codificador alvo EMA $E_\xi$ com gradiente interrompido (stop-gradient) e uma perda de correspondência no espaço latente avaliada após o traço parcial para o fator do sistema e uma camada de projeção.
• Alvos Contrafactuais: Para garantir que o preditor aprenda a usar o canal de ação, o alvo de treinamento é gerado por uma simulação de rolagem sob uma sequência de ações contrafactuais $\tilde{a}$, em vez da trajetória observada forçada pelo professor.

3. Garantia Teórica

A afirmação teórica central é o não-esquecimento do estado conjunto. Como a conjugação unitária é uma isometria, ela preserva o espectro (autovalores), a pureza e a entropia de von Neumann da matriz de densidade do sistema-ambiente conjunto exatamente durante a rolagem cega.

• Limitação: Essa garantia aplica-se ao estado conjunto completo. O estado reduzido do sistema (obtido via traço parcial) não é unitário, o que significa que a informação pode migrar para correlações sistema-ambiente e tornar-se inacessível a sondas lineares, mesmo que o estado conjunto permaneça não dissipativo.

Resultados Chave

1. Sensibilidade à Ação Requer Alvos Contrafactuais

• Falha Forçada pelo Professor: Quando treinado com alvos forçados pelo professor (correspondendo ao futuro observado), o canal de ação torna-se inerte. A razão entre o termo dependente da ação e o termo de dinâmica básica colapsa para $\|H_1\|/\|H_0\| \approx 0.03$. O preditor aprende a corresponder ao alvo independentemente da ação porque o alvo já contém o efeito da ação.
• Sucesso Contrafactual: Substituir o alvo por uma simulação de rolagem sob ações contrafactuais amostradas restaura a sensibilidade à ação, com $\|H_1\|/\|H_0\| = 1.00 \pm 0.15$.
• Controles de Perturbação: O UWM-JEPA treinado com alvos contrafactuais (UWM-JEPA-CF) mostra degradação monótona na precisão da predição à medida que as ações são perturbadas (invertidas, embaralhadas, negadas), enquanto um LSTM-JEPA-CF com parâmetros correspondentes colapsa para a precisão da classe majoritária (0,53) sob todas as condições.

2. Rolagem Cega e Retenção de Curto Prazo

• Indicador de Velocidade Oculta: Em uma tarefa exigindo simulação para frente de 5 passos para prever o sinal da velocidade oculta (com a observação alvo mascarada), o UWM-JEPA-CF alcança 0,77 de precisão. O LSTM-JEPA-CF correspondente falha, prevendo a classe majoritária para todos os exemplos.
• Retenção de Informação ($\Delta R^2$): Medindo a perda de informação relativa a uma sonda do professor durante a rolagem cega:
  • Em horizontes curtos ($k=1, 3$), o UWM-JEPA perde menos de 10 pontos de $R^2$ (0,023 e 0,069, respectivamente).
  • A linha de base LSTM-JEPA perde significativamente mais (0,413 e 0,678).
  • No horizonte treinado ($k=5$), o LSTM se realinha, mas o UWM mantém desempenho superior em etapas intermediárias.
  • Em horizontes longos ($k>10$), ambos os modelos degradam, indicando que a retenção de longo prazo permanece um desafio aberto.

3. A Representação de Contexto Não é o Motor

• Equivalência do Codificador: Uma sonda linear mantida fora do conjunto de dados sobre as representações do codificador congelado mostra que o UWM-JEPA e o LSTM-JEPA performam identicamente ($R^2 = 0,901 \pm 0,032$ vs $0,894 \pm 0,021$, $p=0,70$).
• Conclusão: A lacuna de desempenho não se deve ao UWM-JEPA possuir um codificador de contexto superior. A separação surge inteiramente da capacidade do preditor de manter a estrutura do estado de crença durante a rolagem.

Significado e Alegações

O artigo apresenta o UWM-JEPA como o primeiro modelo de mundo da família JEPA onde o latente é tratado explicitamente como uma crença sobre continuação ocultas (uma matriz de densidade) e o preditor é restrito a evoluir essa crença sem dissipar o espectro do estado conjunto.

• Escopo Moderado: Os autores afirmam explicitamente que este não é um planejador de propósito geral, um codificador universalmente mais forte ou uma solução competitiva para planejamento de longo prazo. A propriedade de "não-esquecimento" é exata apenas para o estado conjunto, não necessariamente para o estado reduzido usado em tarefas subsequentes.
• Contribuição Central: O trabalho demonstra que, para que modelos de mundo JEPA imaginem sob observabilidade parcial, a geometria latente e as dinâmicas do preditor importam mais do que a capacidade de codificação de contexto congelada isoladamente.
• Implicação Prática: O estudo destaca que o treinamento contra alvos contrafactuais é essencial para aprender preditores sensíveis à ação, uma descoberta que se estende além da parametrização unitária específica.

Os resultados são validados em ambientes estruturados parcialmente observáveis de osciladores e estilo CartPole. Os autores reconhecem que a abordagem atual requer acesso ao simulador para treinamento contrafactual e ainda não se estende à predição visual não estruturada ou cadeias de Markov ocultas.