PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Gonzalo Ariel Meyoyan, Luciano Del Corro

arXiv:2601.13288v2cs.CL

TL;DR

Production LLM systems often add separate classifiers for safety and related tasks, increasing deployment cost. This paper trains lightweight probes over the serving model’s full token-layer hidden states, using two-stage aggregation in the same forward pass. Across safety and sentiment tasks, the probes improve over logit-only reuse and remain competitive with larger baselines across multiple backbones.

  • Problem

    Production LLM orchestration commonly adds auxiliary classifiers that increase training, deployment, and inference costs, motivating reuse of the serving model’s computation.

  • Method

    The paper trains lightweight aggregators on frozen serving-LLM hidden states, selecting and combining information across tokens and layers through two-stage aggregation.

  • Results

    Across safety and sentiment benchmarks, the probes match or surpass strong baselines, while the MHA probe reaches 84.51 F1 and 0.898 AUPRC on ToxicChat.

  • Takeaways & Limitations

    Hidden-state reuse enables single-pass classification alongside generation, with competitive performance and substantially lower parameter and deployment overhead than separate guard-model pipelines.

  • Takeaways & Limitations

    Sentiment and emotion were evaluated only on Llama-3.2-3B-Instruct, and substantially larger active capacities remain untested.

Abstract

from arXiv · show

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.

1 Introduction

Production LLM deployments add auxiliary classifiers that increase compute and operational complexity, motivating reuse of the serving model’s hidden states. The paper frames classification as selecting and aggregating token- and layer-level representations with lightweight probes.

  • Auxiliary safety and classification components add training, deployment, and inference costs, often requiring an extra model invocation per request.
  • Prior reuse methods read classification signals from output distributions or fixed latent representations, but transformers expose depth-indexed hidden states with different abstractions.
  • The paper treats moderation and NLU classification as representation selection over the full L×T×d hidden-state tensor.
  • A two-stage aggregator summarizes tokens within each layer and then combines layer summaries into one classification representation.
  • Across safety and sentiment benchmarks, probes match or surpass strong baselines, generalize across three backbones, and classify alongside generation in one model call.

2 Related Work

Related work reduces orchestration cost by extracting classification signals from the serving LLM, while BERTology motivates using information distributed across transformer depth. This paper extends those directions with joint token×layer readouts attached to a frozen serving model.

  • Guard models improve safety but add latency, memory, and deployment complexity compared with reuse of the serving model’s forward pass.
  • Prior approaches use first-token logits, fixed latent representations, final-layer states, or single-layer probes for classification.
  • Transformer layers encode different abstractions, and prior analyses suggest that combining layers can outperform selecting one layer.
  • The proposed readout learns jointly over tokens and layers instead of fixing a depth or position, while running on a frozen serving LLM.
  • Figure 1 depicts the two-stage architecture: token summaries are formed within layers, then aggregated across layers into a classification vector.

3 Probe Architecture

The probe consumes hidden states from every layer and token of a frozen decoder-only LLM, compressing them through token-level and layer-level aggregation before classification. Direct pooling, scoring attention, and downcast MHA provide increasing expressiveness under an overhead constraint.

  • Probe Architecture: A frozen decoder-only LLM produces h^(l) for T tokens at each layer l, and only the probe parameters θ are trained.
  • Two-Stage Aggregation Architecture: The architecture reduces the L×T×d hidden-state tensor to L+1 layer summaries, then to one d-dimensional vector for a linear classification head.
  • Two-Stage Aggregation Architecture: Stage 1 aggregates tokens within each layer, while Stage 2 aggregates across layers and learns task-specific layer weighting instead of using manual selection.
  • Classification Head: The probe emits a label during the same forward pass that initiates generation, allowing unsafe inputs to halt generation without another model invocation.
  • Aggregation Mechanisms: The aggregation operator must balance expressiveness with parameter and inference overhead.
  • Aggregation Mechanisms: The three mechanisms range from fixed mean or max pooling to learned scoring gates and downcast multi-head self-attention.
  • Aggregation Mechanisms: Scoring attention assigns normalized importance weights to valid tokens or layers, while MHA projects into a smaller dimension before pooling.
  • Aggregation Mechanisms: The scoring gate uses (L+2)d trainable parameters, whereas the MHA design uses L+2 modules with downcast projections.

4 Experiments and Results

The experiments evaluate frozen-backbone probes for safety and sentiment classification across datasets and architectures, emphasizing reuse of the serving model’s single forward pass. Learned token–layer aggregation consistently improves over pooling and logit-only reuse while remaining competitive with standalone classifiers.

  • Safety classification: 84.51 F1 and 0.898 AUPRC: the MHA probe surpasses MULI and ToxicChat-T5 on in-distribution ToxicChat without an extra model invocation.MULI reaches 77.8 F1 / 0.829 AUPRC, while ToxicChat-T5 reaches 82.2 F1 / 0.885 AUPRC.
  • Safety classification: 86.17 F1 and 0.915 AUPRC: GPT-OSS-20B gives the strongest MHA result among the three ToxicChat backbones.The pooling < scoring-attention < MHA ranking is preserved across Llama-3.2-3B, GPT-OSS-20B, and Qwen3-30B-A3B.
  • Safety classification: 72.88 F1 / 0.798 AUPRC: MHA leads under WildGuardMix-to-ToxicChat distribution shift, while scoring attention reaches 64.81 F1 / 0.706 AUPRC.The results show transfer remains possible without deploying a separate guard LLM, although performance drops relative to in-distribution evaluation.
  • Safety classification: 88.55 F1: MHA approaches WildGuard’s 88.9 F1 on WildGuardMix while training 35M parameters and reusing the serving pass.Scoring attention reaches 85.98 F1 with 0.10M added parameters, compared with 83.79 F1 for MULI.
  • Sentiment and emotion classification: Learned aggregation outperforms direct pooling across sentiment and emotion tasks, with MHA providing the best overall reuse performance.Pooling underperforms significantly on Emotion at 69.18, while the scoring-attention gate uses only 0.10–0.11M added parameters and MHA is especially beneficial for Emotion.
  • Sentiment and emotion classification: Reuse probes substantially outperform prompting and remain competitive with task-specific classifiers on IMDB, SST-2, and Emotion in a single serving pass.Compared with MULI, the probes are especially stronger on Emotion, where the reported scores are 87.68 versus 64.05.

5 Layer Attention Analysis

Layer-attention analysis shows that safety-relevant evidence is distributed across depth rather than localized to one layer. The probe’s layer-weight profiles vary by content class and prediction correctness, with errors often resembling the predicted class.

  • Layer Attention Patterns: No single layer dominates across all groups; correctly classified toxic prompts distribute attention across later layers, especially L17–L28.These layers receive weight above the uniform baseline, indicating that nearby intermediate-to-late representations contribute to toxic classification.
  • Layer Attention Patterns: Correctly classified non-toxic prompts concentrate most attention on final layers L27–L28, with a smaller contribution from embedding layer L0.
  • Interpretation: The class-conditional profiles support learned multi-layer aggregation because fixed single-layer readouts can discard evidence from nearby layers.The analysis extends layer-wise representation findings to decoder-only LLMs; similar patterns also appear in sentiment classification.
  • Misclassification Profiles: Misclassified toxic prompts predicted as non-toxic concentrate attention on L28, while misclassified non-toxic prompts show the distributed profile typical of toxic predictions.Errors therefore align more with the predicted class than the ground-truth label, possibly reflecting atypical representations, label noise, or ambiguity.
  • Benchmark Setup: Table 5 benchmarks latency across probe rows reusing the serving forward pass and guard-model rows running a separate model before the serving LLM.The benchmark uses 2,000 ToxicChat prompts with maximum length 512, batch size 1, and max_new_tokens=1.

6 Inference Efficiency Analysis

Single-pass probes reduce deployment overhead by deriving classifications from the serving LLM’s existing forward pass instead of invoking a separate guard model. They remain near the serving memory footprint, with pooling and scoring attention favoring efficiency and MHA providing higher accuracy at moderate overhead.

  • Latency and Memory: The MHA probe is over 2× faster than ToxicChat-T5 and uses 1 GB less peak memory, despite adding attention computation.ToxicChat-T5 is the smallest standalone guard baseline at 780M parameters.
  • Latency–Accuracy Trade-off: Pooling adds minimal overhead, scoring attention remains close to serving-only latency, and MHA is substantially faster than any guard-then-serve pipeline.
  • Memory Footprint: Probe variants remain near the serving footprint at 6.5–7.0 GB, while guard-then-serve pipelines range from 8.0 GB to 22.8 GB.The guard-model range depends on the standalone model, spanning ToxicChat-T5 and Llama Guard 3.
  • Latency–Accuracy Trade-off: Within the probe family, pooling and scoring attention offer the best latency–accuracy trade-off, whereas MHA provides higher accuracy at moderate single-pass overhead.

7 Conclusions and Future Work

The paper concludes that lightweight probes on frozen serving-LLM hidden states can support moderation and sentiment or emotion classification without a separate classifier model. Full token–layer aggregation preserves parameter efficiency while avoiding commitment to a single readout location.

  • Conclusions: Lightweight probes on frozen serving-LLM hidden states support moderation and sentiment or emotion classification without deploying a separate classifier model.
  • Conclusions: Treating classification as representation selection over the full L × T × d tensor enables two-stage token–layer aggregation instead of a fixed layer or token.
  • Results: Across ToxicChat and WildGuardMix, the probes approach or match strong guard-model baselines with far fewer trainable parameters.

Limitations

The study’s limitations concern generalization, sequence length, training-data requirements, response generation, and robustness under distribution shift. These boundaries leave larger active capacities, long inputs, limited labels, conditional refusal strategies, and evolving harm categories insufficiently evaluated.

  • Generalization Across Models and Architectures: Sentiment and emotion benchmarks use only Llama-3.2-3B-Instruct, while all three tested backbones have approximately 3B active parameters.The effect of substantially larger active capacity remains untested, and different architectures may require architecture-specific aggregation tuning.
  • Sequence Length Constraints: Longer inputs such as full documents and multi-turn conversations may create memory challenges requiring optimization techniques.
  • Training Data Requirements: The minimum dataset size needed to train these probes effectively remains unclear, particularly for parameter-intensive multi-head attention variants.The smallest dataset used contained 7,000 SST-2 examples.
  • Response Generation Limitation: The approach detects harmful content and interrupts generation but does not itself generate explanatory refusals; conditional reprompting was not evaluated.
  • Robustness and Failure Scenarios: Distribution shift degrades probe performance, and probes may fail silently on novel harm categories or evolving jailbreak strategies without retraining.In safety-critical settings requiring broader or changing coverage, a dedicated guard model may remain necessary despite higher computational cost.

Ethical Considerations

The paper identifies misuse, bias, and privacy risks associated with probing and caching LLM hidden states. It also notes that lightweight probes make extensive experimentation feasible compared with fine-tuning large guard models.

  • Hidden-state safety signals could inform adversarial attacks designed to evade detection, although the authors judge the efficiency benefits to outweigh this risk.
  • Probes inherit biases and representation gaps from the frozen backbone, which can produce disparate error rates across demographic groups.
  • Cached hidden representations may contain sensitive prompt information and should receive access controls and limited retention.
  • Approximately 100 configurations per dataset enable systematic sensitivity analysis that would be prohibitive when fine-tuning multi-billion-parameter guard models.

A.3 Sensitivity Analysis

Sensitivity varies substantially by aggregation mechanism. Self-Attention is comparatively robust across hyperparameter settings, whereas Pooling and Scoring Attention are especially sensitive to learning rate; downcasting also has dataset-dependent optima.

  • Self-Attention sustains high PR-AUC of ∼0.75–0.90 with low variance across approximately 100 ToxicChat configurations.
  • Pooling and Scoring Attention span approximately 0.2–0.9 PR-AUC across settings and are highly sensitive to learning rate.
  • Downcast 32 is optimal for ToxicChat, whereas downcast 8 achieves the best results on SST-2.

B Ablation Study on Attention Downcasting

The downcasting ablation evaluates reduced-dimensional QKV projections under otherwise fixed settings and shows that the preferred downcasting factor depends on the dataset. Layer-wise sentiment attention patterns further support aggregating information across depths.

  • Ablation setup: The ablation varies the downcasting factor for MHA on ToxicChat and SST-2 while keeping the main-results hyperparameter configuration fixed.
  • Ablation setup: The downcast factor sets the QKV inner dimension to d/factor, with d = 3,072 for Llama-3.2-3B.
  • Layer-wise patterns: Correctly classified sentiment samples show class-conditional attention patterns distributed across embedding, intermediate, and final layers.
Loading 2601.13288v2…