Patronus 公开统一安全分类器：单编码器多任务架构实践

Patronus 团队近日在 Reddit r/MachineLearning 发布了一篇工程总结，详细介绍了他们如何将原先 7 个独立的序列分类任务合并到一个共享编码器 + 多任务头（multi-head）的统一模型中，并在 HuggingFace 上以「patronus-studio」的名义公开了权重与逐任务头（per-head）评估指标。

统一架构：一个编码器，七个任务头

整个模型以 mmBERT-small 作为共享编码器（encoder），在它之上挂了 7 个任务头，分别对应不同的安全分类需求：

• 二元提示注入检测（binary injection，BCE）；
• 文档类别（7-way classification）；
• 工具类型（14-way）；
• 工具操作（6-way）；
• 工具数据流标签（3 × BCE，多标签）；
• 意图路由（5-way）；
• 威胁类型（7-way）。

这种「one encoder, seven heads」的设计意味着推理时只需一次编码器前向，而不是最多 7 次。

训练的关键：缺失标签的掩码损失

由于每条训练样本只标注了部分任务的标签，作者在损失计算中把「缺失任务」整段掩码掉。但这一步非常容易出错，于是他们写了一个 self-test，断言「缺失任务的梯度必须严格为 0」，结果在开发过程中直接捕获到了两个隐蔽 bug。作者建议任何做类似 masking 的团队都该跑一遍这种 sanity check。

为了让各个任务头能联合学习，他们额外整理了约 5k 条合成 + 真实混合的多任务样本用于共训练；评估时则坚持使用 100 % 真实数据。

各任务头表现：路由仍是短板

在留出测试集上，7 个头的 F1 表现如下：

• 提示注入：0.962
• 文档类别：0.980
• 工具类型：0.957
• 工具操作：0.945
• 工具数据流标签：0.958
• 意图路由：0.916
• 威胁类型：0.952

其中「意图路由」是明显弱项。作者分析认为这并非模型问题，而是数据本身存在语义重叠——例如「写一段代码来分析我的数据」到底该归到 code 还是 analytics，类目边界本身模糊，仅靠重新标注未必能解决。

量化与部署：96 MB 的边缘版本

为了便于部署，团队同时发布了 ONNX 格式的 INT8 主干 + INT4 embedding 的 -edge 版本，起始体积约 96 MB。实测显示，量化后最差的任务头相比 FP32 也仅下降 0.012，整体性能基本与浮点版本持平，量化指标也一并放在仓库中供对比。

与 7 个独立模型的取舍

作者也一并放出了 7 个单任务专用模型供直接对比。结果并不令人意外：专用模型在多数任务上仍略胜一筹，但统一模型只需一次编码器前向，在吞吐和部署复杂度上更友好。这是一种典型的「专用模型 vs. 通用模型」权衡，对实际工程取舍有参考价值。

权重与逐任务指标已发布在 HuggingFace 的 patronus-studio 主页，有兴趣复现或对比的读者可以直接拉取。