怎樣使用FHE實(shí)現(xiàn)加密大語(yǔ)言模型？

近來(lái)，大語(yǔ)言模型 (LLM) 已被證明是提高編程、內(nèi)容生成、文本分析、網(wǎng)絡(luò)搜索及遠(yuǎn)程學(xué)習(xí)等諸多領(lǐng)域生產(chǎn)力的可靠工具。

大語(yǔ)言模型對(duì)用戶隱私的影響

盡管 LLM 很有吸引力，但如何保護(hù)好 輸入給這些模型的用戶查詢中的隱私 這一問(wèn)題仍然存在。一方面，我們想充分利用 LLM 的力量，但另一方面，存在向 LLM 服務(wù)提供商泄露敏感信息的風(fēng)險(xiǎn)。在某些領(lǐng)域，例如醫(yī)療保健、金融或法律，這種隱私風(fēng)險(xiǎn)甚至有一票否決權(quán)。

一種備選解決方案是本地化部署，LLM 所有者將其模型部署在客戶的計(jì)算機(jī)上。然而，這不是最佳解決方案，因?yàn)闃?gòu)建 LLM 可能需要花費(fèi)數(shù)百萬(wàn)美元 (GPT3 為 460 萬(wàn)美元)，而本地部署有泄露模型知識(shí)產(chǎn)權(quán) (intellectual property, IP) 的風(fēng)險(xiǎn)。

Zama 相信有兩全其美之法: 我們的目標(biāo)是同時(shí)保護(hù)用戶的隱私和模型的 IP。通過(guò)本文，你將了解如何利用 Hugging Face transformers 庫(kù)并讓這些模型的某些部分在加密數(shù)據(jù)上運(yùn)行。完整代碼見(jiàn) 此處。

全同態(tài)加密 (Fully Homomorphic Encryption，F(xiàn)HE) 可以解決 LLM 隱私挑戰(zhàn)

針對(duì) LLM 部署的隱私挑戰(zhàn)，Zama 的解決方案是使用全同態(tài)加密 (FHE)，在加密數(shù)據(jù)上執(zhí)行函數(shù)。這種做法可以實(shí)現(xiàn)兩難自解，既可以保護(hù)模型所有者知識(shí)產(chǎn)權(quán)，同時(shí)又能維護(hù)用戶的數(shù)據(jù)隱私。我們的演示表明，在 FHE 中實(shí)現(xiàn)的 LLM 模型保持了原始模型的預(yù)測(cè)質(zhì)量。為此，我們需要調(diào)整 Hugging Face transformers 庫(kù) 中的 GPT2 實(shí)現(xiàn)，使用 Concrete-Python 對(duì)推理部分進(jìn)行改造，這樣就可以將 Python 函數(shù)轉(zhuǎn)換為其 FHE 等效函數(shù)。

圖 1. GPT2 架構(gòu);

圖 1 展示了由多個(gè) transformer block 堆疊而成的 GPT2 架構(gòu): 其中最主要的是多頭注意力 (multi-head attention，MHA) 層。每個(gè) MHA 層使用模型權(quán)重來(lái)對(duì)輸入進(jìn)行投影，然后各自計(jì)算注意力，并將注意力的輸出重新投影到新的張量中。

在 TFHE 中，模型權(quán)重和激活均用整數(shù)表示。非線性函數(shù)必須通過(guò)可編程自舉 (Programmable Bootstrapping，PBS) 操作來(lái)實(shí)現(xiàn)。PBS 對(duì)加密數(shù)據(jù)實(shí)施查表 (table lookup，TLU) 操作，同時(shí)刷新密文以支持 任意計(jì)算。不好的一面是，此時(shí) PBS 的計(jì)算時(shí)間在線性運(yùn)算中占主導(dǎo)地位。利用這兩種類型的運(yùn)算，你可以在 FHE 中表達(dá)任何子模型的計(jì)算，甚至完整的 LLM 計(jì)算。

使用 FHE 實(shí)現(xiàn) LLM 的一層

接下來(lái)，你將了解如何加密多頭注意力 (MHA) 中的一個(gè)注意力頭。你可以在 此處 找到完整的 MHA 實(shí)現(xiàn)代碼。

圖 2. 在 FHE 中運(yùn)行 LLM 模型的某些部分

圖 2 概述了一個(gè)簡(jiǎn)化的底層實(shí)現(xiàn)。在這個(gè)方案中，模型權(quán)重會(huì)被分成兩個(gè)部分，分別存儲(chǔ)在客戶端和服務(wù)端。首先，客戶端在本地開(kāi)始推理，直至遇到已第一個(gè)不在本地的層。用戶將中間結(jié)果加密并發(fā)送給服務(wù)端。服務(wù)端對(duì)其執(zhí)行相應(yīng)的注意力機(jī)制計(jì)算，然后將結(jié)果返回給客戶端，客戶端對(duì)結(jié)果進(jìn)行解密并繼續(xù)在本地推理。

量化

首先，為了對(duì)加密值進(jìn)行模型推理，模型的權(quán)重和激活必須被量化并轉(zhuǎn)換為整數(shù)。理想情況是使用 訓(xùn)練后量化，這樣就不需要重新訓(xùn)練模型了。這里，我們使用整數(shù)和 PBS 來(lái)實(shí)現(xiàn) FHE 兼容的注意力機(jī)制，并檢查其對(duì) LLM 準(zhǔn)確率的影響。

要評(píng)估量化的影響，我們運(yùn)行完整的 GPT2 模型，并讓其中的一個(gè) LLM 頭進(jìn)行密態(tài)計(jì)算。然后我們基于此評(píng)估權(quán)重和激活的量化比特?cái)?shù)對(duì)準(zhǔn)確率的影響。

單注意力頭量化的平均 top-k 準(zhǔn)確率

上圖表明 4 比特量化保持了原始精度的 96%。該實(shí)驗(yàn)基于含有約 80 個(gè)句子的數(shù)據(jù)集，并通過(guò)將原始模型的 logits 預(yù)測(cè)與帶有量化注意力頭的模型的 logits 預(yù)測(cè)進(jìn)行比較來(lái)計(jì)算最終指標(biāo)。

在 Hugging Face GPT2 模型中使用 FHE

我們需要在 Hugging Face 的 transformers 庫(kù)的基礎(chǔ)上重寫(xiě)加密模塊的前向傳播，以使其包含量化算子。首先通過(guò)加載 GPT2LMHeadModel 構(gòu)建一個(gè) SingleHeadQGPT2Model 實(shí)例，然后手動(dòng)使用 QGPT2SingleHeadAttention 替換第一個(gè)多頭注意力模塊，代碼如下。你可以在 這里 找到模型的完整實(shí)現(xiàn)。

self.transformer.h[0].attn=QGPT2SingleHeadAttention(config,n_bits=n_bits)

至此，前向傳播已被重載成用 FHE 算子去執(zhí)行多頭注意力的第一個(gè)頭，包括構(gòu)建查詢、鍵和值矩陣的投影。以下代碼中的 QGPT2 模塊的代碼見(jiàn) 此處。

classSingleHeadAttention(QGPT2):
"""Classrepresentingasingleattentionheadimplementedwithquantizationmethods."""


defrun_numpy(self,q_hidden_states:np.ndarray):

#ConverttheinputtoaDualArrayinstance
q_x=DualArray(
float_array=self.x_calib,
int_array=q_hidden_states,
quantizer=self.quantizer
)

#Extracttheattentionbasemodulename
mha_weights_name=f"transformer.h.{self.layer}.attn."

#Extractthequery,keyandvalueweightandbiasvaluesusingtheproperindices
head_0_indices=[
list(range(i*self.n_embd,i*self.n_embd+self.head_dim))
foriinrange(3)
]
q_qkv_weights=...
q_qkv_bias=...

#ApplythefirstprojectioninordertoextractQ,KandVasasinglearray
q_qkv=q_x.linear(
weight=q_qkv_weights,
bias=q_qkv_bias,
key=f"attention_qkv_proj_layer_{self.layer}",
)

#Extractthequeries,keysandvales
q_qkv=q_qkv.expand_dims(axis=1,key=f"unsqueeze_{self.layer}")
q_q,q_k,q_v=q_qkv.enc_split(
3,
axis=-1,
key=f"qkv_split_layer_{self.layer}"
)

#Computeattentionmechanism
q_y=self.attention(q_q,q_k,q_v)

returnself.finalize(q_y)

模型中的其他計(jì)算仍以浮點(diǎn)形式進(jìn)行，未加密，并由客戶端在本地執(zhí)行。

將預(yù)訓(xùn)練的權(quán)重加載到修改后的 GPT2 模型中，然后調(diào)用 generate 方法:

qgpt2_model=SingleHeadQGPT2Model.from_pretrained(
"gpt2_model",n_bits=4,use_cache=False
)

output_ids=qgpt2_model.generate(input_ids)

舉個(gè)例子，你可以要求量化模型補(bǔ)全短語(yǔ) “Cryptography is a” 。在 FHE 中運(yùn)行模型時(shí)，如果量化精度足夠，生成的輸出為:

“Cryptography is a very important part of the security of your computer”

當(dāng)量化精度太低時(shí)，您會(huì)得到:

“Cryptography is a great way to learn about the world around you”

編譯為 FHE

現(xiàn)在，你可以使用以下 Concrete-ML 代碼編譯注意力頭:

circuit_head=qgpt2_model.compile(input_ids)

運(yùn)行此代碼，你將看到以下打印輸出: “Circuit compiled with 8 bit-width”。該配置與 FHE 兼容，顯示了在 FHE 中執(zhí)行的操作所需的最大位寬。

復(fù)雜度

在 transformer 模型中，計(jì)算量最大的操作是注意力機(jī)制，它將查詢、鍵和值相乘。在 FHE 中，加密域中乘法的特殊性加劇了成本。此外，隨著序列長(zhǎng)度的增加，這些乘法的數(shù)量還會(huì)呈二次方增長(zhǎng)。

而就加密注意力頭而言，長(zhǎng)度為 6 的序列需要 11622 次 PBS 操作。我們目前的實(shí)驗(yàn)還很初步，尚未對(duì)性能進(jìn)行優(yōu)化。雖然可以在幾秒鐘內(nèi)運(yùn)行，但不可否認(rèn)它需要相當(dāng)多的計(jì)算能力。幸運(yùn)的是，我們預(yù)期，幾年后，硬件會(huì)將延遲提高 1000 倍到 10000 倍，使原來(lái)在 CPU 上需要幾分鐘的操作縮短到 ASIC 上的低于 100 毫秒。有關(guān)這些估算的更多信息，請(qǐng)參閱 此博文。

總結(jié)

大語(yǔ)言模型有望使能大量應(yīng)用場(chǎng)景，但其實(shí)現(xiàn)引發(fā)了用戶隱私的重大關(guān)切。在本文中，我們朝著密態(tài) LLM 邁出了第一步，我們的最終愿景是讓整個(gè)模型完全在云上運(yùn)行，同時(shí)用戶的隱私還能得到充分尊重。

當(dāng)前的做法包括將 GPT2 等模型中的特定部分轉(zhuǎn)換至 FHE 域。我們的實(shí)現(xiàn)利用了 transformers 庫(kù)，用戶還能評(píng)估模型的一部分在加密數(shù)據(jù)上運(yùn)行時(shí)對(duì)準(zhǔn)確率的影響。除了保護(hù)用戶隱私之外，這種方法還允許模型所有者對(duì)其模型的主要部分保密。你可在 此處 找到完整代碼。

Zama 庫(kù) Concrete 和 Concrete-ML (別忘了給我們的 github 代碼庫(kù)點(diǎn)個(gè)星星 ) 允許直接構(gòu)建 ML 模型并將其轉(zhuǎn)換至等價(jià)的 FHE 域，從而使之能夠?qū)用軘?shù)據(jù)進(jìn)行計(jì)算和預(yù)測(cè)。

審核編輯：劉清