Resumen Técnico: Aprendizaje en un Espacio de Pesos Curvo (SymExpLin)

1. Planteamiento del Problema

El entrenamiento estándar de redes neuronales se basa en actualizaciones aditivas en un espacio de parámetros lineal. Sin embargo, muchas operaciones internas de las redes son inherentemente multiplicativas (por ejemplo, el escalado de amplitudes de características, ganancias de normalización o valores singulares). En un espacio lineal, los cambios relativos de igual magnitud (por ejemplo, duplicar frente a reducir a la mitad un peso) requieren distancias de optimización desiguales. Por ejemplo, mover un peso de 1.0 a 2.0 requiere un paso de tamaño 1.0, mientras que moverlo de 1.0 a 0.5 requiere un paso de 0.5.

Este desajuste crea desafíos de optimización:

• Distribuciones de Cola Pesada: Los modelos preentrenados suelen exhibir distribuciones de pesos de cola pesada donde una pequeña fracción de los parámetros crece significativamente más de su escala de inicialización (por ejemplo, 30 veces más grande), mientras que la mayoría permanece cerca de su inicialización.
• Invarianza del Optimizador: Los optimizadores adaptativos como Adam normalizan las actualizaciones por coordenada, lo que hace que los tamaños de paso sean invariantes al la magnitud del parámetro. En consecuencia, un parámetro pequeño (por ejemplo, 0.004) y un parámetro grande (por ejemplo, 0.4) reciben actualizaciones absolutas similares. Esto resulta en una diferencia de 100 veces en la perturbación relativa, lo que puede causar que los parámetros pequeños sean sobre-perturbados mientras que los parámetros grandes y críticos son sub-perturbados.
• Escalado Estático vs. Dinámico: Las soluciones existentes como la Parametrización de Actualización Máxima (µP) abordan los desajustes de escala estáticamente basándose en el ancho del tensor en la inicialización, pero no se adaptan al crecimiento dinámico de los parámetros individuales durante el entrenamiento.

2. Metodología: SymExpLin (SEL)

Los autores proponen SymExpLin (SEL), una reparametrización de pesos no lineal que descompone los pesos efectivos en una vía simétrica-exponencial y una vía lineal. Esto crea una geometría de espacio de parámetros curva donde las actualizaciones aditivas en el espacio bruto se mapean a cambios proporcionales a la magnitud en el espacio de pesos efectivos.

2.1 La Transformación

El peso efectivo $w_{eff}$ se calcula a partir de un peso bruto $w$ utilizando los siguientes componentes:

1. Vía Simétrica-Exponencial ($E_\theta$): Una función aprendible y preservadora del signo basada en la exponencial simétrica (symexp). Para $|w|$ pequeño, se comporta linealmente; para $|w|$ grande, crece exponencialmente. Esto permite que los grandes pasos brutos produzcan mayores desplazamientos efectivos.
   • $E_\theta(|w|) = e_w (\exp(\beta |w| \cdot m - n) - \exp(-n))$
   • Parámetros: $\beta$ (curvatura), $m$ (multiplicador interno), $n$ (desplazamiento para la pendiente en el origen), $e_w$ (escala de salida).
2. Vía Lineal ($L_\theta$): Una vía de identidad similar $L_\theta(w) = w \cdot l_w$ que proporciona una ruta de gradiente directa, la cual se hipotetiza que estabiliza la optimización.
3. Envoltura de Signo (Sign Wrapping): El signo del peso bruto se aplica a la suma de las vías: $w_{eff} = \text{sign}(w) \cdot (E_\theta(|w|) + L_\theta(w)) + \Delta_\phi$.
4. Residual ($\Delta_\phi$): Una corrección aditiva opcional de bajo rango (estilo LoRA).

2.2 Decisiones de Diseño Clave

• Inicialización Desajustada (Mismatched Initialization): Los autores identifican una asimetría beneficiosa donde los pesos brutos se inicializan invirtiendo una regla de combinación "congruente" (donde las partes lineal y exponencial se suman fuera de la envoltura de signo), pero el paso hacia adelante utiliza una regla de "desajuste" (donde ambas están envueltas dentro de la envoltura de signo). Esto resulta en que los pesos negativos se inicialicen con una magnitud menor que los pesos positivos, actuando como un mecanismo temprano de ruptura de simetría.
• Parámetros de Escala Estructurados: Los parámetros de escala aprendibles ($e_w, l_w, m, n$) se parametrizan utilizando patrones multiplicativos de fila-columna (vectores de rango 1) para equilibrar la expresividad y el costo computacional.
• Recocido de la Tasa de Aprendizaje (Learning Rate Annealing): Se utiliza un esquema dirigido para los parámetros de escala. La escala externa $e_w$ comienza con una tasa de aprendizaje alta para comprometerse rápidamente con un signo, mientras que el multiplicador interno $m$ comienza bajo y aumenta mediante recocido para refinar la curvatura una vez que la escala se estabiliza.
• Calibración de la Curvatura: El parámetro de curvatura $\beta$ se establece heurísticamente basado en el ancho del modelo para asegurar que el régimen exponencial se active cerca de las magnitudes típicas de inicialización (que disminuyen a medida que el ancho aumenta).

3. Contribuciones Clave

• SymExpLin (SEL): Una reparametrización de pesos aplicable a cualquier matriz o sesgo (bias) aprendible, descomponiendo los pesos en componentes simétricos-exponenciales y lineales con escalas estructuradas independientemente aprendibles.
• Inicialización Desajustada: El descubrimiento de que usar diferentes reglas de combinación para la inversión de la inicialización frente al paso hacia adelante crea un sesgo asimétrico beneficioso que mejora la optimización temprana.
• Escala Estructurada y Recocido: Introducción de parámetros de escala multiplicativos de fila-columna y un esquema específico de recocido de la tasa de aprendizaje para gestionar los controles exponenciales redundantes.
• Practicidad: El método es amigable para la implementación. La transformación de entrenamiento es elemento a elemento, y tras el entrenamiento, los pesos efectivos se precomputan y se integran en las capas lineales estándar, lo que resulta en costo cero de inferencia.
• Complementariedad: SEL opera en un eje dinámico por parámetro, complementando los métodos estáticos por tensor como µP.

4. Resultados Experimentales

El método fue evaluado en el modelado de lenguaje autorregresivo usando OpenWebText en nueve configuraciones de ancho $\{1024, 2048, 3072\}$ y profundidad $\{12, 24, 36\}$.

• Descenso de la Pérdida: SEL alcanzó consistentemente una pérdida de validación más baja que el baseline (Adam con parametrización estándar) en todas las configuraciones.
• Eficiencia de Pasos: SEL logró la pérdida de validación igualada en 1.32× a 1.49× menos pasos de entrenamiento. Las ganancias fueron más pronunciadas en modelos de mayor ancho (por ejemplo, una aceleración de 1.49× en 3072×24).
• Aceleración de Tiempo de Ejecución (Wall-Clock Speedup): A pesar de un overhead por paso del 1.7% (modelo pequeño) al 5.5% (modelo grande) debido a la transformación, la reducción en el total de pasos resultó en una aceleración de tiempo de ejecución estimada de 1.29× a 1.37×.
• Ablaciones:
  • Vía Lineal: Eliminar la vía lineal resultó en un caso de "fallo de aprendizaje", confirmando su necesidad para la estabilidad.
  • Desajuste vs. Congruente: La inicialización desajustada superó consistentemente a la inicialización congruente, particularmente en modelos más pequeños.
  • Recocido: El esquema de tasa de aprendizaje con recocido para los parámetros de escala proporcionó mejoras consistentes, aunque pequeñas, sobre los multiplicadores fijos.

5. Significado y Reivindicaciones

El artículo afirma que SEL aborda el desajuste fundamental entre la naturaleza multiplicativa de muchas computaciones neuronales y la naturaleza aditiva de la optimización estándar. Al curvar el espacio de parámetros, SEL permite que el optimizador tome pasos que son naturalmente proporcionales a la magnitud de los pesos, proporcionando efectivamente un escalado relativo sin requerir el ajuste de la tasa de aprendizaje por parámetro.

Los autores enfatizan que SEL es un reemplazo práctico y directo para las capas lineales estándar que:

1. Acelera la convergencia, particularmente para modelos con distribuciones de pesos de cola pesada.
2. No requiere cambios arquitectónicos durante la inferencia.
3. Es invertible, lo que permite que los pesos preentrenados estándar se conviertan de nuevo al espacio SEL para continuar el entrenamiento o el ajuste fino (fine-tuning).

El trabajo sugiere que la parametrización de "espacio logarítmico suave" ayuda a los pesos a alcanzar los objetivos de gran magnitud requeridos por la computación de la red, los cuales suelen verse desfavorecidos por las actualizaciones aditivas lineales.