Resumen Técnico: Transferencia de Hiperparámetros con Capas de Mezcla de Expertos

1. Planteamiento del Problema

Las capas de Mezcla de Expertos (MoE) se han convertido en un mecanismo crítico para escalar las redes neuronales modernas, desacoplando el número total de parámetros entrenables de los parámetros activos durante el paso hacia adelante. Sin embargo, el entrenamiento de modelos MoE dispersos introduce una complejidad significativa en la selección de hiperparámetros (HP). A diferencia de los modelos densos, las arquitecturas MoE introducen nuevos parámetros entrenables (pesos del enrutador) y nuevas dimensiones arquitectónicas (número de expertos, tamaño del experto) que requieren un ajuste cuidadoso.

Ajustar directamente los hiperparámetros (como la tasa de aprendizaje, la escala de inicialización y la decaimiento de peso) a gran escala es computacionalmente prohibitivo. Aunque existen técnicas de transferencia de HP para transformadores densos —que permiten extrapolar los hiperparámetros óptimos encontrados en modelos pequeños a modelos más grandes—, estos métodos no se han adaptado rigurosamente a las dinámicas de escalado específicas de las capas MoE dispersas. El desafío central consiste en determinar cómo escalar los hiperparámetros al aumentar simultáneamente el ancho, la profundidad, la cantidad de expertos y el tamaño del experto, sin reajustar en cada escala, garantizando al mismo tiempo la estabilidad y el rendimiento del entrenamiento.

2. Metodología

2.1. Parametrización Propuesta

Los autores proponen una nueva parametrización para modelos transformadores con capas MoE, extendiendo la parametrización CompleteP (desarrollada previamente para transformadores densos) para incluir reglas de escalado específicas de MoE. El objetivo es definir un conjunto de reglas que predigan cómo deben cambiar los valores brutos de los hiperparámetros (desviación estándar de inicialización $\sigma$ y tasa de aprendizaje $\eta$) a medida que escalan las dimensiones del modelo, de modo que las dinámicas de entrenamiento permanezcan consistentes.

La derivación se basa en el principio de parametrización de actualización máxima ($\mu$P), que requiere que los componentes de la red (pre-activaciones y actualizaciones residuales) permanezcan en $O(1)$ en la inicialización y reciban actualizaciones de $\Theta(1)$ por paso de entrenamiento. Los autores extienden esto a los MoE exigiendo que las condiciones de actualización máxima se cumplan no solo para la salida de la capa, sino también para los componentes individuales de los expertos (coeficientes de mezcla y salidas de los expertos).

Las reglas de escalado clave derivadas para el módulo MoE incluyen:

• Pesos del Enrutador: La tasa de aprendizaje escala como $\eta \propto n_{embd}^{-1}$. La inicialización escala como $n_{embd}^{-\gamma}$ (con $\gamma \geq 0.5$).
• Sesgos de los Expertos: Inicializados en cero con una tasa de aprendizaje constante $\eta \propto 1$ (independiente de la cantidad de expertos), siempre que la dispersión se mantenga fija.
• Pesos MLP de los Expertos:
  • Proyección ascendente ($W_{up}$): $\sigma_{init} \propto n_{embd}^{-1/2}$, $\eta \propto n_{embd}^{-1}$.
  • Proyección descendente ($W_{down}$): $\sigma_{init} \propto \alpha_{ffn}^{-1} n_{embd}^{-1/2}$, $\eta \propto \alpha_{ffn}^{-1} n_{embd}^{-1}$.
  • Aquí, $\alpha_{ffn}$ representa el multiplicador del tamaño oculto del experto en relación con la dimensión de incrustación.

Crucialmente, los autores fijan la tasa de dispersión $\kappa = n_{act}/n_{exp}$ (la fracción de expertos activos) como una constante mientras escalan el número total de expertos ($n_{exp}$) y los expertos activos ($n_{act}$). Esto contrasta con enfoques que fijan el número de expertos activos mientras aumentan el grupo total.

2.2. Fundamentación Teórica: Teoría de Campo Medio Dinámico (DMFT)

Para justificar estas reglas de escalado heurísticas, los autores emplean la Teoría de Campo Medio Dinámico (DMFT). Analizan las dinámicas de entrenamiento de redes residuales con capas MoE en el límite simultáneo de ancho infinito ($n_{embd}$), profundidad ($L$), tamaño del experto ($n_{hid}$) y cantidad de expertos ($n_{exp}$), manteniendo una dispersión de activación constante $\kappa$.

El análisis revela una nueva jerarquía de campo medio de tres niveles:

1. Representaciones del Flujo Residual: Campo medio sobre las salidas de los expertos.
2. Salidas de los Expertos: Campo medio sobre las neuronas individuales de cada experto.
3. Neuronas Individuales: Campo medio dentro de cada experto.

El análisis DMFT demuestra que, bajo la parametrización propuesta, las dinámicas de entrenamiento límite son:

• Independientes de la relación FFN ($\alpha_{ffn}$): Las dinámicas no dependen del tamaño específico de los expertos en relación con la dimensión de incrustación, siempre que se tome el límite de escalado conjunto.
• Dependientes únicamente de la dispersión $\kappa$: Las dinámicas son consistentes a través de todos los parámetros de escalado siempre que la tasa de dispersión se mantenga fija.
• Invariancia de Escala: La evolución de las estadísticas resumidas de la red (por ejemplo, kernels de características por capa) es consistente a través de las escalas, garantizando teóricamente que los hiperparámetros se transfieran de manera fiable.

2.3. Configuración Experimental

Los autores validan empíricamente su parametrización utilizando modelos de lenguaje Transformer solo de decodificador en los conjuntos de datos FineWeb y C4.

• Modelos Base: Ajustados en modelos pequeños con aproximadamente 38M de parámetros activados.
• Escalado: Escalado hasta modelos con hasta 2B de parámetros totales, variando ancho, profundidad, cantidad de expertos y tamaño del experto.
• Presupuesto de Tokens: Los experimentos se realizaron con un presupuesto fijo de 1B de tokens (2000 pasos) para aislar las dinámicas de entrenamiento temprano, así como horizontes más largos (hasta 7.5B de tokens).
• Optimizador: Optimizador Adam estándar.
• Equilibrio de Carga: Se utilizó una estrategia libre de pérdida auxiliar, actualizando los sesgos de los expertos directamente para fomentar el equilibrio de carga sin añadir un término de regularización a la función de pérdida.

3. Contribuciones Clave

1. Parametrización de MoE: El artículo extiende la parametrización CompleteP a modelos MoE dispersos, proporcionando reglas de escalado explícitas para los pesos del enrutador, los sesgos de los expertos y los pesos MLP de los expertos a través del ancho, la profundidad, la cantidad de expertos y el tamaño del experto.
2. Justificación Teórica mediante DMFT: Los autores proporcionan una base teórica rigurosa para su parametrización utilizando DMFT. Derivan una descripción explícita de las dinámicas de entrenamiento en el límite de ancho/profundidad infinitos, demostrando que las dinámicas convergen a un sistema estable e invariante de escala que depende únicamente de la dispersión, y no del escalado específico de las dimensiones de los expertos.
3. Validación Empírica de la Transferencia de HP: El estudio demuestra que los hiperparámetros óptimos (tasa de aprendizaje y escala de inicialización) identificados en modelos base pequeños (38M de parámetros activos) se transfieren de manera fiable a modelos mucho más grandes (hasta 2B de parámetros totales) a través de diversas dimensiones arquitectónicas.
4. Insights Arquitectónicos: Los autores verifican empíricamente que, bajo su parametrización, aumentar el número de expertos (manteniendo fijo el número total de parámetros) produce un mejor rendimiento que aumentar el tamaño de los expertos individuales. Este hallazgo se alinea con la literatura reciente, pero se logra aquí sin la necesidad de costosos barridos de hiperparámetros en cada escala.

4. Resultados

• Transferencia Confiable de HP: Bajo las reglas de escalado propuestas, las tasas de aprendizaje base óptimas y las desviaciones estándar de inicialización se transfieren eficazmente a través de modelos que van desde 51M hasta 2B de parámetros totales. Las curvas de pérdida de los modelos escalados colapsan con las del modelo base en las primeras iteraciones antes de divergir (con modelos más grandes logrando una pérdida menor).
• Estabilidad: La parametrización garantiza dinámicas de entrenamiento estables, incluido un equilibrio de carga uniforme de los expertos, incluso al escalar el número de expertos. Los autores señalan que el pre-entrenamiento de MoE es particularmente sensible a hiperparámetros de escala constante (multiplicadores tratados como $\Theta(1)$), y ajustar estos es crucial para la estabilidad.
• Rendimiento: Los modelos entrenados utilizando hiperparámetros de cero-shot (transferidos desde modelos pequeños) logran un rendimiento competitivo frente a baselines densos (por ejemplo, GPT-2 pequeño/mediano) cuando se igualan en conteo de parámetros activos.
• Cantidad de Expertos vs. Tamaño: Los experimentos confirman que aumentar el número de expertos es más eficiente en parámetros que aumentar el tamaño de los expertos con un conteo fijo de parámetros. Este beneficio persiste incluso en horizontes de entrenamiento más largos (5B de tokens).
• Sensibilidad a la Dispersión: El estudio destaca que la transferencia de hiperparámetros es válida únicamente cuando la tasa de dispersión $\kappa$ se mantiene fija. Escalar el número de expertos mientras se fija el número de expertos activos (haciendo que $\kappa \to 0$) rompe la transferibilidad de los hiperparámetros óptimos.

5. Significado y Afirmaciones

El artículo afirma proporcionar un marco práctico pero riguroso para escalar modelos MoE. Al combinar la parametrización heurística con el análisis DMFT, los autores ofrecen un método para:

• Reducir los Costos de Entrenamiento: Permitir la selección de hiperparámetros óptimos para modelos MoE a gran escala ajustando únicamente modelos base pequeños y económicos.
• Garantizar Estabilidad: Proporcionar reglas que previenen la inestabilidad del entrenamiento (por ejemplo, colapso o divergencia de expertos), algo común en el entrenamiento de MoE dispersos.
• Guiar el Diseño Arquitectónico: Ofrecer evidencia teórica y empírica de que aumentar la cantidad de expertos es preferible a aumentar el tamaño de los expertos bajo dispersión fija, ayudando en el diseño de modelos grandes y eficientes.

Los autores reconocen limitaciones, señalando que el trabajo actual se centra en la tasa de aprendizaje y la inicialización, dejando otros hiperparámetros (tamaño de lote, decaimiento de peso, programas de LR) para futuras investigaciones. También notan que, aunque el análisis DMFT apoya el límite de ancho infinito, el comportamiento de las transferencias de ancho pequeño requiere un estudio teórico adicional. Además, el artículo no afirma resolver las leyes de escalado "óptimas en cómputo" para MoE (por ejemplo, exponentes de Chinchilla), ya que el compromiso FLOP-rendimiento en MoE difiere significativamente de los modelos densos debido a las limitaciones de hardware inducidas por la dispersión.

En resumen, este trabajo establece una parametrización fundamental que permite la extrapolación fiable de hiperparámetros desde modelos MoE pequeños a grandes, fundamentada en una nueva teoría de campo medio de tres niveles de las dinámicas de entrenamiento.