Resumo Técnico: Quando o LeJEPA Aprende um Modelo de Mundo?

Declaração do Problema

O artigo aborda uma lacuna fundamental na aprendizagem auto-supervisionada (SSL): enquanto Arquiteturas de Predição de Incorporação Conjunta (JEPAs) demonstraram sucesso empírico na aprendizagem de representações para visão, vídeo e planejamento, não há garantia teórica de que essas representações constituam um "Modelo de Mundo". Especificamente, permanece incerto sob quais condições uma representação aprendida recupera fielmente a estrutura latente subjacente do mundo (por exemplo, posição, velocidade, cor) em vez de entrelaçar esses fatores.

Uma representação que embaralha os graus de liberdade verdadeiros (por exemplo, misturando a posição de um objeto com sua cor) pode performar bem em tarefas restritas, mas falha quando o mundo muda ou quando usada para generalização composicional. A questão central é: Quando uma representação aprendida recupera linearmente as variáveis latentes do mundo? Os autores definem essa propriedade como identificabilidade linear, uma condição necessária para sondagem linear confiável e planejamento ótimo no espaço latente.

Metodologia e Estrutura Teórica

O Modelo de Mundo

Os autores formalizam o "Mundo" como um sistema com variáveis latentes $z \in \mathbb{R}^n$ que não são observadas. Os dados observados $x$ são gerados por uma função de mistura desconhecida, potencialmente altamente não linear $g$, tal que $x = g(z)$. O objetivo de aprendizagem envolve pares positivos $(z, z')$ gerados por um processo estacionário com transições independentes e aditivas de ruído.

O artigo foca em uma classe específica de mundos onde:

1. Independência: As variáveis latentes e suas transições são independentes entre as dimensões.
2. Estacionariedade: A distribuição marginal das variáveis latentes permanece constante entre as visualizações.
3. Ruído Aditivo: As transições seguem $z' = m(z) + \eta$, onde $\eta$ é ruído independente.

Crucialmente, os autores especializam essa classe para o Mundo Gaussiano, onde $z \sim \mathcal{N}(0, I_n)$. Sob essas suposições, as dinâmicas de transição correspondem a um processo de Ornstein–Uhlenbeck (OU): $z' = \rho z + \sqrt{1-\rho^2}\eta$.

O Aprendiz: LeJEPA

A arquitetura de aprendizagem é o LeJEPA (JEPA de LeCun), que treina um codificador $f$ para produzir incorporações $h(z) = f(g(z))$ minimizando uma função de perda composta por dois termos:

1. Perda de Alinhamento: Minimiza a distância entre as incorporações de pares positivos, $E[\|h(z') - h(z)\|^2]$. Isso incentiva a representação a capturar características lentas e previsíveis.
2. SIGReg (Regularização Isotrópica Gaussiana Esboçada): Regulariza explicitamente a distribuição da incorporação para ser uma Gaussiana isotrópica, $h(z) \sim \mathcal{N}(0, I_n)$, prevenindo o colapso da representação.

O problema de otimização é:
$$\min_h E[\|h(z') - h(z)\|^2] \quad \text{s.t.} \quad h(z) \sim \mathcal{N}(0, I_n)$$

Abordagem Teórica

Os autores empregam análise espectral do operador de transição $T$, definido como $(Th)(z) = E[h(z')|z]$.

• Expansão de Hermite: Para latentes Gaussianos, as autofunções do operador de transição são os polinômios de Hermite. Os autovalores decaem geometricamente como $\rho^d$, onde $d$ é o grau do polinômio.
• Penalidade Espectral: A perda de alinhamento penaliza componentes de grau mais alto (não lineares) mais severamente do que componentes lineares, pois eles se descorrelacionam mais rapidamente sob a transição.
• Teoria de Sturm-Liouville: Para distribuições gerais (não Gaussianas), as autofunções são caracterizadas pela teoria de Sturm-Liouville. A primeira autofunção não constante é sempre monótona, mas a linearidade é garantida apenas se a distribuição for Gaussiana.

Principais Contribuições

1. Direção Direta: Identificabilidade Linear para Mundos Gaussianos

Teorema 1 prova que, se as variáveis latentes são Gaussianas e a transição é um processo OU, qualquer minimizador do objetivo LeJEPA (sujeito a restrições de Gaussianidade) deve ser uma transformação ortogonal linear das verdadeiras latentes: $h(z) = Qz$ para alguma matriz ortogonal $Q$.

• Mecanismo: A prova baseia-se em uma decomposição espectral onde qualquer não linearidade reduz estritamente a correlação entre pares positivos. Como a perda de alinhamento busca maximizar essa correlação, o ótimo global força a representação a ser puramente linear (polinômios de Hermite de grau 1).
• Resultado: A representação aprendida recupera as variáveis latentes até uma rotação global, e as dinâmicas de transição são recuperadas exatamente no espaço latente.

2. Direção Conversa: Unicidade da Gaussiana

Teorema 2 estabelece que a distribuição Gaussiana é a única distribuição latente para a qual a identificabilidade linear vale nesta classe de mundos.

• Mecanismo: Usando a teoria de Sturm-Liouville, os autores mostram que exigir que a primeira autofunção do operador de transição seja afim (linear) força a função de pontuação da distribuição latente a ser linear. Esta condição caracteriza unicamente a distribuição Gaussiana.
• Implicação: Se as variáveis latentes são não Gaussianas, o LeJEPA não pode garantir identificabilidade linear; a representação aprendida será necessariamente uma transformação não linear das latentes.

3. Identificabilidade Aproximada

Teorema 3 fornece um limite quantitativo sobre o erro de recuperação quando os objetivos são apenas aproximadamente satisfeitos (ou seja, alinhamento imperfeito e branqueamento imperfeito).

• O erro escala de forma graciosa com a lacuna de alinhamento $\delta$ e o erro de branqueamento $\epsilon$.
• O limite destaca que a lacuna de alinhamento é o fator dominante no erro de recuperação, enquanto o branqueamento aproximado é relativamente "gratuito" de alcançar.

4. Planejamento Latente Ótimo

Teorema 4 conecta a identificabilidade linear à teoria de controle. Prova que, se o codificador é linearmente identificável ($h(z) = Qz$), então qualquer problema de controle ótimo com custos invariantes sob transformações ortogonais (por exemplo, distância Euclidiana até um objetivo) pode ser resolvido optimalemente no espaço latente aprendido. As ações ótimas e funções de valor no espaço latente mapeiam diretamente para o mundo verdadeiro, não exigindo modificação da função de custo.

Resultados Experimentais

Os autores validam sua teoria em vários cenários experimentais:

1. Simulações 2D: Usando várias funções de mistura não lineares (espiral, cisalhamento, RealNVP), o LeJEPA recupera com sucesso as variáveis latentes até uma rotação quando as latentes são Gaussianas, confirmando o Teorema 1.
2. Escalabilidade de Alta Dimensão: Experimentos até $N=1024$ dimensões mostram que regularizadores de estatísticas de lote (SIGReg, VICReg) mantêm identificabilidade linear quase perfeita ($R^2 > 0.999$), enquanto métodos contrastivos (InfoNCE) degradam em escala devido a incompatibilidades na largura do kernel.
3. Ablações Distribucionais: Varredura da forma da distribuição latente (de caudas pesadas a uniforme) demonstra que a recuperação linear atinge um pico agudo na Gaussiana ($\alpha=2$) e degrada para distribuições não Gaussianas, validando o Teorema 2.
4. Controle Robótico Baseado em Pixels: Usando o DeepMind Control Suite Reacher, os autores comparam o treinamento em transições Gaussianas sintéticas (OU) versus trajetórias reais de RL.
   • Condição OU: Alcança alta identificabilidade linear ($R^2 \approx 0.95$).
   • Condição Trajetória: Políticas do mundo real induzem marginais não Gaussianas e anisotrópicas, resultando em identificabilidade significativamente menor ($R^2 < 0.5$).
5. Validação de Planejamento: Na tarefa Reacher, a interpolação em linha reta no espaço latente aprendido do codificador treinado com Gaussiana decodifica para caminhos quase retos no espaço de juntas (planejamento ótimo). Em contraste, o codificador treinado com trajetórias produz caminhos curvos e subótimos, confirmando o Teorema 4.

Significado e Alegações

O artigo alega fornecer o primeiro resultado de identificabilidade para JEPAs, transformando uma receita empiricamente bem-sucedida em uma garantia matemática.

• Garantia Fundamental: O trabalho estabelece que o LeJEPA recupera provavelmentemente a estrutura do mundo (variáveis latentes e dinâmicas) sob condições específicas e bem definidas (latentes Gaussianas, ruído aditivo estacionário).
• Papel da Gaussianidade: O artigo inverte a narrativa clássica da ICA. Na ICA linear, a Gaussianidade é a distribuição onde a separação de fontes falha. Neste cenário não linear e auto-supervisionado, a Gaussianidade é a única distribuição que permite a identificabilidade linear.
• Implicações Práticas: A teoria explica por que estratégias de exploração que aproximam passeios aleatórios isotrópicos (mantendo os dados no regime Gaussiano) são cruciais para aprender modelos de mundo úteis, enquanto políticas direcionadas a objetivos (que colapsam dados em regiões de baixa entropia) podem quebrar a identificabilidade.
• Fundação para Planejamento: Ao provar que a identificabilidade linear permite planejamento ótimo com custos invariantes à rotação, o artigo fornece uma justificação teórica para o uso de espaços latentes aprendidos para controle sem a necessidade de mapear de volta para as coordenadas físicas verdadeiras.

Os autores notam limitações: a teoria assume que a dimensão do codificador corresponde à dimensão latente e depende da otimalidade ao nível da população. No entanto, eles demonstram que as garantias degradam-se de forma graciosa com amostras finitas e erros de otimização, e verificam formalmente todos os teoremas principais usando o provador de teoremas Lean 4.