近期，基于 TensorFlow 的隐私 AI 框架 Rosetta 正式开源（https://github.com/LatticeX-F...）。借助 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）等抽象代数结构上的操作，而具体的在代码中则落地到对于大整数、多项式等数据结构上的处理，所以框架设计上要在三个方面达成平衡：
1、尽可能的实现用户对于内部密文数据结构的透明无感；
2、TensorFlow 的 DAG 图构建、自动求导等核心功能仍需要无缝的支持；
3、支持不同的 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 等框架的一些设计方案。后续的文章中，我们会就核心的关键模块进行进一步的介绍。
本文转自 公众号：AI前线 ，作者 Rosetta 技术团队，点击阅读原文
