一个编码器，七个头：使用掩码损失训练统一安全分类器的经验 [P]

Reddit r/MachineLearning 模型
multi-head security-classifier masked-loss quantization onnx mmbert

摘要

他们发布了一个统一的多头安全分类器（mmBERT-small编码器，七个头），使用掩码损失处理缺失任务标签，在七个任务上取得了高F1分数，同时还发布了量化的ONNX INT8变体。权重和指标已公开。

过去几个月，我们将七个独立的序列分类器整合为一个多头模型，可以说是我们的巅峰模型。现在权重已公开，我想分享哪些方法有效以及哪些地方让我们感到惊讶。设置：共享的mmBERT-small编码器，带有七个任务头，即二进制注入（BCE）、文档分类（7类）、工具类型（14类）、工具操作（6类）、工具数据流标签（3× BCE，多标签）、意图路由（5类）和威胁类型（7类）。需要注意的部分是：我们的训练行只携带了部分任务的标签，因此缺失的任务标签被完全从损失中屏蔽掉。我们最终编写了一个自测程序，断言缺失任务的梯度绝对为零，这发现了两个细微的bug，我建议任何做类似屏蔽的人采用这种方法。大约5k条合成/真实的多任务行有助于各头共同训练；测试集保持100%真实数据。每个头在保留集上的结果：注入F1 0.962，文档0.980，工具类型0.957，工具操作0.945，工具标签0.958，路由0.916，威胁0.952。量化：统一模型和专用的单任务变体都发布了量化边缘构建版本（ONNX INT8 + INT4嵌入，从96 MB开始），并在仓库中提供了对等的基准测试，最差情况下头部损失0.012（相对于FP32）。与七个专用模型相比是否值得？我们发布了两种变体，您可以自行判断，专用模型在大多数任务上得分略高，但统一模型只需进行一次编码器前向传播，而非最多七次。我们的薄弱点：路由，得分为0.916。意图类别在语义上有重叠（"编写代码分析我的数据"——这是代码还是分析？），我怀疑这种歧义确实存在于数据中。如果您有除了重新标注之外的想法，请告知我 :) 权重和每个头的指标：https://huggingface.co/patronus-studio
查看原文

相似文章

@maximelabonne：像2020年那样训练编码器！

X AI KOLs Following

Liquid AI 对其 LFM2.5-Encoder 模型（230M 和 350M）进行了微调，使其能够在单次前向传播中执行多标签分类，无需解码循环或解析。这展示了在 NLP 任务中进行高效标签评分的潜力。