Résumé Technique : FPTQuant : Transformations Préservant la Fonction pour la Quantification des LLM

Problématique

Les grands modèles de langage (LLM) engendrent des coûts de calcul et d'énergie significatifs lors de l'inférence, limitant leur déploiement sur les appareils de bord aux ressources contraintes. Bien que la quantification (la mise en correspondance de valeurs continues avec une grille entière uniforme) soit la technique la plus efficace pour améliorer l'efficacité, la quantification naïve des LLM entraîne souvent une dégradation sévère des performances. Cela est principalement dû aux « outliers » (valeurs aberrantes) — des activations et des poids ayant de grandes magnitudes qui perturbent la distribution uniforme requise pour une quantification à faible nombre de bits.

Le défi fondamental réside dans le compromis entre plage et précision :

1. Augmenter la plage pour capturer les outliers entraîne une perte de précision pour la majorité des valeurs regroupées autour de zéro.
2. Maintenir la précision en tronquant les outliers entraîne une perte d'information et un effondrement du modèle.

Les travaux antérieurs ont exploré les Transformations Préservant la Fonction (FPT — Function-Preserving Transforms) — des opérations telles que le redimensionnement ou les rotations qui lissent les outliers sans altérer la sortie du modèle en l'absence de quantification. Cependant, les méthodes existantes font souvent face à un compromis entre expressivité (la capacité à réduire efficacement l'erreur de quantification) et surcoût de calcul (coût d'inférence). De nombreuses transformations expressives sont non fusionnables, nécessitant de nouvelles opérations dans le graphe de calcul et entraînant une latence importante.

Méthodologie : FPTQuant

Le document présente FPTQuant, un cadre utilisant trois nouvelles FPT légères et expressives conçues spécifiquement pour les transformeurs. La philosophie centrale est d'exploiter les équivariances et les indépendances inhérentes aux opérations canoniques des transformeurs pour concevoir des transformations qui sont à la fois fusionnables (intégrables dans les poids existants) et expressives.

Les trois principales transformations sont :

1. Transformation Pre-RoPE Fusionnable pour les Queries et les Keys ($T_k, \bar{T}_k$)

• Contexte : Les LLM modernes utilisent les plongements positionnels rotatifs (RoPE), ce qui empêche la fusion naïve de transformations dans les poids de la requête ($W_q$) et de la clé ($W_k$) car le RoPE ne commute pas avec des transformations linéaires arbitraires.
• Mécanisme : Les auteurs introduisent une paire de transformations appliquées avant le RoPE. $T_k$ est appliquée aux clés, et $\bar{T}_k$ (une inverse approximative) est appliquée aux requêtes. Celles-ci consistent en des matrices de rotation $2 \times 2$ mises à l'échelle.
• Propriété : En raison de la structure diagonale par blocs du RoPE, ces matrices de rotation/redimensionnement spécifiques commutent avec l'opération RoPE. Cela permet de fusionner directement les transformations dans les poids $W_q$ et $W_k$ après l'entraînement, n'entraînant aucun surcoût d'inférence.
• Applicabilité : Valide pour les modèles avec RoPE (ex: série Llama) et les modèles sans plongements positionnels.

2. Transformation de Valeur Multi-tête Fusionnable ($T_v, \bar{T}_v$)

• Contexte : Dans le mécanisme d'attention, la multiplication matricielle par lots (BMM) opère sur les valeurs ($V$) et la projection de sortie ($W_o$). La dimension $d$ (dimension de la valeur) est indépendante du calcul de la probabilité d'attention.
• Mécanisme : Une transformation invertible $T_v$ est appliquée aux poids de projection de la valeur ($W_v$), et son inverse $\bar{T}_v$ est appliquée aux poids de la projection de sortie ($W_o$).
• Propriété : Puisque les différentes têtes d'attention sont indépendantes, une matrice inversible unique peut être optimisée pour chaque tête. Cette transformation est fusionnée dans $W_v$ et $W_o$.
• Avantage : Cela offre une haute expressivité (matrices inversibles complètes par tête) tout en restant fusionnable, contrairement aux méthodes de rotation précédentes qui partageaient souvent une seule rotation entre les têtes ou n'étaient pas fusionnables.

3. Mise à l'échelle Résiduelle Pseudo-Dynamique ($S_n$)

• Contexte : Les connexions résiduelles dans les transformeurs contournent souvent les couches de normalisation, ce qui conduit à des jetons ayant des échelles très différentes. Cela provoque des outliers dans les blocs suivants (ex: la sortie de SwiGLU).
• Mécanisme : Les auteurs proposent de déplacer l'RMSNorm pour qu'elle s'applique également au flux résiduel. Un facteur de mise à l'échelle par jeton $S_n$ (basé sur la racine carrée de la moyenne du carré — RMS — du résiduel) est appliqué à l'entrée résiduelle. Pour préserver la sortie, les sorties des blocs d'attention et MLP ($Y_n$) sont redimensionnées par la même valeur $S_n$ avant d'être ajoutées au résiduel.
• Propriété : La multiplication matricielle et les couches linéaires sans biais commutent avec la mise à l'échelle par jeton. Cela permet d'appliquer la mise à l'échelle profondément dans les blocs d'attention et de MLP pour réduire les outliers intra-blocs.
• Implémentation : Les échelles sont calculées de manière récursive en utilisant la sortie normalisée du bloc précédent. En pratique, la division finale par l'échelle est inutile car la tête du LM commence par un RMSNorm.

Stratégie d'Optimisation

FPTQuant emploie un processus d'optimisation en deux étapes :

1. Optimisation Locale : Les transformations sont optimisées localement et indépendamment pour minimiser la norme $L_p$ (spécifiquement $L_4$) des poids fusionnés. Cela réduit les outliers et fournit une initialisation stable.
2. Entraînement de Bout en Bout (Étudiant-Professeur) : Un modèle étudiant quantifié (avec les transformations) est entraîné pour approximer un modèle professeur en pleine précision en utilisant une perte de divergence de Jensen-Shannon (JSD).
   • Distinction Cruciale : Contraurs aux méthodes utilisant la perte de prédiction du prochain jeton, l'approche étudiant-professeur empêche le surapprentissage sur le jeu de données d'entraînement (ex: Wikitext) en faisant correspondre la distribution de probabilité complète du professeur. Cela améliore la généralisation aux tâches en aval.

Contributions Clés

1. Nouvelles FPT : Introduction de trois transformations spécifiques (Pre-RoPE, Valeur et Mise à l'échelle Résiduelle) qui sont conçues pour être à la fois hautement expressives et fusionnables, éliminant ainsi le surcoût d'inférence.
2. Quantification INT4 Statique : Démonstration que FPTQuant permet une quantification INT4 statique (poids, activations et KV cache) avec un surcoût minimal, atteignant jusqu'à 3,9× d'accélération par rapport au FP16 sur les tâches de prefill.
3. Compromis Précision-Vitesse : FPTQuant atteint des performances égales ou supérieures aux méthodes de pointe (SOTA) comme FlatQuant et SpinQuant, tout en étant nettement plus rapide (jusqu'à 29 % plus rapide que FlatQuant) et ne nécessitant aucun noyau (kernel) personnalisé.

Résultats Expérimentaux

Les auteurs ont évalué FPTQuant sur divers modèles (Llama-2/3, Ministral, Qwen) à travers plusieurs configurations de quantification (W4A4KV4, W4A8KV8, etc.).

• Vitesse : FPTQuant atteint systématiquement une accélération de prefill de 2,8× à 3,9× par rapport au FP16. Il est plus rapide que SpinQuant et systématiquement plus rapide que FlatQuant (accélération relative de 15 à 29 %).
• Précision :
  • Dans les configurations de déploiement "difficiles" (quantification de toutes les activations intermédiaires), FPTQuant surpasse les bases de référence comme SpinQuant et FlatQuant.
  • Dans les configurations standards (Linéaire+KV), FPTQuant est comparable à FlatQuant, qui est actuellement la SOTA pour la précision.
  • Sur le cas exigeant W4A4KV4, FPTQuant montre une perplexité et des scores de raisonnement zero-shot compétitifs, surpassant souvent SpinQuant et QuaRot.
• Stabilité : L'analyse de sensibilité montre que FPTQuant est une fonction préservée même sous des mises à jour de paramètres bruitées, alors que des méthodes comme FlatQuant (qui reposent sur des inversions en virgule flottante) et OSTQuant (non préservant la fonction) présentent une instabilité ou une dégradation avec le bruit.
• Coût d'Entraînement : L'entraînement de FPTQuant est efficace. Bien que légèrement plus lent que RTN-opt, il est généralement plus rapide que FlatQuant car ses transformations sont fusionnables, évitant les passages avant (forward passes) en ligne requis par les transformations non fusionnables.

Signification et Revendications

L'article affirme que FPTQuant représente une étape significative pour rendre les LLM viables pour le déploiement en périphérie (edge). En exploitant explicitement les propriétés mathématiques des opérations des transformeurs (équivariances), les auteurs ont créé des transformations qui ne sacrifient pas la vitesse d'inférence pour la précision.

• Impact Pratique : FPTQuant permet l'utilisation de largeurs de bits plus petites (INT4) sans la lourde pénalité de calcul de la quantification dynamique ou des transformations non fusionnables. Cela réduit l'empreinte mémoire et la consommation d'énergie des LLM, démocratisant potentiellement l'accès à l'IA sur les smartphones et les appareils de bord.
• Contribution Méthodologique : Ce travail souligne que la compréhension des propriétés de commutation des opérations existantes est la clé pour concevoir des stratégies de quantification efficaces. Il plaide contre l'utilisation de la perte de prédiction du prochain jeton pour l'entraînement des FPT en raison des risques de surapprentissage, préconisant plutôt l'entraînement étudiant-professeur par correspondance de distribution.
• Limites : Les auteurs notent modestement que bien que l'approche soit robuste à travers différentes familles de modèles, elle n'est pas garantie de se traduire de la même manière pour tous les LLM (ex: Mixture-of-Experts). De plus, les limitations matérielles actuelles (absence de support natif INT4 sur de nombreux GPU) signifient que les tests reposent sur l'émulation, et que les pleins bénéfices pourraient ne se réaliser que sur le futur matériel ou les accélérateurs spécifiques aux appareils de bord.