在 Elastic 中充分利用 Transformers 模型

序言

在 8.3 版本中，我们的 Elastic Stack 机器学习团队引入了一种将第三方自然语言处理 (NLP) 模型导入 Elastic 的方法。作为安全研究人员，我们必须在安全数据集上试用它。因此，我们决定通过微调 Hugging Face 模型中心上现有的预训练模型来构建一个模型，以识别恶意命令行。

在发现微调后的模型表现（令人惊讶地！）良好后，我们想看看它是否可以替代或与我们之前的基于树的模型结合使用，以检测“活在土地上”（LotL）攻击。但首先，我们必须确保这个新模型的吞吐量和延迟足以进行实时推理。这导致了一系列的实验，我们将在本博客中详细介绍这些实验的结果。

在本博客中，我们将简要介绍如何微调一个用于掩码语言建模 (MLM) 任务的 Transformer 模型，使其适用于分类任务。我们还将了解如何将自定义模型导入 Elastic。最后，我们将深入探讨我们围绕使用微调模型进行实时推理所做的所有实验。

用于命令行分类的 NLP

在开始构建 NLP 模型之前，务必了解NLP模型是否适合当前的任务。在我们的案例中，我们希望将命令行分类为恶意或良性。命令行是用户通过计算机终端提供的一组命令。以下是一个命令行示例

**move test.txt C:\**

上面的命令将文件 test.txt 移动到 C: 目录的根目录。

命令行中的参数是相关的，某些值的共现可能表明存在恶意活动。NLP 模型值得在此处探索，因为这些模型旨在理解和解释自然（人类）语言中的关系，并且由于命令行经常使用一些自然语言。

微调 Hugging Face 模型

Hugging Face 是一个数据科学平台，为机器学习 (ML) 爱好者提供使用开源代码和技术构建、训练和部署 ML 模型的工具。它的模型中心拥有大量针对各种 NLP 任务训练的模型。你可以按原样使用这些预训练模型来对你的数据进行预测，或者在特定于你的 NLP 任务的数据集上微调模型。

微调的第一步是使用特定模型的模型配置和预训练权重来实例化模型。随机权重将分配给基本模型中可能不存在的任何特定于任务的层。初始化后，可以训练模型以学习特定于任务的层的权重，从而针对你的任务进行微调。Hugging Face 有一个名为 from_pretrained 的方法，允许你从预训练的模型配置实例化模型。

对于我们的命令行分类模型，我们创建了一个 RoBERTa 模型实例，其编码器权重从 roberta-base 模型复制，并在编码器之上随机初始化了一个序列分类头

model = RobertaForSequenceClassification.from_pretrained('roberta-base', num_labels=2)

Hugging Face 配备了一个 Tokenizers 库，其中包含当今最常用的一些标记器。对于我们的模型，我们使用了 RobertaTokenizer，它使用字节对编码 (BPE) 来创建标记。此标记化方案非常适合来自与标记化语料库（英语文本）不同的域（命令行）的数据。可以在此处找到我们如何使用 RobertaTokenizer 标记化数据集的代码片段。然后，我们使用 Hugging Face 的 Trainer API 来训练模型，代码片段可以在此处找到。

ML 模型不理解原始文本。在将文本数据用作模型的输入之前，需要将其转换为数字。标记器将大量文本分组为较小的语义上有用的单元（例如但不限于单词、字符或子词 — 称为标记），然后可以使用不同的编码技术将其转换为数字。

• 查看此视频（从 2:57 开始），以查看在标记化之后可能根据你的数据集需要的其他预处理步骤。
• 可以在此处找到有关如何微调预训练的 Hugging Face 模型的完整教程。

将自定义模型导入 Elastic

一旦你拥有一个你满意的训练模型，就可以将其导入 Elastic 了。这是使用 Eland 完成的，Eland 是用于 Elasticsearch 中机器学习的 Python 客户端和工具包。可以在此处找到我们如何使用 Eland 将模型导入 Elastic 的代码片段。
你可以通过 Kibana 中的机器学习 UI 导航到模型管理 > 训练的模型来验证模型是否已成功导入

使用 Transformer 模型进行推理 — 一系列实验

我们进行了一系列实验，以评估我们的 Transformer 模型是否可用于实时推理。对于这些实验，我们使用了一个包含约 66k 个命令行的数据集。

我们使用微调后的 RoBERTa 模型进行的第一次推理运行在测试数据集上花费了约 4 个小时。首先，这比我们试图击败的基于树的模型要慢得多，后者在整个数据集上花费约 3 分钟。很明显，我们需要提高 PyTorch 模型的吞吐量和延迟，使其适合实时推理，因此我们进行了一些实验

使用多个节点和线程

以上延迟数字是在模型在单个节点上的单个线程上运行时观察到的。如果你的 Elastic 部署关联了多个机器学习 (ML) 节点，你可以在多个节点上运行推理，也可以在每个节点的多个线程上运行推理。这可以显著提高模型的吞吐量和延迟。

你可以在通过 API 启动训练的模型部署时更改这些参数

**POST \_ml/trained\_models/\\<model\_id\\>/deployment/\_start?number\_of\_allocations=2&threa ds\_per\_allocation=4**

number_of_allocations 允许你设置跨机器学习节点的模型的总分配数量，并可用于调整模型吞吐量。threads_per_allocation 允许你设置每个模型分配在推理期间使用的线程数，并可用于调整模型延迟。请参阅 API 文档，了解有关设置这些参数的最佳实践。

在我们的案例中，我们将 number_of_allocations 设置为 2，因为我们的集群有两个 ML 节点，并将 threads_per_allocation 设置为 4，因为每个节点都有四个分配的处理器。

使用这些设置运行推理使原始推理时间加快了 2.7 倍。

动态量化

量化是提高模型计算成本的最有效方法之一，同时还能减小模型大小。这里的想法是为权重和/或激活使用降低精度的整数表示形式。虽然有很多方法可以在模型开发期间权衡模型精度以提高吞吐量，但动态量化有助于在事后实现类似的权衡，从而节省在模型训练迭代上花费的时间和资源。

Eland 提供了一种在将模型导入 Elastic 之前动态量化模型的方法。要做到这一点，只需在创建 TransformerModel 对象时传入 quantize=True 作为参数（请参阅导入模型的代码片段），如下所示

**# Load the custom model**
**tm = TransformerModel("model", "text\_classification", quantize=True)**

对于我们的命令行分类模型，我们观察到在动态量化后，模型大小从 499 MB 降至 242 MB。使用此模型在我们的测试数据集上运行推理，速度比原始推理时间提高了 1.6 倍，模型灵敏度略有下降（具体数字见下节）。

知识蒸馏

知识蒸馏是一种通过将知识从大型（教师）模型转移到较小（学生）模型，同时保持有效性来实现模型压缩的方法。从高层次上讲，这是通过使用教师模型在每一层的输出，来反向传播通过学生模型的误差来实现的。这样，学生模型就可以学习复制教师模型的行为。模型压缩是通过减少参数数量来实现的，而参数数量与模型的延迟直接相关。

为了研究知识蒸馏对我们模型性能的影响，我们为我们的命令行分类任务微调了一个 distilroberta-base 模型（遵循微调部分中描述的相同步骤），并将其导入到 Elastic 中。distilroberta-base 有 8200 万个参数，而其教师模型 roberta-base 有 1.25 亿个参数。微调后的 DistilRoBERTa 模型的模型大小为 329 MB，低于 RoBERTa 模型的 499 MB。

在使用此模型运行推理时，我们观察到速度比原始推理时间提高了 1.5 倍，并且模型灵敏度（具体数字见下节）略好于微调后的 roberta-base 模型。

动态量化和知识蒸馏

我们观察到，动态量化和模型蒸馏都显著提高了原始推理时间的速度。因此，我们最终的实验涉及使用微调后的 DistilRoBERTa 模型的量化版本运行推理。

我们发现这使速度比原始推理时间提高了 2.6 倍，并且模型灵敏度略好（具体数字见下节）。我们还观察到量化后模型大小从 329 MB 降至 199 MB。

整合所有内容

根据我们的实验，动态量化和模型蒸馏显著提高了推理速度。将这些改进与分布式并行计算相结合，我们进一步将测试集上的总推理时间从四小时缩短到 35 分钟。然而，即使我们最快的 Transformer 模型，尽管使用了明显更多的 CPU 资源，速度仍然比基于树的模型慢几个数量级。

Elastic 的机器学习团队正在 Elastic Stack 的 8.4 版本中引入推理缓存机制，以节省在重复样本上执行推理所花费的时间。这些在真实环境中很常见，尤其是在安全方面。有了这种优化，我们乐观地认为，未来我们将能够在基于树的模型旁边使用 Transformer 模型。

对我们的基于树的模型和 Transformer 模型的灵敏度（真阳性率）和特异性（真阴性率）的比较表明，两者的集成可能会产生性能更高的模型

如上所示，基于树的模型更适合对良性数据进行分类，而 Transformer 模型在恶意样本上表现更好，因此加权平均或投票集成可以很好地通过平均两个模型的预测来减少总误差。

下一步

我们计划在后续博客中介绍我们从推理缓存和模型集成中获得的研究成果。请继续关注！

与此同时，我们很乐意了解您正在为 Elastic 中的推理构建的模型。如果您想分享您正在做的事情或在此过程中遇到任何问题，请通过我们的社区 Slack 频道和讨论论坛与我们联系。祝您实验愉快！