一、MPCViT
1.1 引言
- 隐私保护需求:随着隐私保护的重要性日益增加,MPC(安全多方计算)成为实现隐私保护计算的关键技术。在视觉任务中,Transformer模型(如ViT)因其强大的表现力和灵活性受到广泛关注。
- 研究动机:ViT模型在MPC环境下的计算延迟和通信成本较高,尤其是在多头注意力机制和MLP块的GeLU激活函数方面。
1.2 MPC下的ViT隐私推理框架
- 半诚实威胁模型:假设参与方诚实但好奇,严格遵循通信协议但试图尽可能多地学习信息。
- 秘密分享方案:将秘密数据拆分成多个份额,每个份额由不同参与方保存,单个参与方无法恢复秘密,只有若干个参与方协作才能恢复秘密信息。
1.3 ViT模型架构
- ViT架构:基于Transformer的编码器架构,包括输入嵌入层、多个Transformer块和分类层。每个Transformer块包含多头注意力层和MLP块,MLP块包含两个线性层和一个GeLU激活函数。
1.4 研究动机
- 延迟和通信瓶颈:研究发现,ViT模型的主要通信瓶颈在于多头注意力机制和MLP中的GeLU激活函数。
- 不同注意力机制的对比:研究对比了不同注意力机制的准确率和延迟,发现并不是所有的注意力机制都同等重要。
二、MPCViT的整体算法流程
2.1 搜索空间设计
- 粒度划分:设计了三种不同粒度的搜索空间:
- 粗粒度:Transformer层级粒度。
- 中粒度:注意力头级粒度。
- 细粒度:注意力行级粒度(token级)。
2.2 算法步骤
- MPC感知神经架构搜索:通过可微分搜索和基于延迟限制的架构参数二值化,实现灵活适应的架构设计。
- 异构注意力ViT重训练:借助知识蒸馏技术提升模型性能。传统的基于logits的蒸馏和基于特征的蒸馏相结合,提高了异构注意力机制ViT的准确率。
三、基于SecretFlow搭建Vision Transformer框架
3.1 SecretFlow架构
- 前端:高层次的模型编写。
- 编译器:针对MPC设计的优化。
- 运行时:映射到底层MPC协议。
3.2 搭建流程
- ViT模型组件:
- Patch Embedding:使用卷积实现图像Patch嵌入。
- 注意力机制:构建多头注意力层。
- MLP模块:通常设置较小架构,GeLU激活函数的维度较大。
- Transformer模块:完整的Transformer层构建。
3.3 ViT模型隐私推理
- 隐私推理协议及配置:模拟WAN和LAN网络环境,配置SPU运行时。
- 执行流程:初始化SPU环境,设置输入变量维度,创建计算对象,在密文环境中计算并输出结果。
四、实验结果
4.1 实验对比
- 与相关工作的对比:在不同数据集上,与Linformer、THE-X和MPCFormer等基线模型相比,MPCViT在延迟和准确率方面均有显著提升。
- CIFAR-10、CIFAR-100、Tiny-ImageNet:在Tiny-ImageNet数据集上,MPCViT比基线模型在延迟方面降低了6.2×、2.9×和1.9×,准确率提高了1.9%、1.3%和3.6%。
4.2 消融实验
- 知识蒸馏的重要性:实验表明,基于logits和基于特征的知识蒸馏显著提高了ViT的性能,特别是在大数据集上。
4.3 神经架构搜索的一致性和扩展性
- 一致性:在不同λ、数据集和注意力头数量的设置下,搜索结构表现出一致性。
- 可视化:架构参数可视化显示中间靠前的注意力更倾向于保留,符合层级推理延迟和通信的可视化结果。
进一步阅读