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

AI summary

Overview

• Research area: Natural Language Processing — efficient inference and LLM orchestration; safety moderation and text classification via hidden-state probing.
• Technical level: Intermediate. The paper assumes familiarity with transformer hidden states, pooling, and attention, but the core idea (reading a label out of a frozen model's internal representations) is explained clearly.
• Scope: The paper proposes and evaluates two-stage, token- and layer-selective probes that attach to a frozen serving LLM and classify safety and sentiment/emotion labels in the same forward pass used for generation.

What This Paper Is About

Production LLM systems usually bolt on separate guard models or classifiers for moderation, jailbreak detection, and filtering, which adds latency, VRAM, and deployment complexity. This paper asks whether those labels can instead be read out of computations the serving LLM already performs, by training small probes on its hidden states. Rather than fixing a single readout position (such as first-token logits or final-layer pooling), the authors treat classification as representation selection over the full layer × token × hidden-dimension tensor and learn which tokens and layers matter.

Key Contributions

1. A two-stage aggregation probe architecture that first summarizes tokens within each layer (producing L+1 layer summaries) and then aggregates across layers into a single vector fed to a linear classification head, keeping the serving LLM frozen.
2. Three instantiations of the aggregation template spanning a simple-to-expressive spectrum: direct pooling (mean or max), a ~100K-parameter scoring-attention gate, and a downcast multi-head self-attention probe with up to 35M trainable parameters.
3. Multi-backbone validation on one dense model (Llama-3.2-3B-Instruct) and two mixture-of-experts models (GPT-OSS-20B, Qwen3-30B-A3B, both ~3B active parameters), showing that the method ranking (pooling < scoring attention < MHA) holds across architectures and a 10× total parameter range.
4. An inference-efficiency benchmark and a layer-attention analysis quantifying the latency/VRAM profile of probes versus guard-then-serve pipelines, and showing class-conditional differences in which layers the probe attends to.

Main Findings

• Aggregation beats fixed readouts on ToxicChat (in-distribution): On the Llama-3.2-3B backbone, direct pooling reaches 73.53 F1, the scoring-attention gate 80.49 F1 (0.10M added parameters), and the MHA probe 84.51 F1 with 0.898 AUPRC (35M added parameters). These beat the logit-reuse baseline MULI (77.8 F1 / 0.829 AUPRC) and the standalone ToxicChat-T5-large classifier (82.2 F1 / 0.885 AUPRC), the latter of which requires an extra model call.
• Cross-backbone consistency: On the same ToxicChat task, GPT-OSS-20B gives 77.36 (pooling), 79.23 (scoring), and 86.17 F1 (MHA, 0.915 AUPRC); Qwen3-30B-A3B gives 73.24, 80.94, and 83.76 F1 (0.905 AUPRC). All MHA probes match or exceed the 82.2 F1 standalone baseline, and GPT-OSS-20B achieves the highest F1 (86.17) and AUPRC (0.915).
• Cross-dataset transfer: Trained on WildGuardMix and evaluated on ToxicChat, pooling drops to 53.33 F1, scoring attention reaches 64.81 F1 / 0.706 AUPRC, and MHA reaches 72.88 F1 / 0.798 AUPRC — outperforming several guard-model and API baselines listed in the paper without an extra model call. The authors note these baselines differ in architecture and training data, so they treat the cross-dataset comparison as indicative rather than a controlled head-to-head match.
• WildGuardMix: Direct pooling reaches 82.84 F1, the scoring gate 85.98 F1 (0.10M parameters), and MHA 88.55 F1, approaching the strongest standalone guard baseline in the table (WildGuard at 88.9 F1) while training 35M parameters. MULI scores 83.79 F1.
• Sentiment and emotion: On IMDB/SST-2/Emotion, the MHA probe scores 95.15/95.39/87.68, versus MULI's 86.50/93.19/64.05, DeBERTa V3 Large's 95.34/90.38/87.65, RoBERTa Large's 94.30/95.99/84.16, and SentriLlama 3.2 (3B) Instruct's 95.79/95.94/82.20. Prompting the same backbone with zero-shot, few-shot, and chain-of-thought is markedly weaker (for example, chain-of-thought reaches 91.54/93.06/56.05).
• Learned aggregation wins within the probe family: The ordering pooling < scoring attention < MHA is consistent across all three backbones and all datasets. The authors attribute the gains mainly to where the readout occurs rather than large added capacity, since the probes use aggressive attention downcasting.
• Layer attention is class-conditional: The scoring-attention gate on ToxicChat attends to later layers (L17–L28) for toxic prompts, while non-toxic prompts concentrate on layers L0 and L27–L28 (uniform baseline is 1/(L+1) ≈ 0.034). Misclassified examples show weight profiles resembling the predicted class rather than the ground-truth label.
• Efficiency: In a latency benchmark (2000 samples, max length 512, batch size 1, max_new_tokens=1), the base Llama-3.2-3B runs at 37.84 samples/s, 26.43 ms/sample, 6497.63 MB peak GPU. Adding the pooling probe gives 33.72 samples/s and 29.66 ms/sample at unchanged memory; scoring gives 32.36 samples/s and 30.90 ms/sample at 6749.00 MB; MHA gives 24.83 samples/s and 40.27 ms/sample at 6968.92 MB. Guard-then-serve pipelines are far slower: ToxicChat-T5 (780M) 11.32 samples/s and 88.33 ms/sample at 7992 MB, WildGuard (7B) 9.97 and 100.32 at 21084 MB, Llama Guard 3 (8B) 8.12 and 123.21 at 22769 MB. Even against the smallest standalone baseline, the MHA probe is over 2× faster and uses 1 GB less peak memory.
• Hyperparameter sensitivity: About 100 configurations per dataset were explored. Self-attention sustained high PR-AUC (~0.75–0.90) with low variance across settings, while pooling and scoring attention were highly sensitive to learning rate (spanning ~0.2–0.9).

Methodology in Plain English

The researchers keep the serving LLM completely frozen and attach a small trainable classifier to it. During the model's normal forward pass, they collect the hidden states at every layer for every token, forming a three-dimensional tensor of shape layers × tokens × hidden dimension. Instead of flattening or hand-picking from this tensor, they use a two-stage aggregation: first, within each layer, tokens are combined into a single layer summary vector; second, those layer summaries are combined into one vector that a simple linear head turns into class logits. Both stages use the same aggregation operator, which the authors compare in three forms — fixed mean/max pooling, a tiny learned gate that scores each token or layer with a single linear projection and softmax weights, and a downcast multi-head self-attention block that shrinks the query/key/value dimension (e.g., by a factor of 4 to 64) to keep cost low and can use PyTorch's scaled dot-product attention with FlashAttention when available.

Training uses cross-entropy with standard hyperparameter search: learning rates from 10⁻⁵ to 10⁻³, batch sizes of 8/16/32/64, up to 10 epochs with early stopping, AdamW (β₁ = 0.9, β₂ = 0.999), weight decay in [0, 0.05], cosine annealing, and mean-versus-max pooling choices. To fit training in memory, hidden states are optionally precomputed and cached before probe training, decoupling backbone inference from classifier training. Probe training ran on a single NVIDIA RTX 3090 (24GB VRAM, 96GB system RAM); hidden states for the two additional backbones were pre-extracted on an NVIDIA A100 80GB GPU.

Why This Matters

• For research: The paper reframes moderation and NLU classification as a search over where in a transformer's depth and token sequence a signal is most separable, rather than as a question of whether the serving model's computation can be reused. It extends classical BERTology findings about layer-wise abstraction to decoder-only and mixture-of-experts LLMs.
• Real-world applications:
  • Content safety moderation in chat products, screening prompts for toxicity and jailbreak-style inputs before generation is streamed.
  • Sentiment and emotion tagging over user reviews and feedback, run alongside generation in a single pass.
  • Policy compliance and retrieval filtering inside orchestration pipelines, where an extra classifier invocation per request is costly.
  • Latency- or memory-constrained deployments (the reported probe footprint stays near 6.5–7.0 GB, versus 8.0–22.8 GB for guard-then-serve pipelines).
• Industry relevance: The approach claims a path to additional safety without training or serving a separate multi-billion-parameter guard model, requiring only up to 35M trainable parameters beyond the backbone. It also integrates directly into existing orchestration: if a prompt is classified unsafe, the orchestration layer can halt generation before tokens are streamed and return a templated refusal or re-prompt the same model, with no additional model invocation. The authors frame this against the "alignment tax" of safety-tuning the main model, and the paper reports an open-source code repository.

Future Directions

• Broaden cross-backbone evaluation: Sentiment and emotion benchmarks were run only on Llama-3.2-3B-Instruct, and all three tested backbones share roughly 3B active parameters, leaving the effect of substantially larger active capacity untested.
• Handle longer sequences: VRAM constraints limited longer inputs with larger batches; full documents and multi-turn conversations may require further memory optimization.
• Determine minimum data requirements: The smallest dataset used was SST-2 with 7,000 examples, and the minimum dataset size needed — especially for parameter-heavier MHA probes — remains unclear, raising the question of few-shot or transfer-learning approaches.
• Integrate conditional refusal generation and new tasks: The authors did not evaluate conditional re-prompting strategies that would let the serving model produce a contextual rejection after probe detection, and they propose extending multi-backbone evaluation to additional tasks and fine-grained safety taxonomies.
• Address robustness under distribution shift: Out-of-distribution results show degradation, and the authors warn that probes may fail silently on novel harm categories or evolving jailbreak strategies, since they cannot adapt without retraining.

Target Audience

This paper is most useful for applied ML engineers and researchers working on LLM serving, orchestration, and safety infrastructure who need efficient alternatives to separate guard-model pipelines. It also suits interpretability researchers interested in how task-relevant information is distributed across layers and tokens, and practitioners evaluating the accuracy-versus-latency-versus-VRAM trade-offs of attaching lightweight classifiers to frozen backbones. Readers should be comfortable with transformer internals such as hidden states, attention, and pooling, and with standard classification metrics (F1, AUPRC, accuracy, ROC AUC).

Note: the provided paper content is truncated. The appendix table reporting the downcast-factor ablation is cut off mid-row (showing ToxicChat results for downcast factors 4, 8, and a partial row for 16), and figures 3 and 4 are referenced but only Figure 2's content is described in the text.

Authors’ 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. 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.