Een BERTologie-perspectief op LLM-orchestraties: Token- en laagselectieve probes voor efficiënte classificatie in één doorloop

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

January 19, 2026

Auteurs: Gonzalo Ariel Meyoyan, Luciano Del Corro

cs.AI

Samenvatting

Productie-LLM-systemen vertrouwen vaak op aparte modellen voor veiligheid en andere classificatie-intensieve stappen, wat de latentie, VRAM-behoefte en operationele complexiteit verhoogt. Wij hergebruiken in plaats daarvan reeds betaalde rekencapaciteit van het servende LLM: we trainen lichtgewicht probes op zijn verborgen toestanden en voorspellen labels in dezelfde voorwaartse pass die voor generatie wordt gebruikt. We benaderen classificatie als representatieselectie over de volledige token-laag verborgen-toestand tensor, in plaats van ons te committeren aan een vaste token of vaste laag (bijv. eerste-token logits of final-layer pooling). Om dit te implementeren, introduceren we een tweetraps-aggregator die (i) tokens binnen elke laag samenvat en (ii) over de laagsamenvattingen aggregeert om één enkele representatie voor classificatie te vormen. We concretiseren deze template met direct pooling, een 100K-parameter scoring-attention gate, en een downcast multi-head self-attention (MHA) probe met tot 35M trainbare parameters. Op veiligheids- en sentimentbenchmarks overtreffen onze probes logit-only hergebruik (bijv. MULI) en zijn ze concurrerend met aanzienlijk grotere taakspecifieke baselines, terwijl ze de latentie bijna gelijk aan serven behouden en de VRAM- en latentiekosten van een apart guard-model pipeline vermijden.

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.