基于隐语的VisionTransformer框架

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框架搭建流程。