Resumen Técnico: GeoNorm: Unificar Pre-Norm y Post-Norm mediante Optimización Geodésica

Declaración del Problema

La ubicación de las capas de normalización dentro de las arquitecturas Transformer —específicamente la elección entre Pre-Norm y Post-Norm— sigue siendo una decisión de diseño crítica que influye en la estabilidad del entrenamiento, el rendimiento del modelo y el flujo de gradientes. Si bien Pre-Norm se ha convertido en la opción predominante en los modelos de lenguaje de gran tamaño (LLMs) modernos debido a su estabilidad, introduce gradientes desequilibrados donde las capas inferiores reciben actualizaciones más grandes que las superiores. Por el contrario, Post-Norm, diseñado originalmente para facilitar el entrenamiento de redes profundas, es propenso a la inestabilidad en el entrenamiento y a picos de pérdida. Las soluciones existentes como DeepNorm y SandwichNorm mejoran la estabilidad mediante la inicialización de parámetros o capas de normalización adicionales, pero carecen de un marco teórico unificado para analizar estos mecanismos.

Metodología

Los autores proponen GeoNorm, un nuevo método de normalización basado en la optimización de variedades (manifolds). La idea teórica central es reinterpretar las operaciones de actualización estándar de Transformer ($x + \text{FFN}(x)$ y $x + \text{Attention}(x)$) como pasos de optimización en una esfera. En esta visión:

1. Los embeddings de los tokens actúan como el estado inicial en una variedad.
2. Los módulos de Atención y FFN generan direcciones de actualización (gradientes estimados).
3. La normalización se ve como una proyección de vuelta a un conjunto de restricciones (la esfera).

La normalización estándar (Post-Norm) realiza una actualización no restringida seguida de una proyección euclidiana de vuelta a la esfera. Los autores argumentan que esta proyección descarta información geométrica y puede distorsionar las direcciones de actualización. En su lugar, GeoNorm reemplaza el paso de proyección con una actualización geodésica utilizando el mapa exponencial.

Para un punto $x$ en una esfera y un vector tangente $v$ (la dirección de actualización proyectada), la actualización se realiza intrínsecamente a lo largo de la geodésica:
$$\text{exp}_x(v) = \cos\left(\frac{\|v\|}{\|x\|}\right)x + \|x\|\sin\left(\frac{\|v\|}{\|x\|}\right)\frac{v}{\|v\|}$$

Esta formulación asegura que la representación se mueva suavemente a lo largo de la variedad, preservando la geometría intrínseca del espacio de características. Además, los autores introducen un mecanismo de decaimiento de actualización por capa para el tamaño del paso ($\alpha_k$), análogo a los programas de tasa de aprendizaje (learning rate schedules). Evalúan tres estrategias de decaimiento:

• Armónico: $\alpha_k = \alpha / k$
• Raíz Cuadrada: $\alpha_k = \alpha / \sqrt{k}$
• Lineal: $\alpha_k = \alpha (T - k) / T$

El artículo también demuestra teóricamente que Pre-Norm puede interpretarse como un caso especial de GeoNorm con ángulos dependientes de la profundidad y no adaptativos, unificando así ambos enfoques bajo un único marco geomético.

Contribuciones Clave

1. Marco Teórico: Los autores proporcionan una perspectiva teórica unificada que interpreta las capas de Transformer como pasos de optimización iterativos en una variedad esférica, cerrando la brecha entre Pre-Norm y Post-Norm.
2. Algoritmo GeoNorm: Introducen GeoNorm, que reemplaza la normalización basada en proyección estándar con actualizaciones geodésicas mediante el mapa exponencial. Este método no requiere parámetros entrenables adicionales y añade una sobrecarga computacional insignificante.
3. Decaimiento por Capa: Una propuesta novedosa de decaimiento de actualización por capa para los componentes de atención y FFN para estabilizar la dinámica del entrenamiento.
4. Validación Empírica: Experimentos extensos a través de varios tamaños de modelo (de 125M a 1.3B), conjuntos de datos (Arxiv, Books3, FinWeb-Edu) y longitudes de entrenamiento demuestran la eficacia del método.

Resultados Experimentales

El artículo reporta que GeoNorm supera consistentemente a los métodos de normalización existentes (Pre-Norm, Post-Norm, DeepNorm, SandwichNorm) en todas las configuraciones probadas:

• Rendimiento: En los conjuntos de datos Arxiv y Books3 con un modelo de 125M, GeoNorm logró la pérdida de validación más baja (1.8792 en Arxiv; 3.4040 en Books3) en comparación con todos los baselines.
• Escalabilidad: La ventaja de rendimiento persiste y se amplía a medida que aumenta el tamaño del modelo. Para un modelo de 1.3B, GeoNorm logró una pérdida de 3.1759, superando al mejor baseline (Pre-Norm con 3.2067).
• Estabilidad: GeoNorm demostró una estabilidad de entrenamiento superior, particularmente en ejecuciones de entrenamiento largas (hasta 4096 de longitud de secuencia). Mientras que Post-Norm exhibió picos de pérdida y colapso del entrenamiento en longitudes mayores, GeoNorm mantuvo una convergencia estable.
• Mecanismo de Decaimiento: Entre las estrategias de decaimiento, el decaimiento armónico ($\alpha/k$) produjo el mejor rendimiento, estableciendo nuevos récords en los conjuntos de datos probados.

Significado y Reivindicaciones

El artículo afirma que GeoNorm ofrece una integración fluida en las arquitecturas Transformer estándar, ya que no requiere modificaciones en los bloques centrales (Atención o FFN) y no introduce parámetros entrenables adicionales. Los autores sostienen que, al aprovechar la optimización geodésica, GeoNorm permite que las representaciones se muevan de manera más coherente a lo largo de la variedad esférica, lo que conduce a una mejor propagación de la información y estabilidad en el entrenamiento.

Este trabajo se posiciona como un paso fundamental para comprender los mecanismos de normalización a través de la lente de la geometría de Riemann. Sugiere que el éxito empírico de Pre-Norm puede explicarse por su decaimiento implícito del tamaño del paso, un mecanismo que GeoNorm hace explícito y adaptativo. Los autores concluyen que esta visión unificada proporciona una base robusta para futuras innovaciones en el diseño de arquitecturas Transformer, particularmente para escalar modelos hacia mayores profundidades y complejidades.