G.O.S.S.I.P 阅读推荐 2024-01-29 TEESLICE

今天要为大家介绍北京大学计算机学院操作系统实验室关于使用可信执行环境（Trusted Execution Environments, TEEs）保护深度学习模型的工作——No Privacy Left Outside: On the (In-)Security of TEE-Shielded DNN Partition for On-Device ML ，这项工作已经被IEEE S&P 2024接收。该工作深入研究了面向TEE的模型划分技术（TEE-Shielded DNN Partition，TSDP），在评估并总结现有模型划分方案局限性的基础上，提出了一种新型模型划分方法——TEESLICE。

该论文针对的是在用户侧部署深度学习模型时的安全性和隐私问题。当把模型部署到用户端后，攻击者可以从系统中获取模型的白盒信息（如模型参数），并使用白盒的模型信息进行模型窃取攻击（Model Stealing）以及成员推理攻击（Membership Inference），从而威胁模型拥有者的安全及隐私。为了解决该问题，研究人员提出使用TEE保护设备上的DNN模型，从而把准确、高效的白盒攻击降低为不准确、低效的黑盒攻击（即攻击者只能通过模型输入输出获取信息）。但是，由于因为现有TEE设备计算资源有限，且难以使用外部GPU进行计算加速，因此简单地把全部DNN模型放到TEE中会导致计算速度过低的问题。为了解决这个问题，研究人员提出了多种面向TEE的模型划分（TEE-Shielded DNN Partition，TSDP）方法。这些方法把DNN模型分成两部分，一部分包含模型的少量参数，这些参数与模型隐私以及模型功能关联性较强；另一部分包含模型的大量参数，这些参数与模型功能关联性较弱。当使用TEE保护模型时，就可以把模型中较轻量的部分部署在TEE中，而把模型的大部分部署在GPU上。这种方法既使用TEE保护了关键模型参数，又使得模型中的很大一部分可以使用GPU进行加速运算。

现有框架安全性评估

本文首先对现有TSDP方法进行归纳总结及安全评估。作者调研了近五年发表在计算机领域顶级会议的文章，以及被这些顶会文章引用的文章。除去一些针对不同场景的方案，作者把现有保护模型推理阶段的TSDP方法归为五类：保护网络深层，保护网络浅层，保护大幅权重，保护中间层和保护非线性层。

在调研的过程中，作者发现已有的模型划分方法的安全性都是建立在经验性的观察总结基础上，并且暴露在GPU上的模型参数仍然有可能泄露部分隐私信息，因此有必要对这些方案进行系统的安全性评估。本文使用模型窃取攻击和成员推理攻击衡量模型功能和隐私的泄露程度，并使用一个较强的攻击者攻击这些模型划分方法。这个攻击者可以使用互联网上的公开模型和数据尝试分析并破解这些保护机制，从而提高模型窃取和成员推理的攻击效果。实验设置包括三个模型结构（AlexNet、ResNet18和VGG16_BN）和四个不同的数据集（CIFAR10, CIFAR100, STL10和UTKFace）。

文章首先从每类模型划分方法中选择一个具体的方案进行评估。实验结果表明，现有的TSDP方案在受到更强的攻击者攻击时难以提供黑盒级别的安全保障，在某些场景下，攻击者可以达到与白盒攻击类似的效果。这是因为现有模型划分方案的配置都倾向于使用TEE保护非常少量的模型参数，而把大量“看似”与模型功能无关的参数暴露在GPU上。当攻击者使用更多信息进行攻击时，暴露在GPU上的看似与隐私无关的参数仍然会泄露大量的模型信息。

在此基础上，作者进一步探究这个安全泄露问题的根本原因是什么。因为每种模型划分方法都可以有不同的配置（即在TEE中保护的参数量），那么是否是因为所选择方案的配置问题而导致了信息泄露呢？为了探究这个问题，作者遍历了每种划分方法的不同配置方案，并衡量了这些方案的防御效果，结果如图所示：

图中每个子图代表在一个数据集和模型上的攻击效果，子图中的每一条曲线代表一种划分方法在不同配置下的保护效果和计算量的对应关系，横轴代表计算量，纵轴代表攻击准确率。可以看到在不同场景下，不同划分方法都存在性能和安全性的权衡问题。同时，这些方法在不同场景下的最优配置是不同的，即难以找到一个系统的方法自动确定一个安全性和计算量权衡之下的最优配置。由此说明，现有划分方法仍有一定缺陷，防御者仍然面临如何快速找到最优配置来划分模型的问题。

TEESLICE框架

文章认为导致现有划分方法缺陷的原因是它们均使用“先训练再划分”的模式，即先对DNN模型进行端到端的训练，再对训练好的模型划分。但是因为在训练过程中，因为全部模型参数都被隐私数据更新过，所以难以在训练后的模型上实现隐私隔离。针对这个问题，本文提出了使用“先划分后训练”的模式，即在训练前划分隐私模块与非隐私模块，然后在训练过程中仅使用隐私数据更新隐私模块，不更新非隐私模块，从而在训练后实现隐私划分。

基于“先划分后训练”的模式，本文提出了TEESlice。该方法使用公开预训练模型作为模型骨干网络（Backbone），并在骨干网络中插入多个与隐私相关的层作为模型切片（Model Slice），这些切片包含了私有信息。骨干和模型切片组合成一个混合模型（Hybrid Model），可以模拟原始隐私模型的功能。TEESLICE的工作流程主要包括两个阶段：模型切片提取（训练阶段）和混合模型部署（推断阶段）。在模型切片提取阶段，TEESlice首先使用目标隐私模型监督训练具有稠密切片的混合模型，然后使用自适应的迭代剪枝策略减少模型切片的数量，从而降低TEE内部的计算量。在模型部署阶段，TEESlice把模型切片和主干网络轻量的非线性层部署在TEE中，把主干网络的线性层部署在GPU上，在保护隐私信息的前提下使用GPU加速计算。同时，TEESLICE使用One-Time Pad和Freivalds算法来保护TEE-GPU之间的数据传输信息和GPU计算结果的完整性。

实验与结论

文章使用前文所述的评估方法对TEESlice的安全性和性能进行了评估，实验结果表明，TEESlice可以在达到黑盒保护的前提下，大幅降低TEE内部需要执行的计算量。文章还探究了TEESlice是否可能在其他方面造成性能损失，包括模型推断准确率是否会降低以及暴露在GPU上的主干网络是否会增加攻击成功率，结果表明在这两方面TEESlice都没有额外损耗。同时，作者还使用Intel SGX SDK和PyTorch实现了一个面向TEE的深度学习模型可信异构部署框架 TAOISM（A TEE-based Confidential Heterogeneous Framework for DNN Models），以评估在真实设备上的性能。实验结果表明，在真实设备上TEESlice也可以大幅提高推理性能。

未来展望

基于模型划分的可信计算是近几年新起的一个方向，该方法通过分离DNN模型的关键模块，并使用可信计算方法（如TEE及密码学方法）保护关键模块，从而降低性能损耗。最近，工业界和学术界都提出了许多方案提高TEE的计算能力，如NVDIA H100具有可信计算的功能，Intel TDX支持虚拟机级别的隔离并可以支持多达1GB的加密内存，以及一些把GPU集成到ARM CCA的工作。作者认为这些硬件的解决方案并不能替代基于模型划分的方法，因为这些方案还没有达到商用的规模，高性能的H100更是一卡难求。基于模型划分的方案作为一种软件解决方案，可以大幅降低所依赖的硬件门槛（低端GPU也可以部署），并且可以支持没有TEE功能的异构硬件，具有更广泛的应用领域，因此在未来仍然是一个重要的研究方向。

论文链接：https://arxiv.org/pdf/2310.07152.pdf

代码链接：

• S&P文章的artifact：https://github.com/ziqi-zhang/TEESlice-artifact

• 面向模型划分的异构可信部署框架TAOISM：https://github.com/ziqi-zhang/TAOISM

投稿作者介绍：

张子祺，北京大学博雅博士后，于2023年在北京大学计算机学院取得博士学位，师从陈向群教授、郭耀教授和李锭助理教授，研究领域包括可信人工智能和智能软件工程，已在IEEE S&P、CCS、ICSE、ESEC/FSE、ISSTA、Ubicomp和WWW等多个计算机领域国际顶级会议与期刊发表多篇CCF-A类学术论文。

个人主页：https://ziqi-zhang.github.io/

原文始发于微信公众号（安全研究GoSSIP）：G.O.S.S.I.P 阅读推荐 2024-01-29 TEESLICE