🤖 machine learning

MESH: Memory-Efficient Sinkhorn Optimization for Mixture-of-Experts Training

Este artículo presenta MESH, un método de optimización de Sinkhorn eficiente en memoria que incorpora actualizaciones de momento oculto y precondicionamiento de bloques opcional para permitir con éxito el entrenamiento de modelos de Mezcla de Expertos con bajo consumo de memoria al abordar la inestabilidad temporal de los gradientes de los expertos enrutados que causa el fallo de los métodos de Sinkhorn sin estado estándar.

Autores originales: Masato Fujitake

Publicado 2026-08-06

📖 4 min de lectura☕ Lectura para el café

CC BY 4.0

Autores originales: Masato Fujitake

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que estás intentando enseñarle a un robot gigante y superinteligente a escribir historias. Para que este robot sea realmente bueno, necesitas un cerebro enorme hecho de miles de millones de diminutos interruptores. Pero aquí está el truco: el cerebro del robot es tan grande que apenas cabe en la memoria de la computadora; es como intentar reproducir una película de gran éxito en un teléfono con una batería diminuta; el teléfono se calienta, se ralentiza y podría colapsar.

Para solucionar esto, los científicos usan "entrenadores" especiales llamados optimizadores. El entrenador más famoso, llamado AdamW, es como un entrenador muy cuidadoso que lleva un cuaderno detallado para cada uno de los interruptores del cerebro del robot, recordando exactamente cómo darle un empujoncito la próxima vez. Esto funciona de maravilla, pero esos cuadernos ocupan un espacio enorme. Así que, los investigadores inventaron un entrenador más ligero llamado Sinkhorn. Es como un entrenador que no lleva ningún cuaderno; en su lugar, simplemente observa el error actual y realiza un ajuste rápido y sin memoria. Esto ahorra un espacio masivo, pero resulta que para un tipo específico de cerebro de robot llamado "Mezcla de Expertos" (Mixture-of-Experts o MoE), el entrenador sin memoria se confunde y el robot deja de aprender bien. La pregunta que los científicos se hacen es: ¿Podemos mantener los beneficios de ahorro de memoria del entrenador rápido sin perder la capacidad del robot para aprender?

Este artículo presenta un nuevo método llamado MESH (Sinkhorn de Memoria Eficiente para Expertos) para resolver exactamente ese problema. Los investigadores descubrieron que el fallo del entrenador rápido y sin memoria ocurre porque los "expertos" en el cerebro del robot no siempre llegan a trabajar. En un modelo de Mezcla de Expertos, el robot solo elige a algunos expertos específicos para manejar cada frase, algo así como un restaurante donde solo se llama a unos pocos chefs para un pedido específico. Debido a que los chefs cambian de un pedido a otro, el entrenador rápido ve una señal desordenada y ruidosa y se pierde.

El artículo sugiere que la solución no es darle al entrenador un cuaderno completo de nuevo, sino darle una "memoria oculta". En lugar de mirar solo el pedido actual, el entrenador ahora recuerda el promedio de los últimos pedidos antes de tomar una decisión. El autor llama a esto "momento oculto" (hidden momentum). Lo probaron en un modelo pequeño de 110 millones de parámetros (un "nanowhale") y descubrieron que este truco simple —suavizar el ruido antes de ajustar los pesos— restaura la capacidad de aprendizaje del robot.

Esto es lo que encontraron y lo que descartaron:

La solución principal: El artículo muestra que el fallo ocurre porque los "expertos enrutados" (los chefs que solo trabajan a veces) necesitan que sus señales se suavicen a lo largo del tiempo antes de que ocurra el ajuste sin memoria. Al promediar el historial reciente de estos expertos (usando un búfer oculto que no cuenta como memoria permanente), el nuevo método MESH funciona mucho mejor.
Lo que no funcionó: Los investigadores probaron muchas otras ideas para solucionar el problema, pero las descartaron. Descubrieron que simplemente hacer los ajustes más pequeños, cambiar la forma en que se agrupan los expertos o usar solo información de "signo" (saber solo si el error subió o bajó, no cuánto) no solucionó el problema. La clave era específicamente el suavizado temporal, no solo cualquier tipo de dato adicional.
¿Qué tan seguros están?: Los resultados se basan en experimentos con un modelo pequeño específico. En estas pruebas, el nuevo método redujo la memoria necesaria para el estado del optimizador en un 62.5% (bajando de 0.883 GB a 0.331 GB) y redujo el uso de memoria máxima de la computadora en aproximadamente un 12.6%. Sin embargo, la puntuación de la prueba final del robot (pérdida de evaluación) fue ligeramente peor que la del entrenador pesado que guarda cuadernos (AdamW), con puntuaciones alrededor de 3.64 para el nuevo método frente a 3.59 para el antiguo. El autor sugiere que esta pequeña brecha puede deberse a que otras partes del cerebro del robot (como el vocabulario) aún necesitan al entrenador pesado, pero están seguros de que el "momento oculto" es la pieza que falta para los expertos.

En resumen, el artículo argumenta que no necesitas un cuaderno completo para entrenar estos cerebros de robot especiales; solo necesitas una forma de recordar el pasado reciente sin escribirlo permanentemente. Este enfoque de "momento oculto", llamado MESH, te lleva casi a la meta, ahorrando mucha memoria mientras mantiene al robot lo suficientemente inteligente como para aprender.

Resumen Técnico: MESH – Optimización de Sinkhorn con Eficiencia de Memoria para el Entrenamiento de Mixture-of-Experts

1. Planteamiento del Problema

El artículo aborda el desafío de optimizar modelos de lenguaje Mixture-of-Experts (MoE) con eficiencia de memoria. Aunque AdamW es el optimizador estándar para el preentrenamiento debido a su combinación de suavizado del primer momento temporal y escalado adaptativo por coordenada, incurre en un alto costo de memoria al almacenar dos tensores de estado completos por parámetro.

Las alternativas con eficiencia de memoria, como los optimizadores de matrices estilo Sinkhorn y los híbridos SAGE/Sinkhorn, reducen este costo eliminando los estados de momento para matrices densas. Sin embargo, el autor demuestra que aplicar directamente estos métodos sin estado a los expertos de MoE falla. En un entorno controlado de "nanowhale" tipo DeepSeek de 110 millones de parámetros, un híbrido SAGE/Sinkhorn estándar redujo significativamente el estado del optimizador (de 0.883 GB a 0.331 GB) pero resultó en una degradación severa en la pérdida de evaluación (3.8265 frente al baseline de AdamW de 3.58–3.64).

El fallo central se identifica como un problema de ordenamiento específico de los expertos enrutados:

Los Gradientes Enrutados son Condicionales: A diferencia de los gradientes de las MLP densas, los gradientes de los expertos enrutados son muestras condicionales basadas en las decisiones del router y los pesos de las compuertas (gates), que varían temporalmente.
No Linealidad de Sinkhorn: La normalización de Sinkhorn es una operación no lineal. Aplicar Sinkhorn a un gradiente enrutado instantáneo y de alta varianza ( $Sinkhorn(G_t)$ ) se comporta de manera diferente que aplicarlo a una señal promediada temporalmente ( $Sinkhorn(E[G_t])$ ).
Amplificación de Ruido: El Sinkhorn sin estado normaliza el ruido inducido por el enrutamiento en cada paso, mientras que el fallo de los métodos sin estado proviene de la falta de suavizado temporal antes de que ocurra la normalización de la matriz.

2. Metodología: MESH y MESH-B

El autor propone MESH (Memory-Efficient Sinkhorn for Experts), un optimizador híbrido diseñado específicamente para arquitecturas MoE.

Mecanismo Central: Momento Oculto (Hidden Momentum)

MESH reemplaza la actualización de Sinkhorn sin estado para los expertos MoE con un enfoque de momento oculto: en lugar de almacenar un tensor explícito del primer momento en el estado del optimizador (lo que anularía el objetivo de ahorro de memoria), MESH mantiene el primer momento dentro del ciclo de vida del búfer de gradientes:

Acumulación: Se mantiene un acumulador de momento de experto no normalizado: $H_{e,t} = \beta_1 H_{e,t-1} + G_{e,t}$ .
Ciclo de Vida del Búfer: El acumulador $H_{e,t}$ se almacena en el búfer de gradientes en lugar de en el diccionario de estado del optimizador.
Ciclo de Actualización: Antes del paso hacia atrás (backward pass), el búfer se vuelve a adjuntar como el gradiente ( $p.grad \leftarrow \beta_1 H_{e,t}$ ). El paso hacia atrás suma el gradiente actual, actualizando efectivamente el acumulador para el siguiente paso.
Normalización: La actualización de Sinkhorn se computa sobre la señal suavizada: $U_{e,t} = Sinkhorn(H_{e,t})$ .

Este enfoque restaura la señal del primer momento temporal requerida para los expertos enrutados sin necesidad de almacenar un tensor de estado de optimizador explícito para ellos.

Extensión Opcional: MESH-B

El artículo también evalúa MESH-B, que añade un multiplicador de inverso de RMS por bloque/neurona grueso a la dirección de Sinkhorn. Esto actúa como un precondicionador a nivel de bloque. Las estadísticas de bloque se computan a partir de la fuente suavizada ( $H_{e,t}$ ) en lugar de un flujo de gradiente bruto separado para mantener la simplicidad de la implementación.

Estrategia de Asignación de Roles

MESH mantiene una estrategia híbrida para diferentes clases de parámetros:

Embeddings / LM Head: Manejados por SAGE (actualizaciones de estado ligero).
Matrices Densas No-Expert: Manejadas por Sinkhorn sin estado.
Expertos Enrutados/Compartidos: Manejados por MESH (Sinkhorn con momento oculto).
Normas / Biases: Manejados por SAGE.

3. Contribuciones Clave

Identificación del Punto de Fallo: El artículo identifica las matrices de expertos MoE enrutados como el punto de fallo específico de la optimización directa de SAGE/Sinkhorn en el preentrenamiento, distinguiéndolos de las matrices densas y los embeddings.
Algoritmo MESH: Introducción de una actualización de Sinkhorn con momento oculto que restaura una señal de primer momento para los expertos sin añadir un estado de optimizador explícito, utilizando el ciclo de vida del búfer de gradientes.
Análisis de Ablación: Ablaciones negativas sistemáticas que muestran que las intervenciones de solo segundo momento, solo signo/momento, solo confianza y solo escala no explican la brecha de rendimiento. El ingrediente causal primario se identifica como el suavizado temporal antes de la normalización de la matriz.
Evaluación de MESH-B: Evaluación de una variante precondicionada por bloques, mostrando que, si bien puede mejorar la frontera de memoria-calidad en semillas específicas, no es universalmente necesaria comparada con el suavizado temporal proporcionado por MESH.

4. Resultados Experimentales

Los experimentos se realizaron en un modelo MoE tipo DeepSeek de 110 millones de parámetros entrenado en FineWeb-Edu durante 5,000 pasos.

Eficiencia de Memoria:
- AdamW: 0.883 GB de estado del optimizador.
- MESH / MESH-B: 0.331 GB de estado del optimizador (una reducción del 62.5% respecto a AdamW).
- Asignación CUDA Pico: Reducida aproximadamente en un 12.6% (de ~4.37 GB a ~3.82 GB).
Rendimiento (Pérdida de Evaluación):
- Baseline AdamW: ~3.5877 (promedio entre semillas 100/512).
- SAGE/Sinkhorn sin estado: 3.8265 (Degradación significativa).
- MESH: 3.6430 (Cercano a AdamW, pero ligeramente superior).
- MESH-B: 3.6365 (Ligeramente mejor que MESH, pero aún por detrás de AdamW).
Hallazgos Diagnósticos (Semilla 42):
- Las variantes de estado completo (MESH-FULL, MESH-B-FULL) que almacenan explícitamente el primer momento recuperan una calidad similar a AdamW (por ejemplo, MESH-B-FULL con V(H) alcanzó 3.6102) con aproximadamente la mitad del estado del optimizador de AdamW.
- Esto confirma que el suavizado temporal es el ingrediente faltante crítico, mientras que la implementación de momento oculto aproxima esto con éxito sin el costo de un estado completo.

5. Significancia y Reclamaciones

El artículo reclama modestamente que MESH proporciona un punto de estado sustancialmente menor en la frontera de memoria-calidad para el entrenamiento de MoE.

Conclusión Primaria: Los expertos de MoE requieren suavizado temporal antes de la normalización de la matriz para funcionar eficazmente con optimizadores tipo Sinkhorn. Este suavizado es el factor causal dominante para cerrar la brecha de rendimiento con AdamW.
Conclusión Secundaria: El precondicionamiento por bloque/neurona (MESH-B) puede mejorar los resultados en configuraciones específicas, pero no se establece como un componente universalmente necesario.
Limitaciones: El autor reconoce una brecha de rendimiento residual entre MESH/MESH-B y AdamW. Hipotetizan que esto se debe a los roles no expertos (específicamente las matrices de vocabulario y las matrices densas no expertas), que permanecen en SAGE o Sinkhorn sin estado en su configuración, mientras que AdamW las adapta con momentos completos.
Alcance: El estudio es intencionalmente estrecho, centrándose en una arquitectura de 110 millones de parámetros y ejecuciones de 5,000 pasos. El autor no afirma que MESH supere a AdamW en todas las métricas, sino que recupera la mayor parte de la calidad perdida por los métodos sin estado mientras reduce drásticamente la sobrecarga de memoria.

En resumen, MESH demuestra que el fallo de los optimizadores con eficiencia de memoria en el entrenamiento de MoE no es inherente a la estructura de MoE en sí, sino al la falta de suavizado temporal para los gradientes enrutados. Al implementar el "momento oculto", MESH resuelve este problema de ordenamiento, ofreciendo un camino viable hacia el preentrenamiento de MoE con eficiencia de memoria.