基于 TensorFlow 的隱私 AI 框架 Rosetta 正式開源（https://github.com/LatticeX-F...）。借助 Rosetta，AI 開發(fā)者不需要了解密碼學等隱私保護技術(shù)，只需要改動幾行代碼，就能賦予自己的程序以保護數(shù)據(jù)隱私的能力。本專欄將通過多篇獨家技術(shù)稿件，深度披露 Rosetta 的整體框架設(shè)計、TensorFlow 的定制化改造最佳實踐、將密碼學理論算法進行高效工程化落地等內(nèi)容。通過本系列專題，希望能讓更多開發(fā)者了解隱私 AI 框架的技術(shù)挑戰(zhàn)，同時為密碼協(xié)議算法工程化、AI 框架深度定制等相關(guān)方向的開發(fā)者提供一些經(jīng)驗參考。

數(shù)據(jù)是 AI 技術(shù)的“燃料”已經(jīng)成為業(yè)界的共識，更多的數(shù)據(jù)往往意味著可以訓練出更準確的模型。但無論是在公司內(nèi)部還是多個企業(yè)之間，為了對用戶數(shù)據(jù)負責、合法合規(guī)，在數(shù)據(jù)的分享使用時，必須注意對于原始明文數(shù)據(jù)的保護。

傳統(tǒng)的面向靜態(tài)數(shù)據(jù)保護的安全手段無法解決數(shù)據(jù)的動態(tài)使用、分享中的隱私泄露問題，而正是這一實際需求催生出了隱私計算（在 AI 場景下，更進一步的可以稱之為隱私 AI）這一新的交叉技術(shù)。隱私計算技術(shù)融合在數(shù)據(jù)的使用過程之中，保障的是計算過程本身（廣義的講，還包括計算結(jié)果）不會泄露原始明文數(shù)據(jù)本身的信息。

目前實現(xiàn)隱私計算的途徑可以分為密碼學、聯(lián)邦學習和硬件可信執(zhí)行環(huán)境（TEE）等幾大類。而其中以密碼學理論為基礎(chǔ)的 MPC（Multi-Party Computation，安全多方計算）是最有安全保障的技術(shù)路線，其秉持的基本理念是信任計算復雜度理論、代碼，而不是信任人、硬件，而聯(lián)邦學習和 TEE 目前還很難講清楚安全性，經(jīng)常被發(fā)現(xiàn)新的安全漏洞。并且，聯(lián)邦學習中核心部分也往往需要使用同態(tài)加密等密碼學手段進行強安全性的保障。從工程技術(shù)的角度上看，聯(lián)邦學習是分布式機器學習技術(shù)的延伸，主要的挑戰(zhàn)是訓練過程中如何進行多異構(gòu)終端的同步更新等 [1]，很多傳統(tǒng)分布式系統(tǒng)開發(fā)經(jīng)驗仍然適用。而以 MPC 為代表的密碼學途徑則帶來了一些全新的挑戰(zhàn)。

其中最核心的困難是，密碼學屬于計算機理論領(lǐng)域，很多概念、算法協(xié)議都需要有長期的專業(yè)知識積累才能理解，而業(yè)務(wù)落地中的典型 AI 方向，無論是計算機視覺、文本挖掘還是用戶行為建模等都更加面向?qū)嶋H場景。如何打通以密碼學為典型代表的隱私保護技術(shù)與 AI 技術(shù)之間的壁壘？這是開發(fā)者在實際構(gòu)建一個通用的、易用的隱私計算框架時需要解決的核心問題。圍繞著這個核心問題，又有一系列具體的工程技術(shù)挑戰(zhàn)：

如何實現(xiàn)系統(tǒng)的易用性？AI 開發(fā)者不會愿意，也不應(yīng)該為了在業(yè)務(wù)中引入數(shù)據(jù)隱私保護能力而費時耗力地學習各種復雜、抽象的密碼算法。一個好的隱私 AI 框架應(yīng)該是易上手的，便于 AI 開發(fā)者使用自己熟悉的方式快速解決自己的數(shù)據(jù)隱私問題。

如何實現(xiàn)系統(tǒng)的高效執(zhí)行？這包括單機和多機兩個層面。密碼學的計算大部分都是在大隨機數(shù)的密文上進行，為此常常需要使用專用的硬件指令、SIMD（Single Instruction/Multiple Data）等技術(shù)來進行單機并行化的加速，這些優(yōu)化實現(xiàn)需要對于密碼學基礎(chǔ)庫有深入的了解，并往往需要根據(jù)協(xié)議算法做進一步的并行優(yōu)化。而在多機層面上，則需要考慮如何和很多 AI 框架自身的并行優(yōu)化技術(shù)兼容。

如何實現(xiàn) MPC 多方之間的高效通信？在 MPC 中，多方之間需要進行大量的同步通信，而且信道上的內(nèi)容大都是無規(guī)律的、不可壓縮的一次性使用的隨機數(shù)，這就需要在保證安全性的同時，根據(jù)具體的計算邏輯進行很多工程優(yōu)化以減少通信量和通信次數(shù)。

如何保障隱私保護技術(shù)的可擴展性？MPC 等隱私計算技術(shù)還在不斷發(fā)展之中，也是學術(shù)研究上的熱點問題，所以一個好的隱私 AI 框架，需要能夠支持研究者簡單快速地將新的算法協(xié)議集成進來。

針對這些問題，業(yè)界已經(jīng)有一些探索，下面我們結(jié)合 Rosetta 來具體談一談在隱私 AI 框架的設(shè)計和實現(xiàn)中如何克服這些挑戰(zhàn)。篇幅所限，本文主要整體性地先介紹宏觀設(shè)計，后續(xù)系列文章會進一步剖析一些技術(shù)細節(jié)。

如同其他的隱私 AI 框架一樣，Rosetta 仍然處于發(fā)展的早期，尚有一些不完善之處。我們在此以 Rosetta 為例是希望能具體化地講清楚這一領(lǐng)域中的細節(jié)挑戰(zhàn)，也希望激發(fā)更多開發(fā)者參與到未來的隱私 AI 系統(tǒng)設(shè)計中來。

隱私 AI 框架整體設(shè)計思路

目前業(yè)界尚沒有大規(guī)模落地的、成熟完善的隱私 AI 框架，但是已經(jīng)有一些探索性的開源隱私 AI 框架，比如 PySyft、TF Encrypted 和 CrypTen。

從整體上看，這些框架都是在 TensorFlow 或 PyTorch 的前端 Python 層進行封裝集成的。這樣做的好處是可以直接使用這些 AI 框架的上層接口來實現(xiàn)隱私計算算法，而且天然的可以直接調(diào)用框架自身封裝好的一些高層次 API 高功能。這對于聯(lián)邦學習這種本身就源自分布式機器學習的技術(shù)來說是較為適合的，但是對于密碼學的 MPC 來說會有一些不足：

首先，單機性能無法得到充分的提升。用 Python 來實現(xiàn)各種復雜的密碼學計算、多方之間的通信無法充分利用底層操作系統(tǒng)、硬件層的并行優(yōu)化。而且更現(xiàn)實的是，大部分高性能的密碼學庫提供的是 C/C++ 的接口，如果在 AI 框架的前端上層實現(xiàn) MPC 等密碼技術(shù)，則很難復用這些業(yè)界長時間積累（同時也還在不斷發(fā)展中）下來的成果。

其次，密碼協(xié)議等隱私技術(shù)的實現(xiàn)和 AI 框架本身耦合過深，不利于擴展。由于這些 AI 框架本身提供的對外 API 接口本就是面向 AI 需求的，在實現(xiàn)較為復雜的 MPC 等密碼協(xié)議時不僅需要熟練使用這些框架的 API，還往往需要大量直接使用numpy等庫來實現(xiàn)復雜的計算邏輯。這一方面破壞了對 AI 框架自身使用上的自洽，不再能將全部的計算邏輯完全承載在 AI 框架的邏輯執(zhí)行圖上，另一方面也使得每一次引入新的后端密碼協(xié)議時都需要重新基于 AI 框架進行實現(xiàn)，這對于密碼協(xié)議開發(fā)者來說成本很高。

基于上述認識，Rosetta 在現(xiàn)階段首先以 TensorFlow 這一流行 AI 框架為基礎(chǔ)，深度改造其前端 Python 入口和后端 kernel 實現(xiàn)，并封裝可插拔的 MPC 算法協(xié)議作為“隱私保護引擎”來驅(qū)動整個計算過程中數(shù)據(jù)的安全流動。

Why TensorFlow？

TensorFlow 和 PyTorch 是目前工業(yè)界最主流使用的開源 AI 框架。雖然很多公司在內(nèi)部可能也會根據(jù)自身需求定制化改造一些組件，或者推出各具特色的新框架以在易用性、高效性、完備性等不同的維度上進行進一步突破，但是整體上看，這些框架基本的設(shè)計范式是較為相似的。大多是通過豐富的接口 API 讓用戶以有向無環(huán)圖 DAG 的形式表達上層計算邏輯，而框架本身則會在實際調(diào)度執(zhí)行這些計算任務(wù)時進行一系列的優(yōu)化。TensorFlow 雖然在用戶友好性上略遜色于 PyTorch，常受開發(fā)者詬病，但是其在可擴展性、高效性上、分布式部署等方面確實是更加均衡、全面（當然這也意味著 TensorFlow 更加復雜，對其改造會更加的具有挑戰(zhàn)性）。所以綜合考慮下來，Rosetta 在當前版本中選擇 TensorFlow 作為基本的底層承載體，在設(shè)計開發(fā)的過程中，一方面會充分利用 TensorFlow 內(nèi)在的計算圖并行執(zhí)行優(yōu)化等功能以提升效率，另一方面也會盡量克制，主要是利用其作為深度學習框架通用性的一些接口特性，而不會過于依賴其獨有的一些組件。

Rosetta 框架核心設(shè)計思想

隱私算子（SecureOp）作為核心抽象接口連接 AI 框架和隱私計算技術(shù)。TensorFlow 在不同的層次上提供了多樣的擴展方式，Rosetta 選擇后端算子（Operation）層作為核心切入點，TensorFlow 在執(zhí)行算子時會動態(tài)的綁定到具體 MPC 協(xié)議中的 SecureOp 實現(xiàn)中。通過這樣的抽象，密碼協(xié)議開發(fā)者可以不需要了解 AI 框架，只需要用 C++ 實現(xiàn)滿足接口定義的各個功能函數(shù)即可，而 AI 開發(fā)者也不需要深入了解 MPC 等技術(shù)的實現(xiàn)細節(jié)，而只需要在現(xiàn)有算子的基礎(chǔ)上進一步封裝自己想要的上層高級功能即可。

基于優(yōu)化遍（Pass）的分階段轉(zhuǎn)換。為了盡可能給 AI 開發(fā)者提供易用的接口，減少給線上 AI 程序賦予數(shù)據(jù)隱私保護能力時的改造成本，Rosetta 在整體的設(shè)計中借鑒了程序編譯器領(lǐng)域的核心概念：Pass。Pass 是編譯器中常用的技術(shù)，主要用作將源碼文件一步步轉(zhuǎn)變?yōu)闄C器碼過程中的多輪轉(zhuǎn)化和優(yōu)化。在 Rosetta 中，用戶使用原生 TensorFlow 接口編寫的 DAG（有向無環(huán)圖）形式的邏輯計算圖會被分階段的轉(zhuǎn)換、替換為多方協(xié)作執(zhí)行的 MPC 程序，這樣可以實現(xiàn)對于用戶 API 層最少的改動。

具體的，在 Rosetta 中，有兩個階段的 Pass，一個在前端 Python 層的全局 DAG 構(gòu)建過程中生效的 Static Pass，會將原生Tensor轉(zhuǎn)換為支持自定義密文類型的RttTensor，將原生Operation轉(zhuǎn)換為支持tf.string格式輸入輸出的RttOp，并最終在圖開始啟動時進一步的轉(zhuǎn)換為承載實際 MPC 操作的SecureOp。

Static Pass

另一個是在SecurOp的實際執(zhí)行時所進行的 Dynamic Pass 處理，會動態(tài)的根據(jù)當前用戶選擇的協(xié)議選擇對應(yīng)的實際算子實現(xiàn)去執(zhí)行，同時可以在此時嵌入基于執(zhí)行上下文的優(yōu)化處理。

Dynamic Pass

融合 MPC 技術(shù)的分布式隱私 AI 架構(gòu)

理清楚整體的分布式結(jié)構(gòu)對于了解一個系統(tǒng)的架構(gòu)大有裨益。整個隱私 AI 系統(tǒng)對外接口會涉及三個方面，如何指定物理部署上的網(wǎng)絡(luò)拓撲？數(shù)據(jù)在整個計算的過程中是如何安全輸入、流動、輸出的？隱私計算邏輯要如何表達？

多方網(wǎng)絡(luò)的建立

MPC 技術(shù)本身就是要求多方（multi-party）參與的，一般稱他們?yōu)?Player，不同的 MPC 算法協(xié)議會有不同個數(shù)的參與方。以 Rosetta 中目前實現(xiàn)的三方協(xié)議 SecureNN[2] 為例，系統(tǒng)中存在三個邏輯參與方，P0、P1 和 P2。

在 v0.2.1 版本中，在這一方面的用戶接口層次上，為了保障對外的靈活性，目前支持用戶通過配置文件來一次性指定多機之間的網(wǎng)絡(luò)關(guān)系，也支持調(diào)用接口動態(tài)的激活、解除多方之間的網(wǎng)絡(luò)拓撲：

# 調(diào)用 activate 接口會根據(jù)配置參數(shù)或配置文件建立起網(wǎng)絡(luò) rtt.activate(protocol_nam, protocol_config_str=None) # 調(diào)用 deactivate 接口會釋放網(wǎng)絡(luò)鏈接等資源 rtt.deactivate()

在內(nèi)部實現(xiàn)中，每一個參與方都會監(jiān)聽一個本地的 server 端口，而同時分別建立到另外兩方之間的 client 網(wǎng)絡(luò)鏈接。這樣的好處是相互之間的網(wǎng)絡(luò)鏈接關(guān)系簡單清晰，當然也需要解決隨之而來的SecureOp并發(fā)同步執(zhí)行時的一致性問題，這一點我們會在后續(xù)文章中討論。

一些注意點

熟悉 TensorFlow 的讀者可能會疑惑，這種多方基于不同數(shù)據(jù)跑相同程序的模式，不就是 TensorFlow 分布式執(zhí)行下對數(shù)據(jù)并行進行支持的 In-graph replication 和 Between-graph replication 嗎？并不是這樣，實際上它們是不同層次的結(jié)構(gòu)，這里講的是上層邏輯視角的 MPC 各參與方，在實際中，你甚至可以進一步的將各方內(nèi)部執(zhí)行的這一 task 按照 TensorFlow 的分布式規(guī)范進行集群部署，而將集群中的 "server“作為統(tǒng)一的對外邏輯代表。

上面一直講的是“邏輯上”的三方，那么在實際的業(yè)務(wù)場景中，可能是 2 個、4 個或以上公司之間的數(shù)據(jù)合作，是不是就不能用這些架構(gòu)了嗎？其實不然，我們完全可以在上層進行一層映射，以 Privacy-as-a-Service 的形式提供對上層的服務(wù)，關(guān)于這一點，后續(xù)文章也會進一步介紹。

隱私數(shù)據(jù)的流動

每一個邏輯參與方都可以有自己私有的明文輸入數(shù)據(jù)，也可以繼續(xù)處理上一次任務(wù)輸出的密文結(jié)果。在整個程序的運行過程中，只有開始和結(jié)束時數(shù)據(jù)才會以明文的形式存在：開始時對于私有數(shù)據(jù)的引入，以及最后可配置是否將結(jié)果以明文的形式恢復出來加以輸出。而在中間各個算子的計算過程中，數(shù)據(jù)總是以密文形式在本地的邏輯上下文、多方之間進行交互。

對外接口方面，在實際的業(yè)務(wù)中多方數(shù)據(jù)之間是需要關(guān)聯(lián)對齊的，目前 Rosetta 提供常見的兩種數(shù)據(jù)集處理方法，一是對應(yīng)于整體上數(shù)據(jù)集是在各方之間“水平劃分”的場景，即各方擁有不同樣本 id 的全部特征屬性值；另一種對應(yīng)于整體數(shù)據(jù)集是在各方之間“垂直劃分”的場景，即各方之間擁有相同樣本 id 的部分特征屬性值。這些都可以調(diào)用PrivateDataset類的load_data等接口方便的處理。而在輸出階段，提供了如下兩個接口：

# 將一個密文形式的 cipher_tensor 恢復為明文， receive_party 參數(shù)指定 3 方中哪幾方可以獲得明文結(jié)果 rtt.SecureReveal(cipher_tensor, receive_party=0b111) # 與原生 TensorFlow 中模型保存接口 SaveV2 具有一樣的函數(shù)原型，可通過配置文件指定哪幾方可以獲得明文模型文件 rtt.SecureSaveV2(prefix, tensor_names, shape_and_slices, tensors, name=None)

隱私集合求交 PSI（Private Set Intersection）技術(shù)

在實際場景中還存在一個很現(xiàn)實的問題，就是多方之間樣本的對齊問題，比如將 A 方的樣本 id 所指向的樣本和 B 方此樣本 id 對應(yīng)的屬性信息給對應(yīng)起來。PSI 技術(shù)可以安全的解決上述問題，目前各個開源框架中還沒有將這一功能很好的集成進來，Rosetta 目前正在集成這一功能，將在近期版本中發(fā)布。

在內(nèi)部實現(xiàn)中，密碼學中的很多運算是在空間較大的環(huán)（Ring）、域（Field）等抽象代數(shù)結(jié)構(gòu)上的操作，而具體的在代碼中則落地到對于大整數(shù)、多項式等數(shù)據(jù)結(jié)構(gòu)上的處理，所以框架設(shè)計上要在三個方面達成平衡：

1、盡可能的實現(xiàn)用戶對于內(nèi)部密文數(shù)據(jù)結(jié)構(gòu)的透明無感；

2、TensorFlow 的 DAG 圖構(gòu)建、自動求導等核心功能仍需要無縫的支持；

3、支持不同的 MPC 協(xié)議使用自定義的具體數(shù)據(jù)結(jié)構(gòu)對象以便于擴展。

為同時實現(xiàn)這些目標，Rosetta 基于tf.string這一 TensorFlow 原生數(shù)據(jù)結(jié)構(gòu)來承載各協(xié)議自定義的密文數(shù)據(jù)，然后通過對 TensorFlow 源碼代碼進行深度的 hook 改造使得 DAG 圖構(gòu)建、自動求導等功能仍然可用。

DAG 的執(zhí)行

如上述網(wǎng)絡(luò)結(jié)構(gòu)圖所展示的那樣，各 Player 運行的是同一份基于 TensorFlow 編寫的 AI 二進制代碼，比如訓練一個簡單神經(jīng)網(wǎng)絡(luò)模型的程序。用戶直接的使用 TensorFlow 中原生的算子 API 來構(gòu)建邏輯計算圖 DAG，Rosetta 內(nèi)部會在圖開始執(zhí)行時完成到隱私算子SecureOp的轉(zhuǎn)換。與其他隱私計算框架相比，這樣的切換成本是最低的。

在執(zhí)行過程中，各 Player 自身都是在按照這個 DAG 圖在運行，特殊的地方在于在各個算子的內(nèi)部執(zhí)行過程中，各個 Player 會根據(jù)自己所屬的角色遵循 MPC 協(xié)議執(zhí)行不同的操作，這些操作即包括本地的在密文上的處理，也包括在多方之間進行強同步的通訊交互，傳輸大量隨機數(shù)形式的密文。

小 結(jié)

在本篇文章里，我們結(jié)合 Rosetta 框架，從整體上介紹了隱私 AI 框架在工程落地時所需要面對的挑戰(zhàn)，以及 Rosetta 等框架的一些設(shè)計方案。后續(xù)的文章中，我們會就核心的關(guān)鍵模塊進行進一步的介紹。

文章來源：AI前線 ，作者 Rosetta 技術(shù)團隊

審核編輯 黃昊宇