Accueil › Glossaire › Auxiliary Loss

Auxiliary Loss

L’auxiliary loss (perte auxiliaire) est un terme de régularisation ajouté à la fonction de perte principale d’un modèle Mixture-of-Experts (MoE) pour forcer une répartition équitable des tokens entre les différents experts, évitant ainsi l’effondrement du routage et le gaspillage de capacité.

Le problème que l’auxiliary loss résout

Dans un modèle MoE, un réseau de routage (le router) décide quels experts traitent chaque token. Sans contrainte, ce mécanisme converge naturellement vers un état où il sélectionne toujours les mêmes quelques experts pour chaque entrée. C’est une boucle auto-renforçante : un expert fréquemment activé reçoit plus de gradients, s’améliore plus vite, et devient donc encore plus attractif pour le routeur. Les experts restants stagnent et deviennent inutiles.

Ce phénomène s’appelle le routing collapse (effondrement du routage). Il a été observé dès les premiers travaux sur les MoE. Dans l’article fondateur de Jacobs et al. (1991), sur 8 experts entraînés, seuls 2 ou 3 finissaient réellement actifs. Pour un modèle sparse moderne qui déploie des dizaines voire des centaines d’experts dans le but de réduire le coût de calcul, c’est un problème critique : vous payez le stockage mémoire de 64 experts mais vous n’en utilisez que 5.

L’auxiliary loss intervient comme un signal d’entraînement supplémentaire qui pousse le routeur à distribuer les tokens de façon plus homogène entre tous les experts. Concrètement, elle pénalise le déséquilibre dans la répartition de la charge.

Formule mathématique

La formulation la plus répandue est celle introduite par le Switch Transformer (Fedus, Zoph et Shazeer, 2022), qui simplifie les deux pertes distinctes (load-balancing loss + importance-weighting loss) du papier original de Shazeer et al. (2017) en une seule expression.

Pour une couche MoE donnée avec N experts et un batch B contenant T tokens :

Plus précisément, f_i se calcule en comptant combien de tokens du batch ont eu l’expert i comme choix principal (argmax des probabilités), divisé par le nombre total de tokens :

f_i = (1/T) · Σ(x ∈ B) 𝟙{argmax p(x) = i}

Et P_i est la moyenne des probabilités assignées par le routeur à l’expert i sur l’ensemble du batch :

P_i = (1/T) · Σ(x ∈ B) p_i(x)

L’objectif idéal est que chaque vecteur f et P ait des valeurs uniformes de 1/N pour chaque expert. Dans ce cas, chaque expert traite le même nombre de tokens et reçoit la même confiance du routeur.

Pourquoi un produit scalaire ?

Le choix du produit f_i · P_i n’est pas anodin. Le vecteur f (tokens effectivement routés) n’est pas directement différentiable : c’est le résultat d’un argmax, une opération discrète. En revanche, P (probabilités moyennes) est différentiable puisqu’il passe par la fonction softmax du routeur. En multipliant les deux, la loss génère un gradient sur P qui est proportionnel à f : un expert qui reçoit beaucoup de tokens (f_i élevé) se verra appliquer un gradient plus fort pour réduire sa probabilité, ce qui redistribuera les tokens vers les experts sous-utilisés.

Le rôle critique de α

L’hyperparamètre α contrôle l’importance relative de l’auxiliary loss par rapport à la perte principale (cross-entropy pour un modèle de langage). C’est un réglage délicat :

Le Switch Transformer a testé des valeurs de α allant de 10⁻¹ à 10⁻⁵ par puissances de 10 et a retenu α = 10⁻² (soit 0,01) comme le meilleur compromis. Cette valeur permet un équilibrage rapide sans interférer significativement avec la convergence sur la tâche principale. La valeur optimale varie toutefois selon l’architecture, la taille du modèle et les données d’entraînement.

Fonctionnement dans la boucle d’entraînement

L’auxiliary loss s’intègre dans la perte totale du modèle de façon additive. Pour chaque couche MoE du réseau, on calcule la perte auxiliaire, puis on l’ajoute à la loss de la tâche :

L_total = L_tâche + Σ(couches MoE) L_aux

En pratique, dans des frameworks comme HuggingFace Transformers, le paramètre aux_loss est exposé et la loss est calculée automatiquement pour chaque couche Switch lors du forward pass. Voici un exemple simplifié en PyTorch illustrant le calcul :

# Calcul de l'auxiliary loss pour une couche MoE
# router_probs : [batch_size, seq_len, num_experts] - probabilités du routeur
# expert_indices : [batch_size, seq_len] - expert choisi par token

import torch
import torch.nn.functional as F

num_experts = 8
alpha = 0.01

# f_i : fraction de tokens routés vers chaque expert
expert_mask = F.one_hot(expert_indices, num_classes=num_experts).float()
f = expert_mask.mean(dim=[0, 1])  # [num_experts]

# P_i : probabilité moyenne allouée à chaque expert
P = router_probs.mean(dim=[0, 1])  # [num_experts]

# Auxiliary loss
aux_loss = alpha * num_experts * torch.sum(f * P)

# Perte totale
total_loss = cross_entropy_loss + aux_loss

Le calcul est léger en termes de compute : il s’agit essentiellement de moyennes et d’un produit scalaire. Le surcoût est négligeable comparé au coût des couches MoE elles-mêmes.

Les variantes de l’auxiliary loss

La loss d’importance (Shazeer et al., 2017)

Le papier original « Outrageously Large Neural Networks » de Shazeer et al. (2017) utilisait en réalité deux pertes auxiliaires distinctes. La première, la load-balancing loss, pénalisait le déséquilibre du nombre de tokens par expert. La seconde, l’importance loss, ciblait la variance du « score d’importance » de chaque expert (la somme des poids de gating sur un batch). Le Switch Transformer a ensuite unifié ces deux pertes en une seule formulation plus simple, qui est devenue le standard de facto.

La Router Z-Loss (ST-MoE)

Proposée dans le papier ST-MoE (Zoph et al., 2022), la router z-loss est une perte auxiliaire complémentaire qui ne cible pas directement l’équilibrage mais la stabilité numérique. Elle contraint la magnitude des logits (avant softmax) produits par le routeur. Sans cette contrainte, ces logits peuvent devenir très grands pendant l’entraînement, provoquant des erreurs d’arrondi dans la fonction exponentielle du softmax, même en précision float32. La formulation est :

L_z = (1/T) · Σ(x ∈ B) (log Σ(i=1 à N) exp(z_i(x)))²

Le coefficient de pondération recommandé pour la z-loss est environ 10⁻³ (un ordre de grandeur plus faible que la load-balancing loss). En pratique, les implémentations modernes utilisent souvent les deux losses conjointement : la load-balancing loss pour l’équilibre, et la z-loss pour la stabilité.

Pertes au niveau device (DeepSeek-V2)

DeepSeek-V2 a introduit des pertes auxiliaires supplémentaires spécifiques à l’entraînement distribué. Quand les experts sont répartis sur différents GPU (ce qu’on appelle l’expert parallelism), il ne suffit pas d’équilibrer les tokens par expert : il faut aussi équilibrer la charge par device. DeepSeek-V2 ajoute donc deux termes supplémentaires qui encouragent un routage équilibré au niveau des nœuds de calcul, garantissant que la communication inter-GPU et la computation restent synchronisées.

Limites de l’auxiliary loss classique

Malgré son efficacité, l’auxiliary loss classique présente plusieurs problèmes reconnus dans la littérature :

Interférence des gradients. L’auxiliary loss ajoute un signal de gradient qui pousse le routeur à distribuer uniformément, même quand la distribution optimale pour la tâche serait non-uniforme. Certains tokens complexes bénéficieraient d’être envoyés vers des experts spécialisés, mais la perte auxiliaire force une dispersion artificielle.

Sensibilité à α. La valeur optimale de α n’est pas universelle. Elle dépend de la taille du modèle, du nombre d’experts, du dataset, et même de la phase d’entraînement. Des travaux récents suggèrent que la valeur idéale évolue au cours de l’entraînement : plus forte au début (pour éviter l’effondrement précoce), puis réduite progressivement (pour laisser le modèle optimiser la tâche).

Réduction de la spécialisation. L’un des avantages clés du MoE est la spécialisation des experts. Une perte auxiliaire trop forte pousse chaque expert à traiter un nombre identique de tokens, ce qui peut diluer cette spécialisation et réduire les bénéfices de l’architecture.

Problème résiduel avec le capacity factor. L’auxiliary loss ne garantit pas à elle seule un équilibrage parfait. En pratique, les implémentations combinent souvent la loss avec un mécanisme de capacité maximale par expert. Les tokens excédentaires sont soit envoyés vers un autre expert, soit « droppés » (ignorés et passés à la couche suivante via la connexion résiduelle). Ce token dropping dégrade la qualité du modèle.

L’alternative de DeepSeek : auxiliary-loss-free balancing

Face à ces limites, DeepSeek a proposé une approche radicalement différente dans un article dédié (Wang et al., 2024), ensuite déployée dans DeepSeek-V3. L’idée : supprimer complètement la perte auxiliaire et utiliser un mécanisme de biais dynamique pour équilibrer la charge.

Le mécanisme de biais dynamique

Au lieu d’ajouter une loss, DeepSeek ajoute un terme de biais b_i aux scores d’affinité de chaque expert lors de la sélection top-K. Ce biais est ajusté dynamiquement à chaque pas d’entraînement :

C’est la distinction clé : dans l’approche classique, l’auxiliary loss modifie les gradients de la backpropagation. Dans l’approche DeepSeek, l’ajustement du biais opère en dehors du graphe de calcul des gradients. Le modèle optimise sa tâche principale sans perturbation, tandis qu’un mécanisme séparé s’occupe de l’équilibrage.

Résultats expérimentaux

Les ablations publiées dans le rapport technique de DeepSeek-V3 montrent que la stratégie auxiliary-loss-free surpasse systématiquement l’approche classique à la fois sur des modèles de 15,7 milliards et de 228,7 milliards de paramètres totaux. La méthode a été validée sur le même jeu de données et la même architecture, seule la stratégie d’équilibrage changeant.

Fait notable : malgré le nom « auxiliary-loss-free », DeepSeek-V3 conserve en réalité une petite perte auxiliaire complémentaire au niveau séquence (avec un coefficient très faible) pour gérer les cas extrêmes où une séquence individuelle favoriserait massivement certains experts. L’équilibrage au niveau batch ne garantit pas l’équilibre au niveau de chaque séquence, surtout en inférence où le contenu des séquences est imprévisible.

Quels modèles utilisent l’auxiliary loss ?

Le cas de Mixtral est intéressant : Mistral AI a montré qu’avec un routage top-2, une fonction softmax standard et des données suffisamment diverses, l’équilibrage peut émerger naturellement sans aucune perte auxiliaire explicite. Cela simplifie l’implémentation mais ne fonctionne pas systématiquement sur toutes les configurations.

Implémentation pratique

Dans HuggingFace Transformers

La bibliothèque Transformers de HuggingFace expose le paramètre router_aux_loss_coef dans la configuration des modèles MoE comme le Switch Transformer. Ce paramètre correspond directement au coefficient α. La loss est calculée automatiquement dans le forward pass et ajoutée au output du modèle :

from transformers import SwitchTransformersConfig

config = SwitchTransformersConfig(
    num_experts=8,
    router_aux_loss_coef=0.01,  # α = 10⁻²
    expert_capacity=1.25,       # capacity factor
)

# Lors de l'entraînement, model(**inputs) retourne
# un objet avec .loss qui inclut déjà la perte auxiliaire

Conseils de réglage

Commencer par α = 0,01. C’est la valeur validée par le Switch Transformer et utilisée dans de nombreuses implémentations de référence. C’est votre point de départ.

Monitorer la distribution des tokens par expert. Utilisez TensorBoard ou Weights & Biases pour tracer, à chaque étape d’entraînement, le nombre de tokens traités par chaque expert. Si certains experts reçoivent systématiquement moins de tokens que la moyenne, augmentez α progressivement.

Surveiller la loss principale en parallèle. Si la validation loss stagne ou remonte alors que l’équilibrage s’améliore, vous avez poussé α trop loin.

Combiner avec la router z-loss. Si vous observez de l’instabilité en entraînement (spikes de loss), ajoutez une z-loss avec un coefficient d’environ 10⁻³. La z-loss cible un problème différent (magnitude des logits) et ne crée pas de redondance avec la load-balancing loss.

Envisager le capacity factor. Un capacity factor de 1,25 en entraînement (2,0 en évaluation) est un bon point de départ. Ce mécanisme de hard cap complète l’auxiliary loss en imposant une limite physique au nombre de tokens par expert.

Auxiliary loss et fine-tuning : un cas particulier

Le fine-tuning des modèles MoE pose un défi spécifique. Les travaux sur ST-MoE ont montré que les modèles sparse bénéficient de taux d’apprentissage plus élevés et de batch sizes plus petits lors du fine-tuning, deux choix qui augmentent le bruit et améliorent la généralisation.

Fait surprenant : on peut souvent désactiver l’auxiliary loss pendant le fine-tuning sans dégradation notable. Les modèles sparse sont robustes au dropping de 10 à 15 % des tokens à cette étape. L’explication probable est que la distribution apprise pendant le pré-entraînement est suffisamment stable pour que le routeur ne s’effondre pas sur les quelques milliers de steps du fine-tuning.

En revanche, d’autres travaux (MoEs Meets Instruction Tuning, 2023) montrent que lors d’un instruction tuning, le fait de conserver l’auxiliary loss est bénéfique car il prévient le surapprentissage sur les données d’instruction, qui sont typiquement moins volumineuses que les données de pré-entraînement.

Verdict

L’auxiliary loss reste, à ce jour, la technique standard pour l’équilibrage des modèles MoE. Elle est simple à implémenter, bien comprise théoriquement, et utilisée dans la majorité des architectures MoE en production. Son principal défaut, l’interférence des gradients avec la tâche principale, est réel mais gérable avec un réglage soigné de α.

L’approche auxiliary-loss-free de DeepSeek est une avancée significative qui devrait se généraliser dans les futures architectures. Elle élimine le compromis performance/équilibrage en séparant complètement les deux mécanismes. Mais elle n’est pas encore le standard par défaut : la plupart des frameworks et tutoriels utilisent encore l’auxiliary loss classique, et son implémentation est plus directe.

Si vous concevez un modèle MoE aujourd’hui : commencez par l’auxiliary loss classique avec α = 0,01. Si vous observez une dégradation de performance liée à l’équilibrage, passez à une stratégie de biais dynamique inspirée de DeepSeek. Et dans tous les cas, monitorez activement la distribution des tokens par expert tout au long de l’entraînement.