技术前沿｜NeurIPS 2024 中稿快讯！两方密态推理框架 Nimbus，大模型隐私推理场景的重要技术

2024-11-14
浙江
本文字数：1817 字
阅读完需：约 6 分钟

“隐语” 是开源的可信隐私计算框架，内置 MPC、TEE、同态等多种密态计算虚拟设备供灵活选择，提供丰富的联邦学习算法和差分隐私机制

开源项目：

github.com/secretflow

gitee.com/secretflow

导语：NeurIPS（ Conference on Neural Information Processing Systems）是人工智能领域的顶级会议之一，每年吸引全球众多学者和研究人员参与。NeurIPS 2024 将于当地时间 2024 年 12 月 9 日至 15 日在加拿大温哥华举行。蚂蚁密算科技隐语团队与上海交通大学合作论文《Nimbus: Secure and Efficient Two-Party Inference for Transformers》在本次投稿的 15,671 篇有效论文投稿且最终录取率为 25.8% 中脱颖而出，顺利中选。

本论文提出了一种新的两方密态推理框架 Nimbus，专门用于 Transformer 神经网络的隐私保护推理，该框架在保护模型和用户数据隐私的前提下，实现了高效的推理过程，为大模型的隐私推理场景提供了重要技术支持。本文将带来本篇论文的技术解读，一起来关注！

论文亮点

Transformer 神经网络在各领域上展现出了惊人的效果，也是最热门大模型的结构基础，在众多任务上有潜在的实际应用。为了解决伴随而来的隐私问题，本论文提出了新的两方隐私推理框架 Nimbus，为 Transformer 神经网络中线性层的矩阵乘法及非线性层的激活函数提出了针对性的加速。下面我们将分别简单介绍这两部分具体的技术。

线性层：基于外积编码的高效矩阵乘法协议

对于 Transformer 网络的线性层中大量出现的矩阵乘法，现有的两方推理大多采用同态加密进行计算。线性层的时间主要由三部分组成：

加解密以及矩阵乘法本身的计算；
输入密文通信；
输出密文通信。

如何将矩阵编码到同态明文/密文多项式中对计算和通信效率有着重要影响。之前的工作，如 BOLT（S&P 2024）采用的 SIMD 编码和 BumbleBee（NDSS 2024）使用的系数编码，在设计时需要兼顾输入通信，输出通信和计算，三者的互相影响使每一部分都无法做到高效。Nimbus 的矩阵乘法协议重构了两方的乘法模式，利用参数的静态性完全消除了输入密文的通信！在消除输入密文通信的约束后，我们提出了一种新的编码方式，通过外积的方式实现矩阵乘法。进一步达到了更高效的计算和输出密文通信，极大地提升了矩阵乘法的效率。矩阵乘法协议示意图如下：

非线性层：输入分布感知的分段多项式拟合

对于 Transformer 网络中的激活函数 softmax (exponential) 和 GELU，一种主流的密态计算方式是通过分段多项式来拟合其非线性运算。现有的方法通过最小化分段多项式与原始激活函数之间的拟合误差来维持模型精度。这种做法隐含了激活函数的输入分布是均匀的假设。本研究认为这一假设引入了冗余的拟合难度，导致现有工作需要采用更多的分段和更高阶的多项式来避免精度损失。

Nimbus 提出了一项观察，Transformer 网络中激活函数的输入呈现出规律性的非均匀分布。比如下图中，GELU 函数的输入有 85%都是小于零的，exponentia l 函数的输入 75%都集中在[-4,0]区间。因此，在对激活函数进行分段与拟合时，Nimbus 针对那些出现概率较高的区间给予更精细的拟合，而对于低概率出现的区间，则允许其存在一定程度上的拟合误差。用这样的方式，相比于现有方法将输入视为均匀分布来直接拟合激活函数的方法，Nimbus 可以在保持相同模型精度的情况下，有效减少多项式拟合的分段数量和最高阶数，以减少密文运算的数量和通信量。

实验结果

Nimbus 基于 Secretflow-SPU 框架实现了 Transformer 模型的两方密态推理。我们在不同尺寸的 Transformer 网络以及不同输入句长上进行了大量实验。在保持模型精度不变的前提下，Nimbus 相比于最新的工作 BumbleBee（NDSS 2024）可以实现 2.7 到 4.7 倍的加速。其中，矩阵乘法的加速更加显著，为 2.9 到 12.5 倍；激活函数的加速为 2.9 到 4.0 倍。

SPU GitHub：https://github.com/secretflow/spu