SimLens for Early Exit in Large Language Models: Eliciting Accurate Latent Predictions with One More Token

Les grands modèles de langage (LLM) forment leurs prédictions progressivement à travers leurs couches profondes. Bien que les états cachés des couches intermédiaires contiennent des informations sémantiques pertinentes, les décoder avec précision, en particulier dans les premières couches, reste un défi majeur.

Les méthodes existantes de type "lunette" (comme Logit Lens, Tuned Lens ou Future Lens) reposent généralement sur une lecture linéaire directe (une projection linéaire apprise ou partagée) des états cachés vers l'espace du vocabulaire. Bien que simples, ces approches souffrent d'un décalage important par rapport à la prédiction finale du modèle, surtout aux couches précoces. Le goulot d'étranglement n'est pas l'absence d'information latente, mais la difficulté à la lire fidèlement sans les transformations non linéaires des couches supérieures.

L'objectif est de permettre un arrêt anticipé (early exit) fiable : arrêter l'inférence dès qu'une couche intermédiaire contient une prédiction suffisamment précise, réduisant ainsi considérablement le coût de calcul.

2. Méthodologie

L'article propose une approche novatrice qui remplace la lecture linéaire pure par une continuation minimale du modèle.

A. SimLens (Le Décodeur)

Au lieu de projeter linéairement un état intermédiaire, SimLens conserve un état minimal composé de deux tokens :

1. Le token de début <s> : Il agit comme une condition globale et une ancre de contexte/position (s'appuyant sur le comportement des "attention sinks").
2. Le token de réponse candidat <a> : Il sert d'ancre sémantique pour la réponse évaluée.

Fonctionnement :
Pour une couche intermédiaire $l$, SimLens prend l'état de ces deux tokens et les fait passer à travers les couches supérieures restantes ($l+1$ à $L$) en utilisant la dynamique native du modèle (sans apprentissage supplémentaire).
$$[<s>^L, <a>^L] = T_{l}^{L} ([<s>^l, <a>^l])$$
Ensuite, seule la projection linéaire finale est appliquée au token <a> pour obtenir la probabilité de la réponse. Cette "étape de plus" (one more token) permet de récupérer des prédictions latentes beaucoup plus précises que les méthodes linéaires.

B. Linear SimLens (Approximation pour la Confiance)

Pour éviter le coût d'une propagation complète à chaque vérification de confiance lors d'un arrêt anticipé, les auteurs introduisent Linear SimLens.

• C'est une approximation linéaire légère entraînée pour mapper l'espace de représentation d'une couche $l$ vers celui de la couche finale $L$.
• Contrairement aux méthodes classiques qui visent la sortie du modèle original, Linear SimLens est distillé à partir des cibles de SimLens. Cela lui permet de capturer le décalage d'espace de représentation nécessaire sans nécessiter la propagation complète de la séquence pendant le monitoring.

C. SimExit (Mécanisme Hybride d'Arrêt Anticipé)

SimExit combine les deux approches pour un arrêt adaptatif :

1. Surveillance : À intervalles réguliers, Linear SimLens calcule un score de confiance basé sur l'entropie de la distribution des probabilités.
2. Décision : Si l'entropie tombe en dessous d'un seuil prédéfini, l'inférence de la séquence complète est interrompue.
3. Décodage Final : Le modèle ne propage plus que les deux tokens (<s> et <a>) à travers les couches restantes via SimLens pour produire la réponse finale.

3. Contributions Clés

1. Introduction de SimLens : Un décodeur sans entraînement qui utilise une continuation à deux tokens pour extraire des prédictions latentes précises des couches intermédiaires.
2. Rôle des Tokens : Démonstration que <s> (conditionnement global) et <a> (ancre sémantique) jouent des rôles distincts et complémentaires essentiels à la performance.
3. Cadre SimExit : Un mécanisme hybride combinant une estimation de confiance peu coûteuse (Linear SimLens) et un décodage précis réduit (SimLens) pour l'arrêt anticipé.
4. Analyse Équitable : Une évaluation rigoureuse incluant le surcoût de calcul des deux tokens supplémentaires ("fair compute"), prouvant que la méthode reste supérieure même avec ce coût.

4. Résultats Expérimentaux

Les expériences ont été menées sur les modèles LLaMA-7B et Vicuna-7B avec les jeux de données ARC, BoolQ et HeadQA.

• Précision Latente : SimLens dépasse systématiquement les méthodes linéaires (Logit Lens, Tuned Lens) aux couches précoces.
  • Gain moyen de précision (Iso-Compute) de +0,43 sur les six configurations (modèle + jeu de données), même en incluant le coût des deux tokens supplémentaires.
  • Sur ARC, SimLens atteint la précision du modèle complet à la couche 18, contre la couche 28 pour Logit Lens.
• Transfert de Tâche : Linear SimLens entraîné sur ARC fonctionne bien sur BoolQ et HeadQA sans réentraînement, démontrant une bonne généralisation.
• Performance d'Arrêt Anticipé (SimExit) :
  • Accélération : Gain moyen de vitesse de 1,15x aux points de fonctionnement à meilleure précision.
  • Compromis Précision/Vitesse : Avec une tolérance de perte de précision de 1 point de pourcentage, l'accélération atteint 1,40x.
• Ablations :
  • Retirer <s> (SimLens-NoS) entraîne une chute de performance moyenne de 10,68 points, prouvant son rôle de conditionnement global.
  • Remplacer <a> par un token incorrect ou aléatoire fait chuter la précision drastiquement (jusqu'à -36 points), confirmant son rôle d'ancre sémantique.

5. Signification et Impact

Ce travail remet en question la nécessité de décodeurs linéaires complexes pour l'analyse des couches intermédiaires. Il démontre qu'une petite quantité de calcul non linéaire supplémentaire (la propagation de deux tokens à travers les couches supérieures) est bien plus efficace pour transformer les représentations intermédiaires en prédictions utilisables que n'importe quelle projection linéaire.

Implications pratiques :

• Réduction des coûts : Permet d'accélérer l'inférence des LLMs sans sacrifier la précision, en arrêtant le calcul dès que la confiance est suffisante.
• Simplicité : La méthode est sans entraînement (training-free) pour SimLens et utilise une approximation légère pour la surveillance, la rendant facile à déployer.
• Limites et Avenir : L'approche est actuellement optimisée pour les tâches de décision à un seul token (QCM, classification). L'extension à la génération autoregressive longue pose des défis liés à la gestion du cache KV (remplissage des couches sautées), mais l'article suggère que le principe de base reste valide.

En résumé, SimLens offre une nouvelle perspective sur l'interprétabilité et l'efficacité des LLMs, suggérant que la "profondeur" nécessaire pour une prédiction fiable peut être réduite de manière significative grâce à une propagation ciblée et minimale.