Learning the symmetric group: large from small

Resumen Técnico: Aprendizaje del Grupo Simétrico: Grande a partir de Pequeño

Enunciado del Problema
El artículo investiga si las redes neuronales basadas en Transformer pueden aprender la estructura del grupo simétrico $S_n$ entrenándose exclusivamente en subgrupos más pequeños $S_m$ (donde $m < n$) y, posteriormente, generalizar para predecir permutaciones en grupos más grandes sin codificar manualmente las relaciones del grupo. La tarea específica consiste en tomar una palabra $w \in S_n$ expresada como un producto de transposiciones y predecir la permutación resultante del conjunto $\{1, \dots, n\}$. Los autores abordan dos formulaciones distintas de este problema:

1. Transposiciones Generales: Palabras factorizadas utilizando transposiciones arbitrarias $s_{i,j}$.
2. Transposiciones Adyacentes: Palabras factorizadas utilizando únicamente transposiciones adyacentes elementales $s_i = s_{i,i+1}$.

Un desafío clave es manejar longitudes de palabra variables. Los autores señalan que las longitudes de palabras reducidas varían de $0$ (identidad) a $n-1$ (caso general) o $n(n-1)/2$ (caso adyacente). Para gestionar esto, los modelos se entrenan con palabras no reducidas de longitud fija, aumentadas con transposiciones de identidad.

Metodología
Los autores emplean una arquitectura Transformer autoregresiva (Figura 1) que predice tokens de permutación uno por uno. La entrada consiste en una "ventana de contexto" que contiene la secuencia de transposiciones y la permutación parcialmente predicha.

• Generación y Aumento de Datos:

  • Aumento de Identidad: Para manejar longitudes variables, las palabras se rellenan hasta una longitud máxima $N$ con transposiciones que actúan como identidades bajo las relaciones del grupo. El modelo debe aprender implícitamente estas relaciones para ignorar el relleno.
  • Entrenamiento en Subgrupos: El modelo se entrena con datos generados a partir de $S_m$ (específicamente $m=10$), pero se prueba en $S_n$ ($n=25$ para el caso general, $n=16$ para el caso adyacente).
  • Estrategia de Reetiquetado: Para generar datos de entrenamiento para $S_n$ que solo involucren $S_m$, los autores mapean palabras de $S_m$ a $S_n$ mediante una permutación aleatoria $\sigma \in S_n$.
    • Para Transposiciones Generales, se utiliza una reetiquetado aleatorio simple de índices $\{1, \dots, m\}$ a un subconjunto aleatorio de $\{1, \dots, n\}$.
    • Para Transposiciones Adyacentes, un método de "ventana" ingenuo (reetiquetando $s_1, \dots, s_{m-1}$ a un bloque contiguo $s_{k+1}, \dots$) falló porque el modelo aprendió a identificar la ubicación de la ventana en lugar de la estructura del grupo. Para contrarrestar esto, los autores introdujeron un método de Ventana Particionada. Esto implica descomponer $m$ en una composición (por ejemplo, $12 = 3+6+3$), crear múltiples ventanas no contiguas de longitudes variables dentro de $S_n$, y mapear los tokens de entrada a estas ventanas con desplazamientos aleatorios. Esto obliga al modelo a aprender las relaciones locales del grupo en lugar de pistas posicionales globales.

• Detalles de la Arquitectura:

  • Tokenización: Las transposiciones generales se mapean a enteros $x_k = i_k - 1 + n(j_k - 1)$. Las transposiciones adyacentes se mapean directamente a índices $x_k = i_k$.
  • Embeddings: El modelo utiliza embeddings de tokens aprendidos y embeddings posicionales. La entrada se procesa a través de capas estándar de Transformer (atención multi-cabeza, alimentación hacia adelante, normalización) antes de una capa lineal y probabilidades de salida Softmax.

Resultados Clave
Los experimentos demuestran que los modelos alcanzan una precisión cercana al 100% en conjuntos de prueba fuera de distribución (OOD):

• Transposiciones Generales: Un modelo entrenado en $S_{10}$ generalizó a $S_{25}$ con pérdida de entrenamiento cero y rendimiento de prueba casi perfecto (Figura 2).
• Transposiciones Adyacentes: Un modelo entrenado en $S_{10}$ generalizó a $S_{16}$ con un rendimiento alto similar (Figura 3). Los autores señalan que el caso adyacente es teóricamente más complejo debido a las longitudes de palabra más largas y las restricciones de localidad más estrictas, sin embargo, el método de ventana particionada permitió exitosamente la generalización.

Observaciones de Interpretabilidad
Los autores analizaron los mapas de calor de auto-similitud de los embeddings aprendidos:

• Transposiciones Generales: La matriz de embeddings de tokens reveló una estructura de red regular, indicando que el modelo aprendió relaciones de alto nivel entre generadores (por ejemplo, $s_{i,j}$ se relaciona fuertemente con $s_{i,k}$ y $s_{j,i}$). El embedding de posición no mostró un orden intrínseco entre las transposiciones generales.
• Transposiciones Adyacentes: El embedding de posición para transposiciones adyacentes exhibió una sorprendente subestructura de bloques dentro de las primeras $n(n-1)/2$ filas, aunque los autores notan que esto no fue consistente en todos los tamaños de grupo y su significado sigue siendo poco claro.

Significado y Afirmaciones
El artículo afirma demostrar que los Transformers pueden aprender la estructura algorítmica subyacente de la multiplicación de grupos a partir de conjuntos de entrenamiento pequeños y generalizar a grupos significativamente más grandes, incluso cuando la representación de entrada cambia (mediante reetiquetado y particionamiento).

• Escalabilidad: El trabajo propone un método escalable donde los modelos entrenados en estructuras más simples ($S_m$) pueden resolver tareas para grupos más grandes y complejos ($S_n$).
• Aprendizaje Fuera de Distribución: Esto sirve como un caso específico de aprendizaje OOD en matemáticas, mostrando que las características aprendidas en subgrupos pequeños pueden ser lo suficientemente universales para aplicarse a grupos más grandes.
• Direcciones Futuras: Los autores sugieren que este enfoque podría proporcionar información sobre el problema de la palabra para grupos finitamente generados y aplicaciones a problemas NP-difíciles en genómica comparativa. Identifican específicamente al grupo trenzado como un objetivo para trabajos futuros, motivado por su relevancia para el problema del nudillo, señalando que el grupo simétrico está "bien comportado" y que grupos más exóticos presentarán mayores desafíos.

Los autores permanecen modestos respecto a los mecanismos internos, afirmando que es una "cuestión abierta" si los modelos han descubierto un algoritmo general o han aprendido representaciones específicas, y planean realizar un análisis de interpretabilidad mecánica en trabajos futuros.