A BERTology View of LLM Orchestrations: Token- and Layer-Selective Probes for Efficient Single-Pass Classification

1. Énoncé du problème

Les systèmes de production modernes de modèles de langage de grande taille (LLM) reposent généralement sur un modèle d'orchestration où un LLM de service principal est entouré de modèles auxiliaires « gardiens » pour la modération de la sécurité, la conformité aux politiques et l'analyse des sentiments. Cette approche souffre de trois principaux inconvénients :

• Latence : Chaque requête nécessite un appel de modèle supplémentaire (modèle gardien $\to$ LLM de service), doublant le temps d'inférence.
• Surcharge des ressources : Le déploiement de modèles séparés augmente l'empreinte VRAM et la complexité opérationnelle.
• Sous-utilisation : Le LLM de service calcule déjà de riches représentations internes (états cachés) contenant des signaux de sécurité et de classification, pourtant ces derniers sont jetés lorsqu'on utilise uniquement les logits de sortie finale ou un modèle séparé.

Les méthodes de « réutilisation » existantes tentent d'extraire des signaux du modèle de service mais s'engagent souvent sur des points de lecture fixes (par exemple, uniquement les logits du premier token ou l'état caché de la dernière couche), risquant ainsi de manquer des informations discriminatives distribuées à travers différentes couches et positions de tokens.

2. Méthodologie

Les auteurs proposent un cadre de sonde sélective par token et par couche qui traite la classification comme un problème de sélection de représentations sur le tenseur complet des états cachés ($L \times T \times d$, où $L$=couches, $T$=tokens, $d$=dimension). L'innovation centrale est une architecture d'agrégation en deux étapes appliquée à un LLM de service figé.

A. Aperçu de l'architecture

La sonde consiste en un classifieur léger entraîné sur les états cachés générés lors du passage avant du LLM de service. Il ne modifie pas le LLM de base.

1. Étape 1 (Agrégation au niveau des tokens) : Pour chaque couche $l$, les $T$ vecteurs de tokens sont agrégés en un seul vecteur de résumé de couche $v^{(l)}$.
2. Étape 2 (Agrégation au niveau des couches) : Les $L+1$ vecteurs de résumé de couche sont agrégés en un seul vecteur de représentation globale $v$.
3. Tête de classification : Une couche linéaire mappe $v$ vers les logits de classe.

B. Mécanismes d'agrégation

L'article évalue trois mécanismes pour effectuer l'agrégation, équilibrant expressivité et efficacité des paramètres :

1. Regroupement direct (Direct Pooling) : Opérateurs fixes (Max ou Moyenne) sur les positions valides de tokens/couches. (Coût le plus faible, ~0,003M de paramètres).
2. Porte d'attention par notation (Scoring Attention Gate) : Un mécanisme apprenable qui attribue des scores d'importance scalaires aux tokens/couches en utilisant une projection linéaire unique et une normalisation softmax. (~0,1M de paramètres).
3. Attention auto-à-têtes multiples (MHA) : Un module MHA réduit (utilisant une réduction agressive de dimension, par exemple $d/16$) pour capturer des interactions complexes entre tokens et couches. (~35M de paramètres).

C. Flux de travail opérationnel

Pendant l'inférence, la sonde s'exécute en parallèle avec le processus de génération. Si la sonde prédit une étiquette « non sûre », la couche d'orchestration peut immédiatement arrêter la génération avant que des tokens ne soient streamés, renvoyant un refus ou une nouvelle demande au modèle. Cela élimine le besoin d'un deuxième appel de modèle.

3. Contributions clés

• Passage des lectures fixes aux lectures sélectives : L'article va au-delà des lectures de couche ou de token fixes (comme [CLS] de BERT ou les logits du premier token) pour adopter un mécanisme d'apprentissage de sélection qui identifie où dans le tenseur $L \times T$ le signal est le plus séparable.
• Modèle d'agrégation en deux étapes : Une architecture novatrice qui résume d'abord les tokens au sein des couches, puis agrège à travers les couches, permettant au modèle d'apprendre l'importance spécifique à la tâche de la profondeur et des tokens.
• Compromis Efficacité vs Performance : Démontre qu'une classification de haute précision peut être atteinte avec un nombre minimal de paramètres entraînables (jusqu'à 35M) et une latence d'inférence supplémentaire nulle par rapport au seul modèle de service.
• Généralisation : Valide l'approche sur des architectures diverses, incluant des modèles denses (Llama-3.2-3B) et des mélanges d'experts (GPT-OSS-20B, Qwen3-30B-A3B).

4. Résultats expérimentaux

Modération de la sécurité (ToxicChat & WildGuardMix)

• Performance : La sonde MHA a atteint 84,51 F1 sur ToxicChat et 88,55 F1 sur WildGuardMix.
• Comparaison :
  • A surpassé les baselines de réutilisation basées uniquement sur les logits (par exemple, MULI : 77,8 F1).
  • A égalé ou dépassé les modèles gardiens autonomes (par exemple, ToxicChat-T5-large : 82,2 F1 ; WildGuard : 88,9 F1) malgré l'utilisation de plusieurs ordres de grandeur moins de paramètres (35M contre 780M–8B).
  • Latence supplémentaire nulle : Contrairement aux modèles gardiens qui nécessitent un appel API séparé, la sonde s'exécute dans le même passage avant.
• Transversalité des architectures : Le classement des performances (Regroupement < Attention par notation < MHA) est resté cohérent à travers les modèles Llama, GPT-OSS et Qwen.

Sentiment et Émotion (IMDB, SST-2, Emotion)

• Les sondes ont nettement surpassé le prompting zero-shot et few-shot (par exemple, 87,68 % de précision sur Emotion contre 64,05 % pour MULI).
• Les performances étaient compétitives par rapport à de grands classifieurs spécifiques à la tâche (par exemple, DeBERTa, RoBERTa) tout en maintenant un pipeline d'inférence en un seul passage.

Analyse de l'efficacité

• Latence : La sonde MHA n'a ajouté que ~14 ms de latence par échantillon par rapport au modèle de base, tandis qu'un pipeline gardien-then-service ajoutait ~60–100 ms.
• Mémoire : Les variantes de sonde ont maintenu la mémoire GPU de pointe près de l'empreinte du modèle de base (~6,5–7,0 Go), alors que les pipelines gardiens nécessitaient 8,0–22,8 Go.

Interprétabilité (Attention par couche)

• L'analyse de la porte d'attention par notation a révélé que les signaux de sécurité sont distribués. Les prompts toxiques avaient tendance à activer les couches ultérieures (L17–L28), tandis que les prompts non toxiques reposaient sur les couches finales et les embeddings. Les mauvaises classifications montraient souvent des motifs d'attention alignés avec la classe prédite plutôt que la vérité terrain, suggérant que la sonde apprend des variétés de caractéristiques distinctes pour différents types de contenu.

5. Importance et conclusion

Cet article offre une alternative convaincante au modèle d'orchestration standard de « modèle gardien ». En exploitant les principes de la BERTologie (selon lesquels différentes couches encodent différentes abstractions) et en les appliquant aux LLM modernes de type décodeur uniquement, les auteurs montrent que :

1. Les tâches de sécurité et de compréhension du langage naturel (NLU) sont déjà « résolues » en interne par le modèle de service ; le défi consiste à extraire le signal efficacement.
2. La classification en un seul passage est viable pour la production, offrant une réduction massive de la latence et des coûts d'infrastructure sans sacrifier la précision.
3. L'agrégation apprise (spécifiquement MHA) est supérieure au regroupement fixe, prouvant que la « meilleure » couche pour une tâche n'est pas statique mais dépend du contenu de l'entrée.

Ce travail suggère un avenir où l'orchestration des LLM est rationalisée, réduisant le besoin de pipelines de sécurité lourds et séparés, et permettant des systèmes d'IA plus réactifs et rentables.