Resumen Técnico: MapFormers

Enunciado del Problema

Los sistemas de IA basados en Transformadores actuales, a pesar de su impresionante rendimiento en tareas dentro de la distribución, muestran una fragilidad significativa cuando se enfrentan a escenarios fuera de la distribución (OOD). A menudo fallan al generalizar a situaciones novedosas que los humanos y los animales manejan con facilidad. Los autores atribuyen esta brecha a la falta de sesgos estructurales que permiten a los cerebros biológicos construir mapas cognitivos: modelos internos que desentienden la estructura abstracta e invariante de un entorno de su contenido específico y variable.

Los intentos existentes de modelar mapas cognitivos, como la Máquina Tolman-Eichenbaum (TEM) que utiliza matrices dependientes de la acción en redes neuronales recurrentes (RNN), adolecen de dos limitaciones críticas:

1. Falta de Escalabilidad: Requieren multiplicación de matrices secuencial para actualizar posiciones, lo que impide el procesamiento paralelo.
2. Acciones Preespecificadas: A menudo requieren que el conjunto de acciones (por ejemplo, direcciones cardinales) esté codificado de forma rígida en lugar de aprenderse a partir de los datos.

Además, los codificamientos posicionales estándar (como RoPE) son estáticos o dependen de heurísticas que no soportan la composición dinámica de acciones requerida para la integración de trayectorias en entornos desconocidos.

Metodología

El artículo introduce MapFormers, una clase de arquitecturas Transformador diseñadas para aprender mapas cognitivos y realizar integración de trayectorias mediante la predicción del siguiente token sin supervisión. La innovación central reside en el uso de matrices dependientes de la entrada para actualizar los codificamientos posicionales, derivadas de la teoría de grupos de Lie.

Marco Teórico

• Formalismo de Grupo de Lie: Los autores modelan las acciones como elementos de un grupo de Lie (específicamente grupos compactos como $SO(2)$ para rotaciones). En lugar de aprender una matriz separada para cada acción posible, el modelo aprende generadores del álgebra de Lie.
• Integración de Trayectorias mediante Aplicaciones Exponenciales: Las acciones se representan como generadores infinitesimales en el álgebra de Lie. La integración de trayectorias (componer una secuencia de acciones) se realiza sumando estos generadores en el álgebra (acumulación lineal) y luego aplicando una única aplicación exponencial para obtener la transformación finita en el grupo de Lie.
  • Para mapas conmutativos (por ejemplo, navegación en una cuadrícula 2D), esto permite calcular la integración de trayectorias en paralelo mediante suma acumulativa, seguida de una única exponencial de matriz.
  • Este enfoque reduce el sesgo inductivo desde el número de acciones hasta la complejidad estructural del álgebra.

Variantes Arquitectónicas

El artículo propone dos variantes de MapFormers, unificando la codificación posicional absoluta y relativa con modelos de memoria distintos:

1. MapEM (Memoria Episódica):

   • Utiliza incrustaciones posicionales absolutas para representar un mapa cognitivo.
   • Desentiende el contenido (observaciones) y la estructura (posición) en poblaciones neuronales separadas.
   • Calcula la atención sobre la conjunción de contenido y posición ($Q_G = \text{vec}(Q^T \cdot P)$), actuando efectivamente como una máscara de atención donde los codificamientos posicionales se actualizan independientemente del contenido.
   • Adecuado para tareas que requieren un recuerdo robusto de estados específicos.

2. MapWM (Memoria de Trabajo):

   • Utiliza codificación posicional relativa (análogo a RoPE pero dinámico).
   • Representa la estructura implícitamente rotando la actividad neuronal (claves y consultas) hacia subespacios distintos mediante matrices de rotación dependientes de la entrada ($R_{\theta PI}$).
   • El ángulo de rotación $\theta$ se deriva de una proyección de bajo rango del token de entrada, permitiendo que el modelo aprenda qué tokens actúan como acciones (actualizando la posición) y cuáles son observaciones (dejando la posición sin cambios).

Mecanismo Dependiente de la Entrada

Ambas arquitecturas emplean una proyección de dos etapas para determinar la "duración de integración" ($\Delta_t$) para cada token:

1. Proyección Interna: Mapea el token de entrada a una representación de acción de baja dimensión.
2. Proyección Externa: Proyecta esta representación a una duración de integración por bloque.
   El ángulo de rotación final es $\theta_t = \omega \odot \Delta_t$, donde $\omega$ es un vector de velocidad angular aprendible. La integración de trayectorias se calcula como la suma acumulativa de estos ángulos ($\theta_{0 \to t} = \text{cumsum}(\theta_t)$).

Contribuciones Clave

1. Arquitectura MapFormers: Una variante novedosa de Transformador que aprende mapas cognitivos desentendiendo la estructura del contenido mediante codificamientos posicionales dependientes de la entrada basados en generadores del álgebra de Lie.
2. Integración de Trayectorias Paralela: Al aprovechar grupos de Lie conmutativos (por ejemplo, $SO(2)^k$), los modelos logran integración de trayectorias en paralelo, superando el cuello de botella secuencial de los modelos anteriores de mapas cognitivos basados en RNN.
3. Unificación de Modelos de Memoria: El trabajo proporciona un vínculo teórico entre la codificación posicional absoluta/relativa y la memoria episódica/de trabajo, mostrando cómo ambas pueden derivarse de los mismos principios teóricos de Lie.
4. Escalabilidad a Mapas No Conmutativos: El marco se extiende a problemas no conmutativos (por ejemplo, árboles genealógicos) utilizando múltiples generadores, aunque esto requiere integración de trayectorias secuencial.

Resultados Experimentales

Los autores evaluaron MapFormers en tres tareas formales dirigidas a capacidades cognitivas distintas:

1. Copia Selectiva (Puerta Lógica):

   • Tarea: Copiar una secuencia ignorando tokens distractores.
   • Resultado: MapFormers (tanto EM como WM) lograron una generalización OOD perfecta, mientras que las líneas base estándar (RoPE, TAPE) fallaron. Esto demuestra la capacidad de aprender mecanismos de puerta dinámica.

2. Navegación en Cuadrícula (1D y 2D):

   • Tarea: Predecir observaciones en una secuencia de acciones alternas y observaciones aleatorias, requiriendo que el modelo infiera la topología latente e integre la posición.
   • Resultado: MapFormers alcanzaron una precisión casi perfecta en secuencias OOD (variando longitud y densidad), mientras que las líneas base (incluyendo CoPE y PathAtt) fallaron al generalizar.
   • Análisis Neuronal: Los modelos aprendieron a proyectar tokens de observación a rotación cero (identidad) y tokens de acción a direcciones de rotación específicas, desentendiendo efectivamente el contenido de la estructura.

3. Dyck-2 (Jerarquías Anidadas):

   • Tarea: Predecir continuaciones válidas de secuencias de paréntesis, requiriendo el seguimiento de una pila.
   • Resultado: MapFormers generalizaron perfectamente a secuencias más largas y profundas que las vistas en el entrenamiento utilizando solo una capa. Las líneas base requirieron múltiples capas y aún fallaron en datos OOD.
   • Mecanismo: Los modelos incrustaron la pila directamente en el codificamiento posicional utilizando dimensiones ortogonales para diferentes tipos de paréntesis y direcciones opuestas para paréntesis de apertura y cierre.

4. Escalado en Lenguaje Natural:

   • Tarea: Preentrenamiento en OpenWebText.
   • Resultado: Un modelo MapWM de 12 capas mostró mejoras consistentes en la perplejidad sobre una línea base RoPE ($18.79$ vs $19.14$) y una mejor extrapolación de longitud. Sin embargo, en el benchmark sintáctico BLIMP, MapWM no mostró ganancias significativas sobre RoPE, lo que sugiere que la integración de trayectorias conmutativa por sí sola puede no ser suficiente para todo el modelado sintáctico.

Significado y Afirmaciones

El artículo afirma que MapFormers proporcionan un sesgo estructural crítico que permite la generalización sistemática OOD capturando estructuras latentes mediante integración de trayectorias en lugar de depender de heurísticas.

• Robustez: Los modelos demuestran que desentender relaciones abstractas del contenido permite una generalización robusta a longitudes de secuencia, profundidades y densidades ambientales no vistas.
• Escalabilidad: El uso de computación paralela en mapas conmutativos permite que estos principios escalen a dominios reales a gran escala, como lo evidencian los resultados en lenguaje natural.
• Plausibilidad Biológica: El enfoque cierra la brecha entre los mapas cognitivos biológicos (que factorizan el "qué" del "dónde") y el aprendizaje profundo moderno, ofreciendo una vía hacia sistemas de IA que pueden razonar a través de problemas novedosos comprendiendo la geometría subyacente de su entorno.

Los autores se mantienen modestos respecto a las limitaciones, señalando que la integración de trayectorias no conmutativa no es paralelizable y que la restricción conmutativa actual puede limitar el modelado de ciertas estructuras sintácticas complejas en el lenguaje. Sugieren que el trabajo futuro podría explorar arquitecturas híbridas o extensiones no conmutativas para un razonamiento relacional más complejo.