Résumé technique : Compression de poids compensée par le mouvement (MCWC)

Énoncé du problème

Le déploiement des réseaux de neurones est de plus en plus limité non pas par la puissance de calcul, mais par le stockage, la bande passante et la surcharge de démarrage à froid nécessaires pour mettre en place les points de contrôle dans les accélérateurs. Les pipelines de compression existants, dominés par la quantification et l'élagage post-entraînement, traitent généralement chaque couche comme un objet indépendant. Cette approche néglige une redondance inter-couche substantielle induite par des symétries préservant la fonction (par exemple, la capacité de réindexer les unités cachées ou les têtes d'attention sans modifier la fonction entrée-sortie du modèle). Par conséquent, ces méthodes échouent à exploiter la forte similarité et prévisibilité inter-couche inhérentes aux architectures modernes comme les Transformers.

Méthodologie

MCWC propose un codec dédié aux poids qui traite la profondeur du modèle comme un signal séquentiel, analogue aux trames vidéo. L'hypothèse de base est que, en éliminant le « mouvement » induit par les symétries de permutation, l'évolution restante des poids à travers les couches devient hautement prévisible, permettant un codage résiduel efficace. Le pipeline se compose de quatre étapes principales :

1. Compensation de mouvement par alignement fonctionnel

L'encodeur identifie les permutations préservant la fonction pour des blocs de poids (par exemple, unités cachées des FFN, têtes d'attention) qui maximisent la correspondance inter-couche.

• Mécanisme : Pour chaque couche $\ell$, la méthode sélectionne une permutation $\Pi_\ell$ qui aligne les blocs avec leurs homologues de la couche précédente $\ell-1$. Cela est réalisé en résolvant un problème d'affectation linéaire pour maximiser la similarité (un hybride de similarité cosinus basée sur les poids et de similarité basée sur les statistiques d'activation) entre les blocs candidats et la référence précédemment alignée.
• Résultat : Cela transforme la séquence de poids en un « système de coordonnées aligné » où les couches adjacentes sont géométriquement comparables, réduisant considérablement l'énergie des différences entre elles.

2. Codage prédictif séquentiel par couche

Une fois alignées, les poids sont traités comme une séquence temporelle.

• Prédiction : Un prédicteur léger et séquentiel par couche ($g_\theta$) estime les blocs alignés de la couche actuelle en fonction du contexte précédemment décodé (blocs alignés de la couche $\ell-1$), conditionné par des embeddings de couche et de type de bloc.
• Résidus : Le codec encode uniquement les résidus de prédiction quantifiés ($R = \text{CibleAlignée} - \text{Prédiction}$).
• Images clés : Pour prévenir l'accumulation d'erreurs (dérive) sur de longues séquences, des « images clés » périodiques sont encodées sous forme de valeurs absolues plutôt que de résidus. Elles agissent comme des points de synchronisation, permettant au décodeur de réinitialiser son contexte.

3. Quantification apprise et modélisation de l'entropie

• Quantification : Les résidus sont mappés vers des indices de code discrets à l'aide d'un quantificateur scalaire appris avec des pas et des moyennes par canal.
• Codage entropique : Un modèle d'entropie appris attribue des probabilités aux symboles quantifiés, permettant un codage arithmétique. Le modèle est conditionné par l'index de couche, le type de bloc et les statistiques locales dérivées de la sortie du prédicteur.
• Entraînement : L'ensemble du système (prédicteur, quantificateur, modèle d'entropie) est entraîné de bout en bout sous un objectif taux-distorsion ($L = L_{task} + \lambda E[-\log p_\psi(c)]$), où $L_{task}$ est une perte de distillation assurant que les poids reconstruits maintiennent les performances de la tâche.

4. Informations secondaires et décodage

Le flux de bits comprend les résidus/images clés codés par entropie, les informations de permutation secondaires (encodées via des codes de Lehmer codés en delta pour réduire la surcharge) et les paramètres du quantificateur. Le décodeur inverse le processus : il décode par entropie, déquantifie, reconstruit les résidus par prédiction et applique les permutations inverses pour restaurer la structure du tenseur original en vue du déploiement.

Contributions clés

1. Compensation de mouvement dans l'espace des poids : L'introduction de la sélection de permutations par blocs pour augmenter la correspondance inter-couche tout en maintenant la fonction représentée, analogue à l'estimation de mouvement dans la compression vidéo.
2. Codec prédictif pour la profondeur : Une architecture novatrice traitant la profondeur comme un axe temporel, utilisant des images clés explicites et un codage résiduel pour transformer les couches alignées en « trames » prévisibles.
3. Entraînement taux-distorsion de bout en bout : Entraînement conjoint des modèles de quantification et d'entropie pour garantir que le flux de résidus discret est à la fois précis après décodage et de faible entropie sous le modèle de probabilité appris.
4. Validation empirique : Démonstration de frontières de Pareto taux-précision améliorées sur les modèles de langage Transformer (Pythia) et les classificateurs de vision (ViT), avec des études d'ablation confirmant que l'alignement, la prédiction et la modélisation de l'entropie sont tous nécessaires pour obtenir les gains complets.

Résultats

MCWC a été évalué sur Pythia-1.4B (modélisation du langage) et ViT-B/16 (classification de vision) par rapport à des bases de référence solides incluant la quantification uniforme INT8/INT4, SmoothQuant, GPTQ, AWQ et la quantification scalaire apprise.

• Modélisation du langage : À des taux de compression agressifs (par exemple, 2,8 bits/paramètre), MCWC atteint une perplexité de 16,40, surpassant Élagage+Quantification (18,10) et la quantification scalaire apprise (17,30). Dans le régime de taux moyen (4,2 bits/paramètre), il atteint une perplexité de 15,95, dépassant l'INT4 uniforme (16,85) et restant compétitif avec des méthodes sensibles à l'Hessien comme GPTQ (16,40) et AWQ (16,25).
• Classification de vision : À 2,8 bits/paramètre, MCWC conserve une précision top-1 de 80,15 % sur ImageNet-1k, surpassant Élagage+Quantification (77,60 %) et la quantification scalaire apprise (78,70 %). À 4,2 bits/paramètre, il atteint une précision de 80,92 %, dépassant l'INT4 uniforme (79,35 %).
• Efficacité de décodage : Les temps de décodage restent compétitifs par rapport aux bases de référence de quantification standard (par exemple, ~2,5 s pour Pythia-1.4B sur CPU), car la reconstruction prédictive et la permutation inverse sont des opérations de temps linéaire, le coût dominant étant le décodage entropique et la matérialisation des tenseurs.
• Diagnostics : Les études d'ablation confirment que l'alignement fonctionnel augmente la similarité cosinus entre couches adjacentes (de ~0,28 à ~0,67 pour Pythia-1.4B) et réduit l'énergie résiduelle normalisée, validant que l'alignement expose un système de coordonnées plus prévisible.

Importance et revendications

L'article positionne MCWC non pas comme un remplacement des noyaux d'inférence de faible précision (par exemple, multiplication matricielle INT8/INT4), mais comme un codec spécialisé de stockage et de transport de points de contrôle. Sa valeur principale réside dans la réduction de l'encombrement de stockage, de la bande passante réseau et de la latence de démarrage à froid associées à la distribution et à la matérialisation répétée de grands poids de modèles.

Les auteurs affirment que, en exploitant la redondance inter-couche par compensation de mouvement, MCWC adresse un goulot d'étranglement orthogonal à la quantification traditionnelle. La méthode est la plus efficace dans les régimes de déploiement où les points de contrôle sont déplacés, stockés ou répliqués de nombreuses fois (par exemple, démarrages à froid serverless, registres de modèles, distribution en périphérie), permettant d'amortir le coût d'encodage hors ligne unique sur des économies de déploiement répétées. Le travail suggère que traiter les poids neuronaux comme des trajectoires structurées à travers la profondeur, plutôt que comme des tenseurs isolés, offre une nouvelle voie pour une compression efficace des modèles.