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

arXiv: https://arxiv.org/abs/2601.13288 · HTML: https://mlsanity.com/paper/2601.13288/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 16 of its 25 references that are in the map.
- First posted on arXiv 19 Jan 2026.

## Builds on (most cited first, 16 of 16)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [gpt-oss-120b & gpt-oss-20b Model Card](https://mlsanity.com/paper/2508.10925/) (2025)
- [Llama Guard: LLM-based Input-Output Safeguard for Human-AI Conversations](https://mlsanity.com/paper/2312.06674/) (2023)
- [BERT Rediscovers the Classical NLP Pipeline](https://mlsanity.com/paper/1905.05950/) (2019)
- [A Primer in BERTology: What we know about how BERT works](https://mlsanity.com/paper/2002.12327/) (2020)
- [WildGuard: Open One-Stop Moderation Tools for Safety Risks, Jailbreaks, and Refusals of LLMs](https://mlsanity.com/paper/2406.18495/) (2024)
- [EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test](https://mlsanity.com/paper/2503.01840/) (2025)
- [ToxicChat: Unveiling Hidden Challenges of Toxicity Detection in Real-World User-AI Conversation](https://mlsanity.com/paper/2310.17389/) (2023)
- [AEGIS: Online Adaptive AI Content Safety Moderation with Ensemble of LLM Experts](https://mlsanity.com/paper/2404.05993/) (2024)
- [Safety Tax: Safety Alignment Makes Your Large Reasoning Models Less Reasonable](https://mlsanity.com/paper/2503.00555/) (2025)
- [How Alignment and Jailbreak Work: Explain LLM Safety through Intermediate Hidden States](https://mlsanity.com/paper/2406.05644/) (2024)
- [What's so special about BERT's layers? A closer look at the NLP pipeline in monolingual and multilingual models](https://mlsanity.com/paper/2004.06499/) (2020)
- [LLaMAs Have Feelings Too: Unveiling Sentiment and Emotion Representations in LLaMA Models Through Probing](https://mlsanity.com/paper/2505.16491/) (2025)

## Abstract (arXiv)

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.