Résumé technique : Dr.LLM – Routage dynamique des couches dans les LLM

Énoncé du problème

Les modèles de langage larges (LLM) fonctionnent actuellement sous un régime de profondeur statique, traitant chaque token à travers une pile fixe de couches de transformateurs, indépendamment de la complexité de l'entrée. Cette approche entraîne deux inefficacités principales :

1. Calcul gaspillé : Les requêtes simples consomment des ressources de calcul inutiles en traversant toutes les couches.
2. Flexibilité insuffisante : Les tâches de raisonnement complexe nécessitent souvent un traitement plus profond ou un raffinement itératif, que les architectures statiques ne peuvent fournir sans réentraînement ou refonte architecturale.

Les méthodes adaptatives de profondeur existantes (par exemple, sorties anticipées, élagage de couches, mélanges de profondeurs) présentent des limitations significatives : elles sacrifient souvent la précision pour la vitesse, nécessitent un réentraînement étendu ou des modifications architecturales, ou reposent sur une recherche coûteuse au moment de l'inférence qui entrave l'évolutivité. De plus, de nombreuses approches antérieures dégradent la précision du modèle malgré les gains d'efficacité.

Méthodologie : Dr.LLM

Les auteurs proposent Dr.LLM (Routage dynamique des couches pour les LLM), un cadre rétrofitable qui équipe les LLM préentraînés figés de routeurs légers par couche. Ces routeurs décident s'il faut sauter, exécuter ou répéter un bloc de transformateur spécifique pour une entrée donnée.

1. Architecture du routeur

• Structure : Chaque bloc de transformateur $B_\ell$ est apparié à un Perceptron Multicouche (MLP) léger composé de couches Linéaire-GELU-Linéaire.
• Entrée : Le routeur opère sur un résumé compact des états cachés de la couche précédente. Pour assurer la stabilité sur les contextes longs, les auteurs emploient un moyennage par fenêtre, divisant la séquence en fenêtres contiguës et moyennant les états cachés au sein de chaque fenêtre.
• Sortie : Le routeur produit des logits pour trois classes : $\{sauter, exécuter, répéter\}$.
• Surcharge : Les routeurs sont exécutés une fois par séquence d'entrée. La surcharge est négligeable (constante par rapport au nombre de tokens) et entièrement compatible avec la mise en cache KV.

2. Entraînement supervisé via MCTS

Contrairement aux méthodes qui reposent sur l'apprentissage par renforcement ou la recherche au moment de l'inférence, Dr.LLM entraîne les routeurs en utilisant une supervision explicite dérivée d'une recherche arborescente Monte Carlo (MCTS) hors ligne.

• Processus de recherche : Pour chaque entrée, la MCTS explore des chemins d'exécution modifiés (sautant ou répétant des couches) sous un budget de calcul. Elle ne conserve que les chemins qui préservent ou améliorent la précision de la tâche par rapport à l'exécution pleine profondeur par défaut.
• Jeu de données : La recherche génère un jeu de données compact d'environ 4 000 chemins d'exécution optimisés (tuples de question, actions de couche optimales et réponse).
• Objectif d'entraînement : Les routeurs sont entraînés de bout en bout sur ce jeu de données en utilisant une Perte Focale avec des poids d'équilibrage de classe pour gérer le déséquilibre sévère entre la classe dominante « exécuter » et les classes minoritaires « sauter » et « répéter ».
• Gel : Crucialement, les poids du LLM de base restent figés ; seuls les paramètres du routeur sont mis à jour.

3. Inférence

Au moment de l'inférence, le routeur entraîné prend des décisions gloutonnes ($\hat{y}_\ell = \arg\max p_\ell$) pour chaque couche sans aucune recherche.

• Sauter : Contourne la couche ($H^{(\ell)} = H^{(\ell-1)}$).
• Exécuter : Applique la couche une fois.
• Répéter : Applique la couche deux fois de suite.
• Regroupement (Batching) : Pour gérer la synchronisation dans le traitement par lots, les séquences sont regroupées selon leurs décisions de routage à chaque couche. Les auteurs notent qu'un fort consensus de routage (accord de 80 à 89 % entre les séquences) minimise les frais de synchronisation.

Contributions clés

1. Routage dynamique supervisé pour LLM figés : Introduction de routeurs par couche décidant de sauter, d'exécuter ou de répéter, entraînés de bout en bout sur seulement 4 000 chemins dérivés de MCTS sans modifier les poids de base.
2. Supervision de chemins efficace via MCTS : Un algorithme MCTS conscient de la longueur qui découvre hors ligne des configurations de couches préservant ou améliorant la précision, générant un jeu de données de supervision compact sans recherche au moment de l'inférence.
3. Entraînement de routeur léger : Une conception intégrant un moyennage par fenêtre pour la stabilité des contextes longs, une perte focale pour le déséquilibre de classe, et un routeur linéaire compact à deux couches, garantissant une surcharge de calcul négligeable.
4. Gains de précision et d'efficacité : Des améliorations démontrées de la précision du raisonnement accompagnées d'une réduction de l'exécution des couches, surpassant les méthodes de routage antérieures de l'état de l'art.
5. Généralisation robuste : Démonstration que les politiques de routage apprises sur des tâches de raisonnement spécifiques se transfèrent efficacement vers des benchmarks hors domaine avec une dégradation minimale de la précision.

Résultats expérimentaux

Performance en domaine (ARC et DART)

Évalué sur six modèles (variantes LLaMA-3.2 et Qwen-2.5), Dr.LLM a constamment amélioré la précision tout en réduisant le nombre moyen de couches exécutées :

• Précision : Amélioration allant jusqu'à +3,4 points de pourcentage (p.p.) sur DART (mathématiques) et +1,0 p.p. sur ARC (logique). Par exemple, LLaMA-3B-Base a vu un bond de 11,8 % à 15,8 % sur DART.
• Efficacité : Économie moyenne de 5,0 couches par exemple à travers les modèles, avec des cas spécifiques économisant jusqu'à 11,0 couches.
• Latence : Réduction de 15,3 % du temps de génération en temps réel pour des exécutions de 1 000 tokens, la surcharge du routeur contribuant à moins de 1 % de la latence totale.

Généralisation hors domaine

Lorsqu'appliqué à des benchmarks diversifiés (MMLU, GSM8k, AIME, TruthfulQA, SQuADv2, GPQA, AGIEval, PIQA) sans réentraînement :

• Les routeurs ont maintenu l'efficacité tout en subissant une baisse moyenne de précision de 0,85 p.p. sur huit benchmarks.
• Dans certains cas (par exemple, GPQA Diamond avec LLaMA-3B), la précision s'est même améliorée (+2,5 p.p.).
• Cela suggère que les politiques apprises capturent des redondances structurelles générales dans le calcul des transformateurs plutôt qu'un surapprentissage à des motifs de tâches spécifiques.

Comparaison avec l'état de l'art

Dr.LLM a surpassé les méthodes de routage précédentes (par exemple, FlexiDepth, MindSkip, ShortGPT) d'une marge allant jusqu'à +7,7 p.p. en précision sur les benchmarks de raisonnement. Contrairement à FlexiDepth, qui nécessite un entraînement sur des centaines de milliers d'exemples et sacrifie souvent la précision, Dr.LLM atteint une précision supérieure avec une surcharge d'entraînement significativement plus faible (4 000 exemples, un seul GPU) et sans modifications architecturales.

Analyse des motifs de routage

L'analyse des décisions de routage a révélé des comportements structurés :

• Couches précoces : Constamment exécutées (utilisation stable).
• Couches intermédiaires : Fréquemment sautées (économie de calcul).
• Couches tardives : Souvent répétées, en particulier pour les tâches de raisonnement complexe, indiquant un rôle dans le raffinement itératif.
• Apprentissage de caractéristiques : Les routeurs entraînés sur les états cachés de la couche précédente ont nettement surpassé ceux entraînés sur des embeddings d'entrée statiques, confirmant que Dr.LLM apprend à mapper dynamiquement les états internes du modèle aux décisions de routage.

Importance et affirmations

L'article affirme que Dr.LLM démontre que des routeurs explicitement supervisés peuvent rétrofitter des LLM figés pour réaliser une inférence consciente du budget et axée sur la précision. Les points clés à retenir incluent :

• Aucun réentraînement requis : Les poids du modèle de base restent intacts, préservant les capacités originales du modèle.
• Efficacité positive pour la précision : Contrairement à de nombreuses méthodes adaptatives qui sacrifient la précision pour la vitesse, Dr.LLM améliore les deux simultanément sur des tâches riches en raisonnement.
• Déploiement pratique : En éliminant le besoin de recherche au moment de l'inférence et de modifications architecturales, Dr.LLM offre une solution pratique et évolutive pour l'inférence adaptative.
• Généralisabilité : Les politiques de routage apprises se transfèrent de manière robuste à travers les familles de modèles et les domaines de tâches, suggérant que l'allocation adaptative de calcul est une propriété fondamentale et apprenable des architectures de transformateurs.

Les auteurs concluent que Dr.LLM constitue une étape pratique vers une inférence adaptative évolutive, conciliant efficacité, précision et robustesse sans les coûts prohibitifs associés au réentraînement ou à la refonte architecturale.