Une perspective BERTologique sur les orchestration des LLM : sondages sélectifs en tokens et couches pour une classification efficace en passe unique
A BERTology View of LLM Orchestrations: Token- and Layer-Selective Probes for Efficient Single-Pass Classification
January 19, 2026
papers.authors: Gonzalo Ariel Meyoyan, Luciano Del Corro
cs.AI
papers.abstract
Les systèmes de production de LLM reposent souvent sur des modèles distincts pour la sécurité et d'autres étapes à forte composante classification, ce qui augmente la latence, l'empreinte mémoire VRAM et la complexité opérationnelle. Nous réutilisons plutôt le calcul déjà effectué par le LLM de service : nous entraînons des sondes légères sur ses états cachés et prédisons les étiquettes lors de la même passe avant utilisée pour la génération. Nous abordons la classification comme une sélection de représentations sur le tenseur complet des états cachés (couches × tokens), plutôt que de nous restreindre à un token ou une couche fixe (par exemple, les logits du premier token ou le pooling de la dernière couche). Pour implémenter cette approche, nous introduisons un agrégateur à deux étapes qui (i) résume les tokens au sein de chaque couche et (ii) agrège ces résumés de couches pour former une représentation unique destinée à la classification. Nous instancions ce modèle avec un pooling direct, une porte d'attention par scoring de 100K paramètres, et une sonde à auto-attention multi-têtes (MHA) réduite avec jusqu'à 35M paramètres entraînables. Sur des benchmarks de sécurité et de sentiment, nos sondes surpassent la réutilisation des logits seuls (par exemple, MULI) et sont compétitives avec des modèles de référence spécialisés substantiellement plus grands, tout en préservant une latence proche de celle du service et en évitant les coûts VRAM et de latence d'un pipeline séparé de modèle de garde.
English
Production LLM systems often rely on separate models for safety and other classification-heavy steps, increasing latency, VRAM footprint, and operational complexity. We instead reuse computation already paid for by the serving LLM: we train lightweight probes on its hidden states and predict labels in the same forward pass used for generation. We frame classification as representation selection over the full token-layer hidden-state tensor, rather than committing to a fixed token or fixed layer (e.g., first-token logits or final-layer pooling). To implement this, we introduce a two-stage aggregator that (i) summarizes tokens within each layer and (ii) aggregates across layer summaries to form a single representation for classification. We instantiate this template with direct pooling, a 100K-parameter scoring-attention gate, and a downcast multi-head self-attention (MHA) probe with up to 35M trainable parameters. Across safety and sentiment benchmarks our probes improve over logit-only reuse (e.g., MULI) and are competitive with substantially larger task-specific baselines, while preserving near-serving latency and avoiding the VRAM and latency costs of a separate guard-model pipeline.