Resumen Técnico: Diferentes Capas, Diferentes Variedades

Planteamiento del Problema

La geometría del espacio de pesos es un factor crítico en la optimización de redes neuronales; sin embargo, los métodos actuales de entrenamiento con restricciones de variedad suelen aplicar una restricción geométrica única y uniforme (por ejemplo, ortogonalidad o restricciones de Stiefel) a todas las matrices de pesos en un Transformer. Este trabajo investiga si los diferentes módulos de un Transformer —específicamente las capas de atención y los bloques MLP (Perceptrón Multicapa)— prefieren geometrías de espacio de pesos distintas. Los autores postulan que los roles computacionales y los mecanismos no lineales de estos módulos difieren significamente, lo que sugiere que la asignación de una variedad uniforme puede ser subóptima o incluso desestabilizadora.

Metodología

El estudio utiliza Manifold Muon, un optimizador que aplica actualizaciones normalizadas por matriz a los pesos de las capas ocultas mientras los restringe a variedades matriciales específicas. Los autores comparan dos familias principales de restricciones:

1. Restricción de Stiefel: Fuerza la ortonormalidad de las columnas ($W^\top W = I$), controlando estrictamente los valores singulares y proporcionando un fuerte control espectral.
2. Restricción de DGram: Una restricción de espacio de Gram más débil que requiere que las entradas fuera de la diagonal de la matriz de Gram sean cero ($\text{Off}(W^\top W) = 0$). Esto preserva la ortogonalidad de las columnas pero permite que las normas de las columnas (y, por tanto, los valores singulares) varíen, ofreciendo grados de libertad de escala adicionales.

Los autores evalúan cinco estrategias de asignación distintas en un modelo GPT-2 small (124M de parámetros) preentrenado en OpenWebText:

• UNCONSTRAINED (Sin Restricciones): Sin restricciones de variedad.
• ALL-STIEFEL: Restricciones de Stiefel en los pesos de Atención y MLP.
• ALL-DGRAM: Restricciones de DGram en los pesos de Atención y MLP.
• HETERO: Restricciones de Stiefel en la Atención, DGram en el MLP.
• HETERO-INV: Restricciones de DGram en la Atención, Stiefel en el MLP.

Todas las configuraciones comparten las mismas arquitecturas de modelo, cronogramas de entrenamiento e hiperparámetros del optimizador para aislar el efecto de la asignación de la variedad.

Resultados Clave

Los experimentos revelan una clara asimetría en cómo los diferentes módulos responden a las restricciones de variedad:

• Configuración Óptima: La asignación HETERO (Stiefel para la Atención, DGram para el MLP) logró la pérdida de validación estable más baja (3.3544), superando tanto al baseline sin restricciones como a la configuración uniforme ALL-STIEFEL.
• Inestabilidad en DGram-Atención: Las configuraciones que asignan restricciones de DGram a las capas de atención (ALL-DGRAM y HETERO-INV) se volvieron inestables bajo la configuración de hiperparámetros compartidos. Estas ejecuciones exhibieron un crecimiento explosivo en la pérdida de validación y no lograron recuperarse.
• Estabilidad en DGram-MLP: Por el contrario, asignar restricciones de DGram a las capas MLP (cuando se combinan con atención de Stiefel) se mantuvo estable y mejoró el rendimiento en comparación con la restricción uniforme de Stiefel.

Análisis Mecanicista

El artículo proporciona una explicación mecanicista para la inestabilidad observada en la atención con restricción de DGram:

1. Crecimiento de Valores Singulares: A diferencia de las restricciones de Stiefel, la restricción de DGram permite que los valores singulares de las matrices de pesos crezcan arbitrariamente grandes. En las capas de atención, las proyecciones de consulta ($W_Q$) y clave ($W_K$) se multiplican para formar los logits de atención ($Z = X W_Q W_K^\top X^\top / \sqrt{d}$).
2. Amplificación de Logits: El crecimiento incontrolado en las normas espectrales de $W_Q$ y $W_K$ amplifica directamente la magnitud de los logits de atención.
3. Saturación de Softmax: Las magnitudes de los logits grandes empujan la función softmax hacia un régimen de saturación, donde la salida se aproxima a un vector de un solo elemento (one-hot). Esto causa que el Jacobiano de la softmax se desvanezca exponencialmente, provocando la degeneración del gradiente y la inestabilidad del entrenamiento.
4. Tolerancia de MLP: Las capas MLP toleran las restricciones de DGram porque sus funciones de activación (por ejemplo, GELU) operan elemento a elemento y no acoplan los puntajes de los tokens a través de un simplex de probabilidad normalizado. En consecuencia, la amplificación de escala en los pesos de MLP no fuerza a todo el módulo a un régimen de enrutamiento saturado.

Contribuciones

Los autores identifican tres contribaciones principales:

1. Evaluación Empírica: Evalúan las asignaciones de variedad módulo a módulo para Manifold Muon en el preentrenamiento de GPT-2, demostrando que las restricciones uniformes son insuficientes.
2. Descubrimiento de Asimetría: Muestran que las capas de atención y MLP prefieren geometrías de espacio de pesos diferentes, donde la atención requiere un control espectral fuerte (Stiefel) y el MLP se beneficia de una libertad de escala más débil (DGram).
3. Mecanismo de Inestabilidad: Identifican el crecimiento de los valores singulares y la subsiguiente saturación de la softmax como un mecanismo plausible para la inestabilidad de la atención con restricción de DGram.

Significado y Limitaciones

El artículo argumenta que la optimización consciente de la simetría y la geometría para los Transformers debe ser específica para cada módulo en lugar de uniforme. Esto sugiere que la geometría del espacio de pesos es un eje de diseño práctico para mejorar la estabilidad y la dinámica del entrenamiento en las redes modernas.

Los autores mantienen una postura modesta respecto a sus afirmaciones:

• No afirman que la atención de DGram sea inherentemente inestable bajo todas las condiciones, sino que expone una vía inestable bajo la configuración de hiperparámetros compartidos probada sin un control de escala adicional.
• Los experimentos se limitan a GPT-2 small y a un conjunto específico de semillas; la brecha de pérdida de validación entre la mejor configuración heterogénea y la uniforme es modesta.
• El análisis mecanicista sirve como una explicación para la inestabilidad observada más que como una prueba completa de que la atención de DGram no pueda estabilizarse (por ejemplo, mediante el reajuste de tasas de aprendizaje o regulación de escala explícita).

Se sugiere como trabajo futuro el desarrollo de métodos automáticos para asignar restricciones de variedad a través de diferentes módulos, arquitecturas y escalas.