Résumé technique : LEDE – Sorties dynamiques pilotées par l'expérience pour les LLM avec l'apprentissage par renforcement

1. Énoncé du problème

Les grands modèles de langage (LLM) souffrent d'une latence significative lors de l'inférence autoregressive car chaque jeton généré doit traverser l'intégralité de la pile des couches de transformeurs. Bien que la Décodage Auto-Spéculatif (SSD - Self-Speculative Decoding) ait été proposée pour accélérer ce processus en utilisant les propres couches précoces d'un modèle comme un « rédacteur » interne pour générer des jetons candidats en parallèle, les méthodes actuelles de SSD sont fondamentalement limitées par des configurations statiques.

Les approches existantes reposent sur des profondeurs de rédaction fixes et des longueurs de spéculation fixes appliquées uniformément sur l'ensemble du processus de génération. Cette rigidité ne tient pas compte de la rareté contextuelle, où la difficulté de prédire les jetons varie considérablement selon le contexte local. Par conséquent, les méthodes statiques ne peuvent pas équilibrer de manière optimale le compromis entre la qualité de la rédaction (taux d'acceptation) et la vitesse de spéculation, ce qui conduit à une performance d'inférence sous-optimale.

2. Méthodologie : Le cadre LEDE

Les auteurs proposent LEDE (Learning-based Dynamic Exit), un cadre qui reformule le contrôle dynamique du SSD comme un Processus de Décision Markovien (MDP). Au lieu d'utiliser des règles rigides et prédéfinies, LEDE emploie l'apprentissage par renforcement hors ligne (offline RL) pour entraîner une politique qui sélectionne dynamiquement à la fois la couche de sortie optimale et la longueur de spéculation en fonction du contexte local de la séquence générée.

2.1. Formulation du Processus de Décision Markovien (MDP)

Le cadre modélise le processus de prise de décision lors d'une étape de génération unique comme un MDP avec les composants suivants :

• Espace d'états ($S$) : L'agent observe un vecteur d'état $s_l$ à chaque couche candidate $l$. Ce vecteur est dérivé de la distribution interne des jetons du modèle et inclut :
  • Confiance du jeton ($C_k^l$) : La valeur moyenne du logarithme négatif de probabilité des top-$k$ jetons.
  • Entropie du jeton ($H^l$) : L'incertitude de la distribution de prédiction du jeton suivant.
  • Trace moyenne de la confiance du jeton ($C_{trace}$) : La moyenne du logarithme négatif de probabilité du top-1 jeton sur une fenêtre glissante des étapes précédentes.
• Espace d'actions ($A$) : À chaque couche candidate, la politique (un réseau Q) décide entre deux actions :
  • Continuer ($a=0$) : Passer à la couche candidate suivante.
  • Sortir ($a=1$) : Désigner la couche actuelle comme la couche de sortie optimale ($l^*$) pour la phase de rédaction.
• Fonction de récompense ($R$) : Conçue pour optimiser l'efficacité de l'inférence.
  • Récompense positive : Accordée si le jeton rédigé est accepté ($x_t^{l^*} = x_t^L$), pondérée par la profondeur économisée ($1 + \frac{L-l^*}{L}$).
  • Récompense négative : Accordée si le jeton rédigé est rejeté ($-2 + \frac{L-l^*}{L}$).
  • Pénalité : Une petite pénalité ($-0,01$) est appliquée pour la continuation, encourageant les sorties précoces lorsque cela est approprié.

2.2. Optimisation de la politique via l'apprentissage hors ligne (Offline RL)

LEDE utilise un Réseau Q Profond (DQN) entraîné hors ligne en utilisant un tampon de relecture (replay buffer) d'expériences d'inférence passées :

1. Collecte de données : Une phase d'exploration initiale collecte des tuples d'expérience diversifiés $(s_t, a_t, r_t, s_{t+1})$ où l'agent interagit avec l'environnement.
2. Entraînement hors ligne : Le réseau Q est entraîné en minimisant l'erreur de différence temporelle (TD error) en n-étapes. La valeur cible est calculée à l'aide d'un réseau cible stable, équilibrant la récompense immédiate d'une sortie précoce par rapport à la valeur à long terme de l'état.
3. Mécanisme de double contrôle : Lors de l'inférence en ligne, l'agent entraîné effectue deux contrôles adaptatifs :
   • Sélection dynamique de la couche de sortie : Il détermine la profondeur optimale $l^*$ pour commencer la rédaction.
   • Longueur de rédaction adaptative : Après avoir sélectionné $l^*$, l'agent réévalue l'état après chaque jeton rédigé. Une action « sortie » continue la rédaction (indiquant une haute confiance), tandis qu'une action « continuer » (indiquant une faible confiance) interrompt immédiatement la rédaction pour vérification.

3. Contributions clés

Le document présente trois contributions principales :

1. Formulation MDP : C'est le premier travail qui formule le contrôle dynamique de la profondeur de rédaction et de la longueur de spéculation dans le SSD comme un problème de MDP, appliquant l'apprentissage par renforcement pour dériver une politique de contrôle.
2. Cadre LEDE : La conception et l'implémentation d'un cadre qui exploite l'apprentissage par renforcement hors ligne et la relecture d'expérience pour apprendre une politique sensible au contexte, surmontant les limites des configurations statiques et des méthodes simples basées sur des heuristiques.
3. Validation empirique : Des expériences approfondies démontrant des accélérations significatives par rapport à la génération autoregressive et des gains d'efficacité substantiels par rapport aux bases statiques.

4. Résultats expérimentaux

Le cadre a été évalué sur les modèles Llama-2 (7B, 13B) et Llama-3 (1B), ainsi que CodeLLaMA (7B, 34B), à travers diverses tâches incluant le suivi d'instructions, la résumé, la modélisation de langage et la génération de code.

• Performance d'accélération (Speedup) :
  • LEDE atteint une accélération moyenne de 2,32× par rapport au décodage autoregressive standard.
  • Les pics d'accélération atteignent 2,72× (ex: sur Llama-2-7B pour la modélisation de langage).
  • LEDE fournit une accélération supplémentaire de 17 % par rapport à la meilleure base de spéculation statique (LayerSkip).
• Comparaison avec les bases de référence (Baselines) :
  • LayerSkip (Statique) : La performance est très sensible aux configurations fixes et ne généralise pas bien à travers les tâches.
  • Méthodes basées sur des règles (LITE, Draft & Verify) : Contraintes par des seuils prédéfinis, menant à des compromis sous-optimaux entre profondeur et taux d'acceptation.
  • LEDE : En apprenant une politique de contrôle fine, il surpasse systématiquement toutes les bases de référence en termes d'accélération tout en maintenant une qualité de génération comparable (mesurée par Rouge-L).
• Étude d'ablation : La suppression de la sélection dynamique de la couche de sortie ou de la longueur de rédaction adaptative dégrade considérablement les performances (faisant chuter l'accélération à environ 2,0×), confirmant que les deux mécanismes sont synergiques et essentiels.

5. Signification et affirmations

L'article affirme que LEDE représente un changement conceptuel passant d'une inférence rigide pilotée par des heuristiques à une inférence adaptative pilotée par une politique. En apprenant à co-optimiser la profondeur de calcul et la longueur de spéculation en fonction du contexte en temps réel, le cadre permet un décodage plus robuste et plus efficace à travers divers modèles et tâches.

Les auteurs positionnent ce travail comme une étape vers la construction de systèmes de LLM plus intelligents et plus efficaces sur le plan computationnel. Ils notent que les travaux futurs se concentreront sur l'extension de ces avantages aux modèles plus grands (échelle 70B et au-delà), mais la contribution actuelle établit la viabilité de l'utilisation de l'apprentissage par renforcement hors ligne pour gérer dynamiquement le pipeline d'inférence des grands modèles de langage.