Resumen Técnico: Estabilidad de los Transformers bajo la Normalización de Capas

Planteamiento del Problema

A pesar de su dominio en el aprendizaje profundo moderno, el entrenamiento de Transformers profundos sigue siendo propenso a la inestabilidad, particularmente a medida que aumenta la profundidad del modelo. Estas inestabilidades se manifiestan en dos áreas críticas: inestabilidad hacia adelante (crecimiento incontrolado de los estados ocultos que genera problemas numéricos) e inestabilidad hacia atrás (explosión de gradientes durante la retropropagación).

Un componente arquitectónico principal utilizado para mitigar estos problemas es la Normalización de Capas (LN). Sin embargo, la ubicación de la LN dentro de los bloques Transformer ha evolucionado principalmente mediante heurísticas empíricas en lugar de una justificación teórica rigurosa. El artículo identifica tres configuraciones principales:

• Post-LN: LN aplicada después de las conexiones residuales. A menudo requiere un cronograma de optimización delicado y exhibe un rendimiento subóptimo.
• Pre-LN: LN aplicada antes de los módulos de atención y de alimentación hacia adelante (feedforward). Aunque es el estándar por su mejora en la estabilidad de entrenamiento sobre Post-LN, se observa empíricamente que produce estados ocultos excesivamente grandes, lo que conduce a una potencial inestabilidad numérica y explosión de gradientes en redes profundas.
• Peri-LN: Una configuración más reciente que aplica LN tanto a la entrada como a la salida de los módulos. Los estudios empíricos sugieren que ofrece una estabilidad superior y estados ocultos más regulares en comparación con Pre-LN, pero sus propiedades teóricas aún no se comprenden bien.

El problema central abordado es la falta de un marco teórico principista para explicar por qué estas diferentes ubicaciones conducen a distintas dinámicas de entrenamiento y perfiles de estabilidad, y si el entrenamiento impulsa a los modelos hacia soluciones regulares o comportamientos patológicos.

Metodología

Los autores proponen un nuevo marco teórico basado en la teoría de control óptimo de tiempo discreto para analizar la dinámica de entrenamiento de los Transformers.

1. Formulación como Control Óptimo: El problema de entrenamiento del Transformer se plantea como un problema de control de campo medio en tiempo discreto. El índice de la capa $i$ representa el tiempo discreto, y el estado oculto $X_i$ representa el estado del sistema. Las conexiones residuales actúan como la velocidad de la dinámica.
2. Análisis de Hamilton-Jacobi-Bellman (HJB): Los autores analizan las condiciones de optimalidad del problema de entrenamiento utilizando la ecuación diferencial parcial (PDE) de HJB acoplada con ecuaciones de transporte de masa.
   • Examinan el Hamiltoniano, que actúa como un mecanismo de control de retroalimentación que define el campo de velocidad óptimo para los estados ocultos.
   • Para Pre-LN, demuestran que el Hamiltoniano es ilimitado, lo que conduce a condiciones de optimalidad mal planteadas donde los estados ocultos pueden crecer arbitrariamente grandes.
   • Para Peri-LN, la normalización de salida restringe el campo de velocidad a un elipsoide (basado en el Lema 1), asegurando que el Hamiltoniano sea acotado y el problema esté bien planteado.
3. Análisis de Gradiente: El artículo analiza la estabilidad hacia atrás examinando la propagación de gradientes a través del producto de sensibilidades locales (matrices Jacobianas) a través de las capas.
4. Escalamiento de Residuos: El marco se extiende para incluir un parámetro de incremento de tiempo $\Delta t < 1$, escalando efectivamente la magnitud de los pasos residuales para controlar la propagación de la señal.

Contribuciones Clave

1. Marco Teórico para Modelos Entrenados

A diferencia de trabajos previos que a menudo se centran en la inicialización del modelo o dependen de evidencia empírica, este marco analiza modelos entrenados en la convergencia. Proporciona una evaluación sistemática de si las dinámicas de entrenamiento conducen a soluciones regulares o patologías.

2. Análisis Riguroso de las Ubicaciones de LN

• Inestabilidad de Pre-LN: El artículo demuestra que, para Pre-LN, la solución óptima al problema de entrenamiento es ilimitada en magnitud. Incluso con el decaimiento de pesos (weight decay), los estados ocultos pueden crecer exponencialmente ($O(e^D)$), y los gradientes crecen proporcionalmente con las activaciones, lo que provoca inestabilidad en el entrenamiento.
• Estabilidad de Peri-LN: Los autores derivan cotas explícitas que muestran que Peri-LN garantiza un crecimiento lineal para el valor absoluto medio de los estados ocultos ($O(D)$) y un crecimiento cuadrático para la varianza ($O(D^2)$). Crucialmente, se demuestra que la sensibilidad local de Peri-LN es invariante a la magnitud de la activación, evitando la explosión de gradientes incluso cuando las activaciones son grandes.

3. Cuantificación de la Incertidumbre

El análisis establece que Peri-LN proporciona una propagación controlada de la incertidumbre (medida mediante la distancia de Wasserstein) desde la entrada hasta los estados ocultos terminales, mientras que Pre-LN puede amplificar la incertidumbre de la entrada a niveles ilimitados.

4. Escalamiento del Paso Residual

Guiados por el análisis de estabilidad, los autores introducen un mecanismo de escalamiento del paso residual ($\Delta t < 1$). Demuestran teóricamente que este escalamiento:

• Agudiza las cotas sobre el crecimiento de los estados ocultos y la incertidumbre de salida.
• Mitiga la explosión de gradientes escalando la sensibilidad local.
• Mejora tanto la estabilidad hacia adelante como la hacia atrás sin costo computacional adicional.

Resultados Experimentales

Los autores validan sus hallazgos teóricos utilizando variantes de GPT-2 (de 124M a 1.5B de parámetros) entrenadas en el conjunto de datos OpenWebText.

• Estabilidad de Entrenamiento: En experimentos con 5 ensayos, los modelos Pre-LN divergieron en 1 ejecución cuando se habilitó el decaimiento de pesos y en 3 ejecuciones cuando se deshabilitó. En contraste, los modelos Peri-LN se mantuvieron estables en las 5 ejecuciones bajo ambas configuraciones de decaimiento de pesos.
• Rendimiento: Peri-LN logró un rendimiento comparable al de Pre-LN incluso cuando los hiperparámetros se ajustaron para Pre-LN.
• Impacto del Escalamiento: Combinar Peri-LN con el escalamiento de residuo ($\Delta t = 0.1$) produjo consistentemente el mejor rendimiento en todas las métricas (Pérdida de Validación, Perplejidad, ROUGE, puntuaciones BERT) y controló eficazmente el crecimiento de los momentos de los estados ocultos (valor absoluto medio y varianza) a través de las capas, alineándose con las predicciones teóricas.

Significancia y Reivindicaciones

El artículo afirma proporcionar la primera justificación teórica rigurosa de las observaciones empíricas respecto a las ubicaciones de la Normalización de Capas en los Transformers.

• Explicación de la Inestabilidad: Explica por qué Pre-LN conduce a estados ocultos ilimitados y a la explosión de gradientes mediante la falta de idoneidad de la PDE de HJB.
• Justificación de Peri-LN: Proporciona una base teórica para la estabilidad de Peri-LN, mostrando que mantiene un crecimiento controlado y una sensibilidad invariante.
• Guía de Diseño: El marco ofrece un método principista para realizar un "control de calidad" de nuevas modificaciones arquitectónicas antes de un entrenamiento empírico costoso.
• Mejora Práctica: La introducción del escalamiento del paso residual se presenta como una modificación simple y sin costo que mejora significativamente la estabilidad y el rendimiento, validada por experimentos que muestran que Peri-LN con escalamiento supera a las configuraciones estándar.

Los autores enfatizan que su trabajo cierra la brecha entre las heurísticas empíricas y la comprensión teórica, ofreciendo una perspectiva de tiempo continuo que genera tanto conocimientos teóricos como mejoras arquitectónicas prácticas.