Una prospettiva BERTologica sulle orchestrazioni di LLM: Sonde selettive per token e strati per una classificazione efficiente in passaggio singolo

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

January 19, 2026

Autori: Gonzalo Ariel Meyoyan, Luciano Del Corro

cs.AI

Abstract

I sistemi LLM di produzione spesso si basano su modelli separati per la sicurezza e altre fasi ad alta intensità di classificazione, aumentando la latenza, l'occupazione di VRAM e la complessità operativa. Noi, invece, riutilizziamo il calcolo già pagato dall'LLM di servizio: addestriamo "sonde" leggere sui suoi stati nascosti e prevediamo le etichette nello stesso passaggio in avanti utilizzato per la generazione. Inquadriamo la classificazione come una selezione di rappresentazioni sul tensore completo degli stati nascosti token-layer, piuttosto che impegnarci su un token fisso o un layer fisso (ad esempio, i logit del primo token o il pooling dell'ultimo layer). Per implementare ciò, introduciamo un aggregatore a due stadi che (i) riassume i token all'interno di ogni layer e (ii) aggrega i riassunti dei layer per formare una singola rappresentazione per la classificazione. Istanziamo questo modello con pooling diretto, un gate di attenzione di scoring da 100K parametri e una sonda basata su self-attention multi-testina (MHA) a precisione ridotta con fino a 35 milioni di parametri addestrabili. Su benchmark di sicurezza e sentiment, le nostre sonde migliorano rispetto al riutilizzo dei soli logit (ad es., MULI) e sono competitive con baseline specifiche per task sostanzialmente più grandi, preservando una latenza quasi pari a quella di servizio ed evitando i costi in VRAM e latenza di una pipeline separata con modello di guardia.

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.