SecretFlow-SPU 学习笔记
论文链接:https://www.usenix.org/system/files/atc23-ma.pdf
一、引言
1.1 背景与动机
- 数据隐私的重要性:在大数据时代,数据隐私变得至关重要。数据泄露和隐私侵犯事件频发,特别是在金融和医疗等敏感领域。因此,隐私保护成为学术界和工业界的一个重要研究方向。
- 安全多方计算(MPC):MPC是一种密码学技术,可以在不泄露各方数据的情况下进行联合计算。这使得在多方协作时能有效保护数据隐私。
- 挑战:传统机器学习框架与MPC技术之间存在显著差异,导致将MPC应用于机器学习任务变得复杂且具有挑战性。特别是对没有密码学背景的用户来说,这种复杂性显得尤为突出。
1.2 现有解决方案的不足
- API替换和重写成本高:现有PPML解决方案通常需要用户大幅修改现有的机器学习代码,或者完全重写以适应MPC的需求。这种过程既费时又费力。
- 用户友好性差:虽然有些解决方案试图模仿现有机器学习框架的API,但用户仍需掌握一定的密码学知识,无法完全屏蔽底层的复杂性。
1.3 SecretFlow-SPU 的目标
- 高性能与易用性:SecretFlow-SPU的设计目标是提供一个高性能且用户友好的PPML框架,用户只需对现有机器学习程序进行少量修改即可运行。
- 集成MPC与ML:通过前端编译器和后端运行时的合作,将现有的机器学习程序转换为MPC特定的中间表示,并在多个计算节点上高效执行这些表示。
二、系统设计
2.1 威胁模型
- 标准MPC威胁模型:SPU基于标准的MPC威胁模型,运行多个参与方之间的预定义程序,保护输入数据和中间结果,通常只向一些指定方揭示最终结果。
- 配置灵活性:SPU可以根据选定的MPC协议配置系统的威胁模型。例如,使用半诚实(具有诚实多数)的ABY3协议,SPU将继承ABY3的威胁模型。
2.2 架构概述
- 前端编译器:SPU的前端编译器将机器学习程序转换为PPHLO(隐私保护高级操作)中间表示,并进行优化。PPHLO表示包含一系列操作,输入和输出都是张量。
- 后端运行时:后端运行时在多个计算节点上执行PPHLO,支持多种MPC协议。SPU后端运行时由多个连接的计算节点组成,这些节点托管SPU运行时,负责执行MPC协议。
2.3 编程接口
- Python API:SPU提供了简单的Python API,允许开发者通过少量代码修改在SPU上运行机器学习程序。例如,使用装饰器
@ppd.device()
指定数据可见性和受保护函数。 - 示例代码:
pythonCopy codeimport jax.numpy as jnp
import numpy as np
import spu.binding.util.distributed as ppd
# 初始化SPU后端节点
with open("/path/to/config", 'r') as file:
conf = json.load(file)
ppd.init(conf["nodes"], conf["devices"])
# 指定数据可见性
@ppd.device("P1")
def data_from_alice():
return np.random.randint(100, size=(4,))
# 指定数据可见性
@ppd.device("P2")
def data_from_bob():
return np.random.randint(100, size=(4,))
# 指定私有函数
@ppd.device("SPU")
def compare(x, y):
return jnp.maximum(x, y)
# 自动获取x和y作为秘密分享
x = data_from_alice()
y = data_from_bob()
# 私下比较x和y
z = compare(x, y)
# 公开z的真实值
print(f"z = {ppd.get(z)}")
2.4 PPHLO
-
PPHLO 概述:PPHLO是一种基于HLO(高级操作)的定制中间表示,包含一系列操作。每个操作的输入和输出都是张量。PPHLO的主要区别在于张量的类型系统,由<Shape, Data Type, Visibility>三元组表示。Visibility(可见性)决定了张量是否需要保护。
-
类型推导规则
:
- 数据类型提升:如果操作数之一是定点数,则结果也是定点数。
- 可见性缩小:如果操作数之一是秘密,则结果也是秘密。
2.5 前端优化
-
编译优化
:前端编译器接收机器学习程序的HLO和初始数据可见性,通过一系列规则推导出整个计算图的类型和可见性,生成合法的PPHLO表示。实现了多种编译优化技术,以提高性能和效率。
- 混合数据类型乘法融合:将整数与定点数相乘的操作融合为一个操作,减少不必要的截断和转换。
- 混合可见性乘法操作数重排序:重排序操作数,先计算公共定点数的乘积,再与秘密定点数相乘,减少高通信开销的截断操作。
- 反平方根转换:将y/√x+u转换为y*rsqrt(x+eps()),利用快速MPC实现逆平方根操作。
- 选择谓词复用:在多个选择操作中复用谓词,减少冗余转换。
- 最大池化转换:优化前向传播和反向传播中的最大池化操作,减少冗余比较。
2.6 后端优化
- 向量化:通过在一组数据上运行一个操作来减少执行操作的数量,提高并行性和执行效率。
- 流处理:将大张量分块并行处理,以提高网络和计算资源的利用率。
- 并发技术:在PPHLO图的粒度上防止通信操作阻塞计算资源,支持操作的并发执行。
三、实现与评估
3.1 实现
- 实现细节:SPU前端编译器基于MLIR(多层次中间表示)实现,PPHLO作为MPC计算的一个新方言。后端运行时使用现代标准C++开发,计算节点通过高性能RPC库bRPC通信。
3.2 性能评估
- 评估方法:通过与MP-SPDZ、TF Encrypted 和CrypTen等框架在MNIST数据集上的神经网络训练任务进行比较,评估SPU的性能。SPU在大多数设置中表现出色,特别是在广域网环境中。例如,在广域网条件下,SPU比MP-SPDZ快最多4.1倍,比TF Encrypted快最多2.3倍。
3.3 用户友好性评估
- 用户友好性:通过运行JAX程序并进行少量修改,展示了SPU的用户友好性。SPU还支持TensorFlow和PyTorch,展示了其对不同机器学习框架的扩展性。实验表明,SPU可以轻松扩展到其他机器学习模型和框架。
四、局限性和讨论
4.1 局限性
- 定点表示的限制:SPU使用定点表示来编码小数,导致与浮点数相比精度和范围有限。这可能导致某些机器学习程序在SPU上运行时得到不正确的结果。
- 秘密条件的支持:目前SPU不支持一些操作(如while)使用的秘密条件。这些值通常不是需要保护的数据,开发者可以将这些值实现为公共张量。
4.2 讨论
- MPC技术的优势:MPC技术与其他PPML框架(如基于TEE的解决方案和联邦学习)相比具有显著优势,提供端到端的隐私保护。MPC不需要一个独立的服务器负责模型聚合,参与方的数据首先被加密,然后在SPU上进行计算(如梯度更新)并训练模型,最终模型在某些指定方之间重构为明文。