Los modelos mundiales (WM) son un marco central para desarrollar agentes que razonan y planifican en un espacio latente compacto. Sin embargo, entrenar estos modelos directamente a partir de datos de píxeles a menudo conduce a un “colapso de representación”, donde el modelo produce incorporaciones redundantes para satisfacer trivialmente los objetivos de predicción. Los enfoques actuales intentan evitar esto basándose en heurísticas complejas: utilizan actualizaciones de gradiente de parada, promedios móviles exponenciales (EMA) y codificadores congelados previamente entrenados. Un equipo de investigadores que incluye a Yann LeCun y muchos otros (Mila & Université de Montréal, New York University, Samsung SAIL y Brown University) presentaron LeWorldModel (LeWM), el primer JEPA (Joint-Embedding Predictive Architecture) que entrena de forma estable de extremo a extremo a partir de píxeles sin procesar utilizando solo dos términos de pérdida: una pérdida de predicción de la próxima incrustación y un regularizador que aplica incrustaciones latentes distribuidas gaussianas.
Arquitectura Técnica y Objetivo
LeWM consta de dos componentes principales aprendidos conjuntamente: un codificador y un predictor.
Codificador ((zt=encθ (ot)): asigna una observación de píxeles sin procesar a una representación latente compacta y de baja dimensión. La implementación utiliza una arquitectura ViT-Tiny (~5M de parámetros). Predictor (Žt+1=predθ(zt, at)): un transformador (~10M de parámetros) que modela la dinámica del entorno al predecir estados latentes futuros condicionados a las acciones.
El modelo se optimiza utilizando una función objetivo simplificada que consta de solo dos términos de pérdida:
$$\mathcal{L}_{LeWM} \triangleq \mathcal{L}_{pred} + \lambda SIGReg(Z)$$
La pérdida de predicción (Lpred) calcula el error cuadrático medio (MSE) entre las incrustaciones consecutivas previstas y reales. El SIGReg (regularizador gaussiano isotrópico esbozado) es el término anti-colapso que impone la diversidad de características.
Según el artículo de investigación, aplicar una tasa de abandono de 0,1 en el predictor y un paso de proyección específico (MLP de 1 capa con normalización por lotes) después del codificador es fundamental para la estabilidad y el rendimiento posterior.
Eficiencia a través de SIGReg y tokenización dispersa
Evaluar la normalidad en espacios latentes de alta dimensión es un gran desafío de escala. LeWM aborda esto utilizando SIGReg, que aprovecha el teorema de Cramér-Wold: una distribución multivariada coincide con un objetivo (gaussiano isotrópico) si todas sus proyecciones unidimensionales coinciden con ese objetivo.
SIGReg proyecta incrustaciones latentes en M direcciones aleatorias y aplica la estadística de prueba de Epps-Pulley a cada proyección unidimensional resultante. Debido a que el peso de regularización λ es el único hiperparámetro efectivo para ajustar, los investigadores pueden optimizarlo utilizando una búsqueda de bisección con complejidad O (log n), una mejora significativa con respecto a la búsqueda de tiempo polinomial (O (n6)) requerida por modelos anteriores como PLDM.
Puntos de referencia de velocidad
En la configuración reportada, LeWM demuestra una alta eficiencia computacional:
Eficiencia de tokens: LeWM codifica observaciones utilizando ~200 veces menos tokens que DINO-WM. Velocidad de planificación: LeWM logra una planificación hasta 48 veces más rápida que DINO-WM (0,98 s frente a 47 s por ciclo de planificación).
Propiedades del espacio latente y comprensión física
El espacio latente de LeWM admite el sondeo de cantidades físicas y la detección de eventos físicamente inverosímiles.
Violación de expectativas (VoE)
Utilizando un marco VoE, se evaluó la capacidad del modelo para detectar “sorpresas”. Asignó mayor sorpresa a perturbaciones físicas como la teletransportación; las perturbaciones visuales produjeron efectos más débiles y los cambios de color del cubo en OGBench-Cube no fueron significativos.
Enderezamiento de camino emergente
LeWM exhibe el enderezamiento temporal de la ruta latente, donde las trayectorias latentes se vuelven naturalmente más suaves y lineales a lo largo del entrenamiento. En particular, LeWM logra una mayor rectitud temporal que PLDM a pesar de no tener un regularizador explícito que fomente este comportamiento.
Conclusiones clave
Aprendizaje estable de extremo a extremo: LeWM es la primera arquitectura predictiva de incrustación conjunta (JEPA) que entrena de manera estable de extremo a extremo a partir de píxeles sin procesar sin necesidad de heurísticas “de mano” como gradientes de parada, promedios móviles exponenciales (EMA) o codificadores preentrenados congelados. Un objetivo radical de dos términos: el proceso de entrenamiento se simplifica en solo dos términos de pérdida (una pérdida de predicción de la siguiente incorporación y el regularizador SIGReg), lo que reduce la cantidad de hiperparámetros ajustables de seis a uno en comparación con las alternativas de extremo a extremo existentes. Creado para la velocidad en tiempo real: al representar observaciones con aproximadamente 200 veces menos tokens que sus contrapartes basadas en modelos básicos, LeWM planifica hasta 48 veces más rápido, completando optimizaciones completas de trayectoria en menos de un segundo. Anti-colapso demostrable: para evitar que el modelo aprenda representaciones redundantes “basura”, utiliza el regularizador SIGReg; esto utiliza el teorema de Cramér-Wold para garantizar que las incrustaciones latentes de alta dimensión permanezcan diversas y con distribución gaussiana. Lógica física intrínseca: el modelo no solo predice datos; captura una estructura física significativa en su espacio latente, lo que le permite sondear con precisión cantidades físicas y detectar eventos “imposibles” como la teletransportación de objetos a través de un marco de violación de expectativas.
Consulte el documento, el sitio web y el repositorio. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 120.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.