Resumen Técnico: UWM-JEPA

Enunciado del Problema

Las Arquitecturas de Predicción de Incrustación Conjunta (JEPAs) predicen estados futuros en el espacio latente en lugar de reconstruir observaciones crudas, separando eficazmente el aprendizaje de representaciones de la decodificación generativa. Sin embargo, en entornos parcialmente observables, una representación latente vectorial estándar carece de la estructura interna necesaria para transportar la "creencia" sobre continuaciones ocultas. Cuando el historial de observaciones actual no determina unívocamente el estado oculto, el futuro no es un único punto en el espacio de representaciones, sino una distribución sobre estados ocultos compatibles.

Los predictores vectoriales estándar entrenados contra objetivos de Media Móvil Exponencial (EMA) pueden minimizar la pérdida, pero la función de pérdida por sí sola no especifica qué incertidumbres deben preservarse, cuáles deben promediarse o cómo las acciones contrafactuales deben desplazar el estado de creencia. En consecuencia, las JEPAS de latente vectorial a menudo fallan en mantener dinámicas sensibles a la acción o en preservar información relevante para la tarea durante el "despliegue ciego" (imaginar estados futuros sin nuevas observaciones).

Metodología

Los autores introducen UWM-JEPA (Modelo de Mundo Unitario JEPA), una modificación del marco JEPA que trata el latente como un estado de creencia explícito en lugar de un punto determinista.

1. Latente de Estado de Creencia

En lugar de un vector, el codificador en línea $E_\theta$ mapea el contexto observado a una matriz de densidad $\rho_t$ sobre un espacio latente bipartito $\mathcal{H}_T = \mathcal{H}_S \otimes \mathcal{H}_E$ (Sistema $\otimes$ Entorno).

• Dimensiones: El factor del sistema ($d_{sys}=8$) es leído por la pérdida y las sondas; el factor del entorno ($d_{env}=2$) actúa como un registro para transportar modos ocultos y correlaciones sistema-entorno.
• Restricciones: La matriz de densidad está restringida a ser hermítica, semidefinida positiva y de traza unitaria. Estas se imponen por construcción: el estado inicial es máximamente mezclado, las actualizaciones de observación aplican operadores de medición positivos y los pasos de predicción utilizan conjugación unitaria.

2. Predictor Unitario

El predictor avanza el estado latente mediante conjugación unitaria:
$$\hat{\rho}_{t+k|t} = U_k \rho_t (U_k)^\dagger$$
donde $U = \exp(-iH\Delta t)$ y $H$ es un generador hermítico.

• Condicionamiento por Acción: El generador es lineal en la acción $a$: $H(a) = H_0 + a H_1$. Esto permite que la trayectoria unitaria sea dirigida por las acciones.
• Objetivo de Entrenamiento: El modelo utiliza el andamiaje estándar de JEPA: un codificador objetivo EMA $E_\xi$ con gradiente detenido, y una pérdida de coincidencia en el espacio latente evaluada después de la traza parcial al factor del sistema y una cabeza de proyección.
• Objetivos Contrafactuales: Para asegurar que el predictor aprenda a utilizar el canal de acción, el objetivo de entrenamiento se genera mediante un despliegue de simulador bajo una secuencia de acciones contrafactuales $\tilde{a}$, en lugar de la trayectoria observada forzada por el maestro.

3. Garantía Teórica

La afirmación teórica central es el no olvido del estado conjunto. Dado que la conjugación unitaria es una isometría, preserva el espectro (valores propios), la pureza y la entropía de von Neumann de la matriz de densidad del sistema-entorno conjunto exactamente durante el despliegue ciego.

• Limitación: Esta garantía se aplica al estado conjunto completo. El estado reducido del sistema (obtenido mediante traza parcial) no es unitario, lo que significa que la información puede moverse hacia correlaciones sistema-entorno y volverse inaccesible para sondas lineales, incluso si el estado conjunto permanece no disipativo.

Resultados Clave

1. La Sensibilidad a la Acción Requiere Objetivos Contrafactuales

• Fallo Forzado por Maestro: Cuando se entrena con objetivos forzados por el maestro (coincidiendo con el futuro observado), el canal de acción se vuelve inerte. La relación del término dependiente de la acción con el término de dinámicas base colapsa a $\|H_1\|/\|H_0\| \approx 0.03$. El predictor aprende a coincidir con el objetivo independientemente de la acción porque el objetivo ya contiene el efecto de la acción.
• Éxito Contrafactual: Reemplazar el objetivo con un despliegue de simulador bajo acciones contrafactuales muestreadas restaura la sensibilidad a la acción, con $\|H_1\|/\|H_0\| = 1.00 \pm 0.15$.
• Controles de Perturbación: UWM-JEPA entrenado con objetivos contrafactuales (UWM-JEPA-CF) muestra una degradación monótona en la precisión de predicción a medida que las acciones se perturban (invertidas, desordenadas, negadas), mientras que un LSTM-JEPA-CF con parámetros coincidentes colapsa a la precisión de la clase mayoritaria (0.53) bajo todas las condiciones.

2. Despliegue Ciego y Retención a Corto Plazo

• Indicador de Velocidad Oculta: En una tarea que requiere una simulación hacia adelante de 5 pasos para predecir el signo de la velocidad oculta (con la observación objetivo enmascarada), UWM-JEPA-CF logra una precisión de 0.77. El LSTM-JEPA-CF coincidente falla, prediciendo la clase mayoritaria para todos los ejemplos.
• Retención de Información ($\Delta R^2$): Midiendo la pérdida de información relativa a una sonda maestra durante el despliegue ciego:
  • En horizontes cortos ($k=1, 3$), UWM-JEPA pierde menos de 10 puntos de $R^2$ (0.023 y 0.069 respectivamente).
  • La línea base LSTM-JEPA pierde significativamente más (0.413 y 0.678).
  • En el horizonte entrenado ($k=5$), LSTM se realinea, pero UWM mantiene un rendimiento superior en pasos intermedios.
  • En horizontes largos ($k>10$), ambos modelos se degradan, lo que indica que la retención a largo plazo sigue siendo un desafío abierto.

3. La Representación del Contexto No es el Motor

• Equivalencia del Codificador: Una sonda lineal retenida sobre las representaciones del codificador congelado muestra que UWM-JEPA y LSTM-JEPA tienen un rendimiento idéntico ($R^2 = 0.901 \pm 0.032$ vs $0.894 \pm 0.021$, $p=0.70$).
• Conclusión: La brecha de rendimiento no se debe a que UWM-JEPA tenga un codificador de contexto superior. La separación surge enteramente de la capacidad del predictor para mantener la estructura del estado de creencia durante el despliegue.

Significado y Afirmaciones

El artículo presenta UWM-JEPA como el primer modelo de mundo de la familia JEPA donde el latente se trata explícitamente como una creencia sobre continuaciones ocultas (una matriz de densidad) y el predictor está restringido a evolucionar esta creencia sin disipar el espectro del estado conjunto.

• Alcance Modesto: Los autores declaran explícitamente que esto no es un planificador de propósito general, un codificador universalmente más fuerte ni una solución competitiva para la planificación a largo plazo. La propiedad de "no olvido" es exacta solo para el estado conjunto, no necesariamente para el estado reducido utilizado en tareas posteriores.
• Contribución Central: El trabajo demuestra que, para que los modelos de mundo JEPA imaginen bajo observabilidad parcial, la geometría latente y las dinámicas del predictor importan más que la capacidad de codificación de contexto congelada por sí sola.
• Implicación Práctica: El estudio destaca que el entrenamiento contra objetivos contrafactuales es esencial para aprender predictores sensibles a la acción, un hallazgo que se extiende más allá de la parametrización unitaria específica.

Los resultados se validan en osciladores parcialmente observables estructurados y entornos estilo CartPole. Los autores reconocen que el enfoque actual requiere acceso al simulador para el entrenamiento contrafactual y aún no se extiende a la predicción visual no estructurada o cadenas de Markov ocultas.