[隐私计算学习笔记]12——基于隐语的VisionTransformer框架

目录

1. MPCViT：安全且高效的MPC友好型Vision Transformer架构

1.1 研究动机

1.2 整体框架     

1.3 算法流程

1.4 实验结果

2. 基于SecretFlow搭建ViT框架

2.1  SecretFlow Secure Processing Unit（SPU）

2.2  隐私推理协议及通信参数配置（通信协议、通信带宽、传输延迟）

2.3  基于Jax的ViT模型搭建（Embedding\Attention\MLP）

2.4  ViT模型隐私推理

3. 小结

---

1. MPCViT：安全且高效的MPC友好型Vision Transformer架构

1.1 研究动机

• ViT延迟分解及通信瓶颈（在SEMI-2K\Cheetah协议下的延迟分解）

•  不是所有注意力头都同等重要

        参考：

        [1] Less is more: Pay less attention in vision transformers, AAAI 2022

        [2] Q-vit: Fully differentiable quantization for vision transformer, arXiv 2022

• 不同注意力机制变体的准确率和延迟对比

1.2 整体框架     

•  底层协议——基于半诚实威胁模型的2PC秘密分享

        > 参与方诚实且好奇，均严格遵循既定通信协议，但在协议执行中仍会学习尽可能多的信息；

        > 以适当方式拆分秘密，拆分后各部分由不同参与方保存，单个参与者无法恢复秘密信息，只有若干参与者协作才能恢复秘密信息。

•  上层应用——Vision Transformer模型架构

        > ViT是基于Transformer的encoder架构，由输入embedding层、Transformer块以及分类层组成，其中，每个Transformer块包含一个多头注意力层和MLP块，MLP块又包含两个线性层和一个GeLU；

        > 不同于CNN，ViT模型的主要通信瓶颈在于多头注意力机制，以及MLP中的GeLU。

1.3 算法流程

• Step1：设计合适的搜索空间（三种不同粒度）

                 > 粗粒度：Transformer层级粒度

                 > 中粒度：注意力头级粒度

                 > 细粒度：注意力行级粒度（token级）

• Step2：MPC感知神经架构搜索（可微分搜索）

每个注意力的计算方式：

 可微分注意力搜索公式：

• Step3：基于延迟限制的架构参数二值化

• Step4：重训练异构注意力ViT 

        直接训练搜索后的异构注意力机制ViT会导致准确率显著下降，为有效提高异构注意力机制ViT的准确率，采用：

        > 自蒸馏技术：将原始Softmax注意力机制ViT作为教师模型，无需引入任何额外训练和推理开销；

        > 传统的基于logits的蒸馏（软标签蒸馏）；

        > 基于特征的token蒸馏：更加细粒度蒸馏，MPCViT取最后一层特征蒸馏。

1.4 实验结果

• 主要实验结果——与近年来相关工作的对比

主要结论：在Tiny-ImageNet数据集上，和基线ViT、MPCFormer、THE-X相比， MPCViT降低6.2×、2.9×和1.9×延迟，并且提高了1.9%、1.3%和3.6%的准确率。

• 消融实验——知识蒸馏

主要结论：1、与仅使用交叉熵损失函数相比，基于logits和基于特征的知识蒸馏都显著提高了ViT的性能； 2、结合两种知识蒸馏损失函数进一步提高了准确率，特别是在一个更大的数据集上。以上结果表明了知识蒸馏对于ViT训练重要性以及训练损失函数中三部分的不可或缺。

• 神经架构搜索的一致性和可扩展性

注意力头的架构
参数可视化

ViT层级推理延迟和通信可视化

主要结论： 1、在不同的λ、不同的数据集、不同的注意力头数量的设置下，搜索结构都表现出一致性； 2、架构参数可视化与层级可视化一致，中间靠前的注意力更倾向于保留。

2. 基于SecretFlow搭建ViT框架

2.1  SecretFlow Secure Processing Unit（SPU）

详见  SPU框架介绍 

2.2  隐私推理协议及通信参数配置（通信协议、通信带宽、传输延迟）

•  隐私推理协议配置

•  节点配置

•  模拟网络参数配置throttle.sh

2.3  基于Jax的ViT模型搭建（Embedding\Attention\MLP）

• 图像Patch Embedding搭建——用kernel size 等于 patch size 的卷积实现

• 注意力机制搭建

• MLP模块搭建

        GeLU的维度很大，是embedding dim的mlp_ratio倍。通常在较小的架构上，ViT的mlp_ratio取值2。

• Transformer模块搭建

2.4  ViT模型隐私推理

• ViT Benchmark文件

•  ViT隐私推理操作流程

        > 配置Python环境及安装SPU，并激活Python环境

        > 配置并模拟通信网络环境（以WAN设置为例）

 bash ./throttle.sh del

 bash ./throttle.sh wan

        > 模拟MPC环境及协议

python -m spu.binding.util.distributed -c semi2k_spdz_2pc.json up

        > 执行隐私推理：以一个4-head单个注意力层（emb=256）的隐私推理为例

3. 小结

        本篇文章介绍了MPCViT的总体框架、研究动机、算法流程以及实验结果（均优于baseline）。该工作借助SecretFlow-SPU搭建ViT框架，并测试不同模型架构的推理效率。介绍了基于SecretFlow和Jax的ViT框架搭建流程。