Résumé Technique : CLHA – Agrégation de Hessian Cross-Layer pour la Quantification des Transformers à Partage de Poids de type Mixture-of-Experts

1. Énoncé du Problème

Les architectures récentes de Transformers à Mixture-of-Experts (MoE) efficaces, telles que celles de Megrez2, emploient le Partage de Poids Cross-Layer (CLWS). Dans cette conception, des couches adjacentes de transformeurs (par exemple, $\ell$ et $\ell+1$) réutilisent les mêmes tenseurs de poids d'experts routés afin de réduire de moitié l'empreinte des paramètres tout en maintenant le même nombre de paramètres activés par jeton (token).

Cependant, les méthodes existantes de Quantification Post-Entraînement (PTQ), y compris les approches de type GPTQ, traitent ces poids partagés comme s'ils appartenaient à une seule couche. Elles calibrent la quantification en utilisant les statistiques d'activation d'une seule couche « primaire ». Cela crée un décalage fondamental :

• Le flux résiduel (residual stream) évolue entre les couches, ce qui signifie que les distributions de jetons activant l'expert partagé à la couche $\ell_a$ diffèrent de celles de la couche $\ell_b$.
• Les méthodes existantes ignorent les statistiques d'activation de la couche secondaire, ce qui conduit à une minimisation sous-optimale de l'erreur de reconstruction, particulièrement dans les régimes de faible nombre de bits (par exemple, 2 bits).

2. Méthodologie

2.1 Fondement Théorique : Agrégation de Hessian Cross-Layer

Le papier formalise l'objectif de la quantification d'un poids $W$ partagé entre les couches $\ell_1, \dots, \ell_n$. L'objectif correct est de minimiser l'erreur de reconstruction totale à travers toutes les couches utilisant ce poids :
$$L(W_q) = \sum_{i} \alpha_i \text{tr}[(W_q - W)H_{\ell_i}(W_q - W)^\top]$$
où $H_{\ell_i}$ est la Hessian (covariance d'activation) de la couche $\ell_i$, et $\alpha_i$ représente la fréquence de routage (proportion de jetons) pour cette couche.

La Proposition 1 établit que la minimisation de cette erreur combinée est mathématiquement équivalente à l'application de GPTQ avec une Hessian combinée :
$$H_{comb} = \sum_{i} \alpha_i H_{\ell_i}$$
Cela contraste avec la Calibration Indépendante par Couche (PLIC), qui minimise l'erreur en utilisant uniquement la Hessian d'une seule couche ($H_{\ell_a}$).

2.2 Algorithme CLHA

Les auteurs proposent l'algorithme CLHA (Cross-Layer Hessian Aggregation), un algorithme de calibration en trois passes :

1. Collecte : Pour chaque séquence de calibration, accumuler les Hessians par couche et les comptes de jetons pour chaque expert actif.
2. Agrégation : Pour chaque paire d'experts partagés $(\ell_a, \ell_b)$, calculer une Hessian combinée. Le papier évalue deux stratégies de pondération :
   • Pondération par Fréquence de Routage ($\alpha = n_a / (n_a + n_b)$) : Pondère les Hessians par le nombre réel de jetons reçus par chaque couche.
   • Pondération Égale ($\alpha = 0.5$) : Agit comme un régularisateur, utile lorsque les données de calibration sont limitées ou que les estimations de la Hessian de la couche secondaire sont bruitées.
3. Quantification : Appliquer les mises à jour séquentielles par colonnes de GPTQ standard en utilisant la $H_{comb}$ agrégée.

2.3 CLHA-MP (Précision Mixte)

Reconnaissant que les experts partagés toujours actifs reçoivent la totalité des $N$ jetons (contrairement aux experts routés qui ne reçoivent qu'une fraction $K/E$), le papier note que leur sensibilité à la quantification est $E/K$ fois plus élevée. CLHA-MP alloue $b+1$ bits à ces experts partagés tout en maintenant les experts routés à $b$ bits, ce qui n'entraîne qu'un surcoût de stockage minimal (10–20 % des paramètres d'experts).

2.4 Correction d'Implémentation : Ambiguïté des Hooks

Le papier identifie un piège d'ingénierie critique : les hooks (crochets) de propagation directe (forward) standards de PyTorch attachés à un module ExpertFFN partagé s'activent pour n'importe quelle couche appelant ce module. Cela provoque le mélange des activations de différentes couches dans un seul accumulateur de Hessian, corrompant l'attribution par couche.

• Solution : Les auteurs proposent un mécanisme de rappel (callback) sensible à la couche, transmis via la signature MoELayer.forward(). Cela fournit explicitement les indices de couche au rappel, garantissant que des Hessians distinctes sont accumulées pour $H_a$ et $H_b$, même lorsque le module de poids sous-jacent est partagé.

3. Contributions Clés

1. Proposition Théorique : Formalisation du fait que la Hessian optimale pour les poids partagés est la somme pondérée par la fréquence de routage des Hessians par couche.
2. Algorithme CLHA : Une implémentation pratique de la calibration par Hessian agrégée au sein du cadre GPTQ.
3. CLHA-MP : Une extension en précision mixte allouant un bit supplémentaire aux experts partagés hautement sensibles.
4. Correction d'Implémentation : Identification et résolution du problème d'« ambiguïté des hooks » qui corrompt silencieusement la calibration dans les architectures à modules partagés.
5. Validation Empirique : Démonstration que les gains de CLHA sont corrélés à la divergence de routage entre les couches.

4. Résultats Expérimentaux

Les expériences ont été menées sur un modèle CLWS-MoE de 34,5 millions de paramètres entraîné de zéro sur WikiText-2.

• Performance à 2 bits : C'est le régime critique où le mauvais calibrage s'amplifie.
  • CLHA-rt (pondéré par le routage) a réduit la perplexité (PPL) de 4,3 % par rapport à PLIC.
  • CLHA-eq (pondération égale) a réduit la PPL de 5,4 %.
  • CLHA-MP a obtenu une amélioration de 18,6 % par rapport à PLIC.
  • À 4 bits et 3 bits, les différences étaient négligeables (<1 %).
• Calibration vs Erreur de Test : Curieusement, PLIC a obtenu une erreur de Hessian de calibration (CHE) plus faible sur l'ensemble d'entraînement car il optimise directement pour cette distribution. Cependant, CLHA a mieux généralisé sur l'ensemble de test, indiquant que la Hessian combinée capture mieux la véritable distribution du poids partagé qu'une vue d'une seule couche.
• Expérience de Netteté de Routage (Routing Sharpness) : En accentuant artificiellement les logits du routeur (augmentant la divergence entre les couches), les auteurs ont observé :
  • La performance de PLIC se dégrada de 5 % à mesure que la divergence augmentait.
  • CLHA est resté stable (<1 % de dégradation).
  • L'écart de performance entre CLHA et PLIC est passé de +4,3 % (routage naturel) à +8,0 % (haute divergence), confirmant la prédiction théorique que les gains de CLHA croissent avec la divergence cross-layer.

5. Signification et Revendications

Le papier affirme fournir la première analyse de PTQ spécifiquement pour le partage de poids cross-layer dans les modèles MoE. Sa signification réside dans :

• Correction de l'Objectif : Démontrer que traiter les poids partagés comme des entités de couche unique est mathématiquement sous-optimal et que l'agrégation de Hessian est la correction nécessaire.
• Évolutivité des Gains : Montrer qu'à mesure que les modèles deviennent plus grands et que les schémas de routage plus spécialisés (augmentant la divergence cross-layer), les bénéfices de CLHA augmenteront probablement, ce qui en fait essentiel pour les futurs modèles CLWS-MoE en production (ex: Megrez2, DeepSeek, Qwen).
• Nécessité Pratique : Souligner que l'application naïve des outils PTQ existants aux architectures CLWS est fondamentalement erronée en raison de l'ambiguïté des hooks, et fournir la correction nécessaire pour toute méthode de calibration future dans ce domaine.

Les auteurs restent modestes quant à la performance absolue, notant que leur modèle est petit (34,5M) et entraîné pour un nombre limité d'étapes, mais soulignent que les améliorations relatives par rapport à la calibration standard sont robustes et théoriquement fondées.