近期，基于TensorFlow的隐私AI框架Rosetta正式开源。借助Rosetta，AI开发者不需要了解密码学等隐私保护技术，只需要改动几行代码，就能赋予自己的程序以保护数据隐私的能力。本专栏将通过多篇独家技术稿件，深度披露Rosetta的整体框架设计、TensorFlow的定制化改造最佳实践、将密码学理论算法进行高效工程化落地等内容。通过本系列专题，希望能让更多开发者了解隐私AI框架的技术挑战，同时为密码协议算法工程化、AI框架深度定制等相关方向的开发者提供一些经验参考。

数据是AI技术的“燃料”已经成为业界的共识，更多的数据往往意味着可以训练出更准确的模型。但无论是在公司内部还是多个企业之间，为了对用户数据负责、合法合规，在数据的分享使用时，必须注意对于原始明文数据的保护。

传统的面向静态数据保护的安全手段无法解决数据的动态使用、分享中的隐私泄露问题，而正是这一实际需求催生出了隐私计算（在AI场景下，更进一步的可以称之为隐私AI）这一新的交叉技术。隐私计算技术融合在数据的使用过程之中，保障的是计算过程本身（广义的讲，还包括计算结果）不会泄露原始明文数据本身的信息。

目前实现隐私计算的途径可以分为密码学、联邦学习和硬件可信执行环境（TEE）等几大类。而其中以密码学理论为基础的MPC（Multi-Party Computation，安全多方计算）是最有安全保障的技术路线，其秉持的基本理念是信任计算复杂度理论、代码，而不是信任人、硬件，而联邦学习和TEE目前还很难讲清楚安全性，经常被发现新的安全漏洞。并且，联邦学习中核心部分也往往需要使用同态加密等密码学手段进行强安全性的保障。从工程技术的角度上看，联邦学习是分布式机器学习技术的延伸，主要的挑战是训练过程中如何进行多异构终端的同步更新等[1]，很多传统分布式系统开发经验仍然适用。而以MPC为代表的密码学途径则带来了一些全新的挑战。

安全多方计算MPC（图片来自网络）

其中最核心的困难是，密码学属于计算机理论领域，很多概念、算法协议都需要有长期的专业知识积累才能理解，而业务落地中的典型AI方向，无论是计算机视觉、文本挖掘还是用户行为建模等都更加面向实际场景。如何打通以密码学为典型代表的隐私保护技术与AI技术之间的壁垒？这是开发者在实际构建一个通用的、易用的隐私计算框架时需要解决的核心问题。围绕着这个核心问题，又有一系列具体的工程技术挑战：

如何实现系统的易用性？AI开发者不会愿意，也不应该为了在业务中引入数据隐私保护能力而费时耗力地学习各种复杂、抽象的密码算法。一个好的隐私AI框架应该是易上手的，便于AI开发者使用自己熟悉的方式快速解决自己的数据隐私问题。
如何实现系统的高效执行？这包括单机和多机两个层面。密码学的计算大部分都是在大随机数的密文上进行，为此常常需要使用专用的硬件指令、SIMD（Single Instruction/Multiple Data）等技术来进行单机并行化的加速，这些优化实现需要对于密码学基础库有深入的了解，并往往需要根据协议算法做进一步的并行优化。而在多机层面上，则需要考虑如何和很多AI框架自身的并行优化技术兼容。
如何实现MPC多方之间的高效通信？在MPC中，多方之间需要进行大量的同步通信，而且信道上的内容大都是无规律的、不可压缩的一次性使用的随机数，这就需要在保证安全性的同时，根据具体的计算逻辑进行很多工程优化以减少通信量和通信次数。
如何保障隐私保护技术的可扩展性？MPC等隐私计算技术还在不断发展之中，也是学术研究上的热点问题，所以一个好的隐私AI框架，需要能够支持研究者简单快速地将新的算法协议集成进来。

针对这些问题，业界已经有一些探索，下面我们结合Rosetta来具体谈一谈在隐私AI框架的设计和实现中如何克服这些挑战。篇幅所限，本文主要整体性地先介绍宏观设计，后续系列文章会进一步剖析一些技术细节。

如同其他的隐私AI框架一样，Rosetta仍然处于发展的早期，尚有一些不完善之处。我们在此以Rosetta为例是希望能具体化地讲清楚这一领域中的细节挑战，也希望激发更多开发者参与到未来的隐私AI系统设计中来。

隐私AI框架整体设计思路

目前业界尚没有大规模落地的、成熟完善的隐私AI框架，但是已经有一些探索性的开源隐私AI框架，比如PySyft、TF Encrypted和CrypTen。

从整体上看，这些框架都是在TensorFlow或PyTorch的前端Python层进行封装集成的。这样做的好处是可以直接使用这些AI框架的上层接口来实现隐私计算算法，而且天然的可以直接调用框架自身封装好的一些高层次API高功能。这对于联邦学习这种本身就源自分布式机器学习的技术来说是较为适合的，但是对于密码学的MPC来说会有一些不足：

首先，单机性能无法得到充分的提升。用Python来实现各种复杂的密码学计算、多方之间的通信无法充分利用底层操作系统、硬件层的并行优化。而且更现实的是，大部分高性能的密码学库提供的是C/C++的接口，如果在AI框架的前端上层实现MPC等密码技术，则很难复用这些业界长时间积累（同时也还在不断发展中）下来的成果。
其次，密码协议等隐私技术的实现和AI框架本身耦合过深，不利于扩展。由于这些AI框架本身提供的对外API接口本就是面向AI需求的，在实现较为复杂的MPC等密码协议时不仅需要熟练使用这些框架的API，还往往需要大量直接使用numpy等库来实现复杂的计算逻辑。这一方面破坏了对AI框架自身使用上的自洽，不再能将全部的计算逻辑完全承载在AI框架的逻辑执行图上，另一方面也使得每一次引入新的后端密码协议时都需要重新基于AI框架进行实现，这对于密码协议开发者来说成本很高。

基于上述认识，Rosetta在现阶段首先以TensorFlow这一流行AI框架为基础，深度改造其前端Python入口和后端kernel实现，并封装可插拔的MPC算法协议作为“隐私保护引擎”来驱动整个计算过程中数据的安全流动。

Why TensorFlow？

TensorFlow和PyTorch是目前工业界最主流使用的开源AI框架。虽然很多公司在内部可能也会根据自身需求定制化改造一些组件，或者推出各具特色的新框架以在易用性、高效性、完备性等不同的维度上进行进一步突破，但是整体上看，这些框架基本的设计范式是较为相似的。大多是通过丰富的接口API让用户以有向无环图DAG的形式表达上层计算逻辑，而框架本身则会在实际调度执行这些计算任务时进行一系列的优化。TensorFlow虽然在用户友好性上略逊色于PyTorch，常受开发者诟病，但是其在可扩展性、高效性上、分布式部署等方面确实是更加均衡、全面（当然这也意味着TensorFlow更加复杂，对其改造会更加的具有挑战性）。所以综合考虑下来，Rosetta在当前版本中选择TensorFlow作为基本的底层承载体，在设计开发的过程中，一方面会充分利用TensorFlow内在的计算图并行执行优化等功能以提升效率，另一方面也会尽量克制，主要是利用其作为深度学习框架通用性的一些接口特性，而不会过于依赖其独有的一些组件。

Rosetta框架核心设计思想

隐私算子（SecureOp）作为核心抽象接口连接AI框架和隐私计算技术。TensorFlow在不同的层次上提供了多样的扩展方式，Rosetta选择后端算子（Operation）层作为核心切入点，TensorFlow在执行算子时会动态的绑定到具体MPC协议中的SecureOp实现中。通过这样的抽象，密码协议开发者可以不需要了解AI框架，只需要用C++实现满足接口定义的各个功能函数即可，而AI开发者也不需要深入了解MPC等技术的实现细节，而只需要在现有算子的基础上进一步封装自己想要的上层高级功能即可。
基于优化遍（Pass）的分阶段转换。为了尽可能给AI开发者提供易用的接口，减少给线上AI程序赋予数据隐私保护能力时的改造成本，Rosetta在整体的设计中借鉴了程序编译器领域的核心概念：Pass。Pass是编译器中常用的技术，主要用作将源码文件一步步转变为机器码过程中的多轮转化和优化。在Rosetta中，用户使用原生TensorFlow接口编写的DAG（有向无环图）形式的逻辑计算图会被分阶段的转换、替换为多方协作执行的MPC程序，这样可以实现对于用户API层最少的改动。

具体的，在Rosetta中，有两个阶段的Pass，一个在前端Python层的全局DAG构建过程中生效的Static Pass，会将原生Tensor转换为支持自定义密文类型的RttTensor，将原生Operation转换为支持tf.string格式输入输出的RttOp，并最终在图开始启动时进一步的转换为承载实际MPC操作的SecureOp。

Static Pass

另一个是在SecurOp的实际执行时所进行的Dynamic Pass处理，会动态的根据当前用户选择的协议选择对应的实际算子实现去执行，同时可以在此时嵌入基于执行上下文的优化处理。

Dynamic Pass

融合MPC技术的分布式隐私AI架构

理清楚整体的分布式结构对于了解一个系统的架构大有裨益。整个隐私AI系统对外接口会涉及三个方面，如何指定物理部署上的网络拓扑？数据在整个计算的过程中是如何安全输入、流动、输出的？隐私计算逻辑要如何表达？Rosetta的整体逻辑结构如下图所示：

Rosetta 多方网络结构图

多方网络的建立

MPC技术本身就是要求多方（multi-party）参与的，一般称他们为Player，不同的MPC算法协议会有不同个数的参与方。以Rosetta中目前实现的三方协议SecureNN[2]为例，系统中存在三个逻辑参与方，P0、P1和P2。

在v0.2.1版本中，在这一方面的用户接口层次上，为了保障对外的灵活性，目前支持用户通过配置文件来一次性指定多机之间的网络关系，也支持调用接口动态的激活、解除多方之间的网络拓扑：

# 调用activate接口会根据配置参数或配置文件建立起网络
rtt.activate(protocol_name="SecureNN", protocol_config_str=None)
# 调用deactivate接口会释放网络链接等资源
rtt.deactivate()

在内部实现中，每一个参与方都会监听一个本地的server端口，而同时分别建立到另外两方之间的client网络链接。这样的好处是相互之间的网络链接关系简单清晰，当然也需要解决随之而来的SecureOp并发同步执行时的一致性问题，这一点我们会在后续文章中讨论。

一些注意点

熟悉TensorFlow的读者可能会疑惑，这种多方基于不同数据跑相同程序的模式，不就是TensorFlow分布式执行下对数据并行进行支持的In-graph replication和Between-graph replication吗？并不是这样，实际上它们是不同层次的结构，这里讲的是上层逻辑视角的MPC各参与方，在实际中，你甚至可以进一步的将各方内部执行的这一task按照TensorFlow的分布式规范进行集群部署，而将集群中的 "server“作为统一的对外逻辑代表。

上面一直讲的是“逻辑上”的三方，那么在实际的业务场景中，可能是2个、4个或以上公司之间的数据合作，是不是就不能用这些架构了吗？其实不然，我们完全可以在上层进行一层映射，以Privacy-as-a-Service的形式提供对上层的服务，关于这一点，后续文章也会进一步介绍。

隐私数据的流动

每一个逻辑参与方都可以有自己私有的明文输入数据，也可以继续处理上一次任务输出的密文结果。在整个程序的运行过程中，只有开始和结束时数据才会以明文的形式存在：开始时对于私有数据的引入，以及最后可配置是否将结果以明文的形式恢复出来加以输出。而在中间各个算子的计算过程中，数据总是以密文形式在本地的逻辑上下文、多方之间进行交互。

对外接口方面，在实际的业务中多方数据之间是需要关联对齐的，目前Rosetta提供常见的两种数据集处理方法，一是对应于整体上数据集是在各方之间“水平划分”的场景，即各方拥有不同样本id的全部特征属性值；另一种对应于整体数据集是在各方之间“垂直划分”的场景，即各方之间拥有相同样本id的部分特征属性值。这些都可以调用PrivateDataset类的load_data等接口方便的处理。而在输出阶段，提供了如下两个接口：

# 将一个密文形式的cipher_tensor恢复为明文， receive_party参数指定3方中哪几方可以获得明文结果
rtt.SecureReveal(cipher_tensor, receive_party=0b111)
# 与原生TensorFlow中模型保存接口SaveV2具有一样的函数原型，可通过配置文件指定哪几方可以获得明文模型文件
rtt.SecureSaveV2(prefix, tensor_names, shape_and_slices, tensors, name=None)

隐私集合求交PSI（Private Set Intersection）技术

在实际场景中还存在一个很现实的问题，就是多方之间样本的对齐问题，比如将A方的样本id所指向的样本和B方此样本id对应的属性信息给对应起来。PSI技术可以安全的解决上述问题，目前各个开源框架中还没有将这一功能很好的集成进来，Rosetta目前正在集成这一功能，将在近期版本中发布。

在内部实现中，密码学中的很多运算是在空间较大的环（Ring）、域（Field）等抽象代数结构上的操作，而具体的在代码中则落地到对于大整数、多项式等数据结构上的处理，所以框架设计上要在三个方面达成平衡：

尽可能的实现用户对于内部密文数据结构的透明无感；
TensorFlow的DAG图构建、自动求导等核心功能仍需要无缝的支持；
支持不同的MPC协议使用自定义的具体数据结构对象以便于扩展。

为同时实现这些目标，Rosetta基于tf.string这一TensorFlow原生数据结构来承载各协议自定义的密文数据，然后通过对TensorFlow源码代码进行深度的hook改造使得DAG图构建、自动求导等功能仍然可用。

DAG的执行

如上述网络结构图所展示的那样，各Player运行的是同一份基于TensorFlow编写的AI二进制代码，比如训练一个简单神经网络模型的程序。用户直接的使用TensorFlow中原生的算子API来构建逻辑计算图DAG，Rosetta内部会在图开始执行时完成到隐私算子SecureOp的转换。与其他隐私计算框架相比，这样的切换成本是最低的。

在执行过程中，各Player自身都是在按照这个DAG图在运行，特殊的地方在于在各个算子的内部执行过程中，各个Player会根据自己所属的角色遵循MPC协议执行不同的操作，这些操作即包括本地的在密文上的处理，也包括在多方之间进行强同步的通讯交互，传输大量随机数形式的密文。

Rosetta多方协作运行图

小结

在本篇文章里，我们结合Rosetta框架，从整体上介绍了隐私AI框架在工程落地时所需要面对的挑战，以及Rosetta等框架的一些设计方案。后续的文章中，我们会就核心的关键模块进行进一步的介绍。

目前Rosetta框架已经在GitHub开源，欢迎关注。

GitHub地址：
https://github.com/LatticeX-Foundation/Rosetta

作者介绍

Rosetta技术团队，一群专注于技术、玩转算法、追求高效的工程师。Rosetta是一款支持主流深度学习框架TensorFlow的隐私AI开源框架，作为矩阵元公司大规模商业落地的重要引擎，它承载和结合了隐私计算、区块链和AI三种典型技术。

参考资料

Bonawitz, K., Eichner, H., Grieskamp, W., Huba, D., Ingerman, A., Ivanov, V., … & Van Overveldt, T. (2019). Towards federated learning at scale: System design.arXiv preprint arXiv:1902.01046.
Wagh, S., Gupta, D., & Chandran, N. (2018). SecureNN: Efficient and Private Neural Network Training.IACR Cryptol. ePrint Arch.,2018, 442.

创作场景

隐私 AI 工程技术实践指南：整体介绍