担心prompt泄露隐私？这个框架让LLaMA-7B完成安全推理

现如今已有大量提供深度学习服务的供应商，在使用这些服务时，用户需要将自己的信息包含在 prompt 中发送给这些服务商，这会导致隐私泄漏等问题。另一方面，服务商基本不愿意公开自己辛苦训练得到的模型参数。

针对这一问题，蚂蚁集团的一个研究团队提出了 PUMA 框架，可以在不影响模型性能的前提下实现安全的推理。不仅如此，他们也开源了相关代码。

• 论文：https://arxiv.org/abs/2307.12533
• 代码：https://github.com/secretflow/spu/blob/main/examples/python/ml/flax_llama7b/flax_llama7b.py

预训练 Transformer 模型在许多实际任务上都表现优良，也因此受到了很大关注，并且现在已经出现了不少基于这类技术的工具，它们常以「深度学习即服务（DLaaS）」范式提供服务。但是，这些服务常会有隐私问题，比如假如用户要使用 ChatGPT， 要么就需要用户向服务提供商提供自己的私人 prompt，要么就需要服务提供商将自己专有的训练得到的权重配置交给用户。

为了解决 Transformer 模型服务的隐私问题，一种解决方案是安全多方计算（Secure Multi-Party Computation），这可以在推理过程中保证数据和模型权重的安全。但是，多方计算（MPC）中简单基础的 Transformer 推理的时间成本和通信成本都很高，难以用于实际应用。为了取得更好的效率，已经有一些研究成果提出了多种加速 Transformer 模型安全推理的方法，但这些方法仍然存在以下一个或多个缺点：

• 替换很困难。近期一些工作提出，为了降低成本，可使用二次函数和 ReLU 函数等快速近似方法来替代高成本的 GeLU 和 softmax 等函数。但是，如果只是简单替换这些函数，可能会导致 Transformer 模型性能大幅下降（这可能就会需要额外再对模型进行训练，即微调）以及出现部署问题。 
• 推理成本高。有研究提出使用更准确的多项式函数来近似高成本的非线性函数，但其近似方法并未考虑 GeLU 和 Softmax 的特殊性质。因此，使用近似之后，这种方法的成本依然很高。 
• 不容易部署。最近也有些研究提出通过修改 Transformer 的模型架构来加速安全推理，例如分解嵌入过程并重新组织线性层。更糟糕的是，由于 Crypten 框架不支持安全 LayerNorm，因此如果仅使用 BatchNorm 模拟成本，就会导致安全推理得到不正确的结果。这些修改方式与现有的明文 Transformer 系统存在冲突。

综上所述，在 MPC Transformer 推理领域，模型性能和效率难以兼得，而人们可能会有如下问题： 

能否安全又高效地评估预训练大型 transformer 模型，同时无需进一步再训练也能达到与明文模型相近的准确度。 

蚂蚁集团提出的 PUMA 框架正是为了解决这一难题而生，该框架能够安全又准确地执行端到端的安全的 Transformer 推理。这篇论文的主要贡献包括：

• 用于近似非线性函数的新方法。文中提出了更加准确和快速的近似方法，可用于近似 Transformer 模型中高成本的非线性函数（如 GeLU 和 Softmax）。不同于之前的方法，新提出的近似方法基于这些非线性函数的特殊性质，可以兼顾准确度和效率。 
• 更快更准确的安全推理。研究者使用 6 个 transformer 模型和 4 个数据集进行了广泛的实验，结果表明，相比于 MPCFORMER，当使用 PUMA 框架时，准确度在接近明文模型的同时，速度和通信效率都提高了 2 倍左右（并且注意 MPCFORMER 的准确度不及 PUMA）。PUMA 甚至可以在 5 分钟内完成对 LLaMA-7B 的评估，生成一个词。作者表示这是首次采用 MPC 评估如此大的语言模型。 
• 开源的端到端框架。蚂蚁集团的这些研究者成功以 MPC 形式设计并实现了安全的 Embedding 和 LayerNorm 程序。得到的结果是：PUMA 的工作流程遵照明文 Transformer 模型，并未改变任何模型架构，能够轻松地加载和评估预训练的明文 Transformer 模型（比如从 Huggingface 下载的模型）。作者表示这是首个支持预训练 Transformer 模型的准确推理的开源 MPC 解决方案，同时还无需再训练等进一步修改。

PUMA 的安全设计

PUMA 概况

PUMA 的设计目标是让基于 Transformer 的模型能安全地执行计算。为了做到这一点，该系统定义了三个实体：模型所有者、客户端和计算方。模型所有者提供经过训练的 Transformer 模型，客户端负责向系统提供数据和收取推理结果，而计算方（即 P_0、P_1 和 P_2）执行安全计算协议。注意模型所有者和客户端也可以作为计算方，但为了说明方便，这里会将它们区分开。

在安全推理过程中需要保持一个关键的不变量：计算方开始时总是有客户端输入中三分之二的复制的机密份额以及模型的层权重中三分之二的权重，最终计算方也有这些层的输出中三分之二的复制的机密份额。由于这些份额不会向各方泄漏信息，这就能确保这些协议模块能以任意深度按顺序组合起来，从而为任意基于 Transformer 的模型提供安全计算。PUMA 关注的主要问题是降低各计算方之间的运行时间成本和通信成本，同时维持所需的安全级别。通过利用复制的机密份额和新提出的 3PC 协议，PUMA 能在三方设置下让基于 Transformer 的模型实现安全推理。

安全嵌入协议 

当前的安全嵌入（secure embedding）流程需要客户端使用 token id 创建一个 one-hot 向量，这偏离了明文工作流程并会破坏 Transformer 结构。因此，该方法并不容易部署到真实的 Transformer 模型服务应用中。

如此一来，这里的 Π_Embed 就不需要显式地修改 Transformer 模型的工作流程。

安全 GeLU 协议

目前大多数方法都将 GeLU 函数看作是由更小的函数组成的，并会尝试优化其中每一部分，这就让它们错失了从整体上优化私密 GeLU 的机会。给定 GeLU 函数：

过去的一些方法关注的重心要么是 tanh 函数设计高效的协议，要么是将现有的求幂和倒数的 MPC 协议用于 Sigmoid。

但是，这些方法都没有用到这一事实：GeLU 函数两边基本都是线性的，即当 x<−4 时 GeLU (x) ≈ 0，当 x>3 时 GeLU (x) ≈ x。研究者提出在 GeLU 的 [−4,3] 的短区间内，低次多项式的分段近似是一种更高效且更容易实现的安全协议选择。具体来说，这个分段式低次多项式如下 (2) 式所示：

其中多项式 F_0 和 F_1 的计算是通过软件库 numpy.ployfit 实现，如 (3) 式所示。研究者发现，这种多项式拟合虽然简单，但表现却出人意料地好；实验结果的最大误差 < 0.01403，中值误差 < 4.41e−05，平均误差 < 0.00168。

安全 Softmax 协议 

算法 2 给出了新提出的安全协议 Π_Softmax 的详细数学描述，其中提出了两种优化方法：

安全 LayerNorm 协议

实验评估

只需五分钟就能扩展用于 LLaMA-7B。研究者在 3 个阿里云 ecs.r7.32xlarge 服务器上使用 PUMA 评估了大型语言模型 LLaMA-7B，其中每个服务器都有 128 线程和 1 TB RAM，带宽为 20 GB，往返时间为 0.06 ms。如表 5 所示，只需合理的成本，PUMA 就能支持大型语言模型 LLaMA-7B 实现安全推理。举个例子，给定 8 个 token 构成的输入句，PUMA 可以在大约 346.126 秒内以 1.865 GB 的通信成本输出一个 token。研究者表示，这是首次使用 MPC 方案对 LLaMA-7B 实施评估。

PUMA虽然取得了一系列突破，但是它依然是一个学术成果，其推理耗时依然离落地存在一些距离。研究者相信未来与机器学习领域最新的量化技术、硬件领域量新的硬件加速技术相结合之后，真正保护隐私的大模型服务将离我们不再遥远。