TVM如何幫助部署隱私保護和安全的AI應用
導論：保護資料隱私的需求
一直以來，機器學習模型尤其是深度學習模型的成功往往建立在大規模資料的基礎之上。然而在現實生活中，並不是所有的資料都適合被直接被拿來訓練。比方說醫療與金融這類敏感應用，相當一部分資料是個人的隱私資料，這意味著在這樣的資料集上進行訓練存在著洩露患者隱私或者商業機密的風險。迄今為止，已經有相當多的工作關注在如何在這類私密資料上進行機器學習模型的訓練同時保證資料的隱私。比如說安全和隱私保護計算領域最新進展裡的TEE（可信執行環境, Trusted Execution Environments），以及差分隱私（Differential Privacy）等工作，這些工作為互不信任的各方提供了一種有效地訓練機器學習模型而不洩露訓練資料的方法。
來自UC Berkeley的Nick Hynes同學在進行隱私保護方向的研究時，發現相比於其它的機器學習框架，TVM完美適合了在資源有限的TEE環境中部署機器學習模型這一場景。本文則希望向大家分享這一發現。
需求場景
事實上，希望保護訓練資料隱私這樣的場景不在少數：
- 區塊鏈上的機器學習 。目前已有一些區塊鏈專案嘗試去發起使用者來訓練鏈上的機器學習模型，然而需要解決的一個問題是，當用戶使用自己的隱私資料去訓練鏈上模型時，如何保證隱私資料不會被洩露。
- 保證隱私的MLaaS 。雲提供商在使用者的資料上執行他們的機器學習模型。 使用者得到模型輸出且保證自己的資料保持私密，而云提供商則需要確保使用者無法竊取模型。
- 可信的機器學習比賽 。參賽者在比賽資料上訓練模型。 比賽組織者將私密的測試資料傳送到參賽者的模型並獲得可驗證的準確性結果。 在組織者決定採用參賽者的模型之前，參賽者的模型需要保證安全。 並且參賽者也無法獲得測試資料的資訊從而在測試資料上進行訓練來作弊。
- 在一些隱私資料上訓練 。假設研究人員希望在醫院的資料上訓練模型，然而醫院的資料往往無法被直接授予使用，所以需要有一個“受信任的第三方”來訓練一個隱私保護的模型。
TEE是什麼
TEE（可信執行環境, Trusted Execution Environment）本質上允許遠端使用者可以在另一個使用者的機器上執行程式碼而不向硬體提供者透露計算的內容。
技術上來講，TEE提供了一塊包含隔離或加密的儲存器和CPU暫存器的安全區域——enclave，以及可信的隨機來源。TEE還可以傳送已載入程式碼的簽名，以便遠端使用者可以驗證程式碼已正確載入進安全區域enclave中。該過程稱為遠端證明，可用於建立進入enclave的安全通道。在這之後，遠端使用者可以確保安全的為應用配置私鑰，模型引數和訓練資料。
與諸如安全多奇偶校驗計算（MPC）和完全同態加密（FHE）的純加密方法相比，TEE快幾個數量級並且支援通用計算。唯一的缺點也許是硬體信任根（燒入處理器的金鑰）和載入軟體過程中所需要的額外信任假設。
事實上，TEE技術正變得越來越普遍，並且在實際隱私保護方面發揮著重要作用。事實上，通用的TEE技術已經應用在Intel SGX和ARM TrustZone等商用硬體中。此外，完全開源的Keystone Enclave正在開發中。
TVM如何幫助在TEE上部署模型
使用TEE的主要挑戰之一是在TEE內部執行的程式碼無法訪問不受信任的作業系統。這意味著可信軟體無法建立執行緒或者執行IO操作。因而導致了像OpenBLAS這樣的代數庫不能直接在enclave上執行，更不用說其它的深度學習框架如PyTorch和TensorFlow。而相比於其它的機器學習框架，TVM能夠將機器學習模型編譯成非常輕量化，高度優化，且沒有外部依賴的庫，因而完美適合了在資源有限的enclave上部署機器學習模型這樣一類的場景。
TEE實際上有著與資源受限的硬體加速器相類似的程式設計模型，而這恰恰是TVM被創造的原因！在整個工作流程中，使用者首先使用高層的圖表示語言來表示整張訓練的計算圖。 TVM隨後編譯模型並輸出一個包含高度優化運算元的靜態庫，這可以很容易地載入到TEE中。由於運算元是自動生成的並且具有嚴格的邊界檢查，這充分減小了攻擊面。並且它們被輕量且記憶體安全的Rust runtime來裝載執行，也同樣可以輕鬆檢查其安全性和正確性。
當然，效率也是實際應用時需要考量的另一個維度。通過使用runtime來協調執行緒，單個TVM enclave可以充分利用其主機的資源。這使得enclave中的TVM模組具有與基於CPU的本機訓練相當的效能。同時，TVM還提供了對演算法進行更細粒度優化的機會。在差分隱私演算法中，需要裁剪每個訓練sample的梯度並新增一定的噪聲。TVM的排程原語使得實現優化變得簡單與高效，同時，得益於Operator Fusion等內建優化，TVM能將裁剪與噪聲操作進行融合並原地（in-place）計算從而進一步減少延遲與記憶體消耗。這使得我們能夠應用差分隱私演算法且不會引入過高的開銷。
小結
下一代機器學習系統很有可能迎來隱私保護的話題。 隨著TEE技術變得更易於理解和更加廣泛的應用，嘗試使用它去部署隱私保護機器學習模型十分有意義，而TVM則很有可能有助於該方向的研究與開發。歡迎大家嘗試探索以及加入TVM社群。
資源連結
- 英文部落格: ofollow,noindex"> https:// tvm.ai/2018/10/09/ml-in -tees.html
- 技術報告: https:// arxiv.org/abs/1807.0668 9
- TVM in Intel SGX示例: https:// github.com/dmlc/tvm/tre e/master/apps/sgx