💻 computer science

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

This paper proposes an efficient single-pass classification framework that reuses hidden states from production LLMs via lightweight, token- and layer-selective probes, eliminating the latency and resource costs of separate safety or task-specific models while achieving competitive performance across diverse architectures.

Original authors: Gonzalo Ariel Meyoyan, Luciano Del Corro

Published 2026-04-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Gonzalo Ariel Meyoyan, Luciano Del Corro

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you run a busy restaurant (the LLM, or Large Language Model). Every time a customer orders a dish, your head chef (the model) starts cooking. In a typical setup, before the chef even touches the ingredients, you have a separate, expensive security guard standing at the door. This guard checks the order, decides if it's safe, and only then lets the chef start. If the order is bad, the guard stops it.

The Problem:
This "security guard" approach has two big downsides:

It's slow: The customer has to wait for the guard and the chef.
It's expensive: You have to pay for a whole second person (a separate computer model) just to stand there.

The Paper's Idea:
The authors ask: "Why hire a separate guard when the chef already knows if the order is weird?"

They discovered that while the chef is cooking, their brain (the model's internal "hidden states") is processing the order in many different stages. Some stages are thinking about grammar, others about meaning, and others about safety. The paper argues that the "safety signal" isn't just in one specific thought; it's scattered across the whole cooking process.

The Solution: The "Smart Taster"
Instead of hiring a new guard, the authors attach a tiny, lightweight "taster" (a probe) directly to the chef's brain. This taster doesn't stop the chef from cooking; it just listens in on the chef's thoughts as they happen.

Here is how their "Smart Taster" works, using a simple analogy:

The Chef's Thoughts (The Tensor): Imagine the chef's brain activity as a giant grid of notes. It has rows (different stages of cooking/layers) and columns (different words/tokens).
The Old Way (Fixed Reading): Previous methods were like asking the chef, "What did you think at the very start?" or "What did you think at the very end?" They picked just one spot to look.
The New Way (Token- and Layer-Selective): The authors' method is like having a super-smart manager who scans the entire grid of notes. They ask: "Which specific words and which specific stages of cooking hold the most important clues about whether this order is safe?"

How the Taster Scans:
The paper describes a two-stage process:

Stage 1 (Grouping by Word): For every single stage of cooking, the taster groups the notes about each word to create a summary.
Stage 2 (Grouping by Stage): Then, the taster looks at all those summaries from every stage and picks the most important ones to make a final decision.

They tested three types of "tasters":

The Simple Average: Just taking a quick average of all notes (like a basic pool).
The Scoring Gate: A smart filter that learns to give a "score" to which notes matter most, using very few resources.
The Deep Diver (MHA): A more complex, powerful taster that can look for subtle patterns, using a bit more energy but getting the best results.

The Results:

Speed: Because the taster works while the chef is cooking, there is no waiting for a second person. The whole process happens in one go.
Cost: The taster is tiny. It adds almost no extra memory or computing power compared to hiring a whole new "guard model" (which is huge and expensive).
Accuracy: On safety tests (like detecting toxic language) and sentiment tests (like figuring out if a movie review is positive or negative), this tiny taster performed just as well as, or sometimes better than, the expensive separate guards.

Why This Matters:
The paper shows that you don't need a separate security team to keep your AI safe. You can just train a tiny, efficient helper to listen to the AI's own internal thoughts as it works. This makes AI systems faster, cheaper to run, and simpler to manage, without sacrificing safety.

A Note on Limitations:
The authors admit that if the "bad guys" (jailbreakers) change their tactics completely, this taster might miss them because it only knows what it was trained on. Also, if the input is a massive novel instead of a short sentence, the taster might get overwhelmed by the sheer amount of data to scan. But for standard tasks, it's a highly efficient upgrade.

1. Problem Statement

Modern Large Language Model (LLM) production systems typically rely on an orchestration pattern where a primary serving LLM is surrounded by auxiliary "guard" models for safety moderation, policy compliance, and sentiment analysis. This approach suffers from three main drawbacks:

Latency: Each request requires an additional model invocation (guard model $\to$ serving LLM), doubling inference time.
Resource Overhead: Deploying separate models increases VRAM footprint and operational complexity.
Underutilization: The serving LLM already computes rich internal representations (hidden states) that contain safety and classification signals, yet these are discarded when using only the final output logits or a separate model.

Existing "reuse" methods attempt to extract signals from the serving model but often commit to fixed readout points (e.g., only the first token's logits or the final layer's hidden state), potentially missing discriminative information distributed across different layers and token positions.

2. Methodology

The authors propose a Token- and Layer-Selective Probe framework that treats classification as a representation selection problem over the full hidden-state tensor ( $L \times T \times d$ , where $L$ =layers, $T$ =tokens, $d$ =dimension). The core innovation is a Two-Stage Aggregation Architecture applied to a frozen serving LLM.

A. Architecture Overview

The probe consists of a lightweight classifier trained on the hidden states generated during the serving LLM's forward pass. It does not modify the backbone LLM.

Stage 1 (Token-Level Aggregation): For each layer $l$ , the $T$ token vectors are aggregated into a single layer summary vector $v^{(l)}$ .
Stage 2 (Layer-Level Aggregation): The $L+1$ layer summary vectors are aggregated into a single global representation vector $v$ .
Classification Head: A linear layer maps $v$ to class logits.

B. Aggregation Mechanisms

The paper evaluates three mechanisms to perform the aggregation, balancing expressiveness with parameter efficiency:

Direct Pooling: Fixed operators (Max or Mean pooling) over valid token/layer positions. (Lowest cost, ~0.003M params).
Scoring Attention Gate: A learnable mechanism that assigns scalar importance scores to tokens/layers using a single linear projection and softmax normalization. (~0.1M params).
Multi-Head Self-Attention (MHA): A downcast MHA module (using aggressive dimension reduction, e.g., $d/16$ ) to capture complex interactions between tokens and layers. (~35M params).

C. Operational Workflow

During inference, the probe runs in parallel with the generation process. If the probe predicts an "unsafe" label, the orchestration layer can immediately halt generation before any tokens are streamed, returning a refusal or re-prompting the model. This eliminates the need for a second model call.

3. Key Contributions

Shift from Fixed to Selective Readouts: The paper moves beyond fixed-layer or fixed-token readouts (like BERT's [CLS] or first-token logits) to a learned selection mechanism that identifies where in the $L \times T$ tensor the signal is most separable.
Two-Stage Aggregation Template: A novel architecture that first summarizes tokens within layers, then aggregates across layers, allowing the model to learn task-specific depth and token importance.
Efficiency vs. Performance Trade-off: Demonstrates that high-accuracy classification can be achieved with minimal trainable parameters (up to 35M) and zero additional inference latency compared to the serving model alone.
Generalization: Validates the approach across diverse architectures, including dense (Llama-3.2-3B) and Mixture-of-Experts (GPT-OSS-20B, Qwen3-30B-A3B) models.

4. Experimental Results

Safety Moderation (ToxicChat & WildGuardMix)

Performance: The MHA probe achieved 84.51 F1 on ToxicChat and 88.55 F1 on WildGuardMix.
Comparison:
- Outperformed logit-only reuse baselines (e.g., MULI: 77.8 F1).
- Matched or surpassed standalone guard models (e.g., ToxicChat-T5-large: 82.2 F1; WildGuard: 88.9 F1) despite using orders of magnitude fewer parameters (35M vs. 780M–8B).
- Zero Extra Latency: Unlike guard models which require a separate API call, the probe runs in the same forward pass.
Cross-Architecture: The performance ranking (Pooling < Scoring Attention < MHA) remained consistent across Llama, GPT-OSS, and Qwen models.

Sentiment & Emotion (IMDB, SST-2, Emotion)

The probes significantly outperformed zero-shot and few-shot prompting (e.g., 87.68% accuracy on Emotion vs. 64.05% for MULI).
Performance was competitive with large, task-specific classifiers (e.g., DeBERTa, RoBERTa) while maintaining a single-pass inference pipeline.

Efficiency Analysis

Latency: The MHA probe added only ~14ms latency per sample compared to the base model, whereas a guard-then-serve pipeline added ~60–100ms.
Memory: Probe variants kept peak GPU memory near the base model footprint (~6.5–7.0 GB), whereas guard pipelines required 8.0–22.8 GB.

Interpretability (Layer Attention)

Analysis of the Scoring Attention gate revealed that safety signals are distributed. Toxic prompts tended to activate later layers (L17–L28), while non-toxic prompts relied on final layers and embeddings. Misclassifications often showed attention patterns aligning with the predicted class rather than the ground truth, suggesting the probe learns distinct feature manifolds for different content types.

5. Significance and Conclusion

This paper provides a compelling alternative to the standard "guard model" orchestration pattern. By leveraging BERTology principles (that different layers encode different abstractions) and applying them to modern decoder-only LLMs, the authors show that:

Safety and NLU tasks are already "solved" internally by the serving model; the challenge is extracting the signal efficiently.
Single-pass classification is viable for production, offering a massive reduction in latency and infrastructure costs without sacrificing accuracy.
Learned aggregation (specifically MHA) is superior to fixed pooling, proving that the "best" layer for a task is not static but depends on the input content.

The work suggests a future where LLM orchestration is streamlined, reducing the need for heavy, separate safety pipelines and enabling more responsive, cost-effective AI systems.