Unlocking Feature Learning in Gated Delta Networks at Scale

Resumen Técnico: Desbloqueando el Aprendizaje de Características en Redes Delta con Compuerta a Escala

Planteamiento del Problema
El escalado de los Modelos de Lenguaje de Gran Escala (LLMs) enfrenta dos cuellos de botella primarios: el elevado costo computacional de entrenar modelos más grandes y la complejidad cuadrática de la autoatención estándar de los Transformers. Las arquitecturas recurrentes lineales, como las Redes Delta con Compuerta (GDN, por sus siglas en inglés), han surgido como alternativas eficientes que mantienen una inferencia de tiempo lineal mientras capturan dependencias de largo alcance. Sin embargo, existe una brecha crítica para escalar estos modelos de manera efectiva. Aunque la Parametrización de Actualización Máxima (µP) ha permitido con éxito la transferencia de hiperparámetros (específicamente tasas de aprendizaje) zero-shot para Transformers estándar y algunos modelos lineales, su aplicación a las Redes Delta con Compuerta permanece inexplorada. El desafío central reside en la dinámica del estado recurrente de la arquitectura, la cual se actualiza mediante mecanismos de compuerta dependientes de los datos y reglas delta de matriz completa, en lugar de las recurrencias escalares o las estructuras de alimentación hacia adelante (feedforward) estándar asumidas por las derivaciones de µP existentes. Trabajos previos sobre Modelos de Espacio de Estados Estructurados (SSMs) como Mamba han demostrado que el µP convencional falla para los SSMs diagonales, pero el uso de actualizaciones de estado de matriz completa en GDN hace que esas correcciones específicas sean inaplicables.

Metodología
Los autores abordan esta brecha derivando rigurosamente la formulación de µP para las Redes Delta con Compuerta. La metodología consiste en propagar las estimaciones de tamaño de coordenadas a través de todo el paso hacia adelante (forward pass), incluyendo los mecanismos de compuerta únicos y las transiciones de estado recurrente, para determinar las reglas de escalado necesarias para la inicialización, los multiplicadores de avance y las tasas de aprendizaje.

Pasos analíticos clave incluyen:

1. Análisis del Paso hacia Adelante: Los autores derivan los tamaños de coordenadas de las características proyectadas, el estado latente $S_t$ y la lectura (readout). Identifican que la normalización L2 de las consultas (queries) y claves (keys), combinada con las actualizaciones de escritura de rango uno, resulta en escalados de varianza específicos. Crucialmente, proponen insertar un multiplicador $\sqrt{d}$ antes de la capa RMSNorm para asegurar que la entrada a la capa de normalización mantenga un tamaño de coordenada $\Theta(1)$, evitando la interrupción del escalado de gradientes requerido para el aprendizaje de características.
2. Análisis del Mecanismo de Compuerta: El artículo analiza los parámetros de compuerta no lineales $\alpha_t$ (decaimiento) y $\beta_t$ (fuerza de escritura). A diferencia de los pesos estándar, estos se derivan de pre-activaciones escalares que involucran vectores de pesos entrenables ($W_\alpha, W_\beta$) y parámetros escalares ($a_{log}, b$). El análisis revela que estos componentes de compuerta no siguen las leyes de escalado de µP estándar debido a sus roles específicos en la ecuación de actualización de estado.
3. Paso hacia Atrás y Derivación de la Tasa de Aprendizaje:
   • AdamW: Debido a la normalización de los gradientes por parte de AdamW mediante su segundo momento, la magnitud de actualización efectiva es consistente entre las clases de pesos, lo que simplifica los requisitos de la tasa de aprendizaje.
   • SGD: Bajo el uso de SGD simple, las magnitudes de los gradientes brutos varían significativamente entre diferentes clases de pesos. Los autores derivan que las matrices de pesos de compuerta ($W_\alpha, W_\beta$) requieren un escalado de tasa de aprendizaje de $\Theta(1/\sqrt{d})$, y los parámetros de compuerta escalares ($a_{log}, b$) requieren un escalado de $\Theta(\sqrt{d})$. Estos escalados difieren del entorno µP estándar donde las tasas de aprendizaje suelen ser uniformes o siguen una regla simple de $1/n$.

Contribuciones Clave

1. Derivación Teórica: El artículo proporciona la primera formulación completa de µP para las Redes Delta con Compuerta. Deriva las estimaciones de tamaño de coordenadas a través de todo el paso hacia adelante y establece varianzas de inicialización y multiplicadores de avance con principios para todas las clases de pesos.
2. Nuevas Reglas de Escalado: Los autores descubren que los componentes de compuerta requieren escalados de tasa de aprendizaje no estándar bajo SGD:
   • Las matrices de pesos de compuerta ($W_\alpha, W_\beta$) requieren un escalado de $\Theta(1/\sqrt{d})$.
   • Los parámetros de compuerta escalares ($a_{log}, b$) requieren un escalado de $\Theta(\sqrt{d})$.
3. Validación Empírica: Los autores preentrenan modelos de lenguaje GDN a través de múltiples anchuras (que van desde 256 hasta 1563 dimensiones ocultas) utilizando tanto los optimizadores AdamW como SGD.
   • Resultos: Bajo la configuración de µP propuesta, la tasa de aprendizaje óptima se transfiere zero-shot a través de todas las anchuras de modelo para ambos optimizadores. En contraste, la Parametrización Estándar (SP) y la configuración original de µP fallan en la transferencia, con tasas de aprendizaje óptimas que cambian significativamente a medida que aumenta la anchura del modelo.
   • Estabilidad: Las configuraciones propuestas permiten una transferencia estable de la tasa de aprendizaje, confirmando las derivaciones teóricas.

Significancia y Reivindicaciones
El artículo afirma que su análisis resuelve el problema abierto de parametrizar las Redes Delta con Compuerta para el aprendizaje de características a escala. Al demostrar que la formulación de µP derivada permite la transferencia de hiperparámetros zero-shot, el trabajo valida la corrección de la derivación teórica y su utilidad práctica. Los autores enfatizan que este es el primer trabajo que deriva y valida una parametrización consistente con µP específicamente para las Redes Delta con Compuerta, distinguiéndolo del trabajo previo sobre SSMs diagonales. Los hallazgos sugieren que una parametrización adecuada es esencial para escalar modelos recurrentes lineales y que los mecanismos de compuerta específicos en GDN requieren reglas de escalado únicas que difieren tanto de los Transformers estándar como de otras arquitecturas lineales. Los autores esperan que estas derivaciones guíen la investigación futura sobre las leyes de escalado de otras arquitecturas lineales o híbridas.