Résumé technique : Sauter une couche ou la boucler ? Apprentissage de programmes de couches (Program-of-Layers) dans les LLM

1. Énoncé du problème

Les grands modèles de langage (LLM) effectuent actuellement l'inférence via une passe avant à profondeur fixe et ordre fixe, exécutant toutes les couches du transformateur de manière séquentielle pour chaque entrée, quelle que soit la complexité de la tâche. Cette architecture statique ne tient pas compte de la grande variance de la difficulté des entrées, ce qui entraîne deux inefficacités :

1. Sur-calcul (Over-computation) : Les entrées simples peuvent être résolues correctement avec moins de couches, pourtant le modèle exécute la profondeur complète.
2. Sous-calcul (Under-computation) : Les entrées complexes qui échouent lors d'une passe avant standard pourraient être solubles si le modèle pouvait modifier dynamiquement son chemin d'exécution (par exemple, en répétant certaines couches pour le raffinement ou en sautant les couches redondantes).

L'article postule que la passe avant standard ne capture qu'un sous-ensemble étroit de la capacité de raisonnement latent d'un LLM. La question centrale est de savoir si un modèle généraliste peut optimiser son « programme » (la séquence d'exécutions de couches) pour chaque entrée spécifique sans réentraîner le modèle de base.

2. Méthodologie

2.1. Découverte empirique via la recherche arborescente de Monte-Carlo (MCTS)

Avant de proposer une solution basée sur l'apprentissage, les auteurs mènent une étude empirique utilisant la MCTS pour explorer l'espace des Programmes de Couches (POLAR).

• Définition : Les couches sont traitées comme une bibliothèque de fonctions atomiques ($f_0, \dots, f_{D-1}$). Un programme $\pi$ est une séquence d'indices de couches permettant le saut (omission de couches) et la récurrence (répétition de couches).
• Résultats :
  • Optimisation conjointe : Combiner les opérations de saut et de boucle produit une précision nettement supérieure à l'utilisation de l'une ou l'autre seule.
  • Rasoir d'Occam : Pour de nombreuses entrées (y compris celles initialement résolues incorrectement), des programmes d'exécution valides existent qui sont plus courts que la passe avant standard.
  • Mise à l'échelle au moment du test (Test-Time Scaling) : Augmenter le budget de profondeur (via la récurrence) améliore systématiquement la précision sur les tâches plus difficiles, révélant un continuum de calculs latents valides.
  • Biais structurel : Les programmes valides consistent principalement en des segments de couches contigus et nécessitent rarement plus d'une seule récurrence par segment. Les sauts à longue portée ou la réutilisation itérative profonde sont peu communs.

2.2. Le cadre POLAR

Motivés par les découvertes de la MCTS, les auteurs proposent POLAR, un cadre léger qui prédit des programmes d'exécution spécifiques à l'entrée au moment de l'inférence, sans modifier les paramètres du modèle pré-entraîné gelé.

• Représentation du programme :

  • Segmentation : Les $D$ couches sont partitionnées en segments contigus $[s_j, s_{j+1})$.
  • Opérations : Chaque segment se voit assigner l'une des trois opérations : Sauter (omettre), Garder (exécution standard) ou Répéter (exécuter deux fois).
  • Contraintes : Les longueurs de segment sont limitées (fixées empiriquement à $K_{max}=4$) en se basant sur le fait que les programmes valides privilégient les structures locales.

• Architecture du réseau de prédiction :

  • Encodage de l'entrée : L'entrée $x$ est encodée en plongements de jetons (token embeddings) via un modèle d'encodage gelé.
  • Requêtes de couches : Des plongements apprenables représentent chaque indice de couche, agissant comme des requêtes (queries).
  • Attention croisée (Cross-Attention) : Un mécanisme d'attention croisée combine les requêtes de couches avec les plongements de jetons pour générer des représentations conditionnées par l'entrée pour chaque couche.
  • Encodeur inter-couches : Un encodeur transformer léger modélise les dépendances à travers la dimension de profondeur, permettant aux décisions à une couche donnée de dépendre du contexte global.
  • Têtes de prédiction : Deux têtes prédisent (1) les frontières de segmentation (masque binaire) et (2) les étiquettes d'opérations (sauter/garder/répéter) pour chaque début de segment.

• Entraînement et Inférence :

  • Supervision : Le réseau est entraîné sur des programmes de vérité terrain découverts hors ligne via la MCTS.
  • Inférence : Le réseau prédit les logits pour la segmentation et les opérations. Une recherche en faisceau (beam search) est appliquée sur les opérations au niveau du segment pour assurer la cohérence globale avant d'exécuter le LLM gelé.

3. Résultats clés

3.1. Performance en distribution (DART-Math)

Évalué sur des benchmarks de raisonnement mathématique (DM-1 à DM-5) sur quatre modèles (LLaMA-3.2-3B, Qwen1.5-MoE, Qwen2.5-3B, Qwen3-8B) :

• Gains de précision : POLAR surpasse systématiquement l'inférence standard (Base) et les méthodes de profondeur dynamique antérieures (ex: MindSkip, FlexiDepth, DR.LLM).
  • Exemple : Sur LLaMA-3.2-3B, POLAR obtient un gain absolu de +46,8 % sur DM-1 et de +32,0 % sur le plus difficile DM-5 par rapport au modèle de base.
• Mise à l'échelle au moment du test : Augmenter le nombre de programmes candidats ($k$) via la recherche en faisceau produit des améliorations de précision monotones, dépassant de loin les rendements décroissants de l'échantillonnage stochastique standard (Base sampling).
• Efficacité : POLAR parvient souvent à une précision plus élevée tout en exécutant moins de couches uniques en moyenne.

3.2. Généralisation hors distribution (OOD)

Entraîné sur DART-Math, POLAR a été évalué zero-shot sur ASDiv, MAWPS et MMLU-Pro (Mathématiques, Physique, Chimie, Droit, etc.) :

• Transfert de domaine : POLAR surpasse systématiquement la passe avant de base et les autres méthodes dynamiques à travers divers domaines, incluant les problèmes mathématiques textuels et les tâches de culture générale.
• Mécanisme : Les auteurs attribuent cela au fait que la tête de prédiction projette les entrées dans un espace sémantique partagé et à la contrainte de structures simples (segments contigus), ce qui encourage des stratégies réutilisables plutôt que des heuristiques spécifiques aux jeux de données.

3.3. Efficacité et surcharge (Overhead)

• Latence : Le réseau de prédiction introduit une surcharge négligeable (~0,8 % d'une passe avant standard, équivalent à ~0,23 couche).
• Gain net : En sautant des couches sur les entrées faciles, POLAR réduit la latence de bout en bout (ex: 0,83× le temps d'exécution de base sur les entrées faciles) tout en améliorant simultanément la précision.
• Nombre de paramètres : Le prédicteur n'ajoute que ~2,1 millions de paramètres (0,01 % – 0,06 % de la taille du modèle de base).

4. Signification et affirmations

L'article affirme que l'exécution à profondeur fixe est sous-optimale pour les modèles de fondation généralistes. En traitant les couches comme une bibliothèque de fonctions réutilisables, POLAR démontre que :

1. Capacité latente : Les LLM possèdent un vaste espace de calculs alternatifs valides (programmes) capables de corriger des erreurs ou de résoudre des tâches plus efficacement que la passe avant standard.
2. Praticité : Cette capacité latente peut être débloquée via un mécanisme de prédiction léger et sans entraînement (pour le modèle de base) qui remplace la recherche coûteuse par une prédiction en une seule étape.
3. Généralisation : La capacité de reconfigurer dynamiquement la profondeur et l'ordre d'exécution est une compétence transférable qui améliore les performances à travers différents domaines sans redessiner l'architecture ni procéder à un ajustement fin (fine-tuning).

Ce travail suggère que les futures stratégies d'inférence devraient dépasser la profondeur statique, en réallouant la puissance de calcul de manière dynamique pour correspondre à la difficulté de l'entrée, offrant ainsi une voie vers un déploiement plus efficace et plus expressif des modèles de fondation.