使用 FHE 構建加密大型語言模型

大型語言模型 (LLM) 最近被證明是提高許多領域生產力的可靠工具，如程式設計、內容創作、文字分析、網路搜尋和遠端學習。

大型語言模型對使用者隱私的影響

儘管 LLM 很有吸引力，但圍繞著這些模型處理的使用者查詢仍然存在隱私問題。一方面，人們希望利用 LLM 的強大功能，但另一方面，也存在向 LLM 服務提供商洩露敏感資訊的風險。在某些領域，如醫療保健、金融或法律，這種隱私風險是無法接受的。

解決這個問題的一個可能方案是本地部署，即 LLM 所有者將模型部署在客戶的機器上。然而，這並非最佳解決方案，因為構建一個 LLM 可能需要數百萬美元（GPT3 為 460 萬美元），而本地部署存在洩露模型智慧財產權 (IP) 的風險。

Zama 相信你可以兩全其美：我們的目標是既保護使用者隱私，又保護模型的智慧財產權。在這篇部落格中，你將看到如何利用 Hugging Face transformers 庫，並讓這些模型的一部分在加密資料上執行。完整的程式碼可以在這個用例示例中找到。

全同態加密 (FHE) 可解決 LLM 的隱私挑戰

Zama 針對 LLM 部署挑戰的解決方案是使用全同態加密 (FHE)，它能夠在加密資料上執行函式。在保護模型所有者 IP 的同時，仍然維護使用者資料的隱私是完全可以實現的目標。本演示展示了在 FHE 中實現的 LLM 模型能夠保持原始模型預測的質量。為此，有必要修改 Hugging Face transformers 庫中的 GPT2 實現，使用 Concrete-Python 重構推理的部分環節，該工具能將 Python 函式轉換為其 FHE 等效項。

圖 1 展示了 GPT2 架構，它具有一個重複結構：一系列連續應用的多頭注意力 (MHA) 層。每個 MHA 層使用模型權重投射輸入，計算注意力機制，然後將注意力的輸出重新投射到一個新的張量中。

在 TFHE 中，模型權重和啟用值用整數表示。非線性函式必須透過可程式設計自舉 (Programmable Bootstrapping, PBS) 操作來實現。PBS 在加密資料上實現了一個查詢表 (Table Lookup, TLU) 操作，同時重新整理密文以允許任意計算。不利的一面是，PBS 的計算時間遠超線性運算。利用這兩種型別的操作，你可以在 FHE 中表示任何子部分，甚至是完整的 LLM 計算。

使用 FHE 實現 LLM 層

接下來，你將看到如何加密多頭注意力 (MHA) 塊中的單個注意力頭。你也可以在這個用例示例中找到完整 MHA 塊的示例。

圖 2 展示了底層實現的簡化概述。客戶端在本地啟動推理，直到達到已從共享模型中移除的第一個層。使用者加密中間操作並將其傳送到伺服器。伺服器應用部分注意力機制，然後將結果返回給客戶端，客戶端可以解密它們並繼續本地推理。

量化

首先，為了在加密值上執行模型推理，模型的權重和啟用值必須被量化並轉換為整數。理想情況是使用訓練後量化，因為它不需要重新訓練模型。具體過程是實現一個 FHE 相容的注意力機制，使用整數和 PBS，然後檢驗其對 LLM 準確性的影響。

為了評估量化的影響，我們執行一個完整的 GPT2 模型，其中只有一個 LLM 頭在加密資料上執行。然後，我們透過改變權重和啟用值的量化位數來評估獲得的準確性。

此圖顯示，4 位量化保持了 96% 的原始準確率。該實驗使用了約 80 個句子的資料集。指標是透過比較原始模型的 logits 預測與帶有量化頭的模型的預測來計算的。

將 FHE 應用於 Hugging Face GPT2 模型

基於 Hugging Face 的 transformers 庫，我們重寫了想要加密的模組的前向傳播過程，以便包含量化運算元。透過首先載入一個 GPT2LMHeadModel，然後使用一個 QGPT2SingleHeadAttention 模組手動替換第一個多頭注意力模組，來構建一個 SingleHeadQGPT2Model 例項。完整的實現可以在這裡找到。

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

然後，前向傳播過程被重寫，使得多頭注意力機制的第一個頭，包括為構建查詢、鍵和值矩陣所做的投影，都使用 FHE 友好的運算元執行。以下 QGPT2 模組可以在這裡找到。

class SingleHeadAttention(QGPT2):
    """Class representing a single attention head implemented with quantization methods."""


    def run_numpy(self, q_hidden_states: np.ndarray):

        # Convert the input to a DualArray instance
        q_x = DualArray(
            float_array=self.x_calib,
            int_array=q_hidden_states,
            quantizer=self.quantizer
        )

        # Extract the attention base module name
        mha_weights_name = f"transformer.h.{self.layer}.attn."

        # Extract the query, key and value weight and bias values using the proper indices
        head_0_indices = [
            list(range(i * self.n_embd, i * self.n_embd + self.head_dim)) 
            for i in range(3)
        ]
        q_qkv_weights = ...
        q_qkv_bias = ...

        # Apply the first projection in order to extract Q, K and V as a single array
        q_qkv = q_x.linear(
            weight=q_qkv_weights,
            bias=q_qkv_bias,
            key=f"attention_qkv_proj_layer_{self.layer}",
        )

        # Extract the queries, keys and vales
        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}"
        )

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

        return self.finalize(q_y)

模型中的其他計算保持浮點數、非加密狀態，並預計由客戶端在本地執行。

將預訓練的權重載入到以這種方式修改的 GPT2 模型中，你就可以呼叫 generate 方法了。

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


output_ids = qgpt2_model.generate(input_ids)

舉個例子，你可以讓量化模型完成短語“Cryptography is a”。在 FHE 中執行模型時，如果有足夠的量化精度，生成結果是：

“Cryptography is a very important part of the security of your computer”（密碼學是您計算機安全中非常重要的一部分）

當量化精度過低時，你會得到：

“Cryptography is a great way to learn about the world around you”（密碼學是瞭解你周圍世界的好方法）

編譯為 FHE

你現在可以使用以下 Concrete-ML 程式碼編譯註意力頭：

circuit_head = qgpt2_model.compile(input_ids)

執行此程式碼，你將看到以下輸出：“Circuit compiled with 8 bit-width”。這個與 FHE 相容的配置顯示了在 FHE 中執行操作所需的最大位寬。

複雜性

在 Transformer 模型中，計算最密集的操作是注意力機制，它需要將查詢、鍵和值相乘。在 FHE 中，由於加密域中乘法的特殊性，成本會進一步增加。此外，隨著序列長度的增加，這些具有挑戰性的乘法運算的數量會呈二次方增長。

對於加密的注意力頭，一個長度為 6 的序列需要 11,622 次 PBS 操作。這只是一個初步實驗，並未進行效能最佳化。雖然它可以在幾秒鐘內執行，但這需要相當大的計算能力。幸運的是，硬體將使延遲提高 1000 倍到 10000 倍，一旦幾年後 ASIC 可用，CPU 上需要幾分鐘的事情將在 ASIC 上降至 100 毫秒以內。有關這些預測的更多資訊，請參閱這篇部落格文章。

結論

大型語言模型在各種用例中都是出色的輔助工具，但它們的實現引發了嚴重的使用者隱私問題。在這篇部落格中，你看到了讓整個 LLM 在加密資料上工作的第一步，這樣模型就可以完全在雲端執行，同時使用者的隱私得到充分尊重。

這一步包括將像 GPT2 這樣的模型中的特定部分轉換為 FHE 領域。該實現利用了 transformers 庫，並允許你評估當模型的一部分在加密資料上執行時對準確性的影響。除了保護使用者隱私外，這種方法還允許模型所有者對其模型的主要部分保密。完整的程式碼可以在這個用例示例中找到。

Zama 的庫 Concrete 和 Concrete-ML（別忘了在 GitHub 上給這些倉庫點贊 ⭐️💛）允許直接構建 ML 模型並將其轉換為 FHE 等效項，以便能夠在加密資料上進行計算和預測。

希望你喜歡這篇文章；歡迎分享你的想法/反饋！