近期，基于TensorFlow的隱私AI框架Rosetta正式開源（https://github.com/LatticeX-Foundation/Rosetta）。借助Rosetta，AI開發(fā)者不需要了解密碼學(xué)等隱私保護(hù)技術(shù)，只需要改動(dòng)幾行代碼，就能賦予自己的程序以保護(hù)數(shù)據(jù)隱私的能力。本專欄將通過(guò)多篇獨(dú)家技術(shù)稿件，深度披露Rosetta的整體框架設(shè)計(jì)、TensorFlow的定制化改造最佳實(shí)踐、將密碼學(xué)理論算法進(jìn)行高效工程化落地等內(nèi)容。通過(guò)本系列專題，希望能讓更多開發(fā)者了解隱私AI框架的技術(shù)挑戰(zhàn)，同時(shí)為密碼協(xié)議算法工程化、AI框架深度定制等相關(guān)方向的開發(fā)者提供一些經(jīng)驗(yàn)參考。
數(shù)據(jù)是AI技術(shù)的“燃料”已經(jīng)成為業(yè)界的共識(shí)，更多的數(shù)據(jù)往往意味著可以訓(xùn)練出更準(zhǔn)確的模型。但無(wú)論是在公司內(nèi)部還是多個(gè)企業(yè)之間，為了對(duì)用戶數(shù)據(jù)負(fù)責(zé)、合法合規(guī)，在數(shù)據(jù)的分享使用時(shí)，必須注意對(duì)于原始明文數(shù)據(jù)的保護(hù)。
傳統(tǒng)的面向靜態(tài)數(shù)據(jù)保護(hù)的安全手段無(wú)法解決數(shù)據(jù)的動(dòng)態(tài)使用、分享中的隱私泄露問(wèn)題，而正是這一實(shí)際需求催生出了隱私計(jì)算（在AI場(chǎng)景下，更進(jìn)一步的可以稱之為隱私AI）這一新的交叉技術(shù)。隱私計(jì)算技術(shù)融合在數(shù)據(jù)的使用過(guò)程之中，保障的是計(jì)算過(guò)程本身（廣義的講，還包括計(jì)算結(jié)果）不會(huì)泄露原始明文數(shù)據(jù)本身的信息。
目前實(shí)現(xiàn)隱私計(jì)算的途徑可以分為密碼學(xué)、聯(lián)邦學(xué)習(xí)和硬件可信執(zhí)行環(huán)境（TEE）等幾大類。而其中以密碼學(xué)理論為基礎(chǔ)的MPC（Multi-Party Computation，安全多方計(jì)算）是最有安全保障的技術(shù)路線，其秉持的基本理念是信任計(jì)算復(fù)雜度理論、代碼，而不是信任人、硬件，而聯(lián)邦學(xué)習(xí)和TEE目前還很難講清楚安全性，經(jīng)常被發(fā)現(xiàn)新的安全漏洞。并且，聯(lián)邦學(xué)習(xí)中核心部分也往往需要使用同態(tài)加密等密碼學(xué)手段進(jìn)行強(qiáng)安全性的保障。從工程技術(shù)的角度上看，聯(lián)邦學(xué)習(xí)是分布式機(jī)器學(xué)習(xí)技術(shù)的延伸，主要的挑戰(zhàn)是訓(xùn)練過(guò)程中如何進(jìn)行多異構(gòu)終端的同步更新等[1]，很多傳統(tǒng)分布式系統(tǒng)開發(fā)經(jīng)驗(yàn)仍然適用。而以MPC為代表的密碼學(xué)途徑則帶來(lái)了一些全新的挑戰(zhàn)。
安全多方計(jì)算MPC（圖片來(lái)自網(wǎng)絡(luò)）
其中最核心的困難是，密碼學(xué)屬于計(jì)算機(jī)理論領(lǐng)域，很多概念、算法協(xié)議都需要有長(zhǎng)期的專業(yè)知識(shí)積累才能理解，而業(yè)務(wù)落地中的典型AI方向，無(wú)論是計(jì)算機(jī)視覺(jué)、文本挖掘還是用戶行為建模等都更加面向?qū)嶋H場(chǎng)景。如何打通以密碼學(xué)為典型代表的隱私保護(hù)技術(shù)與AI技術(shù)之間的壁壘？這是開發(fā)者在實(shí)際構(gòu)建一個(gè)通用的、易用的隱私計(jì)算框架時(shí)需要解決的核心問(wèn)題。圍繞著這個(gè)核心問(wèn)題，又有一系列具體的工程技術(shù)挑戰(zhàn)：
如何實(shí)現(xiàn)系統(tǒng)的易用性？AI開發(fā)者不會(huì)愿意，也不應(yīng)該為了在業(yè)務(wù)中引入數(shù)據(jù)隱私保護(hù)能力而費(fèi)時(shí)耗力地學(xué)習(xí)各種復(fù)雜、抽象的密碼算法。一個(gè)好的隱私AI框架應(yīng)該是易上手的，便于AI開發(fā)者使用自己熟悉的方式快速解決自己的數(shù)據(jù)隱私問(wèn)題。
如何實(shí)現(xiàn)系統(tǒng)的高效執(zhí)行？這包括單機(jī)和多機(jī)兩個(gè)層面。密碼學(xué)的計(jì)算大部分都是在大隨機(jī)數(shù)的密文上進(jìn)行，為此常常需要使用專用的硬件指令、SIMD（Single Instruction/Multiple Data）等技術(shù)來(lái)進(jìn)行單機(jī)并行化的加速，這些優(yōu)化實(shí)現(xiàn)需要對(duì)于密碼學(xué)基礎(chǔ)庫(kù)有深入的了解，并往往需要根據(jù)協(xié)議算法做進(jìn)一步的并行優(yōu)化。而在多機(jī)層面上，則需要考慮如何和很多AI框架自身的并行優(yōu)化技術(shù)兼容。
如何實(shí)現(xiàn)MPC多方之間的高效通信？在MPC中，多方之間需要進(jìn)行大量的同步通信，而且信道上的內(nèi)容大都是無(wú)規(guī)律的、不可壓縮的一次性使用的隨機(jī)數(shù)，這就需要在保證安全性的同時(shí)，根據(jù)具體的計(jì)算邏輯進(jìn)行很多工程優(yōu)化以減少通信量和通信次數(shù)。
如何保障隱私保護(hù)技術(shù)的可擴(kuò)展性？MPC等隱私計(jì)算技術(shù)還在不斷發(fā)展之中，也是學(xué)術(shù)研究上的熱點(diǎn)問(wèn)題，所以一個(gè)好的隱私AI框架，需要能夠支持研究者簡(jiǎn)單快速地將新的算法協(xié)議集成進(jìn)來(lái)。
針對(duì)這些問(wèn)題，業(yè)界已經(jīng)有一些探索，下面我們結(jié)合Rosetta來(lái)具體談一談在隱私AI框架的設(shè)計(jì)和實(shí)現(xiàn)中如何克服這些挑戰(zhàn)。篇幅所限，本文主要整體性地先介紹宏觀設(shè)計(jì)，后續(xù)系列文章會(huì)進(jìn)一步剖析一些技術(shù)細(xì)節(jié)。
如同其他的隱私AI框架一樣，Rosetta仍然處于發(fā)展的早期，尚有一些不完善之處。我們?cè)诖艘訰osetta為例是希望能具體化地講清楚這一領(lǐng)域中的細(xì)節(jié)挑戰(zhàn)，也希望激發(fā)更多開發(fā)者參與到未來(lái)的隱私AI系統(tǒng)設(shè)計(jì)中來(lái)。
隱私AI框架整體設(shè)計(jì)思路
目前業(yè)界尚沒(méi)有大規(guī)模落地的、成熟完善的隱私AI框架，但是已經(jīng)有一些探索性的開源隱私AI框架，比如PySyft、TF Encrypted和CrypTen。
從整體上看，這些框架都是在TensorFlow或PyTorch的前端Python層進(jìn)行封裝集成的。這樣做的好處是可以直接使用這些AI框架的上層接口來(lái)實(shí)現(xiàn)隱私計(jì)算算法，而且天然的可以直接調(diào)用框架自身封裝好的一些高層次API高功能。這對(duì)于聯(lián)邦學(xué)習(xí)這種本身就源自分布式機(jī)器學(xué)習(xí)的技術(shù)來(lái)說(shuō)是較為適合的，但是對(duì)于密碼學(xué)的MPC來(lái)說(shuō)會(huì)有一些不足：
首先，單機(jī)性能無(wú)法得到充分的提升。用Python來(lái)實(shí)現(xiàn)各種復(fù)雜的密碼學(xué)計(jì)算、多方之間的通信無(wú)法充分利用底層操作系統(tǒng)、硬件層的并行優(yōu)化。而且更現(xiàn)實(shí)的是，大部分高性能的密碼學(xué)庫(kù)提供的是C/C++的接口，如果在AI框架的前端上層實(shí)現(xiàn)MPC等密碼技術(shù)，則很難復(fù)用這些業(yè)界長(zhǎng)時(shí)間積累（同時(shí)也還在不斷發(fā)展中）下來(lái)的成果。
其次，密碼協(xié)議等隱私技術(shù)的實(shí)現(xiàn)和AI框架本身耦合過(guò)深，不利于擴(kuò)展。由于這些AI框架本身提供的對(duì)外API接口本就是面向AI需求的，在實(shí)現(xiàn)較為復(fù)雜的MPC等密碼協(xié)議時(shí)不僅需要熟練使用這些框架的API，還往往需要大量直接使用numpy等庫(kù)來(lái)實(shí)現(xiàn)復(fù)雜的計(jì)算邏輯。這一方面破壞了對(duì)AI框架自身使用上的自洽，不再能將全部的計(jì)算邏輯完全承載在AI框架的邏輯執(zhí)行圖上，另一方面也使得每一次引入新的后端密碼協(xié)議時(shí)都需要重新基于AI框架進(jìn)行實(shí)現(xiàn)，這對(duì)于密碼協(xié)議開發(fā)者來(lái)說(shuō)成本很高。
基于上述認(rèn)識(shí)，Rosetta在現(xiàn)階段首先以TensorFlow這一流行AI框架為基礎(chǔ)，深度改造其前端Python入口和后端kernel實(shí)現(xiàn)，并封裝可插拔的MPC算法協(xié)議作為“隱私保護(hù)引擎”來(lái)驅(qū)動(dòng)整個(gè)計(jì)算過(guò)程中數(shù)據(jù)的安全流動(dòng)。
Why TensorFlow？
TensorFlow和PyTorch是目前工業(yè)界最主流使用的開源AI框架。雖然很多公司在內(nèi)部可能也會(huì)根據(jù)自身需求定制化改造一些組件，或者推出各具特色的新框架以在易用性、高效性、完備性等不同的維度上進(jìn)行進(jìn)一步突破，但是整體上看，這些框架基本的設(shè)計(jì)范式是較為相似的。大多是通過(guò)豐富的接口API讓用戶以有向無(wú)環(huán)圖DAG的形式表達(dá)上層計(jì)算邏輯，而框架本身則會(huì)在實(shí)際調(diào)度執(zhí)行這些計(jì)算任務(wù)時(shí)進(jìn)行一系列的優(yōu)化。TensorFlow雖然在用戶友好性上略遜色于PyTorch，常受開發(fā)者詬病，但是其在可擴(kuò)展性、高效性上、分布式部署等方面確實(shí)是更加均衡、全面（當(dāng)然這也意味著TensorFlow更加復(fù)雜，對(duì)其改造會(huì)更加的具有挑戰(zhàn)性）。所以綜合考慮下來(lái)，Rosetta在當(dāng)前版本中選擇TensorFlow作為基本的底層承載體，在設(shè)計(jì)開發(fā)的過(guò)程中，一方面會(huì)充分利用TensorFlow內(nèi)在的計(jì)算圖并行執(zhí)行優(yōu)化等功能以提升效率，另一方面也會(huì)盡量克制，主要是利用其作為深度學(xué)習(xí)框架通用性的一些接口特性，而不會(huì)過(guò)于依賴其獨(dú)有的一些組件。
Rosetta框架核心設(shè)計(jì)思想
隱私算子（SecureOp）作為核心抽象接口連接AI框架和隱私計(jì)算技術(shù)。TensorFlow在不同的層次上提供了多樣的擴(kuò)展方式，Rosetta選擇后端算子（Operation）層作為核心切入點(diǎn)，TensorFlow在執(zhí)行算子時(shí)會(huì)動(dòng)態(tài)的綁定到具體MPC協(xié)議中的SecureOp實(shí)現(xiàn)中。通過(guò)這樣的抽象，密碼協(xié)議開發(fā)者可以不需要了解AI框架，只需要用C++實(shí)現(xiàn)滿足接口定義的各個(gè)功能函數(shù)即可，而AI開發(fā)者也不需要深入了解MPC等技術(shù)的實(shí)現(xiàn)細(xì)節(jié)，而只需要在現(xiàn)有算子的基礎(chǔ)上進(jìn)一步封裝自己想要的上層高級(jí)功能即可。
基于優(yōu)化遍（Pass）的分階段轉(zhuǎn)換。為了盡可能給AI開發(fā)者提供易用的接口，減少給線上AI程序賦予數(shù)據(jù)隱私保護(hù)能力時(shí)的改造成本，Rosetta在整體的設(shè)計(jì)中借鑒了程序編譯器領(lǐng)域的核心概念：Pass。Pass是編譯器中常用的技術(shù)，主要用作將源碼文件一步步轉(zhuǎn)變?yōu)闄C(jī)器碼過(guò)程中的多輪轉(zhuǎn)化和優(yōu)化。在Rosetta中，用戶使用原生TensorFlow接口編寫的DAG（有向無(wú)環(huán)圖）形式的邏輯計(jì)算圖會(huì)被分階段的轉(zhuǎn)換、替換為多方協(xié)作執(zhí)行的MPC程序，這樣可以實(shí)現(xiàn)對(duì)于用戶API層最少的改動(dòng)。
具體的，在Rosetta中，有兩個(gè)階段的Pass，一個(gè)在前端Python層的全局DAG構(gòu)建過(guò)程中生效的Static Pass，會(huì)將原生Tensor轉(zhuǎn)換為支持自定義密文類型的RttTensor，將原生Operation轉(zhuǎn)換為支持tf.string格式輸入輸出的RttOp，并最終在圖開始啟動(dòng)時(shí)進(jìn)一步的轉(zhuǎn)換為承載實(shí)際MPC操作的SecureOp。
Static Pass
另一個(gè)是在SecurOp的實(shí)際執(zhí)行時(shí)所進(jìn)行的Dynamic Pass處理，會(huì)動(dòng)態(tài)的根據(jù)當(dāng)前用戶選擇的協(xié)議選擇對(duì)應(yīng)的實(shí)際算子實(shí)現(xiàn)去執(zhí)行，同時(shí)可以在此時(shí)嵌入基于執(zhí)行上下文的優(yōu)化處理。
Dynamic Pass
融合MPC技術(shù)的分布式隱私AI架構(gòu)
理清楚整體的分布式結(jié)構(gòu)對(duì)于了解一個(gè)系統(tǒng)的架構(gòu)大有裨益。整個(gè)隱私AI系統(tǒng)對(duì)外接口會(huì)涉及三個(gè)方面，如何指定物理部署上的網(wǎng)絡(luò)拓?fù)洌繑?shù)據(jù)在整個(gè)計(jì)算的過(guò)程中是如何安全輸入、流動(dòng)、輸出的？隱私計(jì)算邏輯要如何表達(dá)？Rosetta的整體邏輯結(jié)構(gòu)如下圖所示：
Rosetta多方網(wǎng)絡(luò)結(jié)構(gòu)圖
多方網(wǎng)絡(luò)的建立
MPC技術(shù)本身就是要求多方（multi-party）參與的，一般稱他們?yōu)镻layer，不同的MPC算法協(xié)議會(huì)有不同個(gè)數(shù)的參與方。以Rosetta中目前實(shí)現(xiàn)的三方協(xié)議SecureNN[2]為例，系統(tǒng)中存在三個(gè)邏輯參與方，P0、P1和P2。
在v0.2.1版本中，在這一方面的用戶接口層次上，為了保障對(duì)外的靈活性，目前支持用戶通過(guò)配置文件來(lái)一次性指定多機(jī)之間的網(wǎng)絡(luò)關(guān)系，也支持調(diào)用接口動(dòng)態(tài)的激活、解除多方之間的網(wǎng)絡(luò)拓?fù)洌?/p>
#調(diào)用activate接口會(huì)根據(jù)配置參數(shù)或配置文件建立起網(wǎng)絡(luò)
rtt.activate(protocol_name="SecureNN",protocol_config_str=None)
#調(diào)用deactivate接口會(huì)釋放網(wǎng)絡(luò)鏈接等資源
rtt.deactivate()
在內(nèi)部實(shí)現(xiàn)中，每一個(gè)參與方都會(huì)監(jiān)聽(tīng)一個(gè)本地的server端口，而同時(shí)分別建立到另外兩方之間的client網(wǎng)絡(luò)鏈接。這樣的好處是相互之間的網(wǎng)絡(luò)鏈接關(guān)系簡(jiǎn)單清晰，當(dāng)然也需要解決隨之而來(lái)的SecureOp并發(fā)同步執(zhí)行時(shí)的一致性問(wèn)題，這一點(diǎn)我們會(huì)在后續(xù)文章中討論。
一些注意點(diǎn)
*熟悉TensorFlow的讀者可能會(huì)疑惑，這種多方基于不同數(shù)據(jù)跑相同程序的模式，不就是TensorFlow分布式執(zhí)行下對(duì)數(shù)據(jù)并行進(jìn)行支持的In-graph replication和Between-graph replication嗎？并不是這樣，實(shí)際上它們是不同層次的結(jié)構(gòu)，這里講的是上層邏輯視角的MPC各參與方，在實(shí)際中，你甚至可以進(jìn)一步的將各方內(nèi)部執(zhí)行的這一task按照TensorFlow的分布式規(guī)范進(jìn)行集群部署，而將集群中的"server“作為統(tǒng)一的對(duì)外邏輯代表。
*上面一直講的是“邏輯上”的三方，那么在實(shí)際的業(yè)務(wù)場(chǎng)景中，可能是2個(gè)、4個(gè)或以上公司之間的數(shù)據(jù)合作，是不是就不能用這些架構(gòu)了嗎？其實(shí)不然，我們完全可以在上層進(jìn)行一層映射，以Privacy-as-a-Service的形式提供對(duì)上層的服務(wù)，關(guān)于這一點(diǎn)，后續(xù)文章也會(huì)進(jìn)一步介紹。
隱私數(shù)據(jù)的流動(dòng)
每一個(gè)邏輯參與方都可以有自己私有的明文輸入數(shù)據(jù)，也可以繼續(xù)處理上一次任務(wù)輸出的密文結(jié)果。在整個(gè)程序的運(yùn)行過(guò)程中，只有開始和結(jié)束時(shí)數(shù)據(jù)才會(huì)以明文的形式存在：開始時(shí)對(duì)于私有數(shù)據(jù)的引入，以及最后可配置是否將結(jié)果以明文的形式恢復(fù)出來(lái)加以輸出。而在中間各個(gè)算子的計(jì)算過(guò)程中，數(shù)據(jù)總是以密文形式在本地的邏輯上下文、多方之間進(jìn)行交互。
對(duì)外接口方面，在實(shí)際的業(yè)務(wù)中多方數(shù)據(jù)之間是需要關(guān)聯(lián)對(duì)齊的，目前Rosetta提供常見(jiàn)的兩種數(shù)據(jù)集處理方法，一是對(duì)應(yīng)于整體上數(shù)據(jù)集是在各方之間“水平劃分”的場(chǎng)景，即各方擁有不同樣本id的全部特征屬性值；另一種對(duì)應(yīng)于整體數(shù)據(jù)集是在各方之間“垂直劃分”的場(chǎng)景，即各方之間擁有相同樣本id的部分特征屬性值。這些都可以調(diào)用PrivateDataset類的load_data等接口方便的處理。而在輸出階段，提供了如下兩個(gè)接口：
#將一個(gè)密文形式的cipher_tensor恢復(fù)為明文，receive_party參數(shù)指定3方中哪幾方可以獲得明文結(jié)果
rtt.SecureReveal(cipher_tensor,receive_party=0b111)
#與原生TensorFlow中模型保存接口SaveV2具有一樣的函數(shù)原型，可通過(guò)配置文件指定哪幾方可以獲得明文模型文件
rtt.SecureSaveV2(prefix,tensor_names,shape_and_slices,tensors,name=None)
隱私集合求交PSI（Private Set Intersection）技術(shù)
在實(shí)際場(chǎng)景中還存在一個(gè)很現(xiàn)實(shí)的問(wèn)題，就是多方之間樣本的對(duì)齊問(wèn)題，比如將A方的樣本id所指向的樣本和B方此樣本id對(duì)應(yīng)的屬性信息給對(duì)應(yīng)起來(lái)。PSI技術(shù)可以安全的解決上述問(wèn)題，目前各個(gè)開源框架中還沒(méi)有將這一功能很好的集成進(jìn)來(lái)，Rosetta目前正在集成這一功能，將在近期版本中發(fā)布。
在內(nèi)部實(shí)現(xiàn)中，密碼學(xué)中的很多運(yùn)算是在空間較大的環(huán)（Ring）、域（Field）等抽象代數(shù)結(jié)構(gòu)上的操作，而具體的在代碼中則落地到對(duì)于大整數(shù)、多項(xiàng)式等數(shù)據(jù)結(jié)構(gòu)上的處理，所以框架設(shè)計(jì)上要在三個(gè)方面達(dá)成平衡：
盡可能的實(shí)現(xiàn)用戶對(duì)于內(nèi)部密文數(shù)據(jù)結(jié)構(gòu)的透明無(wú)感；
TensorFlow的DAG圖構(gòu)建、自動(dòng)求導(dǎo)等核心功能仍需要無(wú)縫的支持；
支持不同的MPC協(xié)議使用自定義的具體數(shù)據(jù)結(jié)構(gòu)對(duì)象以便于擴(kuò)展。
為同時(shí)實(shí)現(xiàn)這些目標(biāo)，Rosetta基于tf.string這一TensorFlow原生數(shù)據(jù)結(jié)構(gòu)來(lái)承載各協(xié)議自定義的密文數(shù)據(jù)，然后通過(guò)對(duì)TensorFlow源碼代碼進(jìn)行深度的hook改造使得DAG圖構(gòu)建、自動(dòng)求導(dǎo)等功能仍然可用。
DAG的執(zhí)行
如上述網(wǎng)絡(luò)結(jié)構(gòu)圖所展示的那樣，各Player運(yùn)行的是同一份基于TensorFlow編寫的AI二進(jìn)制代碼，比如訓(xùn)練一個(gè)簡(jiǎn)單神經(jīng)網(wǎng)絡(luò)模型的程序。用戶直接的使用TensorFlow中原生的算子API來(lái)構(gòu)建邏輯計(jì)算圖DAG，Rosetta內(nèi)部會(huì)在圖開始執(zhí)行時(shí)完成到隱私算子SecureOp的轉(zhuǎn)換。與其他隱私計(jì)算框架相比，這樣的切換成本是最低的。
在執(zhí)行過(guò)程中，各Player自身都是在按照這個(gè)DAG圖在運(yùn)行，特殊的地方在于在各個(gè)算子的內(nèi)部執(zhí)行過(guò)程中，各個(gè)Player會(huì)根據(jù)自己所屬的角色遵循MPC協(xié)議執(zhí)行不同的操作，這些操作即包括本地的在密文上的處理，也包括在多方之間進(jìn)行強(qiáng)同步的通訊交互，傳輸大量隨機(jī)數(shù)形式的密文。
Rosetta多方協(xié)作運(yùn)行圖
小結(jié)
在本篇文章里，我們結(jié)合Rosetta框架，從整體上介紹了隱私AI框架在工程落地時(shí)所需要面對(duì)的挑戰(zhàn)，以及Rosetta等框架的一些設(shè)計(jì)方案。后續(xù)的文章中，我們會(huì)就核心的關(guān)鍵模塊進(jìn)行進(jìn)一步的介紹。
