Resumen Técnico: BHyT – Una alternativa estable y eficiente a la Normalización Pre-Capa (Pre-Layer Normalization)

1. Declaración del Problema

Los Modelos de Lenguaje Extensos (LLMs) dependen fuertemente de la Normalización Pre-Capa (Pre-LN), implementada típicamente mediante RMSNorm, para estabilizar el entrenamiento y permitir un aprendizaje por transferencia efectivo. Sin embargo, a medida que la profundidad del modelo aumenta, el diseño estándar de Pre-LN sufre dos limitaciones críticas:

1. La Maldición de la Profundidad: Análisis recientes indican que la interacción entre las conexiones residuales y la Pre-LN impulsa un crecimiento rápido en la magnitud y la varianza de los estados ocultos a través de las capas. Esta escalada empuja el Jacobiano del bloque hacia la matriz identidad, causando que las capas más profundas se comporten como mapeos de identidad costosos y desestabilizando el entrenamiento.
2. Compromiso entre Eficiencia y Estabilidad: Las soluciones existentes a menudo abordan solo un aspecto. Los métodos orientados a la estabilidad como Peri-LN (normalizando antes y después de las subcapas) o el Escalamiento de CapaNorm (LNS) suprimen el crecimiento de la varianza, pero incurren en una sobrecarga computacional significativa debido a los cálculos estadísticos repetidos. Por el contrario, los métodos orientados a la eficiencia, como la Tanh Dinámica (DyT), eliminan la sobrecarga estadística pero no logran controlar explícitamente el crecimiento residual por profundidad, lo que conduce a una potencial saturación e inestabilidad en redes profundas.

El desafío central es diseñar un mecanismo que combine la baja sobrecarga de las transformaciones acotadas con un control explícito sobre la escala y la varianza del flujo residual para prevenir la deriva por profundidad sin sacrificar el rendimiento (throughput) del entrenamiento.

2. Metodología: Bounded Hyperbolic Tanh (BHyT)

Los autores proponen BHyT, un reemplazo directo para Pre-LN que acopla una no linealidad de tangente hiperbólica con un acotamiento de entrada dependiente de los datos.

Mecanismo Central

A diferencia de DyT, que depende de un escalar global aprendible para reescalar las entradas, BHyT acota explícitamente la entrada de la función $\tanh(\cdot)$ utilizando estadísticas dependientes de los datos.

• Acotamiento Probabilístico de la Entrada: Para evitar la saturación, BHyT restringe el argumento de la pre-$\tanh$ para que se encuentre dentro de un intervalo predefinido $[-\lambda, \lambda]$ con alta probabilidad. Utilizando la desigualdad de Chebyshev, el método deriva una regla de escalamiento basada en la media ($\mu_x$) y la desviación estándar ($s_x$) de la entrada.
• Forma Ideal ($BHyT^*$): Teóricamente, la transformación reescala la entrada $x$ como:
  $$BHyT^*(x) = \gamma \odot \tanh\left(\frac{\lambda}{\kappa s_x + |\mu_x|} x\right)$$
  donde $\gamma$ es un parámetro de escala aprendible, $\lambda$ define el rango objetivo, y $\kappa$ controla el conservadurismo del escalamiento basado en una probabilidad objetivo $p$. Esto asegura que la entrada permanezca en el régimen lineal no saturado de $\tanh$ con alta probabilidad.

Implementación Práctica: Aproximación de Varianza

Calcular la media y la varianza exactas en cada sitio de Pre-LN dentro de un bloque Transformer anularía las ganancias de eficiencia. Para abordar esto, BHyT introduce una aproximación de varianza a nivel de bloque:

1. Primer Sitio (Antes de la Atención): Calcula la varianza exacta por instancia de la entrada.
2. Segundo Sitio (Antes de MLP): En lugar de recalcular la varianza exacta de la salida residual ($x + h_{Attn}$), aproxima la varianza utilizando la varianza de la entrada previamente calculada y una estimación del modelo de la varianza de la salida de atención.
   • La varianza de la salida de atención se aproxima como $\tilde{s}^2_{hAttn} \approx \frac{1}{Td} \|W_V W_O\|_F^2 \cdot \frac{\lambda_{Attn}^2}{\kappa^2}$, asumiendo pesos de atención uniformes y entradas gaussianas de media cero.
   • Esta aproximación permite que la segunda normalización reutilice estadísticas o utilice parámetros del modelo precomputados, evitando una segunda operación de reducción en la ruta crítica de propagación hacia adelante (forward path).

Estabilidad Teórica

El artículo proporciona un límite de propagación de varianza de profundidad finita. El Teorema 3.5 establece que para una red de profundidad $L$, si los hiperparámetros satisfacen $\lambda/\kappa < 1/\sqrt{L}$, la varianza de salida bajo BHyT es estrictamente menor que bajo el Escalamiento de LayerNorm (LNS) para cada capa hasta la profundidad $L$. Esto ofrece una garantía teórica de que BHyT suprime el crecimiento de la varianza por profundidad de manera más efectiva que LNS.

3. Contribuciones Clave

1. Propuesta de BHyT: Un reemplazo directo para Pre-LN que combina una transformación de $\tanh(\cdot)$ acotada con un acotamiento de entrada basado en datos para controlar el crecimiento de las activaciones por profundidad.
2. Eficiencia Teórica y Práctica:
   • Derivación de un límite de variación de profundidad finita que demuestra que BHyT produce un límite de propagación de varianza menor que LNS bajo condiciones específicas.
   • Introducción de una aproximación de varianza a nivel de bloque que reduce los cálculos de varianza exacta repetidos, disminuyendo significativamente la latencia.
3. Validación Empírica: Evaluación exhaustiva en modelos Llama-374M, Llama-1B y Llama-3B. El estudio demuestra que BHyT mejora el compromiso entre estabilidad y eficiencia manteniendo un fuerte rendimiento en tareas posteriores (downstream).

4. Resultados Experimentales

Los experimentos se realizaron en modelos tipo Llama (374M, 1B y 3B) preentrenados en el corpus C4 con presupuestos de tokens proporcionales al tamaño del modelo (aprox. 20 tokens por parámetro).

• Estabilidad: BHyT logró suprimir con éxito el crecimiento por profundidad en la magnitud y varianza de la activación, superando a RMSNorm y DyT. Las estadísticas de activación se mantuvieron controladas a través de las capas, comparables a los métodos orientados a la estabilidad como Peri-LN.
• Rendimiento:
  • Preentrenamiento: BHyT alcanzó la pérdida y la perplejidad de preentrenamiento más bajas en todas las escalas de modelos. También obtuvo la mayor precisión promedio en benchmarks de 3 disparos (ej. ARC, MMLU, Winogrande), superando al segundo mejor método por hasta 0.6 puntos en el modelo de 1B.
  • Post-SFT: Tras el ajuste fino supervisado (SFT) en Lima1K, los modelos BHyT mostraron una menor perplejidad en C4 y WikiText-2 (indicando menos olvido y mejor generalización) y una mayor precisión en tareas posteriores en comparación con RMSNorm, Peri-LN, LNS y DyT.
  • Escala: En una comparación de 20B de tokens en Llama-3B, BHyT superó al fuerte baseline orientado a la estabilidad, Peri-LN, por 3.4 puntos porcentuales en la precisión promedio de tareas posteriores (preentrenamiento) y 2.5 puntos (post-SFT).
• Eficiencia:
  • Rendimiento de Entrenamiento (Throughput): BHyT logró un promedio de 1.6% más de velocidad de entrenamiento en comparación con RMSNorm.
  • Rendimiento de Generación: BHyT alcanzó 1,199.9 tokens/s en Llama-1B, superando a RMSNorm por un 1.77% y a Peri-LN por un 18.0%. Aunque DyT fue el más rápido (debido a la ausencia de estadísticas), BHyT ofreció un compromiso superior entre velocidad y estabilidad.
  • Ablación: El mecanismo de aproximación de varianza proporcionó una aceleración del 14.9% sobre la forma exacta de BHyT* sin degradar materialmente la convergencia o el rendimiento posterior.

5. Significado y Reivindicaciones

El artículo afirma que BHyT ofrece un camino práctico hacia modelos fundacionales más profundos y eficientes. Al unificar la baja sobrecarga de las transformaciones acotadas con el control explícito de la varianza del flujo residual, BHyT aborda la "maldición de la profundidad" sin las penalizaciones computacionales de los esquemas de doble normalización.

Los autores sostienen que el entrenamiento escalable puede mejorar la eficiencia sin sacrificar la estabilidad. BHyT se presenta no solo como una alternativa teórica, sino como un reemplazo robusto y directo para RMSNorm que permite un preentrenamiento estable y un fuerte rendimiento de transferencia a través de diversas escalas de modelos (desde 374M hasta 3B de parámetros), reduciendo simultamente la latencia tanto en el entrenamiento como en la inferencia.