現(xiàn)如今已有大量提供深度學(xué)習(xí)服務(wù)的供應(yīng)商，在使用這些服務(wù)時(shí)，用戶(hù)需要將自己的信息包含在 prompt 中發(fā)送給這些服務(wù)商，這會(huì)導(dǎo)致隱私泄漏等問(wèn)題。另一方面，服務(wù)商基本不愿意公開(kāi)自己辛苦訓(xùn)練得到的模型參數(shù)。

針對(duì)這一問(wèn)題，螞蟻集團(tuán)的一個(gè)研究團(tuán)隊(duì)提出了 PUMA 框架，可以在不影響模型性能的前提下實(shí)現(xiàn)安全的推理。不僅如此，他們也開(kāi)源了相關(guān)代碼。

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

預(yù)訓(xùn)練 Transformer 模型在許多實(shí)際任務(wù)上都表現(xiàn)優(yōu)良，也因此受到了很大關(guān)注，并且現(xiàn)在已經(jīng)出現(xiàn)了不少基于這類(lèi)技術(shù)的工具，它們常以「深度學(xué)習(xí)即服務(wù)（DLaaS）」范式提供服務(wù)。但是，這些服務(wù)常會(huì)有隱私問(wèn)題，比如假如用戶(hù)要使用 ChatGPT， 要么就需要用戶(hù)向服務(wù)提供商提供自己的私人 prompt，要么就需要服務(wù)提供商將自己專(zhuān)有的訓(xùn)練得到的權(quán)重配置交給用戶(hù)。

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

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

綜上所述，在 MPC Transformer 推理領(lǐng)域，模型性能和效率難以兼得，而人們可能會(huì)有如下問(wèn)題： 

能否安全又高效地評(píng)估預(yù)訓(xùn)練大型 transformer 模型，同時(shí)無(wú)需進(jìn)一步再訓(xùn)練也能達(dá)到與明文模型相近的準(zhǔn)確度。 

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

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

PUMA 的安全設(shè)計(jì)

PUMA 概況

PUMA 的設(shè)計(jì)目標(biāo)是讓基于 Transformer 的模型能安全地執(zhí)行計(jì)算。為了做到這一點(diǎn)，該系統(tǒng)定義了三個(gè)實(shí)體：模型所有者、客戶(hù)端和計(jì)算方。模型所有者提供經(jīng)過(guò)訓(xùn)練的 Transformer 模型，客戶(hù)端負(fù)責(zé)向系統(tǒng)提供數(shù)據(jù)和收取推理結(jié)果，而計(jì)算方（即 P_0、P_1 和 P_2）執(zhí)行安全計(jì)算協(xié)議。注意模型所有者和客戶(hù)端也可以作為計(jì)算方，但為了說(shuō)明方便，這里會(huì)將它們區(qū)分開(kāi)。

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

安全嵌入?yún)f(xié)議 

當(dāng)前的安全嵌入（secure embedding）流程需要客戶(hù)端使用 token id 創(chuàng)建一個(gè) one-hot 向量，這偏離了明文工作流程并會(huì)破壞 Transformer 結(jié)構(gòu)。因此，該方法并不容易部署到真實(shí)的 Transformer 模型服務(wù)應(yīng)用中。

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

• 計(jì)算方在接受到來(lái)自客戶(hù)端的 id 向量后，安全地計(jì)算 one-hot 向量。具體來(lái)說(shuō)， 其中 i ∈ [n]. 
• 各計(jì)算方可以通過(guò) 計(jì)算嵌入向量，其中 不需要安全截?cái)啵╯ecure truncation）。

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

安全 GeLU 協(xié)議

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

過(guò)去的一些方法關(guān)注的重心要么是 tanh 函數(shù)設(shè)計(jì)高效的協(xié)議，要么是將現(xiàn)有的求冪和倒數(shù)的 MPC 協(xié)議用于 Sigmoid。

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

其中多項(xiàng)式 F_0 和 F_1 的計(jì)算是通過(guò)軟件庫(kù) numpy.ployfit 實(shí)現(xiàn)，如 (3) 式所示。研究者發(fā)現(xiàn)，這種多項(xiàng)式擬合雖然簡(jiǎn)單，但表現(xiàn)卻出人意料地好；實(shí)驗(yàn)結(jié)果的最大誤差 < 0.01403，中值誤差 < 4.41e?05，平均誤差 < 0.00168。

從數(shù)學(xué)形式上講，給定機(jī)密輸入，新提出的安全 GeLU 協(xié)議的構(gòu)建方式見(jiàn)如下算法 1。

安全 Softmax 協(xié)議 

在函數(shù) 中，關(guān)鍵的挑戰(zhàn)是計(jì)算 Softmax 函數(shù)（其中 M 可被視為一個(gè)偏置矩陣）。為了數(shù)值穩(wěn)定性，可以這樣計(jì)算 Softmax：

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

算法 2 給出了新提出的安全協(xié)議 Π_Softmax 的詳細(xì)數(shù)學(xué)描述，其中提出了兩種優(yōu)化方法：

• 第一種優(yōu)化是將 (4) 式中的 ε 設(shè)置成一個(gè)非常小的正值，比如 ε=10^-6，這樣一來(lái) (4) 式中求冪運(yùn)算的輸入就都是負(fù)值。研究者利用了這些負(fù)操作數(shù)來(lái)提升速度。他們具體通過(guò)簡(jiǎn)單的裁剪使用泰勒級(jí)數(shù)來(lái)計(jì)算其中的冪。

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

安全 LayerNorm 協(xié)議

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

實(shí)驗(yàn)評(píng)估

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

表 1：對(duì)于一個(gè)長(zhǎng)度為 128 的輸入句，Bert-Base、Roberta-Base 和 Bert-Large 的成本。時(shí)間成本以秒計(jì)算，通信成本以 GB 計(jì)算。

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

表 3：對(duì)于 {2, 4, 8, 16} 句子的批次，Bert-Base 和 GPT2-Base 的成本。Bert-Base 和 GPT2-Base 的輸入長(zhǎng)度分別設(shè)定為 128 和 32，GPT2 的數(shù)據(jù)是生成 1 個(gè) token 的成本。

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

圖 2：GPT2-Base 生成不同輸出 token 的成本，輸入長(zhǎng)度為 32。a 是運(yùn)行時(shí)間成本，b 是通信成本。

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

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

PUMA雖然取得了一系列突破，但是它依然是一個(gè)學(xué)術(shù)成果，其推理耗時(shí)依然離落地存在一些距離。研究者相信未來(lái)與機(jī)器學(xué)習(xí)領(lǐng)域最新的量化技術(shù)、硬件領(lǐng)域量新的硬件加速技術(shù)相結(jié)合之后，真正保護(hù)隱私的大模型服務(wù)將離我們不再遙遠(yuǎn)。