论文阅读:SecretFlow-SPU: A Performant and User-Friendly Framework for Privacy-Preserving Machine Learning
动机
开发高效的带有MPC技术的PPML程序对没有密码学背景的用户来说是一个挑战;现有的方法一般是用带有PPML版本的API来机械地替代原有的程序或者是重写整个程序
PPML框架的结构
前端
将输入的程序转化为MPC特定的中间表示intermediate representation(IR);
后端
是一个连接了多个计算节点的虚拟设备,接受IR也就是PPHLO(privacy-preserving high-level operations)为输入,然后把他们变成MPC协议的实现细节,以便于实现ML训练和预测。
文章贡献
- 第一个能够使用MPC的PPML框架,是用来实现ML程序的;显著提升了开发,测试,debug以及部署的速度。
- 设计了一个MPC特定的IR叫做PPHLO,连接了ML和MPC;实现了编译上的优化;开发了一个高性能的运行环境来执行PPHLO。
- 实验结果展现了高性能以及易于使用。
- 代码开源:https://github.com/secretflow
架构总览
威胁模型
SPU底层使用MPC协议是可设置的,威胁模型会根据使用的MPC协议而进行调整。
架构概览
大概流程
一个用JAX编写的ML程序,使用JAX APIs可以将程序转化为HLO和,同时分析数据可见性。框架的前端部分会根据数据可见性和HLO生成SPU特定的中间表示也就是PPHLO,同时,前端也会执行一些MPC特定的优化。优化后的PPHLO被送到后端,(一个连接了多个网络计算节点的虚拟设备)这些节点负责执行MPC操作,节点的数量必须跟MPC协议中支持的参与方一致。
示意图
程序接口
开发者可以使用程序接口定义数据可用性(如@ppd.device(“P1”))和要保护的函数(如@ppd.device(“SPU”)),如图:
隐私保护高级操作
PPHLO中包含了很多操作,这些操作的输入和输出都是张量,PPHLO中的张量类型可以被定义为一个三元组<Shape, Data Type, Visibility>,其中Shape表示张量的维度,数据类型可以是布尔,整数,定点数,可见性可以是公共的,也可以是私密的。
根据输入决定输出的类型:1)当一个操作数是定点数,那么结果也是定点数; 2)当一个操作数是私密的,那么结果也是私密的。
其实这张图有点看不懂哈哈哈哈,不知道pphlo.sec表示的是什么意思。
前端的一些优化操作
混合数据类型乘法融合
一般整数与定点数相乘要先把整数转为定点数,然后再将乘法操作调度到定点数乘法内核中;SPU将两个操作融合到一个操作来减少多余的截断和转化操作。
混合可见性乘法操作重新排序
可以将公共的数先相乘,这样子截断的时候就可以减少开销。
平方根的逆的转化操作
检测到平方根除法就自动转化为逆的乘法(也就是乘以平方根的倒数)。
选择预言的复用
使用同一个预言进行选择,可以将pred(预言)转为算术分享,然后在多个选择门中进行复用
最大池化的转化
前向传播中记下最大值的索引,在反向传播中直接使用索引
前端的优化主要是在操作上的顺序更改,操作的融合以及复用,这些优化能够提高框架的运行速度。
后端
操作调度
SPU运行时会根据数据类型和可见性,将一个操作分解为一些MPC原语函数。
每个MPC原语函数会被调度到MPC层,对应着MPC协议的实现。
新增一个MPC协议只需要实现MPC原语函数集。
运行优化
- 向量化:SIMD操作,单指令多数据(如有两个乘法操作,并行执行)可以减少通信轮次。
- 流化:例如,将张量平铺为子张量,并行处理他们。
- 并行化:SPU支持操作内和操作间的并行化。
对于内部操作的并行化,求平方根的倒数的时候,计算包含了小数部分计算和指数部分计算,可以独立在两个线程中执行。
对于跨操作的并行化,执行PPHLO操作图的时候,SPU会异步展开尽可能多的操作,当一个操作的依赖完成之后,这个操作会被调度执行。
限制和讨论
- 相比起浮点数,定点数的精度和范围都非常有限,可能产生不正确的结果。
- 一些函数实现依赖于浮点数表示,可能会产生意料之外的结果。
- 此外,SPU不支持秘密条件,(条件是用秘密表示的,比如while循环中的条件)
自我感受
这篇文章的主要是搭了一个大概的架构,将程序转化为HLO,然后转化为PPHLO,进行优化之后放到后端进行执行。前端和后端的一些优化能够提高整个框架的执行速度。在实验中,作者发现了最大池化转化和平方根的你能够让性能大大提高。
假如使用中没有使用到池化层,或者很少使用,那么最大池化转化整个操作是不是就不能去起到提升性能的关键作用。此外,能不能有一些其他前端操作上的优化?