Categories
程式開發

隱私AI框架中MPC協議的快速集成


我們在上一篇文章中,介紹了為了實現隱私AI 系統的易用性,如何對TensorFlow 這樣的深度學習框架進行深度的改造。本篇文章進一步進入TensorFlow 後端算子的內部實現,闡述 羅塞塔 中如何通過定義通用的密碼協議抽象接口層實現系統解耦,使得隱私計算研究者可以輕鬆地將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] 來進行整體的設計。整個的密碼協議統一抽象接口層根據功能職責進一步的劃分為三個不同層次,並分別封裝為 ProtocolManagerProtocolBaseProtocolOps 三個類,其中第一個類ProtocolManager 是一個單例(Singleton)類,是上層API、TensorFlow 的後端算子實現中所需要唯一感知的組件。而ProtocolBaseProtocolOps 則是兩個接口類,由它們定義統一的各個後端具體密碼協議所需要實現的功能。這三個類之間的整體關係如下:
內部基礎類UML圖

細心的讀者應該還記得,我們​​在上一篇文章的最後指出,在TensorFlow 後端算子組SecureOpkernel實現中會最終調用這個模塊:

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

這行語句結合上述UML 類圖,可以很清晰地看出各個組件之間的調用鏈關係:通過協議管理組件入口獲取當前上下文的協議對象,協議對象通過算子入口進一步調用具體某一算子函數。

下面就讓我們分別簡要介紹下這三個核心類。

協議管理器

ProtocolManager 是總的入口,負責整體的控制。其內部為了支持多個可選協議,會維護一個協議名到ProtocolBase指針對象的映射表,並據此進行上下文的控制。除了Instance 方法是常規的取得這個Singleton 的對象實例外,它的功能接口可以分為兩大類,一類是面向上層Python SDK 的,一類是面向開發者進行協議擴展時加以支持的。

  • 上層Python 層的一些協議相關的API,如activatedeactivate 等,會內部調用ProtocolManagerActivateProtocolDeactivateProtocol等方法,來實現對當前使用的協議的控制。而這些類成員函數的內部會進一步的調用ProtocolBase 接口基類的InitUninit 等方法。
  • 而當我們需要引入一個新的協議時,在這一層所需要做的僅僅是調用其RegisterProtocol方法來註冊這個新的協議即可。

協議庫

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 中性InitUbinit 方法。

另一個主要的方法GetOps 則會進一步調用對應協議的ProtrocolOps的子類對象來進一步delegate 具體算子的實現。

以Rosetta 中定制化實現的SecureNN 協議為例,我們是通過SnnProtocol 這個子類來具體實現的。其類繼承關係圖如下:
SnnProtocol類關係圖

協議操作

ProtocolOps 用於封裝各個安全協議中具體所需要實現的算子接口。大部分基礎算子的函數原型中,都以字符串向量作為參數類型,並可以通過一個可選的參數傳入相關屬性信息,比如Add的函數原型是:

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

注意: 我們在前面的文章中介紹過,在Rosetta 內部為了支持多種後端協議中自定義的密文格式,我們統一在外層用字符串來封裝密文數據,所以這裡參數的基礎類型都是字符串。

各個具體的協議需要進一步的實現各個算子函數,比如,在Rosetta 中實現的SecureNN 協議中的各個函數的實現是在子類SnnProtocolOps 中加以實現:
SnnProtocolOps類關係圖

在這些具體的各個函數內部實現中,基本的步驟是先將字符串解碼為此協議內部所設定的類型,然後進一步的進行多方之間安全的邏輯計算(這裡一般是需要進行通信交互的) ,最後再將得到的內部計算結果編碼為字符串輸出到出參中。比如下面是 SnnProtocolOps 中矩陣乘法函數 Matmul 的代碼片段:

int SnnProtocolOps::Matmul(const vector& a,
  const vector& b,
  vector& 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  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 out_vec(m * n);
  vector private_a, private_b;
  snn_decode(a, private_a);
  snn_decode(b, private_b);

  std::make_shared(_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)、域(filed)上進行(同時,最新SecureNN 等MPC 協議又為了充分利用基礎硬件的運算加速,已經支持直接在整數環$$ Z_{2^{64}}$$上進行運算處理),所以轉換到大整數上幾乎是所有相關隱私技術必要的內部操作。

示例:Naive 協議的集成

下面,我們結合一個示例協議Naive 來具體演示下如何快速集成MPC 協議到Rosetta中。

注意:

這個Naive 協議是一個不安全的、僅用於演示的協議!不要naive 的在任何生產環境下使用此協議!

這個Naive 協議是一個不安全的、僅用於演示的協議!不要naive 的在任何生產環境下使用此協議!

這個Naive 協議是一個不安全的、僅用於演示的協議!不要naive 的在任何生產環境下使用此協議!

我們僅在Naive 協議中實現最基本的加法和乘法等操作。在這個“協議”中,P0P1 會將自己的私有輸入值平均分為兩份,一份自己持有,另一份發送給對方,作為各自的“密文”。然後在乘法等後續操作中,基於這樣的語義進行對應的操作處理。這個協議顯然是不安全的。

按照上一小節的介紹,我們只需要少量的修改相關代碼即可實現在Rosetta 中使用這個協議,完整的代碼修改可以參考這裡。具體的,類似於上面介紹的SecureNN 中算子的實現,我們在 NaiveOpsImpl 類中實現內部邏輯的處理。其中實現隱私輸入處理和“密文”下乘法操作的部分代碼片段如下:

int NaiveOpsImpl::PrivateInput(int party_id, const vector& in_x, vector& 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 half_share(vec_size, 0.0);
  for(auto i = 0; i 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& a,
                      const vector& b,
                      vector& 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 中,我們直接在“密文”上計算 P0P1 隱私數據的乘積:

#!/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)))

P0P1P2 分別在終端中指定自己的角色後啟動腳本,並根據提示輸入自己的隱私數據,比如P1可以輸入自己的隱私數據為1~6 的整數:
天真演示運行

那麼我們可以得到如下的運行結果(這裡我們假設P0 輸入的也是1~6的整數):
天真的演示結果

bravo!從結果中可以看出,系統通過調用Naive 協議的後端算子來完成TensorFlow 中相關API 的計算,使得隱私輸入、中間計算結果都是以“密文”的形式均分在P0P1 手中。而在最後也可以恢復出“明文”計算結果。

其它相關的get_supported_protocols 等API 此時也可以感知到這個新註冊的後端協議:

樸素的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] 馬丁·羅伯特·C 敏捷軟件開發:原理,模式和實踐。 普倫蒂斯·霍爾(Prentice Hall),2002年。

系列文章:

隱私AI 工程技術實踐指南:整體介紹

面向隱私AI的TensorFlow深度定制化實踐