🤖 machine learning

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

Ce document introduit MESH, une méthode d'optimisation de Sinkhorn efficace en mémoire qui incorpore des mises à jour de moment caché et un préconditionnement par blocs optionnel pour permettre avec succès l'entraînement à faible mémoire des modèles Mixture-of-Experts en traitant l'instabilité temporelle des gradients d'experts routés qui fait échouer les méthodes de Sinkhorn sans état standard.

Auteurs originaux : Masato Fujitake

Publié 2026-08-06

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Masato Fujitake

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez d'apprendre à un robot géant et super intelligent à écrire des histoires. Pour que ce robot devienne vraiment bon, vous avez besoin d'un cerveau massif composé de milliards de minuscules interrupteurs. Mais il y a un piège : le cerveau du robot est si énorme qu'il tient tout juste dans la mémoire de l'ordinateur. C'est comme essayer de faire tourner un film à gros budget sur un téléphone avec une petite batterie ; le téléphone chauffe, ralentit et risque de planter.

Pour résoudre ce problème, les scientifiques utilisent des « entraîneurs » spéciaux appelés optimiseurs. L'entraîneur le plus célèbre, nommé AdamW, est comme un coach très méticuleux qui tient un carnet de notes détaillé pour chaque interrupteur du cerveau du robot, se souvenant exactement de la manière de le pousser la fois suivante. Cela fonctionne très bien, mais ces carnets de notes prennent énormément de place. Ainsi, des chercheurs ont inventé un entraîneur plus léger appelé Sinkhorn. C'est comme un coach qui ne tient aucun carnet de notes ; au lieu de cela, il regarde simplement l'erreur actuelle et effectue un ajustement rapide, sans utiliser de mémoire. Cela économise un espace massif, mais il s'avère que pour un type spécifique de cerveau de robot appelé « Mixture-of-Experts » (MoE), ce coach sans mémoire s'embrouille et le robot arrête d'apprendre correctement. La question que les scientifiques se posent est la suivante : peut-on conserver les avantages de gain de mémoire du coach rapide sans perdre la capacité du robot à apprendre ?

Ce document présente une nouvelle méthode appelée MESH (Memory-Efficient Sinkhorn for Experts) pour résoudre précisément ce problème. Les chercheurs ont découvert que le coach rapide et sans mémoire échoue parce que les « experts » du cerveau du robot ne travaillent pas toujours. Dans un modèle Mixture-of-Experts, le robot ne choisit que quelques experts spécifiques pour gérer chaque phrase, un peu comme un restaurant où seuls quelques chefs sont appelés en cuisine pour une commande spécifique. Comme les chefs changent d'une commande à l'autre, le coach rapide perçoit un signal brouillé et bruyant, et s'y perd.

Le document suggère que la solution n'est pas de redonner un carnet de notes complet au coach, mais de lui donner une « mémoire cachée ». Au lieu de regarder seulement la commande actuelle, le coach se souvient désormais de la moyenne des dernières commandes avant de prendre une décision. L'auteur appelle cela le « momentum caché ». Ils ont testé cela sur un petit modèle de 110 millions de paramètres (un « nanowhale ») et ont constaté que cette astuce simple — lisser le bruit avant d'ajuster les poids — restaure la capacité d'apprentissage du robot.

Voici ce qu'ils ont découvert et ce qu'ils ont écarté :

La solution principale : Le document montre que l'échec se produit parce que les « experts routés » (les chefs qui ne travaillent que parfois) ont besoin que leurs signaux soient lissés au fil du temps avant que l'ajustement sans mémoire ne se produise. En faisant la moyenne de l'historique récent de ces experts (en utilisant un tampon caché qui ne compte pas comme une mémoire permanente), la nouvelle méthode MESH fonctionne beaucoup mieux.
Ce qui n'a pas fonctionné : Les chercheurs ont essayé de nombreuses autres idées pour régler le problème, mais ils les ont écartées. Ils ont découvert que le simple fait de réduire la taille des ajustements, de changer la façon dont les experts sont regroupés ou d'utiliser uniquement des informations de « signe » (savoir seulement si l'erreur a augmenté ou diminué, et non de combien) n'a pas résolu le problème. La clé était spécifiquement le lissage temporel, et non simplement l'ajout de données supplémentaires.
À quel point sont-ils sûrs ? Les résultats sont basés sur des expériences avec un modèle spécifique de petite taille. Dans ces tests, la nouvelle méthode a réduit la mémoire nécessaire pour l'état de l'optimiseur de 62,5 % (passant de 0,883 Go à 0,331 Go) et a abaissé l'utilisation de la mémoire informatique de pointe d'environ 12,6 %. Cependant, le score final du test du robot (perte d'évaluation) est resté légèrement moins bon que celui du coach lourd tenant un carnet de notes (AdamW), avec des scores autour de 3,64 pour la nouvelle méthode contre 3,59 pour l'ancienne. L'auteur suggère que cet écart minime pourrait être dû au fait que d'autres parties du cerveau du robot (comme le vocabulaire) ont toujours besoin du coach lourd, mais il est convaincu que le « momentum caché » est la pièce manquante pour les experts.

En résumé, le document soutient que vous n'avez pas besoin d'un carnet de notes complet pour entraîner ces cerveaux de robots spéciaux ; vous avez juste besoin d'un moyen de se souvenir du passé récent sans l'écrire de façon permanente. Cette approche de « momentum caché », appelée MESH, vous permet d'atteindre presque tous les objectifs, en économisant beaucoup de mémoire tout en gardant le robot assez intelligent pour apprendre.

Résumé Technique : MESH – Optimisation Sinkhorn à efficacité mémoire pour l'entraînement de Mixture-of-Experts

1. Énoncé du Problème

Le papier traite du défi de l'optimisation des modèles de langage Mixture-of-Experts (MoE) avec une efficacité mémoire accrue. Bien que AdamW soit l'optimiseur standard pour le pré-entraînement en raison de sa combinaison de lissage du premier moment temporel et de mise à lance adaptative par coordonnée, il induit un coût mémoire élevé en stockant deux tenseurs d'état complets par paramètre.

Les alternatives économes en mémoire, telles que les optimiseurs matriciels de type Sinkhorn ou les hybrides SAGE/Sinkhorn, réduisent ce coût en supprimant les états de moment pour les matrices denses. Cependant, l'auteur démontre que l'application directe de ces méthodes sans état (stateless) aux experts routés de MoE échoue. Dans un cadre contrôlé de type « nanowhale » DeepSeek de 110 millions de paramètres, un hybride SAGE/Sinkhorn standard a considérablement réduit l'état de l'optimiseur (de 0,883 Go à 0,331 Go) mais a entraîné une dégradation sévère de la perte d'évaluation (3,8265 contre 3,58–3,64 pour la référence AdamW).

L'échec central est identifié comme un problème d'ordonnancement spécifique aux experts routés :

Les gradients routés sont conditionnels : Contrairement aux gradients des MLP denses, les gradients des experts routés sont des échantillons conditionnels basés sur les décisions du routeur et les poids de porte (gate weights), qui varient temporellement.
Non-linéarité de Sinkhorn : La normalisation Sinkhorn est une opération non linéaire. Appliquer une normalisation Sinkhorn à un gradient routé instantané et à haute variance ( $Sinkhorn(G_t)$ ) se comporte différemment de l'application à un signal moyenné temporellement ( $Sinkhorn(E[G_t])$ ).
Amplification du bruit : Le Sinkhorn sans état normalise le bruit induit par le routage à chaque étape, alors que l'échec des méthodes sans état provient de l'absence de lissage temporel avant que la normalisation de la matrice ne se produise.

2. Méthodologie : MESH et MESH-B

L'auteur propose MESH (Memory-Efficient Sinkhorn for Experts), un optimiseur hybride conçu spécifiquement pour les architectures MoE.

Mécanisme Central : Momentum Caché (Hidden Momentum)

MESH remplace la mise à jour Sinkhorn sans état pour les experts MoE par une approche de momentum caché : au lieu de stocker un tenseur explicite du premier moment dans l'état de l'optimiseur (ce qui annulerait l'objectif d'économie de mémoire), MESH maintient le premier moment au sein du cycle de vie du tampon de gradient :

Accumulation : Un accumulateur de momentum d'expert non normalisé est maintenu : $H_{e,t} = \beta_1 H_{e,t-1} + G_{e,t}$ .
Cycle de vie du tampon : L'accumulateur $H_{e,t}$ est stocké dans le tampon de gradient plutôt que dans le dictionnaire d'état de l'optimiseur.
Cycle de mise à jour : Avant la passe arrière (backward pass), le tampon est réattaché comme étant le gradient ( $p.grad \leftarrow \beta_1 H_{e,t}$ ). La passe arrière ajoute le gradient actuel, mettant ainsi à jour l'accumulateur pour l'étape suivante.
Normalisation : La mise à jour Sinkhorn est calculée sur le signal lissé : $U_{e,t} = Sinkhorn(H_{e,t})$ .

Cette approche restaure le signal du premier moment temporel requis pour les experts routés sans stocker de tenseur d'état d'optimiseur explicite pour eux.

Extension Optionnelle : MESH-B

Le papier évalue également MESH-B, qui ajoute un multiplicateur inverse-RMS par bloc/neurone grossier à la direction Sinkhorn. Cela agit comme un préconditionneur au niveau du bloc. Les statistiques de bloc sont calculées à partir de la source lissée ( $H_{e,t}$ ) plutôt qu'à partir d'un flux de gradient brut séparé afin de maintenir la simplicité de l'implémentation.

Stratégie d'Assignation de Rôle

MESH maintient une stratégie hybride pour différentes classes de paramètres :

Embeddings / LM Head : Gérés par SAGE (mises à jour à état léger).
Matrices Denses Non-Experts : Gérées par Sinkhorn sans état.
Experts Routés/Partagés : Gérés par MESH (Sinkhorn à momentum caché).
Normes / Biais : Gérés par SAGE.

3. Contributions Clés

Identification du Point de Défaillance : Le papier identifie les matrices d'experts MoE routés comme le point de défaillance spécifique de l'optimisation SAGE/Sinkhorn directe lors du pré-entraînement, les distinguant des matrices denses et des embeddings.
Algorithme MESH : Introduction d'une mise à jour Sinkhorn à momentum caché qui restaure un signal de premier moment pour les experts sans ajouter d'état d'optimiseur explicite, en utilisant le cycle de vie du tampon de gradient.
Analyse d'Ablation : Ablations négatives systématiques montrant que les interventions basées uniquement sur le second moment, le signe/momentum, la confiance ou l'échelle ne parviennent pas à expliquer l'écart de performance. L'ingrédient causal primaire est identifié comme étant le lissage temporel avant la normalisation de la matrice.
Évaluation de MESH-B : Évaluation d'une variante préconditionnée par bloc, montrant que bien qu'elle puisse améliorer la frontière mémoire-qualité dans certains cas spécifiques, elle n'est pas universellement nécessaire par rapport au lissage temporel fourni par MESH.

4. Résultats Expérimentaux

Les expériences ont été menées sur un modèle MoE de type DeepSeek de 110 millions de paramètres entraîné sur FineWeb-Edu pendant 5 000 étapes.

Efficacité Mémoire :
- AdamW : 0,883 Go d'état d'optimiseur.
- MESH / MESH-B : 0,331 Go d'état d'optimiseur (une réduction de 62,5 % par rapport à AdamW).
- Allocation CUDA de Pointe : Réduite d'environ 12,6 % (de ~4,37 Go à ~3,82 Go).
Performance (Perte d'Évaluation) :
- Référence AdamW : ~3,5877 (moyenne des graines 100/512).
- SAGE/Sinkhorn sans état : 3,8265 (Dégradation significative).
- MESH : 3,6430 (Proche d'AdamW, mais légèrement plus élevé).
- MESH-B : 3,6365 (Légèrement meilleur que MESH, mais reste derrière AdamW).
Résultats de Diagnostic (Graine 42) :
- Les variantes à état complet (MESH-FULL, MESH-B-FULL) qui stockent explicitement le premier moment récupèrent une qualité proche de AdamW (ex: MESH-B-FULL avec V(H) a atteint 3,6102) avec environ la moitié de l'état d'optimiseur d'AdamW.
- Cela confirme que le lissage temporel est l'ingrédient manquant crucial, tandis que l'implémentation du momentum caché réussit à l'approximer sans le coût d'un état complet.

5. Signification et Revendications

Le papier revendique modestement que MESH fournit un point de frontière mémoire-qualité substantiellement plus bas pour l'entraînement de MoE.

Conclusion Primaire : Les experts MoE nécessitent un lissage temporel avant la normalisation de la matrice pour fonctionner efficacement avec les optimiseurs de type Sinkhorn. Ce lissage est le facteur causal dominant pour combler l'écart de performance avec AdamW.
Conclusion Secondaire : Le préconditionnement par bloc/neurone (MESH-B) peut améliorer les résultats dans des contextes spécifiques, mais n'est pas établi comme une composante universellement nécessaire.
Limites : L'auteur reconnaît un écart de performance résiduel entre MESH/MESH-B et AdamW. Il suppose que cela est dû aux rôles non-experts (spécifiquement les matrices de vocabulaire et les matrices denses non-expertes), qui restent sur SAGE ou Sinkhorn sans état dans sa configuration, alors qu'AdamW les adapte avec des moments complets.
Portée : L'étude est intentionnellement étroite, se concentrant sur une architecture spécifique de 110M et des exécutions de 5 000 étapes. L'auteur ne prétend pas que MESH surpasse AdamW dans tous les indicateurs, mais qu'il récupère la majeure partie de la qualité perdue par les méthodes sans état tout en réduisant drastiquement la surcharge mémoire.

En résumé, MESH démontre que l'échec des optimiseurs économes en mémoire dans l'entraînement MoE n'est pas inhérent à la structure MoE elle-même, mais plutôt à l'absence de lissage temporel pour les gradients routés. En implémentant le "momentum caché", MESH résout ce problème d'ordonnancement, offrant une voie viable vers un pré-entraînement MoE efficace en mémoire.