01 隐私保护机器学习背景
随着数据隐私问题日益受到关注,国家对于数据安全重视程度提升,以及监管力度不断加强,机器学习领域同样面临数据安全的问题。数据是机器学习算法预测准确性等效果的基础保证,训练高质量模型需要大量的有效数据,同时训练好的模型对外提供服务,需要用户数据作为输入。这些数据可能包含生物信息、金融信息等敏感数据,需要通过相应的技术手段,达到数据流通的技术信任。
那么如何在发挥数据价值的同时保护数据安全?
解决方案:安全多方计算(MPC)
即多个参与方可以在互不泄露任何信息(除结果外)的情况下协作计算一个函数
上面给出了3个参与方计算一个加法函数的例子。通过分片,在得到加法结果的同时,不会泄露自己的数据,其他参与方得到的只是一个拆分的随机值。实际的建模过程中会涉及到一些更复杂的非线性计算,依赖MPC协议实现。
基于MPC的隐私保护机器学习(PPML)有两大使用场景:隐私训练和隐私推理
隐私训练是在模型训练阶段Alice和Bob在不泄露各自的数据的情况下,可以联合构造一个更高质量的数据集,以提升模型的训练效果。
隐私推理是指在模型的推理阶段,有一个数据提供方Alice和模型提供方Bob。Alice的数据输入Bob的模型后可以得到一个计算的结果,但过程中Alice的数据和Bob的模型参数都不会泄露。
02 SPU架构简介
那么我们是否可以直接以 MPC 的方式高效地运行已有的机器学习程序?
事实上,ML 和MPC领域存在一些差异,如何去跨越这两种不同技术之间的差异? SPU是一种可行的方案。隐私计算实训营第二期第9讲SML入门/基于SPU迁移机器学习算法实践-CSDN博客
SPU的核心系统组件主要分为三块,如图所示:
前端,我们将依赖的AI前端代码翻译成XLA IR
编译器,我们使用MLIR技术栈对HLO进行优化并翻译成PPHLO(SPU字节码)
运行时,我们逐渐将Tensor ops拆解,经过SPU HAL(硬件抽象层,处理fxp/int),最终dispatch到协议层
协议层只需要实现Ring or Field上的基础运算即可
最终,通过编译时和运行时的层层翻译,SPU将AI前端和MPC后端解耦,使得在SPU中扩展的任何安全协议都可以无感的支持多种前端
开发者必看:深度解读隐语密态计算设备SPU · secretflow · Discussion #49 (github.com)
03 NN密态训练/推理示例
以逻辑回归为例,说明四个主要问题:
(1)数据从哪来:数据提供发Alice、Bob分别在本地加载数据。
2)如何加密保护数据:数据方对数据进行加密(数据碎片化),发送给MPC计算方(秘密分享),这里提到采用外包模式(一般数据保护要求高的场景不建议使用,因为很容易发生合谋被窃取数据),好处是可以利用计算方的算力资源。计算方拿到密文(数据碎片)。
(3)如何定义模型计算:使用JAX实现前向和后向传播。
(4)如何定义模型密态计算:计算方以密文数据作为输入,将模型的训练/推理计算图通过SPU编译器转换为相应密态算子计算图,由SPU device按照MPC协议逐个执行。
(5)配置文件说明
对于更复杂模型的密态训练和推理,可以使用stax/flax。
同时也可以复用构建好的模型,例如Huggingface上的模型。(现在很火的大模型的建模是很复杂的,即使使用stax/flax也很难自己实现)
很少的改动量即可实现。
如何支持不同的模型?
04 实践
基于FLAX库实现简单的MLP模型并密态执行对比
SPU 训练神经网络 | SecretFlow v1.7.0b0 | 隐语 SecretFlow
相关代码可以在链接里下载
隐语课程学习笔记10-基于SPU的机器学习建模实操-CSDN博客
这位博主有一些代码的注释,大家可以参考
明文训练结果:
基于SPU将上述明文训练转换为密文训练
import secretflow as sf
# Check the version of your SecretFlow
print('The version of SecretFlow:{}'.format(sf.__version__))
# In case you have a running scretflow runtime already.
sf.shutdown()
sf.init(['alice','bob'],address='local')
alice, bob = sf.PYU('alice'), sf.PYU('bob')
spu = sf.SPU(sf.utils.testing.cluster_def(['alice','bob']))
x1, _ = alice(breast_cancer)(party_id=1, train=True)
x2, y = bob(breast_cancer)(party_id=2, train=True)
init_params=model_init(n_batch)
device=spu
x1_,x2_,y_ = x1.to(device),x2.to(device),y.to(device)
init_params_ = sf.to(alice, init_params).to(device)
params_spu = spu(train_auto_grad, static_argnames=['n_batch','n_epochs','step_size'])(
x1_, x2_, y_, init_params_, n_batch=n_batch, n_epochs=n_epochs, step_size=step_size
)
测试密文训练得到的模型效果
params=sf.reveal(params_spu)
先将密文训练参数恢复成明文,然后利用恢复的明文参数对密文训练模型的预测结果进行评估(生产环境不建议这样使用,仅是为了验证本次实验的准确性)。可以看到密文的评估结果和明文模型基本一致,说明spu可以较好地在保持训练精度下的密文计算。
X_test,y_test=breast_cancer(train=False)
auc=validate_model(params, X_test,y_test)
print(f'auc={auc}')
密文训练的结果:
基于transformer库,调用gpt2,完成明文与密文的推理结果对比
(明文)在CPU上生成文本,可以正常输出下一个词的结果。
(密文)在SPU上生成文本。
import secretflow as sf
# In case you have a running secretflow runtime alLready
sf.shutdown()
sf.init(['alice','bob'],address='local')
alice,bob=sf.PYU('alice'), sf.PYU('bob')
conf = sf.utils.testing.cluster_def(['alice', 'bob'])
conf['runtime_config']['fxp_exp_mode'] = 1
conf['runtime_config']['experimental_disable_mmul_split'] = True
spu = sf.SPU(conf)
def get_model_params():
pretrained_model = FlaxGPT2LMHeadModel.from_pretrained("gpt2")
# 获取预训练模型的参数
params = pretrained_model.params
new_params = {
'transformer': {
'wte': params['transformer']['wte'],
'wpe': params['transformer']['wpe'],
'h': {str(i): params['transformer']['h'][str(i)] for i in range(layers_selected)}, # 仅使用指定层数
'ln_f': params['transformer']['ln_f'],
}}
return new_params
def get_token_ids():
return tokenizer.encode('I enjoy walking with my cute dog', return_tensors='jax')
model_params=alice(get_model_params)()
# input_token_ids=bob(get_token_ids)()
input_token_ids=alice(get_token_ids)()
print("step1")
#device=spu
#model_params_, input_token_ids_ = model_params.to(device), input_token_ids.to(device)
print("step2")
output_token_ids = alice(text_generation)(input_token_ids, model_params)
# output_token_ids = spu(text_generation)(input_token_ids_, model_params_)
如果遇到了内存OOM的问题,可以参照解决。隐语课程学习笔记10-基于SPU的机器学习建模实操-CSDN博客
检查SPU的输出,在SPU上运行GPT-2推理非常简单。可以执行reveal显示SPU生成的文本。
可以发现,SPU生成的文本与CPU生成的文本是完全一致的!