基于CKKS的非交互式安全Transformer推理实现

Secure Transformer Inference Made Non-interactive

本文介绍了如何使用CKKS来计算transformer推理的每个部分。同时给出了一系列优化算法。主要涉及到的计算算法有以下几种: 密文的压缩与分解技术、SIMD槽折叠技术、Sgn()、QuickSum、QuickMax、密文-明文矩阵相乘、密文-密文矩阵相乘法、Softmax算法、归一化、GELU函数、Argmax函数等等。其中密文的压缩与分解技术，和SIMD槽折叠技术是本文的核心创新算法。

Abstract

随着ChatGPT的普及，安全transformer推理已经成为一个突出了研究主题。已有的解决方法通常是交互式的，涉及到客户端和服务端之间大量的通信负载和交互轮次。

本文提出NEXUS，这是第一个用于安全transformer推理的非交互式协议，其中客户端仅需要提交一个加密输入，然后等待来自服务器的加密结果即可。NEXUS的核心是两个创新的技术：SIMD密文压缩和分解技术，以及SIMD槽折叠技术。此外，同24年的另外一个解决方案相比，本方法达到了2.8倍的加速，且减少了368.6倍的带宽消耗。

1 Introduction

Transformers,例如GPT和BERT，已经彻底改变了AI领域。Transformer擅长于广泛领域的应用，比如语言翻译，内容生成以及问题回答。然而这些应用总是涉及到敏感数据，从而导致越来越多地关于用户隐私的担忧。例，OpenAI开发的ChatGPT作为一种在线推理服务，以及为开发人员提供的远程API，其中使用者通过提交prompts或者消息可以很容易地访问这些服务。尽管这些方法是方便的，但是由于使用者提交的数据可能包含敏感信息，故而造成了严重的隐私风险。

Secure inference是一种两方密码协议，该协议使模型推理以如下方式处理运行，即服务器S不会了解到关于客户C提交的输入的任何信息，且C不会了解到关于S的模型的任何信息，仅仅能得到最终的推理结果。

该协议大多被设计于安全CNNs推$[2,27,30,36]$，最近的许多工作也支持基于Transformer的模型$[10,24,26,35,38,40]$​，值得注意的是，这些安全Transformer模型大多都是交互式的，因此会导致巨大的通信开销和交互轮次，这里我们必须强调非交互式安全Transformer推理的重要性。

本文贡献:

本文中，我们提出了NEXUS，第一个secure transformer inference的非交互协议。通过NEXUS，C使用RNS-CKKS加密输入，S对FHE加密数据执行transformer。CKKS的SIMD技术被应用于批处理$N=2^{15}$个数据，多项式近似可以用于处理非线性函数，比如GELU，softmax，层归一化和argmax。

NEXUS不需要对模型进行任何重训练与微调，且为了提高NEXUS的效率，我们提出了两种新颖的且基础的技术。

• SIMD密文压缩与分解：该技术可以将2N个SIMD密文压缩为一个密文，然后可以使用4N个密文—明文乘法和替换将其解压回来。该技术可以大大减少客户端和服务器之间传输的密文数量，而不会为后续计算带来任何额外的开销。

• SIMD槽折叠：在所有SIMD槽中计算关联函数f(),例如sum和max。结果值会自动的填充SIMD密文的槽，允许将其应用于原始密文的每个槽。

本文贡献总结如下：

• secure transformer inference的第一个非交互协议
• 用于密文打包的SIMD密文压缩与分解技术
• SIMD槽折叠技术，以高效操作SIMD密文的槽
• 综合的实现与评估

2 Preliminaries

符号系统描述如下:

Notation	Description	Notation	Description
C	client	S	server
$E(\ast )$	encryption	$\pi (\ast )$	encoding
$Enc(\ast )$	encoding+encryption	$\tilde{a}$	FHE ciphertext
$RotL(\ast )/RotR(\ast )$	左旋转和右旋转	$Subs(\ast )$	替换操作
$Sgn(\ast )$	sign操作	$L$	乘法深度
$N^{\prime }$	CKKS的环维数	$N$	$N=N^{\prime }/2$
$A$	输入矩阵	$W$	权重矩阵

2.1 安全推理和威胁模型

安全推理是一个两方密码学协议，其可以在C和S之间进行模型推理，与此同时还可以保护两个参与方输入隐私。它的正式定义如下:

Definition 1：

针对两方参与者，其中$S$持有模型$M$,且$C$持有输入$A$的协议$\Pi$是安全推理协议，当且仅当以下条件满足时：

(1) 正确性: 该协议的最终输出是正确的推理结果$M(A)$。

(2) 安全性:

$Vie{w}_C^{\Pi }{\approx }_{c}Si{m}_C(A,out)$，其中$Vie{w}_C^{\Pi }$表示协议$\Pi$执行期间$C$的视角，$out$表示推理的结果。

$Vie{w}_S^{\Pi }{\approx }_{S}Si{m}_S(M)$，其中$Vie{w}_S^{\Pi }$表示协议$\Pi$执行期间$S$​的视角。

$Si{m}_{\ast }$可以理解为理想状态下希望实体$\ast$可以得到的信息。

假设$C$或$S$为半诚实对手，其在遵守协议规范的同时也尽可能的在执行过程中手机额外的信息。且假设对手在计算上是有限的。

2.2 Transformer

这里简单介绍一下Transformerd。

图1是transformer的结构与工作流程。它将一个表示为矩阵的嵌入传递给注意层和前馈神经网络，最后根据最终对数最大值输出一个选择向量，且，LayerNorm层被应用于每个块之后。

transformer的结构和工作流程

Attention：

使用三个矩阵($W_Q\in {\mathbb{R}}^{n\times k},W_K\in {\mathbb{R}}^{n\times k},W_V\in {\mathbb{R}}^{n\times k}$)乘嵌入矩阵$A\in {\mathbb{R}}^{m\times n}$，生成一个query矩阵$Q=A\cdot W_Q$，一个key矩阵$K=A\cdot W_K$和一个value矩阵$V=A\cdot W_V$。即对于Attention层的单元，transformer会学习到三个权重矩阵。

attention可以被表示为：
$$Attention(Q,K,V)=Softmax(\frac{Q{K}^T}{\sqrt{k}})\cdot V$$

Layer normalization

该层的输入为$a\in {\mathbb{R}}^n$,均值和标准差分别为$\mu$和$\sigma$，则该层的输出$y\in {\mathbb{R}}^n$可以表示为:
$$y_i=\gamma \cdot \frac{x_i-\mu }{\sigma }+\beta$$
其中，$\gamma ,\beta \in \mathbb{R}$​是两个超参数。

Feed-forward

全连接前馈网络层包含两个线性变换以及一个GELU激活函数：
$$FeedForward(X)=GELU(X{W}_1+b_1)\cdot W_2+b_2$$
其中GELU函数计算如下:
$$GELU(x)=\frac{1}{2}x\cdot (1+erf(\frac{x}{\sqrt{2}}))$$
式中，高斯误差函数为$erf(x)=\frac{2}{\sqrt{\pi }}{\int }_0^x{e}^{-t^2}dt$​。由于其良好的曲率和非单调性，它被用作激活函数。

Argmax

根据最终对数最大值输出一个选择向量

可以看到，只要我们能够使用FHE实现各个层的计算，就可以实现一个安全Transformer。

2.3 Fully Homomorphic Encryption

FHE可以对加密数据执行任意操作，故FHE是使得我们构建非交互式安全transformer推理得主要工具。RNS-CKKS属于级全同态加密，其可以支持L级深度的乘法。RNS-CKKS的明文和密文均是多项式环$R_Q={\mathbb{Z}}_Q[X]/(X^{N^{\prime }}+1)$上的元素。其中$Q={\Pi }_{i=0}^L{q}_i$，且$q_i$​之间互素。若密文的级别变得太低，则可以运行自举操作来刷新密文到高的级别，以允许更多的计算。

简单地说，自举即利用自同构$R_{q_0}\cong R_{q_0}\times R_{q_1}\times ...\times R_{q_L}$，来将密文模从$q_0$提升到$q_L$，以及对密文同态评估解密电路。若自举本身消耗K个级别，则刷新后的密文支持$L-K$个级深度的计算。

RNS-CKKS支持SIMD操作，其可以加密向量$a\in {\mathbb{R}}^N$到一个密文中，且批处理这些加密元素，而不引入其他操作。为了以SIMD格式加密，首先使用编码算法$\pi (\ast )$将向量$a$编码为一个$R_Q$上的多项式，然后使用加密算法$E(\ast )$​加密该多项式。

在整篇文章中，我们使用$E(\ast )$表示加密多项式，使用$Enc(\ast )$表示以SIMD格式加密向量，即$Enc(a)=E(\pi (a))$，其中$a$是一个向量。

---

一个特殊的FHE操作:

$c{t}^{\prime }\leftarrow Subs(ct,k)$:替换操作，该操作以密文$ct=E(p(x))$以及一个奇整数$k$作为输入，然后得到新的密文$c{t}^{\prime }=E(p(x^k))$​​​​。

这里的$Subs(ct,k)$应该是一种密钥交换操作，可以描述如下:

已知密文: $ct=(-a(x)s(x)+e(x)+p(x),a(x))$

将该密文进行自同构操作: ${\kappa }_k(ct)=(-a(x^k)s(x^k)+e(x^k)+p(x^k),a(x^k))$。

然后得到用户提供的交换密钥: $key=(-a(x)s(x)+e(x)+P\cdot s(x^k),a(x))$

然后执行密钥交换操作: $c{t}^{\prime }=({\kappa }_k(ct)[0],0)+(\lfloor P^{-1}\cdot {\kappa }_k(ct)[1]\cdot key\rceil )$

此时新的密文即$c{t}^{\prime }=(-a(x)s(x)+e(x)+p(x^k),a(x))$

注意，这里的$a(x),e(x)$是变化的，也就是不同的密文中，这是不同的。

---

2.4 Homomorphic sign function

由于FHE仅支持线性函数，所以为了实现在FHE下对加密数据的比较，本文需要利用sign函数的多项式近似，即：
$$sign(x)=f^{d_f}(g^{d_g}(x))=\{\begin{array}{ll}-1& (-1\le x\le -2^{-\alpha })\\ 0& (x=0)\\ 1& (2^{-\alpha }\le x\le 1)\end{array}$$
其中，$f(),g()$为两个多项式，$d_f,d_g$为这两个多项式重复的次数。注意，该多项式近似要求输入x取值范围为[-1,1]。因此，对任何输入$a\in [a_{min},a_{max}]$都需要进行归一化处理:
x : = a / m a x { ∣ a m a x ∣ , ∣ a m i n ∣ } x := a/max\{|a_{max}|,|a_{min}|\} x:=a/max{∣amax​∣,∣amin​∣}
这里，我们使用Sgn()表示在SIMD密文上同时运行归一化与sign近似函数：
b ~ ← S g n ( a ~ ) : b i = f d f ( g d g ( a i m a x { ∣ a m a x ∣ , ∣ a m i n ∣ } ) )    ∀ i ∈ [ N ] \widetilde b\leftarrow Sgn(\widetilde a): b_i =f^{d_f}(g^{d_g}({a_i\over{max\{|a_{max}|,|a_{min}|\}}})) \ \ \forall i \in [N] b ←Sgn(a ):bi​=fdf​(gdg​(max{∣amax​∣,∣amin​∣}ai​​))  ∀i∈[N]

在本文的实现中，使用的是9次的$f(\ast )$和$g(\ast )$，且设计$\alpha =16,d_f=2,d_g=2$​​，然后使用BSGS算法来评估多项式。

3 Basic design

本节介绍NEXUS的基础设计，即在不优化的情况下实现上述transformer的每一层计算，在之后的章节中会对本节的算法进行优化。

3.1 Attention

3.1.1 Matrix multiplication(ciphertext-plaintext)

在Attention层的第一个MatMul步骤，我们需要计算三个密文—明文矩阵乘法:
$$\begin{gathered} Q:=A\cdot W_Q; \\ K:=A\cdot W_K; \\ V:=A\cdot W_V; \end{gathered}$$

其中A是我们的输入，$W_Q,W_K,W_V$是三个给定矩阵，下面以$A\cdot W_Q$为例来描述这个密文—明文矩阵乘法，该过程同样适用于$W_K$和$W_V$。

给定矩阵$A\in {\mathbb{R}}^{m\times n}$和矩阵$W_Q\in {\mathbb{R}}^{n\times k}$，计算矩阵$Q:=A\cdot W_Q$。

设$a_{i,j}\in \mathbb{R}$表示矩阵A的第i行第j列的元素，$w_j\in {\mathbb{R}}^k$表示矩阵$W_Q$的第j行的元素向量，$q_i\in {\mathbb{R}}^k$是矩阵$Q$的第i行的元素向量，即:
$$q_i=\sum _{j\in [n]}{a}_{i,j}\cdot w_j$$
因此，上述过程可以描述为，C将A中的每个元素$a_{i,j}$均单独加密为密文发送给S，然后S同态评估MatrixMul,一个演示的示例如下:

图2 SIMD-based matrix multiplication

---

在上述描述中，C需要发送$m\times n$个密文给S，从某一方面来说，这种开销是比较大的，因此本文在第4节提出一种算法可以将如此类型的$m\times n$个密文压缩为$\frac{m\times n}{N^{\prime }}$个密文，即一个密文中存放$N^{\prime }$个元素，随后S可以将压缩后的密文恢复为压缩前的密文形式。

---

3.1.2 Matrix multiplication(ciphertext-ciphertext)

经过上述步骤后，可以获得加密的$(Q,K,V)$，在Attention的第二个MatMul块，S需要计算$Q\cdot K^T$。很明显，现在Q的每一行和$K^T$的每一列已经以SIMD的形式加密为$Enc(q),Enc(k^T)$。如果S可以计算$Enc(q)$和$Enc(k^T)$的内积，则可以获得$Q\cdot K^T$的加密结果。

由于SIMD，S可以很容易的计算得到Enc(u),其中$u=[u_0,...,u_{k-1}]$是q和$k^T$的元素级的乘法，现在为了计算内积，S仅仅需要在SIMD下计算$s:={\sum }_{i=0}^{k-1}{u}_i$。

---

为了计算这个和，我们可以通过k-1次的旋转及加和来计算，从而获得密文Enc([s,s,…,s])，但是本文提出了’QuickSum’算法,该算法仅仅需要logk次旋转就可达到这个目标。'QuickSum’算法在第5节介绍。

---

进一步，S将计算得到每一行的的m个密文组合到单一密文中，计算方法如下:
$$\sum _{i=0}^{m-1}(Enc(s_i,s_i,...,s_i)\cdot b_i)$$
其中$b_i$仅在第i个槽的位置是1，其余槽均为0。

易知，输出矩阵为$A\in {\mathbb{R}}^{m\times m}$，其中A的每行向量以SIMD形式加密，将该结果作为Softmax的输入。

3.1.3 Softmax

Softmax函数需要被应用于A的每一行，该函数评估如下:
$$y_i=\frac{exp(a_i-a_{max})}{{\sum }_{j=0}^{m-1}exp(a_j-a_{max})}\text{\hspace{1em}(1)}$$
其中$a_{max}=max(a_0,...,a_{m-1})$，从而确保指数函数的每个输入$(a_j-a_{max})$​是非正数，保证稳定性。

---

本文提出了’QuickMax’算法，该算法以$Enc([a_0,...,a_{m-1}])$为输入，并输出$Enc([a_{max},...,a_{max}])$​，且，该算法仅需要logm-1次Sgn操作与logm次旋转操作。该算法描述在第5节。

---

给定$Enc([a_0,...,a_{m-1}])$和$Enc([a_{max},...,a_{max}])$。

S进行如下步骤计算:
$$Enc([a_0^{\prime },...,a_{m-1}^{\prime }])=Enc([a_0,...,a_{m-1}])-Enc([a_{max},...,a_{max}])$$
然后根据如下公式计算指数函数，这里使用泰勒展开:
$$exp(x)\approx (1+\frac{x}{2^r}{)}^{2^r},x\le 0$$
其中$r=6$，此时平均误差被限制在$10^{-5}$，即S以SIMD格式计算指数函数:
$$Enc(e_0,...,e_{m-1})=exp(Enc([a_0^{\prime },...,a_{m-1}^{\prime }]))$$
很明显，这里$e_j=exp(a_j^{\prime })$。

接下来，S应用$QuickSum(\ast )$算法来获得$Enc([{\sum }_{j=0}^{m-1}{e}_j,...,{\sum }_{j=0}^{m-1}{e}_j])$​。

进一步的，S使用文献[21,24]中的Goldschmidt除法算法来计算：
$$Enc(y_0,...,y_{m-1})=\frac{Enc(e_0,...,e_{m-1})}{Enc([{\sum }_{j=0}^{m-1}{e}_j,...,{\sum }_{j=0}^{m-1}{e}_j])}$$
Softmax算法的详细描述如算法1所示:

3.1.4 Matrix multiplication(ciphertext-ciphertext)

这里是Attention的最后一个MatMul块，该块的计算原理同3.1.2节完全一致。

3.2 Layer normalization

本文的归一化表示如下(但是不太清楚这个归一化使用的是什么计算公式)：

3.3 Feed forward

前馈网络层涉及到两个矩阵乘法以及一个GELU。矩阵乘法如上文所述来计算。GELU可以使用下述分段多项式来近似，当输入$x\in [-60,60]$，则可以确保误差在$10^{-3}$内。
$$GELU(x)=\in \{\begin{array}{ll}0& (x\le -4)\\ P(x)={\sum }_{i=0}^{i=3}{c}_i{x}^i& (-4<x\le -1.95)\\ Q(x)={\sum }_{i=0}^{i=6}{d}_i{x}^i& (-1.95<x\le 3)\\ x& (x>3)\end{array}$$
首先，使用Sgn操作获得四个加密bit：$b_0,b_1,b_2,b_3$，当且仅当输入x属于第i段时，$b_i=1$，否则$b_i=0$，如此，GELU(x)函数可以表示为:$GELU(x):=b_0\cdot 0+b_1\cdot P(x)+b_2\cdot Q(x)+b_3\cdot x$​。

完整的Secure GELU算法可以表示如下:

3.4 Argmax

transformer最终的输出应该是一个选择向量$Enc([b_0,...,b_{m-1}])$，其中$b_i=1if{a}_i=max(a_0,...,a_{m-1})$，其他情况下$b_i=0$​。

因此，本文的Secure Argmax算法如下：

3.5 Placement of bootstrapping

由于bootstrapping操作是昂贵的，因此合理的放置bootstrapping的位置是至关重要的。

图4 Placement of bootstrapping for a BERT-base transformer

4. SIMD密文的压缩和分解

假设C想要发送N’个密文给S，且每个密文以SIMD方式加密N个相同的值，Enc([$a_0,...,a_0$]),…,Enc([$a_{N^{\prime }-1},...,a_{N^{\prime }-1}$​​])。

---

SIMD密文的压缩算法

C将向量[$a_0,a_1,...,a_{N^{\prime }-1}$]的各个元素打包到一个多项式的系数中,即：
$$p(x)=a_0+a_{1}x+a_2{x}^2+...+a_{N^{\prime }-1}{x}^{N^{\prime }-1}$$
然后将该多项式加密${\tilde{p}}_0=E(p(x))$发送给S。

然后S可以对密文${\tilde{p}}_0$分解从而得到压缩前$N^{\prime }$个SIMD密文。

---

S分解密文${\tilde{p}}_0$过程如下：

---

SIMD密文的分解算法：

（1）执行$Subs({\tilde{p}}_0,N^{\prime }+1)$返回：
$$\begin{gathered} E(a_0+a_1{x}^{N^{\prime }+1}+a_2{x}^{(N^{\prime }+1{)}^2}+...+a_{N^{\prime }-1}{x}^{(N^{\prime }+1{)}^{N^{\prime }-1}}) \\ =E(a_0+a_1(-x)+a_2(-x{)}^2)+...+a_{N^{\prime }-1}(-x{)}^{N^{\prime }-1}) \end{gathered}$$
注意，$x^{N^{\prime }}+1\equiv 0(mod{x}^{N^{\prime }}+1)$,因此$x^{N^{\prime }+1}=x^{N^{\prime }}\ast x=-x(mod{x}^{N^{\prime }}+1)$​,这里的$N’$也就是分圆环的次数。

（2）执行${\tilde{p}}_0+Subs({\tilde{p}}_0,N^{\prime }+1)$​操作，移除p(x)的所有奇数项。
$$\begin{gathered} a_0+a_{1}x+a_2{x}^2+...+a_{N^{\prime }-1}{x}^{N^{\prime }+1} \\ +a_0+a_1(-x)+a_2(-x{)}^2)+...+a_{N^{\prime }-1}(-x{)}^{N^{\prime }-1} \\ =a_0+0x+a_2{x}^2+...+a_{N^{\prime }-2}{x}^{N^{\prime }-2}+0x^{N^{\prime }-1} \end{gathered}$$

（3）通过$log{N}^{\prime }$次$Subs()$操作，S可以提取得到密文：$E(a_0+0x^1+0x^2+...+0x^{N^{\prime }-1})$，实际上，这就是密文Enc([$a_0,a_0,...,a_0$])。完整的操作流程如下：

类似地，为了提取E($a_1+0x^1+...+0x^{N^{\prime }-1}$)，S应该左旋明文多项式p(x)一个单位，通过乘以$x^{-1}$，然后再次执行上述的提取过程。通过执行$N‘$次该提取过程，S可以获得向量[$a_0,a_1,...,a_{N^{\prime }-1}$​]中每个元素的单独SIMD格式加密。

然而上述过程需要执行$(N^{\prime }\cdot log{N}^{\prime })$次$Subs()$操作。对比之下，本文提出一种算法，可以实现相同的目标，但是仅需要$2N^{\prime }$次$Subs()$操作。该算法可以简单地描述如下：

---

算法5是Secure Decompression的详细描述:

下面提供上述分解操作的理论证明：

Theorem 1:

仅有常数项的多项式的加密E($a_s+0x^1+...+0x^{N^{\prime }-1}$)是向量[$a_s,a_s,...,a_s$]的加密Enc([$a_s,a_s,...,a_s$​])。

---

4.1 Application to matrix multiplication

压缩分解技术可以自然地应用于MatrixMul，此外，基于下面的观察结果，本文进一步优化了矩阵乘法，观察到在transformer推理过程中，对于不同输入的矩阵$A\in {\mathbb{R}}^{m\times n}$需要乘以相同的矩阵$W\in {\mathbb{R}}^{n\times k}$​。

设$A=[a_0,...,a_{n-1}]$，其中$a_i\in {\mathbb{R}}^m$表示矩阵$A$的第$i$行。假设S和C需要生成t个响应词，即有t个输入矩阵:
$$\begin{gathered} A_0=[a_{0,0},a_{0,1},...,a_{0,n-1}] \\ {A}_1=[a_{1,0},a_{1,1},...,a_{1,n-1}] \\ ... \\ {A}_0=[a_{t-1,0},a_{t-1,1},...,a_{t-1,n-1}] \end{gathered}$$
令$a_i^{\prime }=\left[\begin{array}{c}{a}_{0,i}\\ a_{1,i}\\ ...\\ a_{t-1,i}\end{array}\right]$和$q_j^{\prime }:={\sum }_{i=0}^{n-1}{a}_i^{\prime }{w}_{i,j}\forall j\in [k]$，则有
$$Q^{\prime }=q_0^{\prime }||q_1^{\prime }||...||q_{k-1}^{\prime }=\left[\begin{array}{c}{A}_{0}W\\ A_{1}W\\ ...\\ A_{t-1}W\end{array}\right]$$

---

---

预计算阶段:

这里，我们引入一个预计算阶段，其中S使用上述提到的密文压缩技术，将压缩后的密文$(En{c}_S([\underset{t\times m}{\underbrace{w_{i,j},w_{i,j},...,w_{i.j}}}])\forall i\in [n],j\in [k])$发送给C。注意，该传输仅只发生一次，除非模型发生改变。接下来，C对压缩的密文执行分解技术，以获得$En{c}_S([\underset{t\times m}{\underbrace{w_{i,j},w_{i,j},...,w_{i.j}}}])\forall i\in [n],j\in [k]$。在预计算阶段C并没有关于输入的信息，采样$U\in {\mathbb{R}}^{(tm)\times n}$，然后计算:
$$En{c}_S(v_j)\leftarrow \sum _{i=0}^{n-1}(u_i\times En{c}_S([w_{i,j},...,w_{i,j}]))\forall j\in [k]$$
其中$u_i$是矩阵$U$的第i列。接下来，C使用自己的密钥来加密$En{c}_S(v_j)$以获得$En{c}_C(En{c}_S(v_j))$，并将其发送给S。注意$En{c}_S(En{c}_C(v_j))=En{c}_C(En{c}_S(v_j))$，故S可以对其进行解密，从而获得$En{c}_C(v_j)$。注意，这里的$v_j$是矩阵$U\cdot W$的第$j$列。

---

切换不同用户加密密钥的过程计算如下:

给定$c{t}_S=(-a{s}_S+m+e)$，

使用$s_C$加密有$c{t}_{C,S}=(-a{s}_S-a{s}_C+m+e+e^{\prime },a)$，

使用$s_S$解密有: $c{t}_C=(-a{s}_S-a{s}_C+m+e+e^{\prime },a)+(a{s}_S,0)=(-a{s}_C+m+e+e^{\prime })$.

---

在线处理阶段:

此时，C知道输入的信息$A^{\prime }=a_0^{\prime }||a_1^{\prime }||...||a_{n-1}^{\prime }$，然后C将明文$(A^{\prime }-U)$发送给S，注意，由于S不知道U的值，故S也不清楚$A^{\prime }$的值，然后S可以计算:
$$\begin{gathered} (A^{\prime }-U)\cdot W+(En{c}_C(v_0)||En{c}_C(v_1)||...||En{c}_C(v_{k-1})) \\ =(A^{\prime }W-V)+(En{c}_C(v_0)||En{c}_C(v_1)||...||En{c}_C(v_{k-1})) \\ =(En{c}_C(q_0^{\prime })||En{c}_C(q_1^{\prime })||...||En{c}_C(q_{k-1}^{\prime })) \end{gathered}$$
其中$q_j^{\prime }$是矩阵$Q^{\prime }$的第$j$列。算法6描述了优化后的矩阵乘法细节:

可以注意到，只需要预计算的过程交互一次，此后C可以直接向S发送明文信息$A^{\prime }-U$，而不会泄露$A^{\prime }$的信息。

5 SIMD槽折叠算法

回想矩阵Q，K，V的行向量是使用SIMD方式加密的。而上述介绍的一系列操作，如内积，Softmax，LayerNorm和Argmax等， 均涉及到利用所有槽元素计算函数$f(\ast )$，并将得到的结果放置到所有槽上。例如给定$Enc([a_0,...,a_{N-1}])$，然后想要获得$Enc([s,...,s])$，其中$s={\sum }_{i=0}^{N-1}{a}_i$，此时$f(\ast )$即求和函数。

本节提供了一种通用的解决方案，只要函数$f(\ast )$满足:
$$f(f(a_0,a_1),a_2)=f(a_0,f(a_1,a_2))$$
算法7描述了槽折叠算法的细节:

这里是一个简单的例子，可以看到算法7的实现流程:

5.1 QuickSum

给定$[a_0,a_1,...,a_{n-1},0,...,0]$，为了获得$[{\sum }_{i=0}^{N-1}{a}_i,...,{\sum }_{i=0}^{N-1}{a}_i,0,...,0]$，可以将算法7的第5行替换为$\tilde{s}\leftarrow \tilde{s}+\tilde{a}$。

5.2 QuickMax

给定$[a_0,a_1,...,a_{n-1},0,...,0]$，为了获得$a_{max},...,a_{max},0,...,0$，其中$a_{max}=max(a_0,a_1,...,a_{n-1})$，很明显$max(a,b)$可以表示为:
$$max(a,b)=\frac{a+b+(a-b)\cdot Sgn(a-b)}{2}$$
因此可以将算法7的第5行替换为:
$$\tilde{s}\leftarrow 0.5\otimes (\tilde{a}\oplus \tilde{s}\oplus (\tilde{a}\ominus \tilde{s})\otimes Sgn(\tilde{a}\ominus \tilde{s}))$$

6. Conclusion

本文提出了NEXUS系统，可以说是第一个不需要客户端和服务器进行交互的安全transformer推理协议。本文提出了适用于RNS-CKKS的一系列新协议，以使得服务器可以高效且精确的在加密数据上计算transformer的每一层。