如何防止模型被窃取?基于TVM的端上模型部署加密方案

AI 2个月前 admin
9 0 0
背景
2023年在AI的发展史上一定是浓墨重彩的一笔,在这一年里出现了百模大战、全民“炼丹”的场面,围绕着各种模型的训练技术和算力需求有很多讨论。随着模型的成熟以及算力市场的发展,7B、13B这类小型号的模型也出现了端上部署的需求,其中以移动设备厂商最为突出。2024年,在端上部署和应用模型也可能会成为各家移动厂商的一个营销热点。
在端上部署推理模型自然就会带来一个安全问题,如何保证模型文件的安全?LLM的架构目前都是Transformer的变体,如果拿到模型文件,借助开源的代码很容易完成适配,然后就可以异地部署起来。LLM现在从零开始预训练、指令微调、强化学习到最后可用,成本并不低,如果很容易被别人窃取到模型文件并被推理使用,会造成不小的损失。
为了解决上面这个安全问题,我们借助TVM编译器帮助我们实现高效模型推理部署,并在TVM中引入Tensor加解密层。这样模型文件在本地是处于加密状态,在推理运行的时候才是解密状态,以此保证模型文件不被轻易窃取。下面介绍技术细节。


TVM简介

TVM全称Tensor Virtual Machine,中文名称张量虚拟机,它主要的用途是将深度学习模型编译成最小可部署的模块,并提供更多的后端自动生成和优化模型的基础设施。架构如下:

如何防止模型被窃取?基于TVM的端上模型部署加密方案

它的特点是:

1.  最佳的性能:通过编译和最小运行时,在现有的硬件上优化机器学习工作负载,进而发挥最佳性能。

2.  丰富的后端:可以在CPU、GPU、浏览器、微控制器、FPGA等硬件上运行,同时可自动在多种后端生成和优化张量算子。

3.  灵活可用:TVM 的灵活设计支持对区块稀疏度、量化(1、2、4、8 位整数,posit)、随机森林/经典 ML、内存规划、MISRA-C 兼容性、Python 原型设计。

4.  简单易用:在 Keras、MXNet、PyTorch、Tensorflow、CoreML、DarkNet 等的深度学习模型上编译。用 Python 借助 TVM 进行编译,即可直接使用 C++、Rust 或 Java 完成推理部署。


TVM工作流程

TVM属于领域编译器,其工作流程也遵循通用CPU编译器的设计规则,将复杂的任务简单化,并进行多层抽象, TVM也是采用多层IR设计,IR的转换如下:

如何防止模型被窃取?基于TVM的端上模型部署加密方案

多层IR转换以后,最后生成指定的后端文件,如上图所示TVM编译会生成三个文件:mod.json是模型的计算图描述文件,mod.so是经过优化生成的算子,mod.param是存储的模型参数文件,其文件格式如下:

如何防止模型被窃取?基于TVM的端上模型部署加密方案

TVM 推理运行如下:

如何防止模型被窃取?基于TVM的端上模型部署加密方案


Safe Tensor

我们提出了Safe Tensor的概念。Safe Tensor的设计原则是简单高效易用,我们放弃了全量加密、mpc等技术,使用了传统的对称加解密算法,在TVM生成mod.param文件的时候,在其保存每层tensor的过程中将里面的数据进行加密。与原始mod.param文件比较,我们只加密了最关键的数据,在加解密速度上保证了最快,也将引入的加解密将损耗降到了最低。Safe Tensor采用了传统的对称加密算法。为了确保业务安全,密钥的保存和使用应该遵循最佳实践。在使用密钥之前,需要使用非对称加密算法进行解密。 

TVM生成模型参数文件格式就变成如下:


如何防止模型被窃取?基于TVM的端上模型部署加密方案

效果如下:

如何防止模型被窃取?基于TVM的端上模型部署加密方案


Safe Tensor的方案只是加密了模型最重要的tensor数据,这对TVM的上层是无感的,运行的方法和之前介绍的方法也相同。

如何防止模型被窃取?基于TVM的端上模型部署加密方案

如何防止模型被窃取?基于TVM的端上模型部署加密方案

总结


Safe Tensor的方案有如下优势:

1.  相对TVM生成的推理程序,加密后业务接入无感,灵活易用。

2.  推理时性能无损失。

3.  模型文件被加密,和推理代码一一绑定,直接拷走模型文件后不可用,从而起到保护作用。

Safe Tensor给大家提供了一个新的模型保护思路,现阶段只是在模型文件生成和加载的时候进行加解密。如果想获得更好的保护效果,还需要进程内存保护的技术保证内存dump相关的攻击。随着未来AI技术的发展,硬件设备厂商也会慢慢关注这块市场,也会推出相关的硬件特性,届时可以使用软硬结合的安全保护,让模型应用既高效又安全。



如何防止模型被窃取?基于TVM的端上模型部署加密方案
END



更多精彩内容点击下方扫码关注哦~



如何防止模型被窃取?基于TVM的端上模型部署加密方案
如何防止模型被窃取?基于TVM的端上模型部署加密方案

   云鼎实验室视频号

  一分钟走进趣味科技

     -扫码关注我们-
如何防止模型被窃取?基于TVM的端上模型部署加密方案

关注云鼎实验室,获取更多安全情报

如何防止模型被窃取?基于TVM的端上模型部署加密方案

原文始发于微信公众号(云鼎实验室):如何防止模型被窃取?基于TVM的端上模型部署加密方案

版权声明:admin 发表于 2024年1月12日 上午11:07。
转载请注明:如何防止模型被窃取?基于TVM的端上模型部署加密方案 | CTF导航

相关文章

暂无评论

您必须登录才能参与评论!
立即登录
暂无评论...