Résumé Technique : RAP (Élagage Adaptatif en Temps d'Exécution pour l'Inférence des LLM)

Énoncé du Problème

Les grands modèles de langage (LLM) rencontrent des obstacles majeurs de déploiement en raison de leurs empreintes computationnelles et mémoire massives. Bien que les techniques de compression de modèles (élagage, quantification, distillation) offrent des solutions, les méthodes existantes souffrent de limitations critiques :

1. Heuristiques Statiques : La plupart des stratégies d'élagage actuelles reposent sur des heuristiques fixes et ponctuelles qui supposent des charges de travail statiques. Elles échouent à s'adapter aux variations de mémoire en temps d'exécution causées par des demandes utilisateurs fluctuantes (par exemple, des tailles de lots et des longueurs de séquence variables).
2. Négligence du Cache KV : L'élagage traditionnel se concentre principalement sur la réduction des poids statiques du modèle (paramètres). Cependant, dans l'inférence autorégressive, la mémoire du cache Clé-Valeur (KV) croît linéairement avec la taille du lot et la longueur de la séquence, dépassant souvent la mémoire des paramètres statiques. Des ratios d'élagage fixes ne peuvent pas accommoder le glissement dynamique des goulots d'étranglement mémoire, passant de la domination des paramètres à celle du cache KV.
3. Variance au Niveau Système : Les environnements réels de périphérie et de cloud connaissent une disponibilité mémoire stochastique due aux applications coexécutées et aux interférences, rendant les hypothèses de budgets fixes invalides.

Le défi central abordé consiste à sélectionner une politique d'élagage optimale pour les LLM qui s'adapte dynamiquement à des charges de travail hétérogènes et variant dans le temps, tout en satisfaisant des budgets mémoire fluctuants sans dégrader la qualité de sortie.

Méthodologie : Le Cadre RAP

Les auteurs proposent RAP (Élagage Adaptatif en Temps d'Exécution), un cadre d'élagage élastique piloté par l'Apprentissage par Renforcement (RL) qui ajuste dynamiquement les stratégies de compression en fonction des états système en temps réel.

1. Importance Séquentielle Gourmande (GSI)

Pour pallier les limites de l'élagage « ponctuel » (qui ignore les dépendances inter-couches et entraîne une dégradation des performances), RAP utilise l'Importance Séquentielle Gourmande (GSI).

• Mécanisme : Au lieu de noter tous les blocs une seule fois, GSI supprime itérativement le bloc dont l'exclusion entraîne la détérioration minimale des performances du modèle (mesurée par la perplexité).
• Recalibrage : Après chaque suppression, les scores d'importance des blocs restants sont réévalués au sein de la nouvelle architecture contractée. Ce processus étape par étape contrôle l'accumulation d'erreurs et prend en compte la sensibilité dynamique des couches à mesure que la structure du modèle change.

2. Formulation par Apprentissage par Renforcement

RAP formule la sélection de la politique d'élagage comme un problème de prise de décision séquentielle (Processus de Décision de Markov) où un agent RL apprend à équilibrer l'efficacité mémoire et la fidélité du modèle.

• État ($S_t$) : L'agent observe un état composite à chaque étape d'inférence $t$ :
  • Caractéristiques de la demande : Taille du lot ($R_{bs}$) et longueur de séquence ($R_{sql}$).
  • Configuration du modèle : Scores d'importance de tous les blocs d'Attention Multi-Têtes (MHA) et de Réseau de Neurones Feed-Forward (FFN) dérivés de GSI.
  • État du système : Mémoire système disponible actuelle ($Sys_{avail}$) et surcharge mémoire estimée ($Sys_{req}$) pour l'action proposée.
• Action ($A_t$) : L'agent sélectionne une politique d'élagage en décidant de conserver ou de supprimer des blocs MHA et FFN spécifiques. Pour gérer l'explosion combinatoire de l'espace d'action ($2^{2N}$), la décision est décomposée en sélections de blocs uniques séquentielles.
• Récompense ($R_t$) : Une fonction de récompense unifiée équilibre deux objectifs :
  • Perplexité ($R_{ppl}$) : Encourage la préservation des performances de modélisation du langage.
  • Pénalité Mémoire ($R_{mem}$) : Pénalise la consommation mémoire maximale durant l'inférence.
  • La récompense est une somme pondérée : $R_t = \sum (At)_i (\alpha R_{ppl, i} - \beta R_{mem, i})$, où les hyperparamètres $\alpha$ et $\beta$ ajustent le compromis précision-mémoire.

3. Flux d'Exécution

1. L'agent RL observe l'environnement en temps réel (longueur d'entrée, lot, disponibilité mémoire).
2. Il sélectionne une action d'élagage (quels blocs supprimer) en fonction de l'état actuel.
3. Le modèle de base exécute l'inférence sur l'architecture élaguée.
4. L'agent reçoit une récompense basée sur l'utilisation mémoire résultante et la perplexité, mettant à jour sa politique pour optimiser les décisions futures.

Contributions Clés

1. Cadre Adaptatif en Temps d'Exécution : RAP est le premier cadre à considérer conjointement les poids du modèle et la mémoire du cache KV « à la volée », ajustant dynamiquement les niveaux de parcimonie en fonction de la charge de travail instantanée et de l'état de l'appareil.
2. Politique Pilotée par RL : Les auteurs formulent l'élagage comme un problème de RL, entraînant un agent à apprendre une politique optimale qui équilibre les économies de mémoire avec la fidélité de génération, surpassant les bases heuristiques statiques.
3. GSI pour l'Élagage Structurel : L'introduction de l'Importance Séquentielle Gourmande résout le problème des dépendances inter-couches dans l'élagage des transformateurs, offrant une méthode plus robuste pour déterminer l'importance des blocs sous une compression agressive.
4. Validation Empirique : Des expériences extensives démontrent que RAP surpasse constamment les bases d'élagage statique de l'état de l'art (par exemple, LLM-Pruner, SliceGPT, ShortGPT) sur divers modèles (Llama2, Llama3, Qwen) et contraintes mémoire.

Résultats Expérimentaux

Les expériences ont été menées sur Llama2-7B, Llama3-8B, Qwen1.5-7B et Qwen2.5-7B sous des budgets mémoire unifiés de 80 % et 60 % (couvrant à la fois les paramètres et le cache KV).

• Capacité de Génération : RAP a présenté la dérive de perplexité la plus faible. Sur Llama2-7B avec un budget de 80 %, RAP a augmenté la perplexité de seulement +6,33 (WikiText2) et +21,47 (PTB), surpassant significativement la méthode suivante par des marges de 22,95 et 146,93 respectivement.
• Tâches en Aval : Dans le raisonnement de bon sens en zéro-shot, RAP a préservé 86 % de la précision du modèle dense sous un budget de 80 %, dépassant les bases leaders de +8,2 % (Llama2-7B) et +11,5 % (Llama3-8B).
• Résilience : Sous des contraintes mémoire sévères de 60 %, RAP était la seule méthode à conserver plus de 60 % des performances du modèle dense, tandis que les autres méthodes subissaient une dégradation catastrophique.
• Études d'Ablation :
  • Le retrait de GSI (utilisant un score ponctuel statique) a causé une baisse de 13,17 % de la précision moyenne en raisonnement de bon sens.
  • Le retrait de l'agent RL (utilisant une suppression aléatoire de blocs) a entraîné une augmentation massive de la perplexité (+301,71 sur WikiText2) par rapport à la politique apprise, soulignant la nécessité d'une prise de décision adaptative.
• Surcharge : Le contrôleur RL ajoute une surcharge temporelle négligeable (<1 % d'augmentation de latence) grâce à son architecture compacte (18K paramètres), tandis que le calcul hors ligne de GSI est un coût unique amorti sur le déploiement.

Signification et Revendications

L'article revendique que RAP comble le fossé entre les techniques de compression statiques et les scénarios de déploiement dynamique du monde réel. En s'éloignant des heuristiques « taille unique », RAP permet :

• Déploiement Évolutif sur Périphérie : Il permet aux LLM de tenir sur l'appareil sous des contraintes mémoire serrées et fluctuantes en élaguant des poids substantiels lorsque nécessaire, tout en conservant une capacité totale lorsque les ressources le permettent.
• Efficacité sans Perte de Qualité : Il réalise des économies de mémoire supérieures et une inférence plus rapide avec un impact minimal sur la qualité de sortie, naviguant efficacement dans les compromis entre les goulots d'étranglement lourds en paramètres et ceux lourds en cache KV.
• Généralisabilité : Le cadre s'avère agnostique à l'architecture, performant sur différentes familles de modèles (Llama, Qwen) et générations.

Les auteurs positionnent RAP comme un mécanisme de compression flexible et conscient du contexte qui rend l'inférence des LLM réalisable dans des environnements hétérogènes où les stratégies fixes échouent.