基于MPC的快速transformer安全推理框架

论文：一种基于安全多方计算的快速Transformer安全推理方案-刘伟欣

摘要

1. 数据隐私泄露问题：当前Transformer推理应用中用户的数据会被泄露给模型提供方
2. 安全推理方法：基于MPC实现Transformer模型的安全推理
   1. 问题：巨大的计算和通信开销
3. 文本贡献：一是针对开销比较的Softmax注意力机制，提出2种安全友好的注意力机制 ；二是更换Transformer框架种的Softmax，并替换激活函数GeLU，集合知识蒸馏技术，提出一个友好安全的Transformer转换框架；三是实验结果表明：本文框架推理在保证安全性的同时，效率也得到提升。

背景

在Transformer应用中，通常有三类参与者：

1. 数据提供方
2. 模型提供方
3. 模型使用方

同时也存在相应的隐私保护问题：

1. 数据提供方：提供的数据可能涉及个人敏感信息等
2. 模型提供方：训练好的模型参数可能会被泄露，且可能根据模型参数复原训练数据等
3. 模型使用方：输入的数据可能包含个人隐私等
   
   前期有研究基础：隐私保护机器学习 ，通过HE、MPC、DP、TEE等技术对模型进行隐私保护，可以保证模型在训练过程中保护数据提供方的原始数据不会泄露给模型提供方，推理过程中模型提供方的模型参数和用户输入输出数据不会泄露。

若对Transformer模型直接使用以上隐私保护方法，开销极大。

1. 如输入长度为512的序列，一个12层的Bert-Base模型在MPC系统中进行安全推理需要59s，而明文下推理只需要不到1s。

机器学习模型可划分为线性函数和非线性函数，其中又以非线性函数处理比较复杂，一般有2个思路：

1. 近似方法。通过函数拟合得到一个分段多项式，将其作为非线性函数的替换，如ABY3等
2. 启发式方法 。使用具有与原函数相似性质的MPC友好函数替换，如使用Quad替换GeLU，如MPCFormer、SecureML等
   1. 此种方式计算效率会更高效，但会导致准确率下降，所以后续需微调，以保证可用。

针对思路2，主要研究方向是将Transformer安全推理中开销比较大的函数替换为MPC友好的函数，再通过知识蒸馏等技术提升准确率。

1. 在Bert-Base模型基于MPC技术的安全推理中，Softmax和GeLU占据大量运行时间。如输入长度为512的序列，Softmax和GeLU分别耗时占总时长的67.8%和18.6%，共占总时长的86.4%。
2. 所以有相关研究：
   1. Softmax采用2ReLU Attention、2Quad Attention、Scaling Attention、Linformer Attention替换；GeLU采用平方函数Quad替换。
3. 但目前的研究提出的替换函数方法没有达到效率和准确率的平衡。
   1. Scaling Attention 等替换方法效率较高，但准确率下降较大；2ReLU Attention 和 2Quad Attention 准确率较高，但同时也导致了较大的开销。

本文在此基础上，一是提出两种新的MPC友好注意力机制，可以在不影响准确率的前提下减少计算量；二是与知识蒸馏结合，提出俺去那友好Transformer框架；

相关工作

函数近似方法。

1. 2022 年 Hao 等人提出的Iron 基于同态加密和秘密分享技术实现各函数的隐私保护协议，从而构建了Transformer 安全推理框架，对于其中的矩阵乘法，在 Cheetah中基于同态加密的方法的基础上进行优化，达到了目前最优的开销。Iron 将非线性函数拆解为若干基础原语，并调用SIRNN进行实现。
2. 2023年 Zheng等人提出的Primer 直接将Transformer 的非线性算子使用混淆电路进行实现。
3. 同年PUMA和Privformerl分别基于ABY3和Falcon框架实现了Transformer 的安全推理，其主要工作是针对 Transformer 中独有的算子利用对应框架提供的协议设计对应的隐私保护协议。
4. 2023年文献对于GeLU进行分段拟合，使用 SIRNN 中的多项技术来优化近似多项式的安全计算。
5. 同年，文献基于函数秘密分享技术构造了GeLU、Softmax 等非线性函数的安全计算协议，大大提升了安全推理的效率。

启发式方法。

1. 2022 年 Chen 等人提出的THE-X 对于非线性部分使用 ReLU函数替换激活函数 GeLU，使用一个3层全连接网络近似 Softmax,但其中 ReLU的计算则由数据持有者在明文上计算,从而导致中间结果泄露给数据持有者，进而可能导致模型持有者的模型参数相关信息遭到泄露．同时THE-X 使用知识蒸馏技术对模型进行了微调，提升了模型准确率。
2. 同年，Li 等人提出 MPCFormer,其主要思想是首先将 Transformer 中开销较大的函数（即Softmax 和GeLU）替换为 MPC 友好的函数 ，之后使用知识蒸馏技术对模型进行微调 ，提升模型准确率，最后使用现有 MIPC 技术实现更新后的模型的安全推理 。
   1. 通过实验，MPCFormer 得出在GLUE数据集下分别将注意力机制和激活函数GeLU 替换为 MPC 友好的 2Quad Attention 注意力机制和平方函数 Quad，在准确率下降幅度不大的前提下获得 2.2× 的安全推理速度提升。
3. MPCViT的思路与 MPCFormer 略有不同，其指出在Vision Transformer（ViT）中并不是所有的注意力机制都同样重要，于是首先分析了多种注意力机制的效率与准确率（效率较高的注意力机制通常具有较低的准确率)，先将所有的 Softmax 注意力机制替换为一种准确率相对较高但效率相对较低的 MPC 友好的注意力机制 2ReLU Attention，之后使用神经架构搜索（neural architecture search,NAS）技术将部分 2ReLUAttention 注意力机制替换为准确率相对较低但效率相对较高的注意力机制 ScalingAttention ，保证在不过多降低准确率的前提下尽可能达到最高的计算效率，最后再使用知识蒸馏技术提升模型准确率。

预备知识

Transformer

Transformer采用的是编码器-解码器结构，Bert模型采用的是编码器结构，下面主要介绍编码器结构的Transformer。

Transformer的输入先经过Embedding层，输出词向量X，之后编码器包含多个Transformer块。其中每个Transformer块包含3部分：

1. 注意力机制（Attention）
2. 前馈神经网络（feed forward network，FFN）
3. 2个归一化层（LN）

注意力机制（Attention）：有助于模型关注输入序列中的不同部分。

多头注意力机制（multi-head attention，MHA）将注意力机制扩展为H个并行的注意力层。

前馈神经网络（FNN）：包含2个线性层和中间的一个激活函数。

归一化层（LN）：需要计算输入的均值和方差。

知识蒸馏

知识蒸馏是模型压缩技术中的一种，用于解决小模型训练困难的问题。

1. 原理是：允许小模型（学生模型）去学习大模型（教师模型）的知识，即学生模型在训练过程中会去模仿教师模型的行为，教师模型的行为能提供比训练数据标签更多的信息，学生模型可以通过学习这些信息获得教师模型的泛化能力。
2. 模型的行为可以用行为函数描述。
3. 在Transformer模型中，FNN层或者中间层都可以看作是行为函数。
   

Transformer安全推理框架

本文使用的安全推理框架，分为3部分。

1. 函数替换
2. 知识蒸馏：以无近似模型为教师模型，替换后的模型为学生模型，使用知识蒸馏技术对替换后的模型进行微调
   1. 前2部分组成了MPC友好的Transformer转换框架
3. 安全推理

MPC友好的Transformer转换框架

概述

本节提出一种MPC友好的Transformer转换框架，该框架可将预训练的Transformer 模型转换为 MPC 友好的 Transformer 模型，同时保证模型准确率不受影响，并且提高安全推理的推理速度，从而可用于 Transformer 安全推理。

框架分为：函数替换和知识蒸馏2个阶段。

1. 由于函数替换改变了注意力机制和激活函数的功能，直接使用原始模型会导致准确率降低。
2. 知识蒸馏以原模型为教师模型，以替换后的模型为学生模型，进行微调 ，使其更新模型参数与其新的注意力机制和激活函数适配，从而提高准确率。
   

函数替换

基于MPC的transformer模型推理计算主要耗时在softmax和GeLU上。

1. softmax注意力机制包含3种线性运算，可以通过MPC协议安全计算
   1. 指数运算
   2. 比较运算
   3. 倒数运算
2. 本文采用2Quad Attention替换方案去除其中的指数运算和比较运算，实验结果证明在SST-2任务上有较好的准确率。
3. 对于倒数运算，考虑使用一个可以在明文下计算的函数计算代替倒数运算，减少倒数运算的开销。【唯一的创新？】
4. 对于激活函数GeLU，直接使用文献【22】现有的替换函数。

知识蒸馏

本文使用无近似模型 M 作为教师模型 ，近似后的模型M'作为学生模型，在下游任务数据集上进行知识蒸馏．初始化时学生模型M'使用与教师模型相同的模型参数，二者仅在使用的函数上不同。

蒸馏任务分为2步：

1. 学生模型主要学习教师模型在Transformer块的行为
2. 学生模型主要学习教师模型在最后预测层的行为

MPC友好的函数替换方案

注意力机制的替换

常用的替换方法是：

1. 将 Softmax 表达式中的指数函数替换为 ReLU的2ReL UAttention
2. 将指数函数替换为平方函数的2Quad Attention
3. 将整个注意力机制替换为矩阵乘法的 Scaling Attention。
   

softmax注意力机制的3个重要特性：

1. 单调性
2. 非负性
3. 求和为1

主要思想是将注意力机制中的指数函数替换为计算较为简单的其他函数。

基于MPC的transformer安全推理的常规套路是：指数函数通常要将其近似为分段多项式，再使用不经意分段多项式求值，需调用多次安全比较协议个安全乘法协议。

而在本文中将指数函数替换为ReLU或者平方函数分别只需要1次安全比较协议和1次乘法协议 ，最后再通过将每一项除以所有项求和的方式满足求和为1的要求。

其中最后一项"安全除以所有项求和"开销比较大，对于规模为n的张量，使用MPC技术计算时，需要执行n次秘密值之间的除法。

解决思路：

1. 由于MPC中的除法计算开销远大于加法运算，所以可以先计算所有项求和的倒数，再计算n次乘法，共计需1次除法和n次乘法的开销。
2. 若用一个常数近似分母（即各项的求和），将注意力机制的第3个特性【求和为1】放宽为求和近似为1，则可在保证准确率的前提下去避免倒数运算，减少开销。
3. 如何选择这个近似的常数？

解决思路：使用预训练的transformer模型对数据集进行多次推理，将每次推理过程中不同输入序列长度下的分母均值进行记录，利用这些数据可以【拟合】得到一个关于输入序列长度的函数f。

1. 即在给定输入x的情况下，使用f(|x|)作为分母近似常数（这不就成明文和密文混合计算了么，分子是密文，分母为明文？）

具体做法为使用预训练的transformer模型对GLUE数据集进行多次推理：

上图表明分母均值并不随输入序列长度的增加而线性变化，本文使用幂函数\(f_{softmax}(x)=ax^b\)对其进行拟合，最终得到\(f_{softmax}(x)=0.9923937666772096x^{0.333253864762922}\)。

用f(x)进行替换分母，得到：

此时映射关系f不依赖于推理阶段的具体输入，可视为模型参数的一部分。

该思想还可以与其他注意力机制结合以获得更好的效果，如2Quad Attention结合，需要做如下修改：

1. 根据预训练的transformer模型在2Quad Attention注意力机制下测试不同输入序列长度时的【分母值】，拟合得到映射函数f(x)。
2. 将映射函与2Quad Attention结合：
   

其中a=84.07373758071103, b=0.7174745255779887,c=5。

下面是本文提出的两种安全友好型的注意力机制在ABY3框架下的计算延迟和加速比：

激活函数GeLU的替换

以往方案采取的替换方式有2种：

1. 将GeLU替换为ReLU
2. 将GeLU替换为平方函数Quad

还有更加精确的方案：对其分段多项式进行精确近似，在准确率几乎不受影响的同时导致较低的计算效率。

作者说无需这么做，所以只选择了以往的2种替换方案。

1. 使用Quad相比于ReLU会导致准确率表现差些

实验与结果

实验环境

1. 显卡P100 16GB
2. MPC协议使用隐语的SPU框架中的半诚实安全的ABY3协议进行安全推理。
3. 模型结构选用Bert-Base-Cased
4. 数据集为GLUE
5. 测试任务：SST-2、STS-B、RTE
6. 输入序列长度为128
7. 知识蒸馏设置
   

实验结果

可以看出，在前面4种注意力机制中，速度提升方面表现最好的是 Scaling Attention,其与激活函数的Quad替换方法组合达到了2.26×的速度提升,但其准确率大幅下降。

而对于本文提出的两种注意力机制，可以看出对于 2Quad freeDiv Attention + Quad的组合在速度提升方面与目前效率提升最明显的Scaling Attention 相当，且在准确率方面表现良好，相比于原模型没有准确率的下降，因此在效率和准确率方面达到了目前所有方案中的最优。

可以看出，使用本文提出的MPC友好的Transformer转换框架转换后的

模型在安全推理速度上达到了 2.26×的提升，没有准确率损失，达到了目前已知方案中的最优。

可以看出，与基准值（Softmax Attention +GeLU）相比，本文提出的freeDiv Attention应用于不同任务上均有较好的准确率表现。

对于单句任务STS-2，freeDiv Attention 至多达到了 91.6%（+0.1%)的准确率；对于相似性任务STS-B，freeDiv Attention的相关系数至多达到了89.0%（+0.3%)；对于释义任务 RTE，freeDiv Attention 至多达到了 61.4%（-1.8%)的准确率，因此本文提出的 freeDiv Attention 在不同任务下具有一定泛化性。

可以看出，对于 SST-2 任务，知识蒸馏过程至多可以提高 42.4%的准确率；对于 STS-B 任务，知识蒸馏过程至多可以提高90.2%的相关系数；对于RTE 任务，知识蒸馏过程至多可以提高12.6%的准确率，因此知识蒸馏过程可以有效提高转换后模型准确率 。同时本文给出了提出的MPC友好的Transformer 转换框架对不同任务的 Transformer 模型进行转换所需时间，即知识蒸馏步骤所需时间。

可以看出，与本文使用同一安全推理框架（函数替换一知识蒸馏一安全推理）的 MPCFormer 方案对比，本文的函数替换方案实现的Transformer安全推理在速度提升上优于MPCFormer。

本文的创新点：提出若用一个常数近似分母（即各项的求和），将注意力机制的第3个特性【求和为1】放宽为求和近似为1，则可在保证准确率的前提下去避免倒数运算，减少开销。