LLM编排的BERTology视角：基于Token和层选择性探针的高效单次分类 A BERTology View of LLM Orchestrations: Token- and Layer-Selective Probes for Efficient Single-Pass Classification
复用LLM隐藏状态训练轻量探针，单次前向传播实现安全分类
前置知识
Transformer隐藏状态（Hidden States）
在Transformer架构中，每一层都会产生一个隐藏状态张量，形状为 $T \times d$（序列长度 $T$ 乘以隐藏维度 $d$）。对于一个 $L$ 层的解码器模型，完整的隐藏状态张量形状为 $L \times T \times d$，即层数×token数×隐藏维度。每一层对输入的表示进行了不同层次的抽象变换：低层更多编码句法模式，高层编码语义信息。传统方法通常只使用最后一层或第一层token的logits，而本文探索如何利用全部 $L \times T \times d$ 的张量信息。
本文的核心就是从这个完整的三维张量中选择和聚合信息，而非固定使用某一位置的表示，理解隐藏状态的多层结构是读懂本文方法的前提。
LLM编排（LLM Orchestration）
在生产环境中，一个服务LLM通常不是孤立运行的，而是被一系列辅助组件包围：安全审核、越狱检测、策略合规、检索过滤等。这种模式被称为LLM编排。每个额外的分类器都会增加训练评估流水线、部署面和推理计算，通常每次请求需要额外调用一次模型。这种模式虽然有效但代价高昂：额外的延迟、显存占用和运维复杂度。
本文的动机就是要解决LLM编排中的计算冗余问题，理解编排模式的具体痛点才能理解本文的实际价值。
BERTology
BERTology是一系列研究Transformer内部工作机制的学术工作总称。经典研究发现Transformer各层像流水线一样工作：不同深度编码不同抽象层次的信息（Tenney et al., 2019; Jawahar et al., 2019）。低层捕获句法模式，高层编码语义。后续分析表明，任务相关信息分布在不同深度，组合多层通常优于单层选择（de Vries et al., 2020; Rogers et al., 2020）。本文将这一来自BERT时代的洞见扩展到仅解码器的大语言模型上。
论文标题直接点明BERTology视角，核心论点是LLM各层信息分布不均，需要学习聚合而非固定读出，这直接源于BERTology研究传统。
探针（Probe）
探针是一种轻量级分类器，附加在冻结的预训练模型上，通过训练探针参数来评估模型内部表示是否包含某类任务的可分离信息。探针通常是一个线性层或小型MLP，参数量极少（本文最小仅3K参数），训练时不更新主模型权重。在本文中，探针附加在服务LLM的隐藏状态上，在与生成相同的前向传播中完成分类预测。
探针是本文方法的核心载体，理解探针的工作方式和设计空间才能理解本文的三种聚合机制及其参数量差异。
Mixture-of-Experts（MoE）
MoE是一种稀疏路由架构，每个token只激活模型参数的一个子集（专家）。例如GPT-OSS-20B总参数200亿但每次仅激活3.6B参数，Qwen3-30B-A3B总参数300亿但仅激活3B参数。本文使用MoE架构验证方法的通用性，因为MoE的稀疏路由机制与密集Transformer有本质区别，如果探针在MoE上同样有效，说明方法不依赖于特定的架构特性。
多骨干网实验是本文验证泛化性的关键设计，涉及密集和稀疏两种架构范式，理解MoE的区别才能理解跨架构验证的意义。
研究动机
现代LLM部署中，服务模型通常需要配合多个辅助分类器使用：安全审核、越狱检测、策略合规、检索过滤等（Inan et al., 2023; Han et al., 2024）。以ToxicChat-T5-large（780M参数）为例，作为独立的安全分类器，它需要一次额外的模型调用，增加约62ms的延迟和1.5GB的显存占用。WildGuard（7B）和Llama Guard 3（8B）则分别增加约74ms和97ms延迟，显存占用分别增加14.6GB和16.3GB。这种guard-then-serve的模式意味着每次请求至少需要两次模型调用：一次由guard模型审核，一次由服务LLM生成。在高并发场景下，这种额外开销直接影响服务成本和用户体验。此外，部署多个模型还带来运维复杂度：每个分类器都有独立的训练评估流水线和部署面。同时，过度的安全对齐（alignment tax）可能损害模型整体能力（Huang et al., 2025），这也是一种隐性代价。
本文的目标是本文的目标是消除LLM编排中的额外模型调用开销。具体而言，作者希望在不部署独立guard模型的前提下，仅通过附加轻量级探针到冻结的服务LLM上，在同一个前向传播中完成分类任务。目标是在保持近似服务延迟和显存占用的同时，实现与独立分类器可比的分类精度。实验覆盖安全审核（ToxicChat、WildGuardMix）和情感分类（IMDB、SST-2、Emotion）两大场景，验证方法的通用性。
与已有工作不同的是，现有的计算复用方法存在一个关键局限：它们固定了读出位置。MULI（Hu et al., 2024）仅使用首token的logits（0.13M参数），LPM（Chrabłaszcz et al., 2025）使用固定潜在表示的距离规则，ShieldHead（Xuan et al., 2025）附加到最后一层隐藏状态，OmniGuard（Verma et al., 2025）探测单一层。这些方法都假设某一个固定位置（特定层或特定token）包含足够的分类信号。然而，BERTology研究早已表明，不同层编码不同抽象层次的信息，任务相关信息分布在多个深度。本文的独特切入角度是：不再假定信号存在于某个固定位置，而是学习一个跨token和层的联合聚合器，在完整的 $L \times T \times d$ 隐藏状态张量上进行表示选择，由数据驱动决定哪些token和哪些层对给定分类目标最具判别力。
核心方法
本文的核心直觉是：Transformer不同层的隐藏状态包含不同抽象层次的信息，任务相关信号分布在多个深度而非固定在某一位置。因此，与其手动选定某一层或某一token位置进行读出，不如学习一个轻量聚合器，让数据决定从哪些位置提取信息。技术路线分为两步：首先对每一层的token维度进行聚合，得到每个层的摘要向量；然后跨层聚合，生成最终的分类表示。这种两阶段设计巧妙地将 $L \times T \times d$ 的高维张量降维到 $d$ 维向量，既避免了将完整张量直接送入分类器的计算爆炸，又保留了跨token和跨层的信息选择能力。整个过程中服务LLM保持冻结，仅训练探针参数（最小仅3K，最大35M），实现了高效的参数复用。
本文的核心创新在于将分类问题重新定义为表示选择问题（representation-selection problem），在完整的 $L \times T \times d$ 隐藏状态张量上学习读出。这与已有方法的本质区别在于三个维度：第一，已有方法（MULI、ShieldHead、OmniGuard）都固定了读出位置——要么是首token的logits，要么是最后一层的隐藏状态，要么是某一选定层——而本文让模型学习哪些位置最具判别力；第二，本文提出了两阶段聚合架构（Stage 1: token级聚合，Stage 2: layer级聚合），用同一机制处理两个维度，保持架构统一性；第三，本文系统比较了从简单到复杂的三种聚合机制（直接池化、评分注意力门控、多头自注意力），揭示了表达能力与参数开销之间的权衡关系。这种设计使得探针可以在与生成相同的前向传播中完成分类，无需额外模型调用，这是guard-then-serve模式无法实现的。
方法步骤详情
方法分为三个步骤。第一步是隐藏状态提取：给定冻结的仅解码器LLM，输入提示词 $x$ 被分词为 $T$ 个token，模型在每一层 $l \in \{0, \ldots, L\}$ 产生隐藏状态 $h^{(l)} \in \mathbb{R}^{T \times d}$，其中 $l=0$ 表示嵌入输出，$d$ 为隐藏维度。第二步是两阶段聚合：Stage 1对每一层应用token级聚合 $A^{(l)}_{\text{token}}$，将 $h^{(l)}$ 压缩为层摘要向量 $v^{(l)} \in \mathbb{R}^d$；Stage 2应用层级聚合 $A_{\text{layer}}$，将 $\{v^{(l)}\}_{l=0}^{L}$ 聚合为单一向量 $v \in \mathbb{R}^d$。第三步是分类头：使用线性层 $\text{logits} = W_{\text{out}} v + b_{\text{out}}$ 产生类别预测，用交叉熵损失训练。三种聚合机制的具体实现：直接池化对有效位置取最大值或均值，零参数开销；评分注意力门控学习标量重要性权重 $s_i = \tanh(w^\top X[i,:])$，用softmax归一化后加权求和，参数量 $(L+2)d$；多头自注意力将QKV投影从 $d$ 降维到 $d_{\text{inner}} < d$（如 $d/16$ 或 $d/32$），拆分为 $H$ 个头计算注意力，参数量 $(L+2) \cdot 4d \cdot d_{\text{inner}}$。
技术新颖性
本文的技术新颖性体现在多个层面。首先，将BERTology关于层信息分布的经典发现系统性地应用于仅解码器LLM的探针设计，这是一个从编码器理解到解码器复用的范式迁移。其次，两阶段聚合架构的设计巧妙地解决了高维张量降维问题，用 $O(L)$ 次 $O(T)$ 的聚合替代直接处理 $L \times T$ 维张量，计算复杂度大幅降低。第三，评分注意力门控机制是一个巧妙的折中：仅用一个线性投影和tanh激活就实现了可学习的注意力权重，参数量仅为 $(L+2)d$（约0.1M），但相比零参数的池化带来了显著的性能提升（ToxicChat上从73.53到80.49 F1）。第四，多头自注意力的激进降维策略（$d/32$）证明了改善主要来自读出位置的选择而非额外容量，因为即使大幅压缩注意力维度，MHA仍然显著优于固定读出方法。最后，多骨干网实验设计（密集Llama、MoE GPT-OSS、MoE Qwen）验证了方法排名的一致性，证明这不是特定架构的巧合。
实验结果
实验结果呈现了清晰的层级关系和多个重要发现。在ToxicChat基准上（表1），Llama-3.2-3B骨干网下，直接池化达到73.53 F1，评分注意力门控提升到80.49 F1（仅0.10M参数），MHA探针达到84.51 F1和0.898 AUPRC（35M参数），显著超过logit复用基线MULI（77.8 F1 / 0.829 AUPRC）和独立分类器ToxicChat-T5-large（82.2 F1 / 0.885 AUPRC）。跨骨干网实验显示，GPT-OSS-20B骨干网下MHA达到86.17 F1和0.915 AUPRC（27M参数），Qwen3-30B-A3B骨干网下达到83.76 F1和0.905 AUPRC（26M参数），三种骨干网下方法排名（池化 < 评分注意力 < MHA）完全一致。在WildGuardMix基准上（表3），MHA达到88.55 F1，接近最强独立基线WildGuard的88.9 F1，但仅需35M额外参数且无需额外模型调用。在情感分类基准上（表4），MHA在IMDB达到95.15%准确率、SST-2达到95.39%、Emotion达到87.68%，与独立分类器DeBERTa V3 Large和RoBERTa Large竞争力相当，远超prompting方法（零样本Llama 3.2在Emotion上仅44.55%）。层注意力分析（图2）揭示了有趣的模式：有毒提示在L17-L28层获得较高注意力权重，无毒提示集中在L0和L27-L28层，错误分类的样本其注意力模式更接近预测类别而非真实类别。推理效率分析（表5）显示，即使MHA探针也比最小的独立基线ToxicChat-T5快2倍以上，显存占用减少约1GB；池化和评分注意力的开销可以忽略不计。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ToxicChat（分布内） | F1 (%) | MHA: 84.51 ± 0.43 | MULI: 77.8, ToxicChat-T5: 82.2 | 超过MULI 6.7个百分点，超过ToxicChat-T5 2.3个百分点 |
| ToxicChat（分布内） | AUPRC | MHA: 0.898 ± 0.005 | MULI: 0.829, ToxicChat-T5: 0.885 | 超过MULI 0.069，超过ToxicChat-T5 0.013 |
| ToxicChat（跨分布，训练于WildGuardMix） | F1 (%) | MHA: 72.88 ± 0.9 | Aegis-Guard-D: 70.0, WildGuard: 70.8, ShieldHead(Llama3.1-8B): 64.3 | 超过WildGuard 2.08个百分点，超过ShieldHead 8.58个百分点 |
| WildGuardMix | F1 (%) | MHA: 88.55 ± 0.30 | MULI: 83.79, WildGuard: 88.9 | 超过MULI 4.76个百分点，接近WildGuard（仅差0.35个百分点） |
| IMDB | Accuracy (%) | MHA: 95.15 ± 0.06 | MULI: 86.50, DeBERTa V3 Large: 95.34 | 超过MULI 8.65个百分点，接近DeBERTa（仅差0.19个百分点） |
| SST-2 | Accuracy (%) | MHA: 95.39 ± 0.32 | MULI: 93.19, RoBERTa Large: 95.99 | 超过MULI 2.20个百分点，接近RoBERTa（仅差0.60个百分点） |
| Emotion | Accuracy (%) | MHA: 87.68 ± 1.09 | MULI: 64.05, DeBERTa V3 Large: 87.65 | 超过MULI 23.63个百分点，与DeBERTa持平 |
| 推理效率（ToxicChat测试） | 延迟 (ms/sample) | 池化: 29.66, 评分注意力: 30.90, MHA: 40.27 | ToxicChat-T5: 88.33, WildGuard: 100.32, Llama Guard 3: 123.21 | MHA比最小独立基线快2.2倍，评分注意力仅增加17%延迟 |
局限与改进
作者在论文中坦诚地指出了多个局限性。首先，跨架构泛化验证仅在ToxicChat上完成，情感和情绪基准仅在Llama-3.2-3B-Instruct上评估，且三个骨干网均约3B活跃参数，更大活跃参数规模对探针精度的影响尚未测试。其次，序列长度受限于VRAM约束，处理更长序列（完整文档、多轮对话）时可能面临内存挑战。第三，最小数据集SST-2包含7,000个样本，探针有效训练的最小数据量尚不明确，尤其是参数量较大的MHA变体。第四，与安全微调或guard模型不同，本文的分类方法只能检测有害内容并中断生成，无法生成解释性拒绝回复。第五，分布外结果表明分布漂移会显著降低探针性能（池化在交叉分布设置下从73.53 F1暴跌至53.33 F1），探针可能在新型攻击策略或训练分布未覆盖的场景下静默失败。此外，本文的探针继承了底层LLM的偏见和表示缺陷，可能导致不同人口群体的差异化错误率。从我个人的观察来看，论文未充分讨论探针在对抗性攻击下的鲁棒性——攻击者如果知道探针的读出机制，可能针对性地构造绕过样本。另外，缓存隐藏状态作为敏感数据的安全处理也未深入讨论。
独立分析的弱点
本文存在几个值得深入探讨的弱点。第一，三种聚合机制在Stage 1和Stage 2使用相同机制，这限制了设计空间——理论上，token级聚合和layer级聚合可能有不同的最优机制，例如token级使用注意力而layer级使用池化。改进方向是探索异构聚合策略的组合。第二，MHA探针虽然性能最佳，但35M参数量相比池化（3K）和评分注意力（0.1M）大了三个数量级，在资源极度受限的边缘部署场景下可能不实际。改进方向是研究知识蒸馏，将MHA的能力迁移到更小的探针中。第三，论文未充分探索探针的持续学习能力——当新的攻击模式出现时，是否可以增量更新探针而非从头重训？第四，层注意力分析虽然直观，但仅在ToxicChat的评分注意力上进行，缺乏对MHA探针（最强变体）的类似分析，也缺乏跨数据集的系统比较。第五，论文声称改善主要来自读出位置选择而非额外容量，但未提供严格的消融实验来证明——例如，将MHA的参数量缩减到与评分注意力相同的量级，性能差异如何？
未来方向
作者提出的关键未来方向包括：将多骨干网评估扩展到更多任务，测试更大活跃参数规模，探索细粒度安全分类体系。基于本文成果，还有几个有前景的延伸方向：第一，将探针方法扩展到多模态LLM，利用视觉-语言模型的多层表示进行多模态安全审核；第二，研究探针的跨任务迁移——在安全任务上训练的探针层权重模式能否指导情感任务的探针初始化？第三，将两阶段聚合架构应用于推测解码中的草稿token预测，借鉴EAGLE-3（Li et al., 2025）的多层特征融合思路；第四，研究探针与条件重提示策略的结合——当探针检测到不安全内容时，如何让服务LLM生成上下文相关的拒绝回复而非简单中断；第五，探索探针在联邦学习场景下的应用，多个用户可以在不共享原始数据的情况下协作训练安全探针。
复现评估
本文的复现条件相对友好。代码已在GitHub开源（https://github.com/gmeyoyan/LLM-Probing-Classifiers），这是一个积极的信号。数据方面，所有五个基准均为公开数据集（ToxicChat、WildGuardMix、IMDB、SST-2、Emotion），且论文引用了具体的子集划分（如IMDB 12,500例、SST-2 7,000例）。算力方面，探针训练仅需单张RTX 3090（24GB VRAM），但MoE骨干网的隐藏状态提取需要A100 80GB GPU。论文报告了约100个超参数配置的搜索范围（表6），复现时需要进行类似的网格搜索。整体而言，复现难度中等——探针训练本身简单高效，但完整的超参数搜索和多骨干网实验需要一定的计算资源和时间投入。论文使用PyTorch的scaled_dot_product_attention实现MHA，并可利用FlashAttention加速，这降低了内存瓶颈。
论文图表