MoP-JEPA: Hard-Assigned Predictor Mixtures for Stochastic JEPA World Models

Resumo Técnico: MoP-JEPA

Definição do Problema
Modelos de mundo JEPA (Arquitetura Preditiva de Embedding Conjunto) tipicamente predizem o próximo estado latente usando um único regressor. Embora eficaz nesta abordagem em ambientes determinísticos, ela falha sob transições estocásticas onde um único contexto (par estado-ação) pode levar a múltiplos sucessores válidos. Objetivos de regressão padrão (perda quadrática ou perda de cosseno) forçam o modelo a prever a média condicional ou a direção normalizada dessa média. Em distribuições multimodais, esse output de "compromisso" frequentemente corresponde a um estado inexistente (por exemplo, um ponto médio inválido entre dois destinos distintos), tornando o preditor inútil para tarefas de planejamento que exigem a enumeração de caminhos futuros válidos. Tentativas existentes de introduzir estocasticidade, como Mixture-of-Experts (MoE) com portões fundidos por uma soma ponderada ou distribuições variacionais de Gauss única, ainda colapsam para um único vetor de saída por contexto, retendo a mesma limitação fundamental.

Metodologia
O artigo propõe o MoP-JEPA (Mixture-of-Predictors JEPA), que substitui o preditor de saída única por uma estrutura capaz de produzir um conjunto finito de sucessores candidatos em uma única passagem direta (forward pass).

• Arquitetura: O modelo mantém o backbone padrão do JEPA (encoder $f_\theta$ e encoder alvo EMA $f_\xi$), mas substitui o preditor de saída única por $K$ cabeças de predição paralelas $\{g_{\phi_k}\}_{k=1}^K$.
• Roteador: Um roteador baseado apenas no contexto $r_\psi$ recebe o contexto $c = [z, a]$ (estado latente e ação) e produz uma distribuição de probabilidade $\pi$ sobre as $K$ cabeças. Crucialmente, o roteador não vê o estado futuro alvo $z^+$, garantindo que ele aprenda a prever sucessores válidos baseando-se apenas no contexto atual.
• Objetivo de Atribuição Rígida (Hard Assignment): O treinamento utiliza um mecanismo "Winner-Take-All" (WTA) inspirado em quantização de vetores condicional. Para cada $z^+_i$ alvo, a cabeça mais próxima $k^*_i$ é identificada. A função de perda consiste em três termos:
  1. Perda de Predição: A distância entre o output da cabeça vencedora e o alvo.
  2. Perda do Roteador: Entropia cruzada para treinar o roteador a prever o índice da cabeça vencedora.
  3. Balanceamento de Carga: Uma penalidade para evitar que o roteador colapse para uma única cabeça, incentivando a utilização de todas as $K$ cabeças.
• Inferência: No tempo de teste, o modelo produz um conjunto de $K$ latentes candidatos $\{u_k\}$ com as probabilidades do roteador associadas $\{\pi_k\}$. Este conjunto finito serve como o grafo de transição para planejadores subsequentes (ex: BFS ou MPC) sem a necessidade de um decoder ou múltiplas execuções (rollouts).

Principais Contribuições

1. Interface de Sucessor Finito para JEPA: O artigo introduz um mecanismo para converter a restrição de saída única do JEPA em um conjunto de candidatos finitos, permitindo que o modelo represente transições multimodais explicitamente.
2. Caracterização Teórica do Colapso: Os autores demonstram formalmente (via Proposições 1–3) que preditores determinísticos e MoEs fundidos otimizam inerentemente para a média condicional, que é frequentemente um estado inválido em ambientes estocásticos. Eles identificam as cabeças de atribuição rígida como uma forma de quantização de vetores condicional que pode cobrir distintos modos.
3. Protocolo de Avaliação de Rota Verificada (Verified-Route): Reconhecendo que uma alta cobertura pode ser alcançada por predição indiscriminada (ex: prever todos os possíveis edges independentemente do contexto), o artigo introduz uma métrica "realroute". Esta métrica avalia se o conjunto de candidatos proposto contém um caminho válido de transições reais do início ao objetivo, distinguindo o aprendizado dependente de contexto de mera cobertura.

Resultos
Os autores avaliam o MoP-JEPA em datasets OGBench (incluindo variantes de PointMaze e AntMaze com teleportação e stitching, que induzem transições multimodais).

• Desempenho de Planejamento: Em transições de teste (held-out) do OGBench, a busca em grafo usando preditores de saída única tem sucesso em apenas 0,02–0,09 das consultas. Em contraste, o MoP-JEPA atinge uma taxa de sucesso de 0,85.
• Recuperação de Modos: Visualizações no espaço latente (Figuras 3 e 4) mostram que preditores densos colapsam para um único ponto médio inválido entre ramificações, enquanto o MoP-JEPA recupera com sucesso picos distintos correspondentes a cada ramificação válida.
• Sucesso de Rota Verificada (Verified-Route): O MoP-JEPA lidera na métrica "realroute" em todos os três ambientes de labirinto. Embora uma Rede de Densidade de Mistura (MDN) atinja alta cobertura bruta, ela prevê muitos edges inexistentes, falhando no teste de rota verificada. O MoP-JEja mantém alta cobertura garantindo que os edges previstos correspondam a transições reais.
• Ablação: O mecanismo de portão (gating) do roteador mostra-se essencial; sem ele, o modelo pode reter cabeças que não correspondem a transições válidas dependentes de contexto.

Significância e Alegações
O artigo alega que a principal limitação dos atuais modelos de mundo JEPA em configurações estocásticas é a restrição de saída única, que força um compromisso que algoritmos de planejamento não podem utilizar. Ao introduzir cabeças de atribuição rígida, o MoP-JEPA permite que o modelo "enumere" modos de sucessores discretos em vez de apenas tirar a média deles. Os autores enfatizam que sua contribuição não é um novo encoder ou maquinaria anti-colapso, mas um mecanismo específico ao nível do preditor que permite evidência de sucesso em nível de planejamento. O trabalho fornece um relato empírico de como a predição de saída única degrada o planejamento sob estocasticidade e oferece um protocolo para distinguir entre sucessores úteis e dependentes de contexto versus cobertura indiscriminada.