A BERTology View of LLM Orchestrations: Token- and Layer-Selective Probes for Efficient Single-Pass Classification
Year: 2026
Type: conference-paper
Abstract: Production LLM systems often rely on separate models for safety and other classificationheavy 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 guardmodel pipeline.Multi-backbone experiments on dense and mixture-of-experts architectures (Llama-3.2-3B,GPT-OSS-20B, Qwen3-30B-A3B) confirm that these findings generalize beyond a single model family.
Authors: Gonzalo Ariel Meyoyan, Luciano Del Corro
Institution: Universidad de Buenos Aires
Language: English
Topic: Topic Modeling
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: hybrid
OpenAlex ID: W7166797743