ModelScan：保护您的机器学习模型免受序列化攻击

ModelScan：为机器学习模型提供全方位安全保护

在人工智能和机器学习日益普及的今天，模型安全已成为不容忽视的重要议题。随着基础模型的兴起，公开的机器学习模型被越来越多地用于进一步训练和微调。这些模型在关键决策和任务关键型应用中发挥着越来越重要的作用。然而，与电子邮件中的PDF文件相比，对模型的安全扫描仍显不足。这种状况亟待改变，而适当的工具则是第一步。

ModelScan：开源的多格式模型扫描工具

ModelScan是由Protect AI公司开发的开源项目，旨在扫描模型以检测是否包含不安全代码。作为业界首个支持多种模型格式的扫描工具，ModelScan目前支持H5、Pickle和SavedModel格式，可为使用PyTorch、TensorFlow、Keras、Sklearn、XGBoost等框架的用户提供保护，未来还将支持更多格式。

为什么需要扫描模型？

模型通常来自自动化流水线或数据科学家的笔记本电脑。无论哪种情况，在使用之前，模型都需要从一台机器转移到另一台机器。将模型保存到磁盘的过程称为序列化。

模型序列化攻击是指在序列化（保存）过程中向模型内容添加恶意代码 - 这是特洛伊木马的现代版本。攻击者通过利用模型的保存和加载过程来实施攻击。当您使用model = torch.load(PATH)加载模型时，PyTorch会打开文件内容并开始运行其中的代码。一旦加载模型，攻击就会立即执行。

模型序列化攻击可能导致以下后果：

• 凭证盗窃（窃取云凭证以访问环境中的其他系统）
• 数据盗窃（窃取发送给模型的请求数据）
• 数据污染（污染模型完成任务后发送的数据）
• 模型污染（篡改模型本身的结果）

ModelScan如何工作？

如果使用机器学习框架加载模型会自动执行攻击，那么ModelScan如何在不加载恶意代码的情况下检查内容呢？

答案很简单：它像读取字符串一样一次读取文件内容的一个字节，寻找不安全的代码签名。这使得扫描过程非常快速，在大多数情况下只需几秒钟即可完成模型扫描。同时，这种方法也非常安全。

ModelScan将不安全代码分为以下几个等级：

• 严重（CRITICAL）
• 高危（HIGH）
• 中危（MEDIUM）
• 低危（LOW）

如果检测到问题，应立即联系模型作者以确定原因。在某些情况下，为了便于数据科学家重现结果，可能会在模型中嵌入代码，但这也会增加被攻击的风险。请根据您的工作负载谨慎判断是否适合使用这种方法。

支持的模型和框架

ModelScan目前支持任何Pickle衍生格式以及许多其他格式：

快速上手

ModelScan作为Python包安装在您的系统上（支持Python 3.9到3.12版本）。您可以通过以下命令安装：

pip install modelscan

安装完成后，可以使用以下命令扫描模型：

modelscan -p /path/to/model_file.pkl

在ML流水线和CI/CD流水线中集成ModelScan

虽然临时扫描是一个很好的第一步，但对于生产级MLOps流程来说还不够。为了全面提高安全性，模型扫描需要在多个阶段进行：

1. 在加载预训练模型进行进一步处理之前进行扫描，以防止受损模型影响您的模型构建或数据科学环境。
2. 训练后扫描所有模型，以检测可能影响新模型的供应链攻击。
3. 在部署到端点之前扫描所有模型，以确保模型在存储后未被篡改。

这些过程同样适用于LLM、基础模型或外部模型的微调或任何修改。如果模型部署在ML流水线之外完成，还应将扫描嵌入到CI/CD系统的部署流程中，以确保使用安全。

深入探索

如果您想更深入地了解模型序列化攻击的原理，可以查看 Model Serialization Attack Explainer。此外，在项目的 notebooks 文件夹中，您可以探索一系列展示如何针对TensorFlow和PyTorch等各种ML框架执行模型序列化攻击的示例。

开源贡献

ModelScan是一个开源项目，欢迎社区贡献。如果您有兴趣为项目做出贡献，请查看 贡献指南 了解详细信息。

通过使用ModelScan，您可以为您的机器学习模型部署增加一层重要的安全保护。在这个AI技术快速发展的时代，确保模型安全不仅是技术问题，更是一种责任。让我们共同努力，为AI的未来构建一个更安全的环境。