Resumen Técnico: ¿Cuándo Aprende LeJEPA un Modelo del Mundo?

Enunciado del Problema

El artículo aborda una brecha fundamental en el aprendizaje auto-supervisado (SSL): si bien las Arquitecturas Predictivas de Incrustación Conjunta (JEPAs) han demostrado éxito empírico en el aprendizaje de representaciones para visión, video y planificación, no existe garantía teórica de que estas representaciones constituyan un "Modelo del Mundo". Específicamente, permanece poco claro bajo qué condiciones una representación aprendida recupera fielmente la estructura latente subyacente del mundo (por ejemplo, posición, velocidad, color) en lugar de entrelazar estos factores.

Una representación que desordene los grados de libertad reales (por ejemplo, mezclando la posición de un objeto con su color) puede rendir bien en tareas estrechas, pero falla cuando el mundo cambia o cuando se utiliza para generalización composicional. La pregunta central es: ¿Cuándo recupera linealmente una representación aprendida las variables latentes del mundo? Los autores definen esta propiedad como identificabilidad lineal, una condición necesaria para la sonda lineal fiable y la planificación óptima en el espacio latente.

Metodología y Marco Teórico

El Modelo del Mundo

Los autores formalizan el "Mundo" como un sistema con variables latentes $z \in \mathbb{R}^n$ que no son observadas. Los datos observados $x$ son generados por una función de mezcla desconocida, potencialmente altamente no lineal $g$, tal que $x = g(z)$. El objetivo de aprendizaje involucra pares positivos $(z, z')$ generados por un proceso estacionario con transiciones aditivas de ruido independiente.

El artículo se centra en una clase específica de mundos donde:

1. Independencia: Las variables latentes y sus transiciones son independientes a través de las dimensiones.
2. Estacionariedad: La distribución marginal de las variables latentes permanece constante entre vistas.
3. Ruido Aditivo: Las transiciones siguen $z' = m(z) + \eta$, donde $\eta$ es ruido independiente.

Crucialmente, los autores especializan esta clase al Mundo Gaussiano, donde $z \sim \mathcal{N}(0, I_n)$. Bajo estas suposiciones, las dinámicas de transición corresponden a un proceso de Ornstein–Uhlenbeck (OU): $z' = \rho z + \sqrt{1-\rho^2}\eta$.

El Aprendiz: LeJEPA

La arquitectura de aprendizaje es LeJEPA (JEPA de LeCun), que entrena un codificador $f$ para producir incrustaciones $h(z) = f(g(z))$ minimizando una función de pérdida compuesta por dos términos:

1. Pérdida de Alineación: Minimiza la distancia entre las incrustaciones de pares positivos, $E[\|h(z') - h(z)\|^2]$. Esto fomenta que la representación capture características lentas y predecibles.
2. SIGReg (Regularización Isotrópica de Gaussiana Esbozada): Regulariza explícitamente la distribución de la incrustación para que sea una Gaussiana isotrópica, $h(z) \sim \mathcal{N}(0, I_n)$, evitando el colapso de la representación.

El problema de optimización es:
$$\min_h E[\|h(z') - h(z)\|^2] \quad \text{sujeto a} \quad h(z) \sim \mathcal{N}(0, I_n)$$

Enfoque Teórico

Los autores emplean análisis espectral del operador de transición $T$, definido como $(Th)(z) = E[h(z')|z]$.

• Expansión de Hermite: Para latentes gaussianos, las autofunciones del operador de transición son los polinomios de Hermite. Los autovalores decaen geométricamente como $\rho^d$, donde $d$ es el grado del polinomio.
• Penalización Espectral: La pérdida de alineación penaliza más severamente los componentes de mayor grado (no lineales) que los componentes lineales, ya que se descorrelacionan más rápido bajo la transición.
• Teoría de Sturm-Liouville: Para distribuciones generales (no gaussianas), las autofunciones se caracterizan por la teoría de Sturm-Liouville. La primera autofunción no constante es siempre monótona, pero la linealidad solo está garantizada si la distribución es Gaussiana.

Contribuciones Clave

1. Dirección Directa: Identificabilidad Lineal para Mundos Gaussianos

El Teorema 1 demuestra que si las variables latentes son gaussianas y la transición es un proceso OU, cualquier minimizador del objetivo de LeJEPA (sujeto a restricciones de gaussianidad) debe ser una transformación ortogonal lineal de los verdaderos latentes: $h(z) = Qz$ para alguna matriz ortogonal $Q$.

• Mecanismo: La prueba se basa en una descomposición espectral donde cualquier no linealidad reduce estrictamente la correlación entre pares positivos. Dado que la pérdida de alineación busca maximizar esta correlación, el óptimo global fuerza a que la representación sea puramente lineal (polinomios de Hermite de grado 1).
• Resultado: La representación aprendida recupera las variables latentes hasta una rotación global, y las dinámicas de transición se recuperan exactamente en el espacio latente.

2. Dirección Converse: Unicidad de la Gaussiana

El Teorema 2 establece que la distribución Gaussiana es la única distribución latente para la cual la identificabilidad lineal se cumple en esta clase de mundos.

• Mecanismo: Utilizando la teoría de Sturm-Liouville, los autores muestran que exigir que la primera autofunción del operador de transición sea afín (lineal) fuerza a que la función de puntuación (score function) de la distribución latente sea lineal. Esta condición caracteriza únicamente a la distribución Gaussiana.
• Implicación: Si las variables latentes no son gaussianas, LeJEPA no puede garantizar la identificabilidad lineal; la representación aprendida será necesariamente una transformación no lineal de los latentes.

3. Identificabilidad Aproximada

El Teorema 3 proporciona un límite cuantitativo sobre el error de recuperación cuando los objetivos solo se satisfacen aproximadamente (es decir, alineación imperfecta y blanqueado imperfecto).

• El error escala de manera suave con la brecha de alineación $\delta$ y el error de blanqueado $\epsilon$.
• El límite destaca que la brecha de alineación es el factor dominante en el error de recuperación, mientras que el blanqueado aproximado es relativamente "gratuito" de lograr.

4. Planificación Latente Óptima

El Teorema 4 conecta la identificabilidad lineal con la teoría de control. Demuestra que si el codificador es linealmente identificable ($h(z) = Qz$), entonces cualquier problema de control óptimo con costos invariantes bajo transformaciones ortogonales (por ejemplo, distancia euclidiana a un objetivo) puede resolverse óptimamente en el espacio latente aprendido. Las acciones óptimas y las funciones de valor en el espacio latente se mapean directamente al mundo real, sin necesidad de modificar la función de costos.

Resultados Experimentales

Los autores validan su teoría en varios entornos experimentales:

1. Simulaciones 2D: Utilizando diversas funciones de mezcla no lineales (espiral, cizalladura, RealNVP), LeJEPA recupera con éxito las variables latentes hasta una rotación cuando los latentes son gaussianos, confirmando el Teorema 1.
2. Escalado a Alta Dimensión: Experimentos hasta $N=1024$ dimensiones muestran que los regularizadores de estadísticas de lote (SIGReg, VICReg) mantienen una identificabilidad lineal casi perfecta ($R^2 > 0.999$), mientras que los métodos contrastivos (InfoNCE) se degradan a escala debido a desajustes en el ancho del kernel.
3. Ablaciones Distribucionales: Recorrer la forma de la distribución latente (desde colas pesadas hasta uniforme) demuestra que la recuperación lineal alcanza un pico agudo en la Gaussiana ($\alpha=2$) y se degrada para distribuciones no gaussianas, validando el Teorema 2.
4. Control Robótico Basado en Píxeles: Utilizando el DeepMind Control Suite Reacher, los autores comparan el entrenamiento con transiciones gaussianas sintéticas (OU) versus trayectorias reales de RL.
   • Condición OU: Logra alta identificabilidad lineal ($R^2 \approx 0.95$).
   • Condición Trayectoria: Las políticas del mundo real inducen marginales no gaussianas y anisotrópicas, resultando en una identificabilidad significativamente menor ($R^2 < 0.5$).
5. Validación de Planificación: En la tarea Reacher, la interpolación en línea recta en el espacio latente aprendido del codificador entrenado con Gaussiana se decodifica en trayectorias casi rectas en el espacio de articulaciones (planificación óptima). En contraste, el codificador entrenado con trayectorias produce trayectorias curvas y subóptimas, confirmando el Teorema 4.

Significado y Afirmaciones

El artículo afirma proporcionar el primer resultado de identificabilidad para JEPAs, convirtiendo una receta empíricamente exitosa en una garantía matemática.

• Garantía Fundacional: El trabajo establece que LeJEPA recupera demostrablemente la estructura del mundo (variables latentes y dinámicas) bajo condiciones específicas y bien definidas (latentes gaussianos, ruido aditivo estacionario).
• Rol de la Gaussianidad: El artículo invierte la narrativa clásica del ICA. En el ICA lineal, la gaussianidad es la distribución donde falla la separación de fuentes. En este entorno no lineal y auto-supervisado, la gaussianidad es la única distribución que permite la identificabilidad lineal.
• Implicaciones Prácticas: La teoría explica por qué las estrategias de exploración que aproximan caminatas aleatorias isotrópicas (manteniendo los datos en el régimen gaussiano) son cruciales para aprender modelos del mundo útiles, mientras que las políticas dirigidas a objetivos (que colapsan los datos en regiones de baja entropía) pueden romper la identificabilidad.
• Fundamento de Planificación: Al demostrar que la identificabilidad lineal permite una planificación óptima con costos invariantes a la rotación, el artículo proporciona una justificación teórica para el uso de espacios latentes aprendidos para el control sin necesidad de mapear de nuevo a las coordenadas físicas reales.

Los autores señalan limitaciones: la teoría asume que la dimensión del codificador coincide con la dimensión latente y se basa en la optimalidad a nivel de población. Sin embargo, demuestran que las garantías se degradan de manera suave con muestras finitas y errores de optimización, y verifican formalmente todos los teoremas principales utilizando el demostrador de teoremas Lean 4.