現如今已有大量提供深度學習服務的供應商，在使用這些服務時，用戶需要將自己的信息包含在 prompt 中發送給這些服務商，這會導致隱私泄漏等問題。另一方面，服務商基本不愿意公開自己辛苦訓練得到的模型參數。

針對這一問題，螞蟻集團的一個研究團隊提出了 PUMA 框架，可以在不影響模型性能的前提下實現安全的推理。不僅如此，他們也開源了相關代碼。

• 論文：https://arxiv.org/abs/2307.12533
• 代碼：https://github.com/secretflow/spu/blob/main/examples/python/ml/flax_llama7b/flax_llama7b.py

預訓練 Transformer 模型在許多實際任務上都表現優良，也因此受到了很大關注，并且現在已經出現了不少基于這類技術的工具，它們常以「深度學習即服務（DLaaS）」范式提供服務。但是，這些服務常會有隱私問題，比如假如用戶要使用 ChatGPT， 要么就需要用戶向服務提供商提供自己的私人 prompt，要么就需要服務提供商將自己專有的訓練得到的權重配置交給用戶。

為了解決 Transformer 模型服務的隱私問題，一種解決方案是安全多方計算（Secure Multi-Party Computation），這可以在推理過程中保證數據和模型權重的安全。但是，多方計算（MPC）中簡單基礎的 Transformer 推理的時間成本和通信成本都很高，難以用于實際應用。為了取得更好的效率，已經有一些研究成果提出了多種加速 Transformer 模型安全推理的方法，但這些方法仍然存在以下一個或多個缺點：

• 替換很困難。近期一些工作提出，為了降低成本，可使用二次函數和 ReLU 函數等快速近似方法來替代高成本的 GeLU 和 softmax 等函數。但是，如果只是簡單替換這些函數，可能會導致 Transformer 模型性能大幅下降（這可能就會需要額外再對模型進行訓練，即微調）以及出現部署問題。 
• 推理成本高。有研究提出使用更準確的多項式函數來近似高成本的非線性函數，但其近似方法并未考慮 GeLU 和 Softmax 的特殊性質。因此，使用近似之后，這種方法的成本依然很高。 
• 不容易部署。最近也有些研究提出通過修改 Transformer 的模型架構來加速安全推理，例如分解嵌入過程并重新組織線性層。更糟糕的是，由于 Crypten 框架不支持安全 LayerNorm，因此如果僅使用 BatchNorm 模擬成本，就會導致安全推理得到不正確的結果。這些修改方式與現有的明文 Transformer 系統存在沖突。

綜上所述，在 MPC Transformer 推理領域，模型性能和效率難以兼得，而人們可能會有如下問題： 

能否安全又高效地評估預訓練大型 transformer 模型，同時無需進一步再訓練也能達到與明文模型相近的準確度。 

螞蟻集團提出的 PUMA 框架正是為了解決這一難題而生，該框架能夠安全又準確地執行端到端的安全的 Transformer 推理。這篇論文的主要貢獻包括：

• 用于近似非線性函數的新方法。文中提出了更加準確和快速的近似方法，可用于近似 Transformer 模型中高成本的非線性函數（如 GeLU 和 Softmax）。不同于之前的方法，新提出的近似方法基于這些非線性函數的特殊性質，可以兼顧準確度和效率。 
• 更快更準確的安全推理。研究者使用 6 個 transformer 模型和 4 個數據集進行了廣泛的實驗，結果表明，相比于 MPCFORMER，當使用 PUMA 框架時，準確度在接近明文模型的同時，速度和通信效率都提高了 2 倍左右（并且注意 MPCFORMER 的準確度不及 PUMA）。PUMA 甚至可以在 5 分鐘內完成對 LLaMA-7B 的評估，生成一個詞。作者表示這是首次采用 MPC 評估如此大的語言模型。 
• 開源的端到端框架。螞蟻集團的這些研究者成功以 MPC 形式設計并實現了安全的 Embedding 和 LayerNorm 程序。得到的結果是：PUMA 的工作流程遵照明文 Transformer 模型，并未改變任何模型架構，能夠輕松地加載和評估預訓練的明文 Transformer 模型（比如從 Huggingface 下載的模型）。作者表示這是首個支持預訓練 Transformer 模型的準確推理的開源 MPC 解決方案，同時還無需再訓練等進一步修改。

PUMA 的安全設計

PUMA 概況

PUMA 的設計目標是讓基于 Transformer 的模型能安全地執行計算。為了做到這一點，該系統定義了三個實體：模型所有者、客戶端和計算方。模型所有者提供經過訓練的 Transformer 模型，客戶端負責向系統提供數據和收取推理結果，而計算方（即 P_0、P_1 和 P_2）執行安全計算協議。注意模型所有者和客戶端也可以作為計算方，但為了說明方便，這里會將它們區分開。

在安全推理過程中需要保持一個關鍵的不變量：計算方開始時總是有客戶端輸入中三分之二的復制的機密份額以及模型的層權重中三分之二的權重，最終計算方也有這些層的輸出中三分之二的復制的機密份額。由于這些份額不會向各方泄漏信息，這就能確保這些協議模塊能以任意深度按順序組合起來，從而為任意基于 Transformer 的模型提供安全計算。PUMA 關注的主要問題是降低各計算方之間的運行時間成本和通信成本，同時維持所需的安全級別。通過利用復制的機密份額和新提出的 3PC 協議，PUMA 能在三方設置下讓基于 Transformer 的模型實現安全推理。

安全嵌入協議 

當前的安全嵌入（secure embedding）流程需要客戶端使用 token id 創建一個 one-hot 向量，這偏離了明文工作流程并會破壞 Transformer 結構。因此，該方法并不容易部署到真實的 Transformer 模型服務應用中。

為了解決這個問題，這里研究者提出了一種新的安全嵌入設計。令 token id ∈ [n] 且所有嵌入向量均表示為 ，則嵌入可以表示為。由于 (id, E) 共享秘密，則新提出的安全嵌入協議的的工作方式如下：

• 計算方在接受到來自客戶端的 id 向量后，安全地計算 one-hot 向量。具體來說， 其中 i ∈ [n]. 
• 各計算方可以通過 計算嵌入向量，其中 不需要安全截斷（secure truncation）。

如此一來，這里的 Π_Embed 就不需要顯式地修改 Transformer 模型的工作流程。

安全 GeLU 協議

目前大多數方法都將 GeLU 函數看作是由更小的函數組成的，并會嘗試優化其中每一部分，這就讓它們錯失了從整體上優化私密 GeLU 的機會。給定 GeLU 函數：

過去的一些方法關注的重心要么是 tanh 函數設計高效的協議，要么是將現有的求冪和倒數的 MPC 協議用于 Sigmoid。

但是，這些方法都沒有用到這一事實：GeLU 函數兩邊基本都是線性的，即當 x<?4 時 GeLU (x) ≈ 0，當 x>3 時 GeLU (x) ≈ x。研究者提出在 GeLU 的 [?4,3] 的短區間內，低次多項式的分段近似是一種更高效且更容易實現的安全協議選擇。具體來說，這個分段式低次多項式如下 (2) 式所示：

其中多項式 F_0 和 F_1 的計算是通過軟件庫 numpy.ployfit 實現，如 (3) 式所示。研究者發現，這種多項式擬合雖然簡單，但表現卻出人意料地好；實驗結果的最大誤差 < 0.01403，中值誤差 < 4.41e?05，平均誤差 < 0.00168。

從數學形式上講，給定機密輸入，新提出的安全 GeLU 協議的構建方式見如下算法 1。

安全 Softmax 協議 

在函數 中，關鍵的挑戰是計算 Softmax 函數（其中 M 可被視為一個偏置矩陣）。為了數值穩定性，可以這樣計算 Softmax：

其中  是輸入向量 x 的最大元素。對于普通的明文 softmax，ε = 0。對于二維矩陣，則是將 (4) 式用于其每個行向量。

算法 2 給出了新提出的安全協議 Π_Softmax 的詳細數學描述，其中提出了兩種優化方法：

• 第一種優化是將 (4) 式中的 ε 設置成一個非常小的正值，比如 ε=10^-6，這樣一來 (4) 式中求冪運算的輸入就都是負值。研究者利用了這些負操作數來提升速度。他們具體通過簡單的裁剪使用泰勒級數來計算其中的冪。

• 研究者提出的第二種優化是降低除法量，這最終能降低計算和通信成本。為此，對于大小為 n 的向量 x，研究者將 Div (x, Broadcast (y)) 運算替換成了 x?Broadcast (1/y)，其中 。這種替換可以有效地將 n 次除法約簡至 1 次倒數運算和 n 次乘法運算。這種優化對于 Softmax 運算尤其有益。在定點值設置下，Softmax 運算中 1/y 依然足夠大，難以維持足夠的準確度。由此，這樣的優化可以在保證準確度的同時顯著降低計算和通信成本。

安全 LayerNorm 協議

回想一下，給定大小為 n 的向量 x，，其中 (γ, β) 是已訓練的參數， 且 。在 MPC 中，關鍵挑戰是評估除以平方根公式 。為了安全地評估這一公式，CrypTen 的做法是按這個順序執行這個 MPC 協議：平方根、倒數和乘法。但螞蟻集團的研究者觀察到 。而在 MPC 方面，計算平方根倒數 σ^{-1/2} 的成本接近平方根運算的成本。此外，受前一小節中第二種優化的啟發，他們還提出首先計算 σ^{-1/2}，然后廣播 Broadcast (σ^{-1/2}) 以支持快速和安全的 LayerNorm (x)。算法 3 給出了 Π_LayerNorm 協議的數學形式。

實驗評估

圖 1：在 GLUE 和 Wikitext-103 V1 基準上的性能表現，模型方面，a 是 Bert-Base，b 是 Roberta-Base，c 是 Bert-Large，d 包括 GPT2-Base、GPT2-Medium、GPT2-Large。

表 1：對于一個長度為 128 的輸入句，Bert-Base、Roberta-Base 和 Bert-Large 的成本。時間成本以秒計算，通信成本以 GB 計算。

表 2：GPT2-Base、GPT2-Medium 和 GPT2-Large 的成本。輸入句的長度為 32，這些是生成 1 個 token 的成本。

表 3：對于 {2, 4, 8, 16} 句子的批次，Bert-Base 和 GPT2-Base 的成本。Bert-Base 和 GPT2-Base 的輸入長度分別設定為 128 和 32，GPT2 的數據是生成 1 個 token 的成本。

表 4：不同輸入長度（#Input）下 Bert-Base 和 GPT2-Base 的成本。Bert-Base 和 GPT2-Base 的輸入長度分別設定為 {64, 128, 256, 512} 和 {16, 32, 64, 128}.GPT2 的數據是生成 1 個 token 的成本。

圖 2：GPT2-Base 生成不同輸出 token 的成本，輸入長度為 32。a 是運行時間成本，b 是通信成本。

表 5：用 LLaMA-7B 執行安全推理的成本，#Input 表示輸入句的長度，#Output 表示所生成的 token 的數量。

只需五分鐘就能擴展用于 LLaMA-7B。研究者在 3 個阿里云 ecs.r7.32xlarge 服務器上使用 PUMA 評估了大型語言模型 LLaMA-7B，其中每個服務器都有 128 線程和 1 TB RAM，帶寬為 20 GB，往返時間為 0.06 ms。如表 5 所示，只需合理的成本，PUMA 就能支持大型語言模型 LLaMA-7B 實現安全推理。舉個例子，給定 8 個 token 構成的輸入句，PUMA 可以在大約 346.126 秒內以 1.865 GB 的通信成本輸出一個 token。研究者表示，這是首次使用 MPC 方案對 LLaMA-7B 實施評估。

PUMA雖然取得了一系列突破，但是它依然是一個學術成果，其推理耗時依然離落地存在一些距離。研究者相信未來與機器學習領域最新的量化技術、硬件領域量新的硬件加速技術相結合之后，真正保護隱私的大模型服務將離我們不再遙遠。