隐私AI框架中MPC协议的快速集成

我们在上一篇文章中，介绍了为了实现隐私 AI 系统的易用性，如何对 TensorFlow 这样的深度学习框架进行深度的改造。 本篇文章进一步进入 TensorFlow 后端算子的内部实现，阐述 Rosetta 中如何通过定义通用的密码协议抽象接口层实现系统解耦，使得隐私计算研究者可以轻松地将 MPC 协议这样的隐私计算技术给集成进来。

在第一篇整体介绍中，我们简要对比过 PySyft [1] 等探索性隐私 AI 框架，它们都是在 PyTorch 等深度学习框架之上，在 Python 接口层利用高层 API 来实现密码学协议的。这种方式虽然具有可以直接复用 AI 框架提供的接口、简化在 Python 进行的并发优化等优点，但也使得密码学专家等隐私计算技术的开发者必须要了解具体的 AI 框架。

此外，由于密码学的基础运算较为耗时，所以实际中为了有更高性能的实现方式，大部分相关的优秀基础库软件都是用 C/C++ 等语言实现，而且还会融合不同底层硬件体系结构下的指令集做进一步的加速。

所以，从便利隐私计算开发者、系统性能提升等角度出发，Rosetta 在后端将隐私计算技术的具体实现给抽象解耦出来，定义了一层较为通用的抽象接口。当开发者需要引入定制化的新隐私算法协议时，只需要参考接口定义规范，自由地按照自己熟悉的方式实现基础功能，就可以很快地将新功能引入进来。而在 Python API 层使用时，通过一个接口的调用就可以完成协议的切换。

接下来，我们首先会整体介绍 Rosetta 中为支持协议扩展所设计的抽象接口层，然后在第二部分结合一个Naive协议示例，来具体说明如何基于这些组件快速的集成一个新的自定义 MPC（Multi-Party Computation，安全多方计算）协议。

注意： 目前，MPC 是基于密码学手段实现隐私计算这个方向上使用的最主要具体技术。所以，下文中除特别指明外，我们所称的“隐私协议”、“密码协议”都是指 MPC 安全协议。 这里的相关介绍仍基于 Rosetta V0.2.1 版本，后续随着项目的迭代优化，可能会有局部调整。

密码协议统一接口模块

为了使得整体架构上足够的灵活、可扩展，Rosetta 的后端 C++ 开发中同样遵循经典的 SOLID 原则 [2] 来进行整体的设计。整个的密码协议统一抽象接口层根据功能职责进一步的划分为三个不同层次，并分别封装为 ProtocolManager、ProtocolBase 和 ProtocolOps 三个类，其中第一个类ProtocolManager 是一个单例（Singleton）类，是上层 API、TensorFlow 的后端算子实现中所需要唯一感知的组件。而ProtocolBase 和 ProtocolOps 则是两个接口类，由它们定义统一的各个后端具体密码协议所需要实现的功能。这三个类之间的整体关系如下：

细心的读者应该还记得，我们在上一篇文章的最后指出，在 TensorFlow 后端算子组SecureOp的 kernel实现中会最终调用这个模块：

// call protocol ops
vector<string> outstr(m*n);
ProtocolManager::Instance()->GetProtocol()->GetOps(msg_id().str())->Matmul(in1, in2, outstr, &attrs_);

这行语句结合上述 UML 类图，可以很清晰地看出各个组件之间的调用链关系：通过协议管理组件入口获取当前上下文的协议对象，协议对象通过算子入口进一步调用具体某一算子函数。

下面就让我们分别简要介绍下这三个核心类。

ProtocolManager

ProtocolManager 是总的入口，负责整体的控制。其内部为了支持多个可选协议，会维护一个协议名到ProtocolBase指针对象的映射表，并据此进行上下文的控制。除了Instance 方法是常规的取得这个 Singleton 的对象实例外，它的功能接口可以分为两大类，一类是面向上层 Python SDK 的，一类是面向开发者进行协议扩展时加以支持的。

• 上层 Python 层的一些协议相关的 API，如activate 和 deactivate 等，会内部调用ProtocolManager 的ActivateProtocol、DeactivateProtocol等方法，来实现对当前使用的协议的控制。而这些类成员函数的内部会进一步的调用ProtocolBase 接口基类的Init 和 Uninit 等方法。
• 而当我们需要引入一个新的协议时，在这一层所需要做的仅仅是调用其RegisterProtocol方法来注册这个新的协议即可。

ProtocolBase

ProtocolBase 是每一个具体的协议都需要最终实现的接口基类。其中Init接口定义如何进行协议的初始化，具体协议中需要在这个接口中根据用户传入的配置信息，实现多方之间网络的建立、本地秘钥和随机数的设置等操作。其函数原型如下：

  /**
   * @desc: to init and activate this protocol. 
   *         Start the underlying network and prepare resources.
   * @param:
   *     config_json_str: a string which contains the protocol-specific config.
   * @return:
   *     0 if success, otherwise some errcode
   * @note:
   *   The partyID for MPC protocol is also included in the config_json_str,
   *   you may need extract it.
   */
  virtual int Init(string config_json_str = "");

  /**
   * @desc: to uninit and deactivate this protocol.
   * @return:
   *     0 if success, otherwise some errcode
   */
  virtual int Uninit();

在 Rosetta 中，为了进一步便于简单协议的集成，我们用一个子类 MpcProtocol 封装了可以复用的一些功能，比如一般 MPC 协议中常用的一个技巧是：多方两两之间通过设定共同的 shared key 来配置伪随机数发生器 PRG，这样可以减少实现协议时多方之间的交互次数和通讯量。这个子类中就基于这些可能可以复用的功能实现了 ProtocolBase 中的Init 和Ubinit 方法。

另一个主要的方法GetOps 则会进一步调用对应协议的ProtrocolOps的子类对象来进一步 delegate 具体算子的实现。

以 Rosetta 中定制化实现的 SecureNN 协议为例，我们是通过SnnProtocol 这个子类来具体实现的。其类继承关系图如下：

ProtocolOps

ProtocolOps 用于封装各个安全协议中具体所需要实现的算子接口。大部分基础算子的函数原型中，都以字符串 向量作为参数类型，并可以通过一个可选的参数传入相关属性信息，比如Add的函数原型是：

# `attr_type` is just an inner alias for `unordered_map<string, string>`
int Add(const vector<string>& a,
    	const vector<string>& b,
    	vector<string>& output,
    	const attr_type* attr_info = nullptr);

注意: 我们在前面的文章中介绍过，在 Rosetta 内部为了支持多种后端协议中自定义的密文格式，我们统一在外层用字符串来封装密文数据，所以这里参数的基础类型都是字符串。

各个具体的协议需要进一步的实现各个算子函数，比如，在 Rosetta 中实现的 SecureNN 协议中的各个函数的实现是在子类SnnProtocolOps 中加以实现：

在这些具体的各个函数内部实现中，基本的步骤是先将字符串解码为此协议内部所设定的类型，然后进一步的进行多方之间安全的逻辑计算（这里一般是需要进行通信交互的），最后再将得到的内部计算结果编码为字符串输出到出参中。比如下面是 SnnProtocolOps 中矩阵乘法函数 Matmul 的代码片段：

int SnnProtocolOps::Matmul(const vector<string>& a,
  							const vector<string>& b,
  							vector<string>& output,
  							const attr_type* attr_info) {
  int m = 0, k = 0, n = 0;
  if (attr_info->count("m") > 0 && attr_info->count("n") > 0 && attr_info->count("k") > 0) {
    m = std::stoi(attr_info->at("m"));
    k = std::stoi(attr_info->at("k"));
    n = std::stoi(attr_info->at("n"));
  } else {
    log_error << "please fill m, k, n for SnnMatmul(x, y, m, n, k, transpose_a, transpose_b) ";
    return -1;
  }

  bool transpose_a = false, transpose_b = false;
  if (attr_info->count("transpose_a") > 0 && attr_info->at("transpose_a") == "1")
    transpose_a = true;
  if (attr_info->count("transpose_b") > 0 && attr_info->at("transpose_b") == "1")
    transpose_b = true;

  vector<mpc_t> out_vec(m * n);
  vector<mpc_t> private_a, private_b;
  snn_decode(a, private_a);
  snn_decode(b, private_b);

  std::make_shared<rosetta::snn::MatMul>(_op_msg_id, net_io_)
    ->Run(private_a, private_b, out_vec, m, k, n, transpose_a, transpose_b);

  snn_encode(out_vec, output);
  return 0;
}

从中可以看出，我们先从属性信息中直接取出矩阵输入参数的 shape 信息，然后将输入数据通过snn_decode 转换为内部的mpc_t 类型。再调用根据 SecureNN 的协议算法实现的多方协同计算的内部函数 MatMul之后就会得到更新之后的结果密文数据，最后通过snn_encode 重新将密文数据由mpc_t 转换为字符串类型加以输出。

• SecureNN 中的mpc_t 类型就是 uint64_t （如果用户配置了使用128位的大整数，则是uint128_t）。因为很多密码学的基础操作都需要在抽象代数的环（ring）、域（field）上进行（同时，最新 SecureNN 等 MPC 协议又为了充分利用基础硬件的运算加速，已经支持直接在整数环$Z_{2^{64}}$上进行运算处理），所以转换到大整数上几乎是所有相关隐私技术必要的内部操作。

示例：Naive 协议的集成

下面，我们结合一个示例协议 Naive 来具体演示下如何快速集成 MPC 协议到 Rosetta中。

注意： 这个 Naive 协议是一个不安全的、仅用于演示的协议！不要 naive 的在任何生产环境下使用此协议！ 这个 Naive 协议是一个不安全的、仅用于演示的协议！不要 naive 的在任何生产环境下使用此协议！ 这个 Naive 协议是一个不安全的、仅用于演示的协议！不要 naive 的在任何生产环境下使用此协议！

我们仅在 Naive 协议中实现最基本的加法和乘法等操作。在这个“协议”中，P0 或 P1 会将自己的私有输入值平均分为两份，一份自己持有，另一份发送给对方，作为各自的“密文”。然后在乘法等后续操作中，基于这样的语义进行对应的操作处理。这个协议显然是不安全的。

按照上一小节的介绍，我们只需要少量的修改相关代码即可实现在 Rosetta 中使用这个协议，完整的代码修改可以参考这里。具体的，类似于上面介绍的 SecureNN 中算子的实现，我们在 NaiveOpsImpl 类中实现内部逻辑的处理。其中实现隐私输入处理和“密文”下乘法操作的部分代码片段如下：

int NaiveOpsImpl::PrivateInput(int party_id, const vector<double>& in_x, vector<string>& out_x) {
  log_info << "calling NaiveOpsImpl::PrivateInput" << endl;
  int vec_size = in_x.size(); 
  out_x.clear();
  out_x.resize(vec_size);
  string my_role = op_config_map["PID"];

  // In this insecure naive protocol, we just half the input as local share.
  vector<double> half_share(vec_size, 0.0);
  for(auto i = 0; i < vec_size; ++i) {
    half_share[i] = in_x[i] / 2.0;
  }

  msg_id_t msgid(_op_msg_id);
  if (my_role == "P0") {
    if (party_id == 0) {
      io->send(1, half_share, vec_size, msgid);
    } else if (party_id == 1) {
      io->recv(1, half_share, vec_size, msgid);
    }
  } else if (my_role == "P1") {
    if (party_id == 0) {
      io->recv(0, half_share, vec_size, msgid);
    } else if (party_id == 1) {
      io->send(0, half_share, vec_size, msgid);
    }
  }
  for(auto i = 0; i < vec_size; ++i) {
    out_x[i] = std::to_string(half_share[i]);
  }
  return 0;
}

int NaiveOpsImpl::Mul(const vector<string>& a,
                      const vector<string>& b,
                      vector<string>& output,
                      const attr_type* attr_info) {
  log_info << "calling NaiveOpsImpl::Mul" << endl;
  int vec_size = a.size();
  output.resize(vec_size);
  for (auto i = 0; i < vec_size; ++i) {
    output[i] = std::to_string((2 * std::stof(a[i])) * (2 * std::stof(b[i])) / 2.0);
  }
  return 0;
}

而在框架集成方面，只需要在ProtocolManger中添加一行协议注册代码：

REGISTER_SECURE_PROTOCOL(NaiveProtocol, "Naive");

在完成上述简单代码修改后，我们重新编译整个 Rosetta 代码库，就可以把这个协议集成进来了！完全不需要修改任何 TensorFlow 相关的代码。下面让我们运行一个上层 demo 验证下效果，在这个 demo 中，我们直接在“密文”上计算 P0、P1 隐私数据的乘积：

#!/usr/bin/env python3

# Import rosetta package
import latticex.rosetta as rtt
import tensorflow as tf

# Attention! 
# This is just for presentation of integrating a new protocol.
# NEVER USE THIS PROTOCOL IN PRODUCTION ENVIRONMENT!
rtt.activate("Naive")

# Get private data from P0 and P1
matrix_a = tf.Variable(rtt.private_console_input(0, shape=(3, 2)))
matrix_b = tf.Variable(rtt.private_console_input(1, shape=(3, 2)))

# Just use the native tf.multiply operation.
cipher_result = tf.multiply(matrix_a, matrix_b)

# Start execution
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    # Take a glance at the ciphertext
    cipher_a = sess.run(matrix_a)
    print('local shared matrix a:\n', cipher_a)
    cipher_result_v = sess.run(cipher_result)
    print('local ciphertext result:\n', cipher_result_v)
    # Get the result of Rosetta multiply
    print('plaintext result:\n', sess.run(rtt.SecureReveal(cipher_result)))

P0、P1 和P2 分别在终端中指定自己的角色后启动脚本，并根据提示输入自己的隐私数据，比如P1可以输入自己的隐私数据为 1~6 的整数：

那么我们可以得到如下的运行结果（这里我们假设P0 输入的也是1~6的整数）：

bravo！从结果中可以看出，系统通过调用 Naive 协议的后端算子来完成 TensorFlow 中相关 API 的计算，使得隐私输入、中间计算结果都是以“密文”的形式均分在P0、P1 手中。而在最后也可以恢复出“明文”计算结果。

其它相关的get_supported_protocols 等 API 此时也可以感知到这个新注册的后端协议：

小结

在本篇文章中，我们介绍了 Rosetta 中是如何通过引入一个中间抽象层组件，来使得后端隐私协议开发完全和上层 AI 框架相解耦的。对于密码学专家等开发者来说，只要参考我们这里介绍的示例协议，在很短的时间内，就可以快速的将自己新设计的安全协议引入到上层 AI 场景应用中来。

本文中介绍的是一个用于协议集成演示的不安全的协议，至于如何真正的集成一个业界前沿的密码学 MPC 协议，并进行面向生产环境落地的高性能改造，以使得用户的隐私数据在整个计算过程中安全的流动，我们会在下一篇文章中具体阐述。stay tuned！

作者介绍：

Rosetta技术团队，一群专注于技术、玩转算法、追求高效的工程师。Rosetta是一款基于主流深度学习框架TensorFlow的隐私AI框架，作为矩阵元公司大规模商业落地的重要引擎，它承载和结合了隐私计算、区块链和AI三种典型技术。目前Rosetta已经在Github开源（https://github.com/LatticeX-Foundation/Rosetta） ，欢迎关注并参与到Rosetta社区中来。

参考资料:

[1] 隐私 AI 框架 PySyft： https://github.com/OpenMined/PySyft

[2] Martin, Robert C. Agile software development: principles, patterns, and practices. Prentice Hall, 2002.

系列文章：

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

面向隐私AI的TensorFlow深度定制化实践