arXiv:2601.13288cs.CL2026-01ACL

用大模型隐状态做轻量分类,不加额外模型却更准更快。

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

  • 在大模型的每一层隐藏状态中选关键信息,统一处理而非固定取某层或某词。
  • 在安全与情感任务上优于仅用最终输出的方法,性能接近专用大模型。
  • 适合想省显存、低延迟部署分类任务的工程团队使用。

生产级大模型系统常需独立模型完成安全检测等分类任务,导致延迟高、显存占用大、运维复杂。本文提出复用生成时已产生的大模型隐藏状态,训练轻量探针,在同一前向传播中完成分类。将分类视为对完整词-层隐藏状态张量的表征选择,而非固定选取某一层或某一词。设计两阶段聚合器:(i) 每层内汇总词信息,(ii) 跨层聚合形成最终分类表示。实现方式包括直接池化、10万参数的评分注意力门控,以及最多3500万可训练参数的下采样多头自注意探针。在安全与情感多个基准测试中,其性能优于仅复用最终输出(如MULI)的方法,且媲美显著更大的专用模型,同时保持接近原服务的延迟,避免额外显存与延迟开销。在密集结构与专家混合架构(Llama-3.2-3B、GPT-OSS-20B、Qwen3-30B-A3B)上的多主干实验表明,该方法具备跨模型泛化能力。

原文摘要 · Abstract (English)

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.

大模型推理轻量探针高效分类

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。