Resumen Técnico: Sub-JEPA

Enunciado del Problema

Las Arquitecturas Predictivas de Incrustación Conjunta (JEPAs) ofrecen un marco para aprender modelos del mundo mediante la predicción de representaciones latentes futuras sin reconstruir píxeles crudos. Sin embargo, entrenar estos modelos de extremo a extremo enfrenta una compensación sesgo-varianza crítica. Sin restricciones suficientes, el codificador colapsa hacia soluciones triviales (mapeando todas las entradas a incrustaciones idénticas); por el contrario, restricciones excesivas suprimen la riqueza de las representaciones aprendidas.

Trabajos recientes, LeWorldModel (LeWM), abordaron esto aplicando una priori gaussiana isotrópica a las incrustaciones latentes, permitiendo un entrenamiento estable de extremo a extremo. No obstante, los autores identifican una limitación fundamental en este enfoque: las tareas de control naturales típicamente residen en variedades de baja dimensión dentro de espacios ambientales de alta dimensión. Imponer una priori gaussiana isotrópica global directamente en el espacio ambiental de alta dimensión introduce un sesgo excesivamente fuerte que no coincide con la geometría intrínseca de la dinámica subyacente, lo que potencialmente obstaculiza el rendimiento en tareas con baja dimensionalidad intrínseca.

Metodología: Sub-JEPA

Los autores proponen Sub-JEPA (Regularización Gaussiana en Subespacio), un método que busca un punto operativo favorable en la frontera sesgo-varianza al reubicar las restricciones gaussianas desde el espacio ambiental hacia múltiples subespacios de baja dimensión.

Mecanismo Central

1. Proyección en Subespacio Ortogonal: En lugar de regularizar el vector latente completo de $D$ dimensiones $z$, el método proyecta $z$ sobre $K$ matrices de proyección aleatorias, fijas y ortonormales por filas $\{P_k\}$. Cada proyección mapea el vector latente a un subespacio de menor dimensión de dimensión $d_s$ (donde $d_s \approx D/K$).
   $$z^{(k)} = P_k z$$
2. Regularización Gaussiana Multi-Subespacio: La restricción gaussiana se aplica independientemente dentro de cada uno de los $K$ subespacios.
   • Dentro de cada subespacio $k$, las incrustaciones proyectadas se proyectan adicionalmente sobre $M$ vectores unitarios aleatorios.
   • Se calcula la estadística de normalidad de Epps–Pulley para estas proyecciones unidimensionales para imponer la gaussianidad.
   • La pérdida total de regularización $L_{reg}$ es el promedio de estas estadísticas a través de todos los subespacios y direcciones aleatorias.
3. Objetivo de Entrenamiento: El modelo se entrena de extremo a extremo utilizando una pérdida combinada:
   $$L_{total} = L_{pred}(\hat{z}_{t+1}, z_{t+1}) + \lambda L_{reg}(Z)$$
   donde $L_{pred}$ es el error de predicción del estado latente (MSE) y $\lambda$ es el peso de regularización.

Decisiones de Diseño

• Proyecciones Fijas: Las matrices de proyección se inicializan ortogonalmente y se mantienen fijas. Esto evita que el regularizador se adapte a la distribución latente en evolución, asegurando que la restricción permanezca estable e independiente.
• Ortogonalidad: El uso de proyecciones ortogonales asegura que los subespacios proporcionen vistas no redundantes y equilibradas de la representación latente, preservando la estructura geométrica.

Contribuciones Clave

1. Identificación del Sesgo: El artículo identifica que la priori gaussiana isotrópica global en LeWM introduce un sesgo excesivo para tareas con baja dimensionalidad intrínseca, no coincidiendo con la geometría subyacente de la tarea.
2. Marco Sub-JEPA: Los autores proponen un método que relaja la priori global aplicando regularización gaussiana sobre múltiples vistas proyectadas aleatorias, de baja dimensión y ortonormales por filas. Esto preserva el efecto anti-colapso mientras permite que el espacio latente se ajuste mejor a la dimensionalidad intrínseca de la dinámica.
3. Validación Empírica: Experimentos extensos en cuatro entornos de control continuo demuestran que Sub-JEPA supera consistentemente a LeWM. Las ganancias de rendimiento se muestran correlacionadas con reducciones en la rango efectivo de las incrustaciones aprendidas, confirmando que el método respeta mejor la dimensionalidad intrínseca.

Resultados Experimentales

El método se evaluó en cuatro puntos de referencia de control continuo: Two-Room (baja dimensionalidad intrínseca), Reacher, PushT y OGB-Cube (3D visualmente rico).

• Rendimiento de Planificación: Sub-JEPA superó consistentemente a LeWM en todos los entornos. La mejora fue más pronunciada en Two-Room, donde la priori gaussiana de espacio completo de LeWM estaba más desacoplada de la baja dimensionalidad de la tarea. Sub-JEPA también superó a PLDM (una línea base de múltiples objetivos) y DINO-WM (una línea base de codificador fijo) en varias tareas sin depender de espaldas visuales preentrenadas.
• Análisis de Rango Efectivo: Se observó una correlación directa entre la reducción del rango efectivo (de LeWM a Sub-JEPA) y el aumento en la tasa de éxito de planificación. Las tareas con compresiones de rango mayores (Two-Room, OGB-Cube) mostraron las mayores ganancias de rendimiento.
• Decodificabilidad Física: Experimentos de sondeo en el entorno PushT mostraron que las representaciones de Sub-JEPA preservaron la información del estado físico (ubicación del agente, ubicación del bloque) tan bien o mejor que LeWM, con una decodificabilidad no lineal mejorada.
• Calidad de la Trayectoria Latente:
  • Coherencia Temporal: Las visualizaciones mostraron que Sub-JEPA produjo trayectorias latentes más coherentes temporalmente en comparación con la evolución menos organizada en LeWM.
  • Enderezamiento de Trayectorias: Sub-JEPA logró una mayor "rectitud" en los vectores de velocidad latente, indicando dinámicas más suaves y que evolucionan de manera más lineal, lo cual es crucial para la planificación.
  • Estabilidad a Largo Plazo: En despliegues en bucle abierto, Sub-JEPA mantuvo la geometría de la escena y exhibió menos deriva espacial a lo largo de horizontes largos en comparación con LeWM.

Significado y Afirmaciones

El artículo afirma que Sub-JEPA ofrece una solución simple pero efectiva a la compensación sesgo-varianza en modelos del mundo basados en JEPA. Al mover la regularización gaussiana a subespacios, el método:

• Relaja la priori isotrópica global sin perder los beneficios anti-colapso.
• Coincide mejor con la dimensionalidad intrínseca de la dinámica subyacente, particularmente para tareas donde el espacio de estados es de baja dimensión.
• Sirve como una línea base sólida para futuras investigaciones sobre modelos del mundo basados en JEPA, demostrando que las restricciones estructurales pueden simplificarse y hacerse más flexibles sin sacrificar la estabilidad.

Los autores enfatizan que su enfoque conserva la simplicidad de la receta de entrenamiento de LeWM (utilizando solo pérdida de predicción y un regularizador) mientras introduce una priori estructural más flexible que produce un rendimiento de planificación superior y una mayor calidad de representación.