Résumé technique : LAMP (Inférence à Précision Mixte Anticipée des Grands Modèles de Langage)

Énoncé du problème

Les stratégies d'inférence à précision mixte actuelles pour les grands modèles de langage (LLM) basés sur les Transformers reposent principalement sur une quantification uniforme des entrées et une accumulation de haute précision des sorties. Bien que efficaces pour les produits matriciels, ces approches traitent uniformément la sortie des calculs intermédiaires, ignorant la nature compositionnelle des réseaux de neurones profonds.

Du point de vue de l'analyse numérique, l'évaluation en virgule flottante (FP) introduit des erreurs d'arrondi à chaque étape. Dans les compositions profondes $f(g(x))$, les erreurs d'arrondi locales se propagent et peuvent être amplifiées par les opérateurs subséquents. Les bornes d'erreur au pire cas existantes pour les réseaux de neurones profonds (DNN) croissent exponentiellement avec la profondeur, pourtant les méthodes standard à précision mixte ne tiennent pas compte des composants spécifiques d'un résultat intermédiaire les plus susceptibles d'être amplifiés par la couche suivante. L'article identifie un écart dans les stratégies d'« accumulation à précision mixte » qui seraient explicitement conscientes des effets compositionnels, visant à contrôler la précision des calculs plutôt qu'à simplement réduire l'empreinte mémoire via la quantification.

Méthodologie : Précision Mixte Anticipée (LAMP)

La contribution centrale est le cadre LAMP, une stratégie adaptative qui sélectionne un sous-ensemble épars de composants au sein d'une fonction intermédiaire $g(x)$ pour être recalculés avec une plus grande précision, tandis que le reste est calculé avec une précision inférieure. Cette sélection repose sur une analyse « anticipée » de la manière dont les erreurs d'arrondi dans $g(x)$ se propageront à travers l'opérateur subséquent $f$.

Fondement théorique

La méthode s'appuie sur une analyse affinée de l'erreur d'arrondi de la composition $f(g(x))$.

1. Dérivation de la borne d'erreur : Les auteurs dérivent une borne pour l'erreur relative de la composition. L'erreur dépend du nombre de conditionnement de la fonction externe $f$ et des erreurs d'arrondi introduites par la fonction interne $g$.
2. Formulation de l'optimisation : L'objectif est de minimiser le nombre de recalculs en haute précision (éparsité) tout en garantissant que l'erreur propagée reste en dessous d'un seuil $\tau$. Cela se formule comme un problème d'optimisation :
   $$\min \|q\|_0 \quad \text{s.t.} \quad \kappa(q) \leq \tau$$
   où $q \in \{0, 1\}^n$ est un vecteur binaire indiquant quels composants de $g(x)$ nécessitent une précision supérieure, et $\kappa(q)$ représente le nombre de conditionnement composant par composant ou normique de la composition.
3. Solutions analytiques pour les Transformers : Une insight théorique clé est que pour des non-linéarités spécifiques des Transformers, le problème d'optimisation admet des solutions analytiques, surmontant la difficulté combinatoire du cas général :
   • Fonctions d'activation : La Jacobienne est diagonale, permettant l'identification immédiate des composants sensibles.
   • Normalisation de couche : Les auteurs prouvent qu'une solution presque la plus éparse possible peut être trouvée de manière gloutonne en triant les magnitudes des entrées.
   • Softmax (Attention) : Pour le mécanisme d'attention, l'approche composant par composant est instable pour des entrées fortement négatives. Les auteurs pivotent vers un objectif $\ell_1$ normique, qui borne le décalage total de la masse de probabilité. La condition optimale pour recalculer un produit scalaire Key-Query (KQ) $y_j$ est dérivée comme suit :
     $$2z_j(1 - z_j)|y_j| > \tau$$
     où $z_j$ est la probabilité softmax résultante. Cela implique que les produits scalaires conduisant à des résultats « confus » (également probables) ou à des motifs de sensibilité spécifiques nécessitent une plus grande précision.

Stratégie d'implémentation

L'algorithme LAMP opère en deux étapes :

1. Évaluation de base : Calculer $g(x)$ en basse précision.
2. Sélection anticipée : Estimer le nombre de conditionnement $\kappa$ basé sur la sortie de base. Résoudre le problème d'optimisation pour identifier l'ensemble éparse d'indices $q$ nécessitant un recalcul.
3. Recalcul adaptatif : Recalculer uniquement les composants sélectionnés de $g(x)$ en plus haute précision (par exemple FP32) avant de les transmettre à $f$.

Contributions clés

1. Nouveau cadre de calcul : L'article introduit LAMP, une approche mathématiquement fondée pour l'inférence à précision mixte qui cible l'accumulation des produits matriciels basée sur la propagation d'erreurs compositionnelles, complétant les stratégies existantes de quantification des opérandes.
2. Faisabilité théorique pour les Transformers : Les auteurs démontrent que le problème combinatoire complexe de sélection des composants à recalculer se simplifie en solutions analytiques pour les blocs de construction standard des Transformers (activations, normalisation de couche, softmax).
3. Preuve de concept : La méthode est validée numériquement sur des modèles GPT-2, ciblant spécifiquement les produits scalaires KQ dans le mécanisme d'attention.

Résultats expérimentaux

Les auteurs ont évalué LAMP sur GPT-2 XL en utilisant le jeu de données OpenWebText, simulant des formats de virgule flottante personnalisés (PS($\mu$)) avec des largeurs de mantisse variables ($\mu$).

• Gains de précision : La méthode atteint des améliorations allant jusqu'à deux ordres de grandeur en précision (mesurée par la divergence de Kullback-Leibler par rapport à une référence FP32 uniforme) avec des taux de recalcul très faibles.
  • Par exemple, avec $\mu=4$ (très basse précision), un taux de recalcul de seulement 0,3 % a réduit la divergence KL de 12 fois.
  • À 1,6 % de recalcul, la réduction était de 83 fois.
  • À 7,6 % de recalcul, la réduction était de 385 fois.
• Sélectivité : Les résultats confirment que la sélection adaptative est cruciale ; un recalcul aléatoire du même nombre de composants ne produit aucune amélioration significative de la précision.
• LAMP assoupli : Pour répondre à la compatibilité avec des noyaux d'inférence efficaces comme FlashAttention (qui nécessitent un softmax en un seul passage), les auteurs proposent un LAMP à « seuil relatif assoupli ». Cette variante supprime la dépendance à la constante de normalisation, ne montrant qu'une perte marginale de performance par rapport à la solution optimale « stricte ».

Importance et affirmations

L'article positionne LAMP comme une preuve de concept et un fondement théorique plutôt que comme un noyau de production entièrement déployé.

• Complémentaire à la quantification : Les auteurs déclarent explicitement que LAMP ne remplace pas la quantification (il ne réduit pas l'empreinte mémoire ni ne modifie les opérandes) mais contrôle la précision computationnelle. Il est conçu pour fonctionner parallèlement aux techniques de quantification.
• Implications matérielles : Les expériences actuelles simulent une accumulation en basse précision. Les auteurs notent que le matériel actuel ne prend pas nativement en charge l'accumulation de produits matriciels en basse précision compactée. Par conséquent, LAMP sert de « blueprint » pour les accélérateurs d'IA de nouvelle génération qui pourraient prendre en charge nativement de telles opérations.
• Limitations :
  • L'algorithme LAMP « strict » pour le softmax est incompatible avec FlashAttention, standard de l'industrie, car il nécessite la matérialisation du vecteur complet des probabilités softmax.
  • La méthode est actuellement validée sur des modèles GPT-2 ; bien que les auteurs soutiennent que les effets numériques s'échelle avec la profondeur du modèle et la concentration des probabilités, la validation sur des modèles plus grands est une direction future.
  • L'accent est strictement mis sur l'inférence ; l'application de LAMP au passage arrière (entraînement) est identifiée comme un travail futur.

En résumé, l'article soutient qu'en analysant mathématiquement la manière dont les erreurs d'arrondi se propagent à travers les compositions de Transformers, on peut atteindre une précision proche du FP32 avec une précision significativement inférieure pour la grande majorité des calculs, à condition qu'un sous-ensemble petit et adaptativement sélectionné soit recalculé avec une plus grande précision.