A Lightweight Library for Energy-Based Joint-Embedding Predictive Architectures

1. El Problema

El campo de los modelos del mundo (world models) y la generación de video ha avanzado significativamente, pero enfrenta desafíos fundamentales:

• Costo Computacional: Los enfoques generativos actuales (como los modelos de difusión) deben modelar todos los píxeles, incluyendo detalles irrelevantes para la tarea, lo que requiere recursos computacionales masivos.
• Barreras de Entrada: Las implementaciones existentes de Arquitecturas de Predicción de Incrustación Conjunta (JEPAs) suelen estar diseñadas para entrenamiento a gran escala, son difíciles de navegar y carecen de ejemplos accesibles que conecten los principios teóricos con la aplicación práctica.
• Complejidad en la Planificación: Aprender a predecir el futuro en un espacio de representación semántico, especialmente cuando se condiciona por acciones (para la planificación), es complejo y a menudo requiere configuraciones específicas (como codificadores preentrenados congelados) que dificultan la investigación educativa y la experimentación rápida.

2. Metodología: EB-JEPA

El artículo presenta EB-JEPA, una biblioteca de código abierto que implementa JEPAs basados en Modelos Basados en Energía (EBMs). En lugar de reconstruir observaciones en el espacio de píxeles, JEPAs aprenden a predecir en un espacio de representación latente.

Marco Unificado

El núcleo de la metodología es minimizar una función de energía que mide la compatibilidad entre entradas y salidas. El objetivo de entrenamiento general es:
$$L = L_{pred}(g_\phi(z, u), z') + \lambda R(z)$$
Donde:

• $z = f_\theta(x)$ es la representación del input.
• $g_\phi$ es el predictor.
• $u$ es la información de condicionamiento (opcional, ej. controles robóticos).
• $R(z)$ es un término de regularización crucial para evitar el colapso de la representación (que el modelo aprenda constantes triviales).

La biblioteca implementa tres instancias progresivamente complejas:

1. Image-JEPA: Aprendizaje de representaciones invariantes a la vista (dos vistas de la misma imagen).
2. Video-JEPA: Predicción temporal en espacio latente (predecir el siguiente frame basado en el contexto).
3. AC-Video-JEPA (Action-Conditioned): Modelado del mundo donde el predictor se condiciona en las acciones ($a_t$) para predecir el siguiente estado, habilitando la planificación.

Regularización y Prevención de Colapso

Para evitar que las representaciones colapsen a valores constantes sin usar gradientes stop-gradient (técnica común en otros métodos), EB-JEPA utiliza regularizaciones explícitas:

• VICReg: Utiliza términos de varianza (para asegurar dispersión de características) y covarianza (para descorrelacionar dimensiones).
• SIGReg: Una alternativa teóricamente fundamentada que fuerza la distribución de las incrustaciones a ser una Gaussiana isotrópica mediante pruebas de normalidad en proyecciones 1D. Es más estable y requiere menos ajuste de hiperparámetros.
• Regularizaciones Adicionales para el Mundo: En entornos aleatorios, se añaden pérdidas de similitud temporal (suavidad de la trayectoria) y modelo de dinámica inversa (IDM) (predecir la acción dada la transición de estados) para evitar correlaciones espurias.

Planificación

Para la planificación, el sistema optimiza secuencias de acciones para alcanzar un objetivo ($x_g$). Utiliza el algoritmo MPPI (Model Predictive Path Integral), que muestrea trayectorias de acciones y las pondera según su energía acumulada (costo) a lo largo del horizonte de planificación, no solo en el estado final.

3. Contribuciones Clave

1. Implementaciones Accesibles: Tres ejemplos progresivos (imagen, video, planificación) entrenables en una sola GPU en pocas horas.
2. Arquitectura Modular: Componentes reutilizables (codificadores, predictores, regularizadores, planificadores) que facilitan la combinación para nuevas aplicaciones.
3. Recurso Educativo: Código bien documentado diseñado para ayudar a los investigadores a comprender los principios de JEPA y el aprendizaje auto-supervisado basado en energía.
4. Análisis Exhaustivo: Ablaciones sistemáticas que demuestran la importancia crítica de cada componente de regularización.

4. Resultados Experimentales

Los experimentos se realizaron en tres tareas de complejidad creciente:

• Representación de Imágenes (CIFAR-10):

  • Lograron una precisión de 91.02% con SIGReg y 90.12% con VICReg mediante linear probing en una red ResNet-18.
  • Se demostró que el uso de un projector (capa de proyección) mejora el rendimiento en ~3 puntos.
  • SIGReg mostró mayor estabilidad frente a variaciones de hiperparámetros en comparación con VICReg.

• Predicción de Video (Moving MNIST):

  • El entrenamiento con rollouts de múltiples pasos (predicción recursiva) mejoró significativamente la precisión promedio en tareas de detección downstream en comparación con la predicción de un solo paso, reduciendo el sesgo de exposición.
  • El modelo mantuvo predicciones coherentes de la trayectoria de dígitos en horizontes extendidos.

• Planificación Condicionada por Acción (Two Rooms):

  • En un entorno de navegación no monótono (donde a veces hay que alejarse del objetivo para llegar a él), el modelo logró una tasa de éxito del 97% utilizando MPPI con costos acumulados.
  • Ablaciones de Regularización:
    • Sin el término IDM, el éxito colapsó al 1% (debido a correlaciones espurias).
    • Sin varianza o covarianza, el éxito cayó a ~46-47%.
    • Sin similitud temporal, cayó al 61%.
  • El uso de un costo acumulado sobre todo el horizonte de planificación superó en un 8% al uso de solo el estado final.

5. Significado e Impacto

El trabajo EB-JEPA es significativo por varias razones:

• Democratización de la Investigación: Reduce la barrera de entrada para estudiar JEPAs y modelos del mundo, permitiendo que investigadores con recursos limitados (una sola GPU) realicen experimentos fundamentales.
• Validación de Principios Teóricos: Confirma empíricamente que la predicción en espacio de representación, combinada con regularizaciones específicas (como IDM y SIGReg), es una vía viable y eficiente para el modelado del mundo y la planificación, evitando los costos de los modelos generativos pixel a pixel.
• Puente entre Teoría y Práctica: Proporciona un "banco de pruebas" ideal para explorar la teoría de la regularización y la dinámica de colapso antes de escalar a configuraciones distribuidas masivas.
• Futuro: Establece una base para futuras investigaciones en modelos jerárquicos del mundo y funciones de valor aprendidas, complementando las bibliotecas de gran escala existentes.

En resumen, EB-JEPA no solo ofrece un rendimiento competitivo en benchmarks estándar, sino que actúa como una herramienta pedagógica y experimental crucial para avanzar en la comprensión de cómo las máquinas pueden aprender modelos internos de su entorno de manera eficiente y semánticamente rica.