使用 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 等效项，以便能够在加密数据上进行计算和预测。

希望你喜欢这篇文章；欢迎分享你的想法/反馈！