朝向使用FHE的加密大型語言模型

TL;DR

Hugging Face 與 Zama 已經展示了一種使用 Fully Homomorphic Encryption (FHE) 在加密資料上運行大型語言模型 (LLM) 部分的方法。此方法使得 LLMs 能在服務提供者看不到原始資料的情況下處理敏感的使用者查詢，同時亦能透過內部部署保護模型擁有者的智慧財產不被洩漏。

使用 Fully Homomorphic Encryption 解決 LLM 隱私問題

Fully Homomorphic Encryption (FHE) 允許直接在加密資料上執行函數，無需在處理前解密資訊。通過適應 Hugging Face transformers 庫中的 GPT-2 實現並使用 Concrete-Python，模型推理過程的特定部分可以轉換為 FHE 等價物。

在 TFHE (Torus Fully Homomorphic Encryption) 方案中，模型權重和激活被表示為整數。非線性函數透過 Programmable Bootstrapping (PBS) 處理，該方法在加密資料上執行表格查詢 (TLU) 並刷新密文以允許任意計算。雖然 PBS 比線性運算計算成本更高，但它允許將 LLM 計算的任何子部分或全部表示為 FHE。

技術實作：加密注意力頭

為實作加密的 LLM 層，系統採用混合方法，其中客戶端先進行本地推理，加密中間運算，然後將其發送到伺服器。伺服器接著在加密資料上應用注意力機制的一部分，並將結果返回給客戶端以解密並繼續本地推理。

量化需求

由於 FHE 僅在整數上運作，模型權重和激活必須進行量化。使用事後訓練量化（這樣可以避免重新訓練的需求），Zama 發現 4-bit 量化能保持原始模型準確度的 96%（基於約 80 句子的數據集）。

與 GPT-2 的整合

該實作涉及重寫特定模組的前向傳遞以包含量化運算子。例如，可以將 GPT2LMHeadModel 的第一個多頭注意力模組替換為 QGPT2SingleHeadAttention 模組來進行修改。

在此設置中，多頭注意力機制的第一個頭（包括查詢、鍵和值矩陣的投影）使用 FHE-friendly 運算子進行計算。其他計算保持為浮點數，並在客戶端以非加密方式執行。

計算複雜度與效能

注意力機制由於查詢、鍵和值的乘法而成為 transformer 模型中計算最密集的部分。在 FHE 中，此成本進一步受到加密域乘法複雜度以及隨著序列長度增長導致運算量平方增加的影響。

主要效能觀察包括：

  • Computational Load: 長度為 6 的序列需要 11,622 次 PBS 運算。
  • Current State: 目前的實作是一個未優化的首次實驗，在 CPU 上運行需要秒級時間，但需要大量計算資源。
  • Future Outlook: Zama 預測未來的 ASIC 硬體可將延遲提升 1,000 倍至 10,000 倍，潛在地將 CPU 上數分鐘的處理時間降低至 100 毫秒以下。

結論

將 FHE 整合到 LLM 中提供了一條通往雲端 AI 服務的道路，在該服務中使用者隱私得到完全尊重且模型智慧財產得到保護。透過利用 transformers 庫以及 Zama 的 Concrete 和 Concrete-ML 庫，開發者可以開始將 ML 模型轉換為 FHE 等價物，以在加密資料上進行預測。

Sources