这篇文章由 Sophos 的 Salma Taoufiq 和 Harini Kannan 共同撰写。
作为下一代网络安全的领导者， Sophos的 致力于保护 500,000 多个国家/地区的 150 多个组织和数百万客户免受不断变化的威胁。 Sophos 由 Sophos X-Ops 的威胁情报、机器学习 (ML) 和人工智能提供支持，提供广泛多样的高级产品和服务组合，以保护和保护用户、网络和端点免受网络钓鱼、勒索软件、恶意软件和范围广泛的网络攻击。
Sophos 人工智能 (AI) 组 (SophosAI) 监督 Sophos 主要机器学习安全技术的开发和维护。 安全是一个大数据问题。 为了逃避检测，网络犯罪分子不断制造新颖的攻击。 这转化为庞大的威胁数据集，该团队必须与之合作以最好地保护客户。 一个值得注意的例子是检测和消除带有恶意软件的文件，其中数据集以 TB 为单位。
在这篇文章中，我们将重点关注 Sophos 针对 PDF 文件格式的恶意软件检测系统。 我们展示了 SophosAI 如何使用 亚马逊SageMaker 使用 TB 级数据进行分布式训练，以训练强大的轻量级 XGBoost（Extreme Gradient Boosting）模型。 这允许他们的团队通过自动超参数调整更快地迭代大型训练数据，而无需管理底层训练基础设施。
该解决方案目前已无缝集成到生产培训管道中，并通过 Sophos 端点服务.
用例上下文
无论您是想分享一份重要的合同，还是保留简历的精美设计，PDF 格式都是最常见的选择。 它的广泛使用以及人们普遍认为此类文件是密封的和静态的，使用户产生了一种虚假的安全感。 因此，PDF 已成为攻击者武器库中的首选感染媒介。 使用 PDF 的恶意行为通常通过嵌入由 PDF 阅读器运行的 JavaScript 有效负载来实现，以从 URI 下载病毒、破坏用户的计算机或窃取敏感信息。
Sophos 使用一组确定性和 ML 模型在攻击的各个点检测恶意 PDF 文件。 下图说明了一种这样的方法，其中恶意 PDF 文件通过电子邮件传递。 一旦进行下载尝试，它就会触发恶意可执行脚本连接到攻击者的命令和控制服务器。 SophosAI 的 PDF 检测器在检测到它是恶意的后会阻止下载尝试。
其他方法包括阻止端点中的 PDF 文件、将恶意文件发送到沙箱（使用多个模型对其进行评分）、将恶意文件提交到评分基础设施并生成安全报告等。
动机
为了构建一个基于树的检测器，该检测器可以高置信度地判定恶意 PDF，同时允许低端点计算功耗和快速推理响应，SophosAI 团队发现 XGBoost 算法是该任务的完美候选者。 出于两个原因，此类研究途径对 Sophos 很重要。 在客户端点级别部署强大而小型的模型对分析师对公司的产品评论有很大影响。 更重要的是，它还提供了更好的整体用户体验。
技术挑战
因为目标是拥有一个内存占用比现有 PDF 恶意软件检测器（磁盘和内存）更小的模型，SophosAI 转而使用 XGBoost，这是一种分类算法，经过证明可以生成比神经网络小得多的模型，同时实现了令人印象深刻的表格数据的表现。 在开始建模 XGBoost 实验之前，一个重要的考虑因素是数据集的绝对大小。 事实上，Sophos 的 PDF 文件核心数据集以 TB 为单位。
因此，主要挑战是使用大型数据集训练模型，而无需进行下采样。 因为对于检测器而言，学会发现任何基于 PDF 的攻击（甚至是大海捞针和完全新颖的攻击以更好地保护 Sophos 客户）至关重要，所以使用所有可用的不同数据集至关重要。
与可以批量训练的神经网络不同，对于 XGBoost，我们需要将整个训练数据集存储在内存中。 该项目最大的训练数据集超过 1 TB，如果不使用分布式训练框架的方法，就无法进行如此规模的训练。
解决方案概述
SageMaker 是一项完全托管的 ML 服务，提供各种工具来构建、训练、优化和部署 ML 模型。 这 SageMaker 内置算法库 由 21 种流行的 ML 算法组成，包括 XGBoost。 （有关详细信息，请参阅 使用XGBoost和Amazon SageMaker简化机器学习.) 使用 XGBoost 内置算法，您可以利用开源 SageMaker XGBoost 容器 通过指定大于 1.0-1 的框架版本，提高了灵活性、可扩展性、可扩展性和托管 Spot 训练，并支持 Parquet 等输入格式，这是用于 PDF 数据集的格式。
SophosAI 选择 SageMaker 的主要原因是能够通过简单地指定多个实例，从多节点 CPU 实例上的完全托管分布式训练中受益。 SageMaker 自动跨节点拆分数据，跨对等节点聚合结果，并生成单个模型。 实例可以是 Spot 实例，从而显着降低训练成本。 随着 XGBoost 的内置算法，您无需任何额外的自定义脚本即可执行此操作。 XGBoost 的分布式版本也以开源形式存在，例如 XGBoost-射线 和 XGBoost4J-火花，但它们的使用需要构建、保护、调整和自我管理分布式计算集群，这代表了科学发展之外的重大努力。
此外， SageMaker自动模型调整，也称为超参数调整，通过运行许多具有您指定的超参数范围的训练作业来找到模型的最佳版本。 然后，它会选择超参数值，从而生成性能最佳的模型，由给定 ML 任务的度量标准来衡量。
下图说明了解决方案体系结构。
值得注意的是，SophosAI 在转向 SageMaker 之前开始 XGBoost 实验时，曾尝试使用大内存 亚马逊弹性计算云 (Amazon EC2) 实例（例如，r5a.24xlarge 和 x1.32xlarge）在尽可能大的数据样本上训练模型。 但是，这些尝试平均花费了 10 多个小时，并且通常由于内存不足而失败。
相比之下，通过使用 SageMaker XGBoost 算法和轻松的分布式训练机制，SophosAI 可以在 20 分钟内在庞大的 PDF 训练数据集上大规模训练一个 booster 模型。 团队只需将数据存储在 亚马逊简单存储服务 (Amazon S3) 作为类似大小的 Parquet 文件，并选择 EC2 实例类型和所需的实例数量，SageMaker 管理底层计算集群基础设施并在集群的多个节点之间进行分布式训练。 在底层，SageMaker 使用 ShardedByS3Key 跨节点拆分数据，以在每个实例之间平均分配文件对象，并使用 XGBoost 实现 兔子协议 （可靠的 AllReduce 和广播接口）启动分布式处理并在主节点和对等节点之间进行通信。 （有关直方图聚合和跨节点广播的更多详细信息，请参阅 XGBoost：可扩展的树增强系统.)
除了使用 SageMaker 训练一个模型之外， XGBoost 超参数调优 通过同时运行不同的实验来微调超参数的最佳组合，它也变得快速而简单。 可调超参数包括特定于助推器的超参数和特定于目标函数的超参数。 两种搜索策略 提供：随机或贝叶斯。 贝叶斯搜索策略已被证明是有价值的，因为它有助于在更少的实验迭代中找到比单纯随机搜索更好的超参数。
数据集信息
SophosAI 的 PDF 恶意软件检测建模依赖于 n-gram 直方图和字节熵特征等多种特征（更多信息，请参阅 MEADE：迈向恶意电子邮件附件检测引擎）。 从收集的 PDF 文件中提取的元数据和特征存储在分布式数据仓库中。 然后计算一个包含 3,500 多个特征的数据集，根据时间进一步拆分为训练和测试集，并以 Parquet 文件的形式分批存储在 Amazon S3 中，以便 SageMaker 轻松访问以进行训练作业。
下表提供了有关训练和测试数据的信息。
| 数据集 | 样品数量 | Parquet 文件数 | 总尺寸 |
| 产品培训 | 70,391,634 | 5,500 | 〜1010 GB |
| 测试 | 1,242,283 | 98 | 〜18 GB |
数据大小已根据以下公式计算：
数据大小 = N × (nF + nL）×4
该公式具有以下参数：
- N 是数据集中的样本数
- nF 是特征的数量，其中 nF = 3585
- nL 是地面实况标签的数量，nL = 1
- 4 是特征数据类型所需的字节数：
float32
此外，以下饼图提供了训练集和测试集的标签分布，引出了 PDF 恶意软件检测任务中面临的类别不平衡。
分布从训练集转移到一个月测试集。 应用基于时间的数据集拆分为训练和测试，以模拟现实生活中的部署场景并避免临时窥探。 例如，这种策略还允许 SophosAI 在面对以前看不见的全新 PDF 攻击时评估模型的真正泛化能力。
实验和结果
为了启动实验，SophosAI 团队使用默认参数训练了一个基线 XGBoost 模型。 然后他们开始使用贝叶斯策略使用 SageMaker 执行超参数微调，这就像指定 超参数 要调整的值和所需的值范围、评估指标（在这种情况下为 ROC（接收器操作特征）AUC）以及训练和验证集。 对于 PDF 恶意软件检测器，SophosAI 优先考虑超参数，包括增强轮数（num_round), 最大树深度 (max_depth), 学习率 (eta)，以及构建树时的列采样率 (colsample_bytree）。 最终，获得了最佳超参数并用于在完整数据集上训练模型，最后在保持测试集上进行评估。
下图显示了目标指标 (ROC AUC) 与在调优作业中运行的 15 个训练作业。 最好的超参数是对应于第九个训练作业的那些。
在 SophosAI 在 SageMaker 上的实验开始时，一个特别重要的问题是：需要什么类型的实例以及需要多少实例才能在手头的数据上训练 XGBoost？ 这一点至关重要，因为使用错误数量或类型的实例可能会浪费时间和金钱； 由于内存不足，训练必然会失败，或者，如果使用太多太大的实例，这可能会变得不必要地昂贵。
XGBoost 是一种内存限制（与计算限制相反）的算法。 因此，通用计算实例（例如，M5）是比计算优化实例（例如，C4）更好的选择。 为了做出明智的决定，有一个简单的 SageMaker 指南来选择在完整数据集上运行训练所需的实例数量：
总训练数据大小 × 安全系数（*） < Instance Count × Instance Type 的总内存
在这种情况下：总训练数据大小 × 安全系数 (12) = 12120 GB
下表总结了所选实例类型为 ml.m5.24xlarge 时的要求。
| 训练规模 × 安全系数 (12) | 实例内存 ml.m5.24xlarge | 训练所需的最小实例数 |
| 12120 GB | 384 GB | 32 |
*由于 XGBoost 分布式训练的性质，需要在训练之前将整个训练数据集加载到 DMatrix 对象中，并且需要额外的可用内存，因此建议使用 10-12 的安全系数。
为了更详细地了解在提供的数据集上对 XGBoost 进行完整 SageMaker 训练的内存利用率，我们提供了从训练中获得的相应图表 亚马逊CloudWatch 监控。 对于这个训练作业，使用了 40 ml.m5.24xlarge 实例，最大内存利用率达到了 62% 左右。
通过将 SageMaker 等托管 ML 服务集成到数据管道中节省的工程成本约为 50%。 使用 Spot 实例进行训练和超参数调整作业的选项将成本额外降低了 63%。
结论
借助 SageMaker，SophosAI 团队可以通过构建轻量级 PDF 恶意软件检测 XGBoost 模型成功解决复杂的高优先级项目，该模型在磁盘上（最多小 25 倍）和内存中（最多小 5 倍）比其探测器的前身。 这是一个小而强大的恶意软件检测器，AUC 约为 0.99，真阳性率为 0.99，假阳性率为 
. 该模型可以快速重新训练，并且随着时间的推移可以轻松监控其性能，因为在超过 20 TB 的数据上训练它只需不到 1 分钟。
您可以利用 SageMaker 内置算法 XGBoost 用于使用您的表格数据大规模构建模型。 此外，您还可以尝试新的内置 Amazon SageMaker 算法 LightGBM、CatBoost、AutoGluon-Tabular 和 Tab Transformer，如本文所述 新闻.
关于作者
萨尔玛·陶菲克 是 Sophos 的高级数据科学家，在机器学习和网络安全的交叉领域工作。 她拥有计算机科学本科背景，毕业于中欧大学，获得理学硕士学位。 在数学及其应用。 在不开发恶意软件检测器时，Salma 是狂热的徒步旅行者、旅行者和惊悚片的消费者。
哈里尼·坎南 是 SophosAI 的数据科学家。 她从事安全数据科学工作约 4 年。 她之前是被 Sophos 收购的 Capsule8 的首席数据科学家。 她曾在 CAMLIS、BlackHat（美国）、Open Data Science Conference（东部）、Data Science Salon、PyData（波士顿）和 Data Connectors 发表演讲。 她的研究领域包括使用性能计数器检测基于硬件的攻击、用户行为分析、可解释的机器学习和无监督异常检测。
哈桑·波纳瓦拉（Hasan Poonawala） 是位于英国伦敦的 AWS 的高级 AI/ML 专家解决方案架构师。 Hasan 帮助客户在 AWS 的生产环境中设计和部署机器学习应用程序。 他作为数据科学家、机器学习从业者和软件开发人员拥有超过 12 年的工作经验。 在业余时间，Hasan 喜欢探索大自然并与朋友和家人共度时光。
迪甘特帕特尔 是 AWS 的企业支持主管。 他与客户合作，在云中大规模设计、部署和运营。 他感兴趣的领域是 MLOps 和 DevOps 实践，以及它如何帮助客户进行云之旅。 工作之余，他喜欢摄影、打排球以及与朋友和家人共度时光。
