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

TL;DR

The paper tackles the latency and complexity of production LLM safety and moderation by reusing computation in a frozen serving model through lightweight probes that read the full hidden-state tensor across all layers and tokens. It proposes a two-stage, token-then-layer aggregation framework with three aggregation mechanisms (direct pooling, scoring attention gate, and downcast multi-head self-attention) to produce a compact representation for a linear classifier, enabling single-pass classification without extra model invocations. Across safety benchmarks (ToxicChat, WildGuardMix) and sentiment tasks (IMDB, SST-2, Emotion), the MHA-based variant often achieves the strongest performance while incurring modest added parameters (~35M) and maintaining near-serve latency, outperforming logit-based reuse and approaching task-specific baselines. The approach demonstrates how task-relevant signals are distributed across depth and tokens, supports deployment in real-time systems, and highlights considerations around generalization, sequence length, data requirements, and privacy for practical adoption.

Abstract

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.

Figures (4)

• Figure 1: View of our probing architecture. The aggregation method featured here is the scoring attention gate
• Figure 2: Attention weights from the scoring-attention probe on ToxicChat, stratified by label and correctness. The dashed line marks uniform $1/L$ ($\approx 0.034$). Toxic prompts attend to later layers (L17--L28), while non-toxic prompts concentrate on layers L0 and L27--L28. Misclassifications resemble the predicted-class profile.
• Figure 3: Sensitivity of PR-AUC to hyperparameter choices on ToxicChat across ${\sim}100$ configurations. Each row corresponds to an aggregation method; horizontal lines denote medians. Self-Attention sustains high PR-AUC (${\sim}0.75$–$0.90$) with low variance across settings. Pooling and Scoring Attention are highly sensitive to learning rate, spanning ${\sim}0.2$–$0.9$.
• Figure 4: Layer-wise post-softmax attention weights on SST-2, stratified by ground-truth label (Positive vs. Negative) and prediction correctness. Similar to toxicity detection (Figure \ref{['fig:layer-attention-tc']}), correctly classified samples exhibit distinct class-conditional attention patterns: positive sentiment concentrates on later intermediate layers (L17--L28), while negative sentiment shows concentrated attention weights on the embedding and final layers (L0, L28).