BRo-JEPA: Learning Modular Arithmetic in Latent Space

Resumen Técnico: BRo-JEPA – Aprendizaje de Aritmética Modular en el Espacio Latente

Planteamiento del Problema
El artículo investiga una pregunta fundamental en el aprendizaje profundo: ¿aprenden las redes neuronales realmente reglas algebraicas abstractas o simplemente memorizan patrones de entrenamiento? Mientras que los modelos actuales sobresalen en tareas matemáticas dentro de rangos vistos, a menudo fallan al extrapolar a rangos numéricos u operaciones no vistas. Los autores postulan que los modelos supervisados estándar y los modelos JEPA (Arquitectura de Predicción de Incrustación Conjunta) existentes dependen de correlaciones estadísticas en lugar de una lógica aritmética subyacente. Específicamente, al realizar tareas de aritmética modular (por e.j., $y = (x + k) \mod 10$) utilizando dígitos de MNIST como estados, los modelos estándar se ajustan a las operaciones vistas pero fallan al generalizar a operaciones no vistas (generalización estricta zero-shot).

Metodología
Los autores utilizan el marco de trabajo JEPA, donde un modelo predice la representación latente de un dígito objetivo dado un dígito de contexto (imagen) y una operación aritmética (acción), en lugar de predecir salidas de píxeles brutos. El estudio compara varias variaciones arquitectónicas:

1. Líneas Base Estándar: Modelos supervisados (MLP y ResNet-18) que utilizan pérdida de entropía cruzada y embeddings de acción composicionales aditivos. Estos embeddings representan las operaciones como múltiplos escalares de un vector base (e.j., $op(k) = k \cdot v$) para imponer una estructura aritmética en el espacio de incrustación.
2. JEPA con Embeddings Aditivos: Un modelo JEPA que utiliza los mismos embeddings aditivos, aumentado con una "Pérdida de Consistencia de Composición". Esta pérdida auxiliar impone que aplicar una operación compuesta (e.j., $+2$) en un solo paso debe producir el mismo resultado latente que aplicar operaciones primitivas secuencialmente (e.j., $+1$ y luego $+1$).
3. BRo-JEPA (JEPA de Rotación de Bloques): La contribución central es un predictor novedoso que reemplaza el MLP estándar no restringido con un predictor de rotación de bloques condicionado por la operación.
   • Mecanismo: Inspirado en los Rotary Position Embeddings (RoPE), este predictor trata la dinámica del espacio latente como una serie de rotaciones planas en 2D. Para un vector latente $z$ y una operación $k$, la predicción es $z_p = R_{block}(k \cdot \theta) \cdot z$.
   • Sesgo Inductivo: Esta arquitectura impone explícitamente la estructura circular de la aritmética módulo-10 sobre el espacio latente. Los ángulos de rotación $\theta$ pueden ser fijos (establecidos en $2\pi/10$) o aprendibles.
   • Variantes: Los autores prueban la Rotación de Frecuencia Única (SFR), donde todos los bloques latentes comparten un ángulo, y la Rotación de Frecuencia Múltiple (MFR), donde cada bloque tiene un ángulo independiente.
   • Regularización: Para evitar el "colapso latente" (donde diferentes entradas mapean a representaciones indistinguibles), el modelo emplea la Regularización de Varianza-Invarianza-Covarianza (VICReg) y un codificador de objetivo de Promedio Móvil Exponencial (EMA).

Contribuciones Clave

1. Definición del Banco de Pruebas: La introducción de la aritmética modular de MNIST como un banco de pruebas controlado para estudiar el aprendizaje de reglas en modelos de mundo latentes.
2. Análisis de Fallos: Demostración de que las líneas base supervisadas y los modelos JEPA con embeddings aditivos, a pesar de ajustar las operaciones vistas, fallan en la extrapolación zero-shot estricta (capturan los signos de las operaciones pero fallan en codificar la magnitud o la geometría composicional).
3. Propuesta de BRo-JEPA: El desarrollo de un modelo de dinámica latente de estilo JEPA que presenta un predictor de rotación de bloques condicionado por la operación que se alinea con la estructura cíclica de la aritmética modular.
4. Análisis Geométrico: Análisis que muestra cómo BRo-JEPA organiza las representaciones latentes en un manifold geométrico estructurado (un manifold circular), permitiendo al modelo interiorizar la topología de la aritmética modular.

Resultos
Los experimentos se realizaron entrenando con las operaciones $\{-1, +1\}$ y evaluando en operaciones no vistas $\{\pm 2, \dots, \pm 9\}$.

• Rendimiento de la Línea Base: Las líneas base supervisadas lograron una alta precisión en las operaciones vistas (~97-98%) pero una precisión cercana a cero en las operaciones no vistas (2.42% para ResNet-18). Las visualizaciones de PCA revelaron que estos modelos comprimen arbitrariamente las operaciones no vistas en el espacio de representación de las vistas.
• JEPA con Embeddings Aditivos: Aunque estos modelos mostraron una alta precisión de rollout (aplicando $+1$ repetidamente para simular $+k$), fallaron en la predicción directa zero-shot de operaciones no vistas (precisión ~10%). La pérdida de consistencia de composición mejoró esto a ~99%, pero se considera un "zero-shot débil" porque el modelo fue expuesto indirectamente a operaciones compuestas durante el entrenamiento.
• Rendimiento de BRo-JEPA:
  • Precisión Zero-Shot: El BRo-JEPA basado en ResNet con Rotación de Frecuencia Múltiple (MFR) logró un 99.46% de precisión zero-shot en operaciones no vistas.
  • Precisión de Rollout: El modelo también logró un 99.46% de precisión en rollouts de múltiples pasos, confirmando la estabilidad de las dinámicas latentes aprendidas.
  • Estructura Latente: Las visualizaciones de PCA confirmaron que BRo-PEP logra coaccionar el espacio latente hacia un manifold elíptico/circular continuo donde las clases de dígitos se organizan secuencialmente. La variante MFR produjo prototipos de clase mejor separados (menor similitud de coseno fuera de la diagonal) en comparación con SFR, lo que condujo a una generalización superior.
  • Ángulos Aprendibles: Cuando los ángulos de rotación eran aprendibles, MFR convergió rápidamente a una generalización efectiva, mientras que SFR exhibió retrasos significativos e inestabilidad, convergiendo a menudo a soluciones armónicas subóptimas.

Significancia y Reivindicaciones
El artículo afirma que los modelos de mundo latentes pueden aprender reglas de transformación simbólica, pero solo cuando el sesgo inductivo de la arquitectura coincide con la estructura del problema. Los resultados sugieren que:

• Los predictores neuronales estándar carecen de las restricciones necesarias para preservar la geometría composicional para operaciones no vistas.
• Modelar explícitamente la naturaleza cíclica de la aritmética modular mediante rotaciones de bloque en el espacio latente permite una fuerte generalización zero-shot sin requerir la exposición a las operaciones específicas durante el entrenamiento.
• Este enfoque ofrece una dirección prometedora para el aprendizaje de representaciones más allá de la aritmética, sugiriendo que incorporar estructuras de transformación apropiadas en los modelos de transición latente es crucial para el aprendizaje de reglas abstractas.

Los autores mantienen la modestia, señalando que sus resultados son específicos para el entorno altamente controlado de la aritmética modular de MNIST con una estructura de grupo cíclico conocida. No afirman que la arquitectura generalice automáticamente a tareas de razonamiento simbólico arbitrarias o grupos de transformación desconocidos sin descriptores estructurales similares.