隐语实训08-SPU论文解读

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（可见性）决定了张量是否需要保护。

• 类型推导规则

  ：

  1. 数据类型提升：如果操作数之一是定点数，则结果也是定点数。
  2. 可见性缩小：如果操作数之一是秘密，则结果也是秘密。

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上进行计算（如梯度更新）并训练模型，最终模型在某些指定方之间重构为明文。