使用 FHE 的加密大型语言模型之路

TL;DR

Hugging Face 和 Zama 已经展示了一种使用全同态加密(FHE)在加密数据上运行大型语言模型(LLM)部分的方法。这种方法使得 LLMs 能够在服务提供商看不到原始数据的情况下处理敏感用户查询,同时通过本地部署防止模型所有者的知识产权泄露。

使用全同态加密解决 LLM 隐私问题

全同态加密(FHE)使得函数可以直接在加密数据上执行,从而无需在处理前解密信息。通过改编自 Hugging Face transformers 库的 GPT-2 实现并使用 Concrete-Python,模型推理过程的特定部分可以转换为 FHE 等价物。 在 TFHE(Torus 全同态加密)方案中,模型权重和激活被表示为整数。非线性函数通过可编程引导(PBS)处理,PBS 在加密数据上执行表查找(TLU)并刷新密文以允许任意计算。虽然 PBS 比线性操作计算成本更高,但它允许将 LLM 计算的任何子部分或全部表示为 FHE。

技术实现:加密注意力头

为了实现加密的 LLM 层,系统采用混合方法,客户端先执行初始本地推理,对中间操作进行加密并发送到服务器。服务器随后在加密数据上应用部分注意力机制,并将结果返回给客户端以进行解密并继续本地推理。

量化要求

由于 FHE 在整数上运行,模型权重和激活必须进行量化。使用后训练量化(这避免了重新训练的需求),Zama 发现 4 位量化保持了原始模型准确率的 96%(基于大约 80 句的数据集)。

与 GPT-2 的集成

该实现涉及重写特定模块的前向传播以包含量化操作。例如,可以通过将第一个多头注意力模块替换为 QGPT2SingleHeadAttention 模块来修改 GPT2LMHeadModel。 在此设置中,多头注意力机制的第一个头(包括查询、键和值矩阵的投影)使用 FHE 友好的操作符执行。其他计算保持为浮点数,并在客户端以非加密方式执行。

计算复杂度与性能

注意力机制由于查询、键和值的乘法而成为 transformer 模型中计算最密集的部分。在 FHE 中,这一成本进一步增加了加密域乘法的复杂性以及随着序列长度增长而导致的操作数的二次增加。 关键性能观察包括:

  • 计算负载: 长度为 6 的序列需要 11,622 次 PBS 操作。
  • 当前状态: 当前实施是一个未优化的首次实验,在 CPU 上运行需要几秒钟,但需要相当大的计算能力。
  • 未来展望: Zama 预测未来的 ASIC 硬件可以将延迟提高 1,000 倍到 10,000 倍,从而可能将 CPU 上几分钟的处理时间降低到低于 100 毫秒。

结论

将 FHE 集成到 LLMs 中为用户隐私得到完全尊重且模型知识产权得到保护的基于云的 AI 服务提供了一条途径。通过利用 transformers 库以及 Zama 的 ConcreteConcrete-ML 库,开发者可以开始将 ML 模型转换为 FHE 等价物,以在加密数据上进行预测。

Sources