基于之前MPC的基础知识,本讲主要内容是MPCViT基于SecretFlow的VisionTransformer框架,主要从神经网络架构,隐私推理框架和实验结果三方面介绍。
一、MPCViT:安全且高效的MPC友好型 Vision Transformer架构
MPCViT隐私推理总体框架中,其底层协议是基于半诚实威胁模型的 2PC的秘密分享方案,简单来说即服务器和客户端相互不能泄漏自己的信息,最终要得到神经网络正确的推理结果。在上层模型方面,VisionTransformer是一种基于Transformer编码器的视觉模型,包含了embedding层,注意力层,非线性GeLU函数等多个不同模块,总体目标是在保护数据和模型隐私的前提下,以更高的推理效率得到最终图像分类的结果。
MPCViT的三大研究动机:ViT延迟分解(ViT在MPC中推理的延迟进行分解分析,发现在SEMI-2K和Cheetah两种协议下,注意力机制中的SoftMax以及MLP中的GeLU函数都占据了绝大部份的延迟)及通信瓶颈、不同注意力机制变体对比,对比了不同注意力的准确率和延迟(进一步分解注意力机制,可以发现求最大值、求倒数、求指数等运算都造成了一部分延迟,能否尽可能去除ViT中这样一些耗时的操作,提高隐私推理的效率,降低推理的延迟)、不是所有的注意力都同等重要(发现一些注意力机制具有很高的效率,另一些具有很好的性能,一个很自然的想法是能否去合理地融合两种不同的注意力机制,使得在隐私推理时同时具有高效性和高准确率)。
MPCViT的整体算法流程(主要用于解决如何权衡模型准确率和推理延迟,以及融合高准确率注意力机制和低延迟注意力机制)可以分为四步:
1、设计合适的搜索空间(三种不同的粒度)
分为粗粒度:Transformer层级粒度,中粒度:注意力头级粒度,细粒度:注意力行级粒度(token级)
2、MPC感知神经架构搜索(可微分搜索)
对于架构搜素算法,给每个注意力都分配架构参数,在搜索过程中,自动决定应该保留哪个注意力,这种搜索算法是可微分的,搜索效率比较高,只需要通过给损失函数添加正则项即可实现。
3、基于延迟限制的架构参数二值化(灵活适应)
在神经架构搜索完成之后,可以根据具体延迟的需求和限制,通过对进行排序,将具有较大注意力机制,设置成高准确率类型,对于较小的,设置成低延迟类型。
4、重训练异构注意力ViT(借助知识蒸馏提升模型性能)
直接训练搜索后的异构注意力机制ViT会导致显著的准确率下降,如何有效提高异构注意力了机制ViT的准确率?可以采用多粒度自蒸馏方式,让异构ViT去学习到原始Softmax ViT的软标签,以及最后一层的特征,这种粗细粒度的蒸馏能够很大程度的帮助ViT架构的学习。
二、搭建基于SecretFlow的Vision Transformer框架
分为以下四部分:
1、SecretFlow Secure Processing Unit (SPU)
前文已经多次介绍。
2、隐私推理协议及通信参数配置
隐私推理协议及节点配置,模拟网络参数配置throttle.sh,包括广域网和局域网。
3、基于Jax的ViT模型搭建
采用Jax语言,其语法与Pytorch有相似之处。具体包括图像Patch Embedding搭建,注意力机制搭建,MLP模块搭建和Transformer模块搭建。
4、ViT模型隐私推理
ViT Benchmark文件,包括初始化SPU环境,设置输入变量维度,模型实例化,以及密文推理过程。
ViT隐私推理操作流程,包括配置Python环境及安装SPU,配置并模拟通信网络环境,模拟MPC环境及协议,以及执行隐私推理。
三、MPCViT主要实验结果
和近年的SOTA模型相比,MPCViT在不同的数据集上,不同的架构上, 都显著的降低了延迟,并且具有更高的准确率。
在消融实验中可以看到,两种粒度的知识蒸馏对于异构ViT模型的训练都是至关重要的。
在搜索算法方面,在不同模型的架构上,在不同的超参数情况下,都有很好的一致性,同时可以扩展到不同规模的架构以及数据集上,具有很好的扩展性。
四、总结
首先介绍了ViT在MPC环境下进行隐私推理的总体框架、研究动机以及算法流程。
MPCViT借助蚂蚁的SecretFlow-SPU去搭建ViT隐私推理框架,并测试了不同模型架构的推理效率。
介绍了SecretFlow-SPU的特点和功能,并展开讲解了基于SecretFlow和Jax的ViT框架搭建流程
最后,呈现了该工作的主要实验结果,效果均优于基线模型。