Comet: A Communication-efficient and Performant Approximation for Private Transformer Inference

## Abstract

Abstract The prevalent use of Transformer-like models, exemplified by ChatGPT in modern language processing applications, underscores the critical need for enabling private inference essential for many cloud-based services reliant on such models. However, current privacy-preserving frameworks impose significant communication burden, especially for non-linear computation in Transformer model. In this paper, we introduce a novel plug-in method Comet to effectively reduce the communication cost without compromising the inference performance. We second introduce an efficient approximation method to eliminate the heavy communication in finding good initial approximation. We evaluate our Comet on Bert and RoBERTa models with GLUE benchmark datasets, showing up to 3.9 × \times × less communication and 3.5 × \times × speedups while keep competitive model performance compared to the prior art.

## 1 Introduction

Leveraging the Transformer-based architecture [], the Generative Pretrained Transformer (GPT) [] is reshaping the global deep learning landscape by demonstrating remarkable proficiency in comprehending human language and generating multifaceted content.
For instance, users can receive instructional responses by sending queries via ChatGPT web portal. While such client-server interaction scheme enhances efficiency and productivity, privacy has emerged as a concern []. More specifically, machine learning applications like ChatGPT require either users provide language prompts or images, which may include confidential information, to the service provider. On the other hand, the server has concerns on exposing trained model weights, which are considered as vital asset, to the clients. Therefore, the gap between privacy requirements and efficient performance motivates our study of private Transformer inference.

To address the privacy concerns of users and protect the model on the server, several privacy-preserving inference frameworks [, , , ] have been proposed for convolutional neural networks via applying secure multi-party computation (MPC) techniques, such as homomorphic encryption (HE) [, ], secret sharing (SS) [], and oblivious transfer (OT) []. However, directly applying such privacy-preserving frameworks to Transformers leads to overwhelming computing and communication cost, because the Transformer-based models usually face complex hybrid protocols for non-linear functions like GeLU, Softmax, and LayerNorm, which have not been sufficiently addressed in previous studies.
To facilitate the widespread of private Transformer inference services, several works [, ] propose customized protocols and fine-tuning model for reducing communication cost. However, these existing works still encounter the challenge of heavy communication required to find good initial approximations or lengthy fine-tuning processes. For example, our investigations have shown that the Look-Up Table (LUT) method, extensively applied in the state-of-the-art work Iron, necessitates heavy communication on capturing good initial approximations [, ]. To alleviate such communication burdens, MPCFormer [] replaces two heavy non-linear functions, namely Softmax and GeLU, with aggressive quadratic polynomials for communication reduction, albeit at the cost of requireing further lengthy fine-tuning and compromises to lower performance.
Based on our preliminary explorations, we gain observations about the empirical characteristics of Transformer model and secret sharing techniques. These findings present us opportunities for designing of communication-efficient and performant private Transformer inference:

Smoothed approximation for Transformer inference: We observe that smoothed approximation functions can maintain or even enhance the Transformer performance across various tasks when replacing GeLU activation function. Additionally, our experimental results indicate that Softmax function replaced by $\frac{ReLU(x)}{\sum ReLU(x)}$ has marginal influence on model accuracy, which echos the finding of secureML []. This provides us an opportunity to unify common non-linear protocols to one function through the inverse square root.

Affinity of exponent of secret shares:
Current protocols on calculating the inverse square root involve high communication costs, often utilizing LUT or (piecewise) polynomial approximation to find good initial approximations for consecutive Newton iterations. We discover that such heavy communication can be totally removed via our novel design protocol. Because the magnitude of activation values is around zero, this provides us a unique opportunity to propose a share flooding technique to ensure our novel protocol working securely in two-party mode.

Based on the above observations, we propose Comet, a communication-friendly and performant private Transformer plug-in approximation method. Comet unifies hybrid complex non-linear functions and designs new specialized protocols to eliminate most of the communication for unified non-linear function. Specifically, in Sec. 3.1, we first endeavor to harmonize non-linear functions that applies hybrid complex protocols, namely GeLU and Softmax, with smoothed maximum unit (SMU) function []. In Sec. 3.2, we present our novel double approximation protocol that removes the communication cost of finding the initial approximation when calculating the inverse square root. To facilitate the proposed double approximation protocol applied in two-party computation scheme, we design a share flooding technique to render the method fully practical, thereby avoiding potential divergence after Newton iterations in Sec. 3.3. We implement our method and conduct extensive evaluations with BERT [] and RoBERTa-base [] models on the GLUE benchmark [] in Sec. 4. Our experiment results show that Comet achieves up to 3.9$\times$ reduction in communication cost and 3.5$\times$ time speedup, compared with LUT method and common Taylor approximation method utilizing the state-of-the-art framework Iron [] and CrypTen [], respectively.

## 2 Preliminaries

### 2.1 Threat Model

Similar to previous works [, , , ], our method follows the two-party semi-honest threat model. Specifically, the client $C$ and the server $S$ follow the protocol but attempt to infer each other’s input, namely the client’s input data and the server’s model parameters, during the inference process.

Figure: Figure 1: Overview of Comet. The client (left), who holds the input, interacts with and receives output from the server, who holds the model, via private Inference engines, e.g., CrypTen and Iron. Comet (right) unifies the non-linear functions into inverse square root and save communication with double approximation and share flooding in two-party mode.
Refer to caption: x1.png

### 2.2 Additive Secret Sharing and Protocols

Given an original message $m$ at party $P\in\{0,1\}$, one of the two Additive Secret Shares (ASS) is constructed by uniformly sampling randomness $r$ and setting $\left\langle m\right\rangle_{P}=r$, while the other share is formed as $\left\langle m\right\rangle_{1-P}=m-r$.
To reconstruct the message, one can simply add two shares $m=\left\langle m\right\rangle_{P}+\left\langle m\right\rangle_{1-P}$.
In this work, we utilize ASS to share the encrypted output of linear functions.
Existing research has developed accurate computation protocols for non-linear function using secret sharing, which protect the privacy of both the client and server. SIRNN [] designs multiple accurate non-linear computation protocols for convolution neural network, extensively using LUT for layer normalization, Softmax, and exponential function. The functionality of LUT takes as input string $i$ (with $\sigma$ bitwidth) and output $T<i>$ (with $\alpha$ bitwidth) where $T$ is a M-entries table. Such functionality can be achieved via $(^{M}_{1})$-OT with $(2^{\sigma}-\log M)$ offline and $(M*\alpha+\sigma)$ online communication bits [, ].
Iron [] improves the communication efficiency of relevant protocols via customized optimizations. Despite efforts to reduce communication cost regarding accurate protocols, current protocols continue to face significant communication burdens. To ease such overhead, (piecewise) polynomial approximation (e.g., Taylor expansion) can greatly reduce the communication overhead by converting complex protocols to multiplications [, ]. Nevertheless, the use of low-degree polynomials leads to a notable loss in inference performance, while employing high-degree polynomials incurs substantial communication costs.

### 2.3 IEEE 754 Float-point Representation

The IEEE 754 is a technical standard for floating-point representation []. Any floating-point number can be represented in the form of $(1+m)*2^{e}$, where $m$ is mantissa ($m\in[0,1)$) and $e$ is exponent number. To store a floating-point number in IEEE 754 representation, taking 32-bit floating-point number as example, 3 basic components should be filled: The highest bit denotes the sign of floating-point number, where
0 represents a positive number while 1 represents a negative number;
The exponent (E), as shown in Fig. 2 (green part), is filled by adding a bias B$=2^{7}-1$ to the actual exponent (e), which means $E=e+B$.
The Normalised Mantissa (m) can be directly filled in the binary form as shown in red part of Fig. 2.
For convenience, we denote M as the integer mantissa, which is generated by moving fraction dot to the end of mantissa ($M=L*m$), where $L=2^{23}$. $M+EL$ denotes the binary content corresponding to a floating-point number.

### 2.4 Newton-Raphson Method

The Newton-Raphson method [] is a root-finding algorithm which produces successively better approximations to the roots of a real-valued function. Given a single-variable function $f$ defined for a real variable x, the derivative $f^{{}^{\prime}}$, and an initial guess $x_{0}$ for a root of $f$, if the function satisfies sufficient assumptions and the initial guess is close, then
$x_{1}=x_{0}-{\frac{f(x_{0})}{f^{\prime}(x_{0})}}$
is a better approximation of the root than $x_{0}$. Geometrically, $(x_{1},0)$ is the intersection of the x-axis and the tangent of the graph of $f$ at $(x_{0},f(x_{0})))$, or in other words, the improved guess is the unique root of the linear approximation at the initial point. The process is repeated as
$x_{n+1}=x_{n}-{\frac{f(x_{n})}{f^{\prime}(x_{n})}}$
until a sufficiently precise value or a predefined number of iterations is reached. In this study, we follow the Newton Iteration equation $x_{n+1}=x_{n}(\frac{3}{2}-\frac{x}{2}x_{n}^{2})$ for inverse square root in previous works [, , , ], where $x$ is the input and $x_{n}$ is the estimation result for each iteration.

Figure: Figure 2: a IEEE 754 Floating-point representation example
Refer to caption: extracted/5839739/fig/754.png

### 2.5 Private Transformer Inference

Comet, like previous related works [, ], considers the scenarios where the server holds Transformer model, while the client holds and sends private input, as shown in the left part of Fig. 1. Our framework enables clients sending inference requests and receives prediction results on its input. To keep the privacy of both client and server, several private Transformer inference frameworks [, , ] are introduced. These frameworks mainly divide the cryptographic operations into two categories– linear and non-linear. For linear computation like matrix multiplication, HE is commonly used. HE ciphertext allows operations like multiplication on ciphertext without decryption. To maintain the correctness of decryption, ciphertexts need to be refreshed within limited number of operations via bootstrapping [] or re-encryption. To produce the linear results of each attention layers, the client encrypts their input embedding vectors, with coefficient encoding technique [, ], into HE ciphertexts and send them to the server. The server multiplies the ciphertexts-plaintext matrix with HE multiplication. To protect the privacy of kernel parameters, the server should generate two shares with random number mask and send one encrypted share to the client for consecutive non-linear functions.

GeLU, LayerNorm, and Softmax are the most common non-linear functions in Transformer-like models, which induce over 80% of communication cost in inference []. For example, GeLU necessitates the Gaussian Error function $erf(x)$, which is approximated using a high-order Taylor expansion. This introduces much more rounds of multiplications compared with linear part of private Transformer inference. Similar challenges are faced in Softmax and LayerNorm, which require to calculate exponential function and inverse square root. Furthermore, the variety of non-linear functions imposes extra difficulties for reducing communication cost, since different customized protocols must be developed for each non-linear function, such as exponential and complex $erf$ function.

## 3 Method

In this section, we present Comet, which unifies the hybrid complex non-linear protocols and removes the significant communication cost for finding good initial approximation. We provide details on how to unify protocols in Sec. 3.1, how to transfer communication to local computation in Sec. 3.2, and how to avoid divergence with share flooding in Sec. 3.3.

### 3.1 Unify Hybrid Complex Protocols

To address the heavy communication issue in private Transformer inference, we first shed light on the
Transformer architecture, which consists of multiple encoder-decoder architecture. The encoder has similar structure with the decoder, hence we focus on encoder blocks. A typical Transformer model with multiple encoder blocks consists of (1) an embedding layer, (2) a stack of encoder blocks, and (3) a prediction layer. One input token maps to a latent representation vector via the embedding layer. The encoder blocks are composed of attention layers and feed-forward layers as illustrated below.

Attention layers. After taking in the token embedding vector, attention function attempts to generate query, key, and values vectors with corresponding weights, denotes as $X_{Q}$, $X_{K}$, and $X_{V}$, respectively. Then, the layer output the attention vector using the function: $Attention(X_{Q},X_{K},X_{V})=Softmax(\frac{X_{Q}X_{K}^{T}}{\sqrt{d}})X_{V}$, where $d$ is the dimension of embedding vector.

Feed-forward Layers. The layer can be represented as follows: $FeedForward(X)=GeLU(XW_{1}+b_{1})W_{2}+b_{2}$, where GeLU is the Gaussian Error Linear Unit function. The feed-forward layer takes the output attention vectors of attention layers as input. The GeLU function requires Gaussian Error function $erf(x)$: $GeLU(x)=0.5x*(1+erf(\frac{x}{\sqrt{2}}))$.

Layer Normalization. Layer Normalization (LayerNorm) is applied after attention layers and feed-forward layers. By calculating the normalization of mini-batch of input, it smooths gradients for better generalization accuracy, shown in following equation:
$y=\frac{x-E(x)}{\sqrt{Var(x)+\epsilon}}*\gamma+\beta$
, where $E(x)$ denotes the mean of input $x$ and $Var(x)$ is the variance of input, $\gamma$ and $\beta$ are learnable parameter during the training.

$e^{x}$ in Softmax, $erf(x)$ in GeLU, and $\frac{1}{\sqrt{x}}$ in LayerNorm require either different Taylor approximation functions or different specialized LUT protocols to calculate their results.
To unify the complex protocols in non-linear functions, we propose to replace the exponential function $e^{x}$ in Softmax with ReLU function as SecureML [], i.e., $Softmax^{*}(x)=\frac{ReLU(x)}{\sum ReLU(x)}$ since it shows competent performance in experiments (see in Sec. 4.2 and Appendix. C).
Since the erf function in $GeLU$ requires high order Taylor expansion in specialized protocol, we leverage the smoothed maximum unit (SMU) to replace the GeLU function.
The SMU function for GeLU derives as $smu(x)=\frac{1+\alpha}{2}*x+\frac{(1-\alpha)*x^{2}+\mu^{2}}{2*\sqrt{(1-\alpha)
*x^{2}+\mu^{2}}}$, where $\alpha$ and $\mu$ are the trainable parameters to control the slope of negative axis and smoothness of function, respectively (refer to Appendix. B and  LABEL: for detailed derivation and explanation).
We set $\alpha=0,\mu=0$ for ReLU replacement in $Softmax^{*}$ function, and $\alpha=0,\mu=\frac{1}{\sqrt{2}}$ for GeLU replacement in $smu(x)$ function to retain the model performance.
Additionally, we test the flexibility of $\alpha$ and $\mu$ by training from scratch, showing a performance boost in some tasks of GLUE benchmark in Appendix. C.
As the LayerNorm only requires inverse square root, we successfully unify all non-linear functions, namely GeLU, Softmax and LayerNorm, into inverse square root.

### 3.2 Double Approximation

Even though we unify complex non-linear protocols to the inverse square root, we still face the high communication challenge when calculating this function. Current works on calculating the inverse square root requires the Newton iterations, or Goldschmidt’s iteration [], to get accurate results with a initial approximation. Due to the local convergence characteristic of Newton method (refer to Appendix. E for details), the initial approximation has to be close enough to the root of function, referred as "the good initial approximation". However, finding the good initial approximation usually requires over $85\%$ communications of total process using LUT method or high order Taylor expansion [].
In this manner, we demonstrate our double approximation method for such initial approximation finding without communication for inverse square root
$y=\frac{1}{\sqrt{x}}$.
First, we take logarithm on both sides of the equation to
$log(y)=-\frac{1}{2}log(x)$.
Such equation can be easily transformed into secret-shared form as equation ( 1), where $x_{c}$ and $x_{s}$ denote as the share of client and server.

$$ $\begin{split}2log(y)&=-log(x)\\ &=-log(x_{c}+x_{s})\end{split}$ (1) $$

Then we replaced the input x and output y with IEEE 754 floating-point representation:

$$ $\begin{split}2log((1+m_{y})*2^{e_{y}})&=-log((1+m_{x})*2^{e_{x}})\\ &=-log((1+m_{x_{c}})*2^{e_{x_{c}}})+(1+m_{x_{s}})*2^{e_{x_{s}}}))\end{split}$ (2) $$

First, let us focus on one-party calculation of inverse square root, which is the upper equation in equation (2). We change the multiplication in logarithm to addition following the logarithm rule.

$$ $\begin{split}2(log(1+m_{y})+e_{y})=-log(1+m_{x})+e_{x}\end{split}$ (3) $$

As the logarithm can be complex to calculate, we first approximate the logarithm with linear function $log(1+m)\approx m+b,m\in[0,1)$, where $b$ is a constant number we can predefined based on logarithm function.
After the replacement of logarithm with first approximation, we have the equation (4) and reorganized to equation (5) to estimate the inverse square root value:

$$ $\begin{split}2(m_{y}+b+e_{y})\approx-(m_{x}+b+e_{x})\end{split}$ (4) $$

$$ $\begin{split}(m_{y}+e_{y})\approx-(m_{x}+e_{x})/2-3b/2\end{split}$ (5) $$

Such approximation can be efficient to calculate the inverse square root with good precision in one-party mode. However, it is challenging for the secret share scheme, as the logarithm cannot be replaced when two shares are added in the lower part of equation (2). Our insight is that two shares’ exponent part can be approximately equal to each other, as the exponent parts for two shares are only 8-bits length. Then we take the second approximation
$e_{x_{c}}\approx e_{x_{s}}$.
We can further transform the secret-shared form equation ( 2) with second approximation as following:

$$ $\begin{split}2log((1+m_{y})*2^{e_{y}})\approx-log((1+m_{x_{c}}+1+m_{x_{s}})*2^ {e_{x_{s}}}))\\ \end{split}$ (6) $$

$$ $\begin{split}2(log(1+m_{y})+e_{y})\approx&-log((1+(m_{x_{c}}+m_{x_{s}})/2)*2^{ e_{x_{s}}+1}))\\ &=-log(1+(m_{x_{c}}+m_{x_{s}})/2)-e_{x_{s}}-1\end{split}$ (7) $$

Then we apply the first approximation to the lower equation ( 7):

$$ $\begin{split}2(m_{y}+b+e_{y})\approx-(m_{x_{c}}+m_{x_{s}})/2-b-e_{x_{s}}-1\end {split}$ (8) $$

We then replace the $m$ and $e$ with $M=L*m$ and $E=e+B$ as stated in Sec. 2.3:

$$ $\begin{split}2\frac{M_{y}}{L}+3b+2(E_{y}-B)\approx-\frac{(M_{x_{c}}+M_{x_{s}}) }{2L}-E_{x_{s}}+B-1\end{split}$ (9) $$

We reorganize (9) to the equation (10) into two shares mode, where $E_{x_{c|s}}$ denotes we replace the $E_{x_{c}}$ with $E_{x_{s}}$ as second approximation:

$$ $\begin{split}M_{y}+LE_{y}\approx\boxed{-\frac{1}{4}{\color[rgb]{1,.5,0} \definecolor[named]{pgfstrokecolor}{rgb}{1,.5,0}(M_{x_{c}}+LE_{x_{c}})}}_{c} \boxed{-\frac{1}{4}{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb} {0,0,1}(M_{x_{s}}+LE_{x_{c|s}})}}_{s}\boxed{+\frac{(3B-3b-1)L}{2}}_{s}\end{split}$ (10) $$

The orange and blue term can be regarded as the integer value of the client and the server share, respectively.
The last black term is a constant that both parties can learn offline. The boxed term with undertext $c$ and $s$ are the output shares of approximated inverse square root result for client and server, respectively.

In this manner, we can get a approximated value of inverse square root without heavy communication between the client and server, as the client and server can only calculate on their own shares and a constant.
The initial approximation can be produced by adding two shares. The precise result can be approached via Newton Method with 3-4 iterations, as shown in Sec. 4.3.
As the main bottleneck of communication lies in finding good initial approximation, we decrease $\mathcal{O}(2^{\sigma})$ LUT communication cost to $\mathcal{O}(1)$ in private Transformer inference.

### 3.3 Share Flooding

In two-party mode, server needs to generate the shares for client and server itself with random number mask, after the linear results are produced in HE ciphertext. However, our second approximation requires the exponent part of two shares are close to each other to remove the communication cost. If the exponent part is not close enough, the result of Newton iteration would be diverge (see Appendix D for more detailed experiments). As shown in Fig. 3, the random number in the gray box which is too close to the input values can make the counterpart share far from each other, e.g., if the input value is 3.1 and random number is 3.098, the other share 0.002 would generate bad initial approximation that leads to divergence, since the large difference in the exponent of two shares breaks second approximation. This imposes a dilemma for generating two shares as one of the share is the uniformly random number mask that is generated offline by server. It is infeasible to always meet the requirements of second approximation as the server cannot learn the convolution results before generating the random number mask offline. This dilemma brings us a new challenge— How can we securely generate shares and efficiently perform double approximation while satisfying privacy needs and approximation assumptions?

To address this challenge, we propose the share flooding technique for private Transformer inference. Our insight is that the absolute magnitude of tensor values is closely surrounded around zero. This is because the input embedding vector is a 0 to 1 valued vector after softmax using word2vec technique []. Applying Softmax and LayerNorm function in attention block and feed-forward layers iterative also confine the activation magnitude into zero to one range. Meanwhile, the training process usually applies L1/L2 normalization in loss function to discourage large weights to fight overfitting issue, which makes the consecutive activation magnitude to be small as well. This is also validated in our preliminary experiments. We refer readers to Appendix A for more details. With our insight, we design to flood the random number mask with a large absolute value. This flood number can drown the exponent of two shares to satisfy the second approximation requirement. For example, if the flood number is 8192 and input message is 3.1, we add the flood number to the random number mask, such as 3.098, results in 8195.098. The other share is -8191.998. Both shares have the same exponent equal to 140 in IEEE 754 float-point representation.
To be more precise, we flood the random mask offline with a large adjustable flooding number to specific task, as one party’s share.
The corresponding other share is generated when online message subtract with flooded random mask, then transfer the floating-point valued share to the integers in corresponding integer field. As equation (10) shows, the flooding number cannot be offset as the blue and orange shares are same sign. We compensate the over-flooded value by adding $\frac{1}{2}(E_{f}-E_{m})L$, where $E_{f}$ is the exponent of flooding number and $E_{m}$ denotes the exponent of most frequent activation value of the distribution learned in similar tasks.
Note that the fixed-point shares can be easily transfer to floating-point shares by dividing the scale. We securely produce shares as there is no information exposed except the sign of two shares, which would not expose any information of original input value. This novel design enable us for addressing the challenge of security and approximation assumption requirement.

Figure: Figure 3: Demonstration of double approximation divergence example. "d" denotes the upper bound of exponent between two shares that would lead to divergence in Newton Method. One share generated, in integer field $\mathbb{Z}_{2^{31}}$, in grey box (within the bound d) means the share is too close to the input X, which makes the two shares exponent out of the convergence bound of Newton Method.
Refer to caption: x2.png

## 4 Experiments

### 4.1 Experimental Settings

We implement Comet within the secure two-party framework Iron, which uses the EMP toolkit [] for implementing non-linear functions, and the CrypTen framework. The experiments are conduct on two servers with an AMD EPYC 7413 24-core Processor 64GB RAM, under the network bandwidth of 200Mbps. We set the flooding number equal to 8192 ($E_{f}=140,M_{f}=0$), as it only floods the exponent $E_{f}$ with mantissa $M_{f}$ equals to zero, and $E_{m}=128$, as it covers most of the activation distribution in Transformer-based model.
We evaluate the model performance based on HuggingFace implementation with the dataset of GLUE benchmarks.
In following sections, we aim to answer three questions to present the benefits of Comet: (1) The model performance of the unified model. (2) The iteration number of the Newton method required and its effect on model performance. (3) The time and communication reduction of Comet.

**Table 1: The model performance of a subset of GLUE benchmark with different combination of smoothed maximum unit replacement for GeLU and Softmax function. "$Softmax^{*}$" stands for ReLU replaced Softmax in Sec. 3.1. "s-" stands for the $smu(x)$ smoothed function in Sec. 3.1. Average Pearson and Spearman correlation is reported for STS-B. Matthews correlation is reported for CoLA. Accuracy is reported for other datasets.**
| Bert-base | RTE | MRPC | STS-B | SST-2 | CoLA |
| --- | --- | --- | --- | --- | --- |
| GeLU + Softmax | 70.8 | 88.97 | 88.6 | 92.7 | 58.9 |
| GeLU + $Softmax^{*}$ | 66.4 | 86.9 | 85.8 | 91.4 | 54.2 |
| GeLU + s-$Softmax^{*}$ | 67.9 | 86.3 | 86.7 | 90.5 | 55.7 |
| s-GeLU + Softmax | 69.4 | 86.7 | 88.3 | 91.1 | 56.1 |
| s-GeLU + $Softmax^{*}$ | 67.2 | 88.1 | 88.1 | 92.1 | 55.7 |
| s-GeLU + s-$Softmax^{*}$ | 71.5 | 88.6 | 88.8 | 92.6 | 57.9 |
| RoBERTa-base |  |  |  |  |  |
| GeLU + Softmax | 74.2 | 92.3 | 89.1 | 92.1 | 55.7 |
| GeLU + $Softmax^{*}$ | 71.8 | 89.6 | 84.8 | 89.9 | 53.8 |
| GeLU + s-$Softmax^{*}$ | 72.6 | 88.9 | 87.2 | 89.1 | 52.9 |
| s-GeLU + Softmax | 72.7 | 89.1 | 88.1 | 87.3 | 54.2 |
| s-GeLU + $Softmax^{*}$ | 73.4 | 91.7 | 86.3 | 88.4 | 54.1 |
| s-GeLU + s-$Softmax^{*}$ | 74.9 | 92.2 | 90.1 | 90.4 | 55.8 |

### 4.2 Unified Model Performance

To present the model performance of the unified model, we evaluate its performance use the Bert-base and RoBERTa-base model and compare it with baseline models. We use the sequence length of 128 for the selected datasets of GLUE benchmark. Baseline models are selected combinations within the set of {GeLU, s-GeLU, Softmax, $Softmax^{*}$, $s-Softmax^{*}$ }, where "s-" denotes the $smu(x)$-replaced ReLU or GeLU in the following function. We train the baseline models with learning rate from 1e-6, 5e-6, 1e-5, and 1e-4, the number of epochs from 10, 30, and 100. We also set the $\alpha=0,\mu=0$ for $s-Softmax^{*}$ and $\alpha=0,\mu=\frac{1}{\sqrt{2}}$ for s-GeLU, as stated in Sec. 3.1. Table. 1 upper part shows the model performance of Bert model on
a subset of GLUE benchmark with different baseline models. Since all non-linear functions are unified to inverse square root with SMU unit, we name our target model with "s-GeLU + s-$Softmax^{*}$", as the unified model. The unified model achieves marginal performance loss with less than $1\%$. It also shows a small accuracy boost of approximately $1\%$ in the relatively small dataset RTE, while other baselines experience larger performance loss. In this manner, our unified model preserves model performance with Bert-base model. To validate our observation, we also evaluate the unified model with RoBERTa-base model architecture with same datasets. As shown in lower part of Table. 1, the target unified model of RoBERTa-base shows consistent model performance among the GLUE benchmark subset.

### 4.3 Newton Iteration Evaluation on Model Performance

In this section, we evaluate the double approximation method of Comet in model performance. To answer the question of how many iterations that the double approximation method requires to recover the performance of unified model, we conduct experiments on the unified model with varying iteration number when calculating the inverse square root with double approximation method. We set $b=0.045$ by inversely solving $\frac{3}{2}L(B-b)=$ 0x5f3759df derived from equation (5) with $m=M/L$ and $e=E-B$ replaced, given the magic number 0x5f3759df as in fast inverse square root algorithm []. Such magic number can be obtained by minimizing the relative error between the approximated results and real results as shown in []. Our double approximation method can recover the model performance of the unified model in 3$\sim$4 iterations with the initial approximation resulting from our method, as shown in Table. 2. Our method requires fewer iterations compared to the CrypTen, where the inverse square root requires 10 iterations by default with a communication-intensive initial approximation function of $2.2*e^{(-0.5x+0.2)}+0.2$. Even though we incur around extra 2 rounds compared with LUT method (which only requires 1$\sim$2 rounds), LUT method necessitates heavy communication for accurate initial approximation, and such communication can grow exponentially with the number of table entries. Thus, our method outperforms the LUT in the total communication and we elaborate the details in Sec. 4.4.

**Table 2: The model performance of double approximation method on how many iterations to recover the best model performance from initial approximation. The unmodified model is the model with unchanged GeLU and Softmax functions. The unified model is the SMU-replaced GeLU and SMU-replaced $Softmax^{*}$ function model.**
| Bert-base | RTE | MRPC | STS-B | SST-2 | CoLA |
| --- | --- | --- | --- | --- | --- |
| Unmodified Model | 70.8 | 88.97 | 88.6 | 92.7 | 58.9 |
| Unified Model | 71.5 | 88.6 | 88.8 | 92.6 | 57.9 |
| Without Newton Iteration | 64.6 | 81.58 | 86.7 | 90.8 | 52.6 |
| 3 iterations | 70.8 | 86.9 | 88.6 | 92.2 | 56.4 |
| 4 iterations | 71.5 | 88.6 | 87.6 | 92.5 | 57.8 |
| RoBERTa-base |  |  |  |  |  |
| Unmodified Model | 74.2 | 92.3 | 89.1 | 92.1 | 55.7 |
| Unified Model | 74.9 | 92.2 | 90.1 | 90.4 | 55.8 |
| Without Newton Iteration | 66.7 | 87.8 | 86.7 | 85.3 | 51.9 |
| 3 iterations | 72.4 | 86.7 | 90.1 | 89.4 | 53.9 |
| 4 iterations | 74.2 | 88.1 | 92.1 | 90.6 | 55.7 |

### 4.4 End-to-End Inference Communication and Time Comparison

In this section, we evaluate the end-to-end inference time and communication cost of Comet implemented within the state-of-the-art privacy-preserving frameworks Iron and CrypTen. We first compare our method with the $2.2*e^{(-0.5x+0.2)}+0.2$ function applied in CrypTen and high/low order of Taylor expansion for finding initial approximation of inverse square root in both communication and time using the unified model. Comet recovers the model performance with setting of 4 Newton iterations as shown in Sec. 4.3, while high and low order Taylor approximation would take 2 and 8 iterations, respectively. We can see from Table. 3 that Comet achieves a 2.5$\times$ speedup compared to the original CrypTen framework, with 2$\times$ to 3$\times$ communication reduction for different non-linear functions.
Such improvement is attributed to the removal of the bottleneck communication involved in finding initial approximation.
We also compare Comet with Taylor approximation with high/low order polynomial, namely $1-\frac{x}{2}+\frac{1}{2}+\frac{3(x-1)^{2}}{8}-\frac{5(x-1)^{3}}{16}+\frac{35(
x-1)^{4}}{128}-\frac{63(x-1)^{5}}{256}+\frac{231(x-1)^{6}}{1024}-\frac{429(x-1
)^{7}}{2048}$ and $1-\frac{x}{2}+\frac{1}{2}+\frac{3(x-1)^{2}}{8}$, showing similar 4.0$\times\sim$ 2.8$\times$ reduction in time and 3.2$\times\sim$ 1.9$\times$ in communication. This is because the high-order Taylor approximation requires significant communication for calculating the initial approximation, whereas the low-order Taylor approximation typically demands more iterations and may diverge due to the initial approximation surpassing the requirement for local convergence of the Newton Method.

**Table 3: The communication (GB) and inference time (Second) comparison on Bert-base and RoBERTa-base model with CrypTen. "H-Taylor P" and "L-Taylor P" denote 7 order Taylor polynomial and 2 order Taylor polynomial approximating inverse square root generated at 1, respectively.**
| Bert-base | Total Time(s) | LayerNorm Time(s) | LayerNorm Comm | Act Time | Act Comm | Softmax Time | Softmax Comm |
| --- | --- | --- | --- | --- | --- | --- | --- |
| CrypTen | 74.8 | 12.19 | 2.14 | 34.1 | 6.28 | 17.8 | 3.28 |
| H-Taylor P | 73.2 | 13.6 | 2.33 | 32.6 | 5.82 | 16.2 | 3.19 |
| L-Taylor P | 70.6 | 10.5 | 2.03 | 29.6 | 5.57 | 15.8 | 3.09 |
| Comet | 29.8(2.5$\times$) | 3.32(3.7$\times$) | 0.94(2.3$\times$) | 7.92(4.3$\times$) | 1.83(3.4$\times$) | 5.7(3.1$\times$) | 1.58(2.07$\times$) |
| RoBERTa |  |  |  |  |  |  |  |
| CrypTen | 79.3 | 14.19 | 2.6 | 36.3 | 7.4 | 18.2 | 3.34 |
| H-Taylor P | 75.7 | 13.9 | 2.32 | 34.7 | 7.31 | 17.6 | 3.19 |
| L-Taylor P | 73.1 | 13.3 | 2.21 | 30.9 | 6.4 | 15.1 | 3.1 |
| Comet | 34.8(2.3$\times$) | 5.81(2.4$\times$) | 1.59(1.6$\times$) | 9.1(4.0$\times$) | 1.88(3.9$\times$) | 6.9(2.63$\times$) | 1.79(1.8$\times$) |

We also compare Comet to LUT-based framework Iron in the same experimental setting as we did with CrypTen. Our method achieves up to 3.48$\times$ for time speedup and 2.4$\times$ to 3.7$\times$ communication reduction in non-linear functions, as shown in Table. 4. This runtime improvement and communication reduction demonstrate the efficiency of the novel design of Comet.

**Table 4: The communication (GB) and inference time (Second) comparison on Bert-base and RoBERTa-base model with Iron [].**
| Bert-base | Total Time | LayerNorm Time | LayerNorm Comm | Act Time | Act Comm | Softmax Time | Softmax Comm |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Iron | 289.6 | 48.2 | 7.45 | 134.3 | 14.9 | 82.6 | 8.85 |
| Comet | 82.4(3.48$\times$) | 15.3(3.17$\times$) | 3.03(2.4$\times$) | 26.8(4.9$\times$) | 4.05(3.7$\times$) | 20.4(4.0$\times$) | 3.56(2.5$\times$) |
| RoBERTa |  |  |  |  |  |  |  |
| Iron | 290.1 | 52.3 | 7.8 | 127.9 | 13.83 | 77.6 | 8.03 |
| Comet | 92.1(3.15$\times$) | 14.2(3.68$\times$) | 2.8(2.8$\times$) | 28.5(4.45$\times$) | 4.3(3.2$\times$) | 18.9(4.12$\times$) | 3.28(2.4$\times$) |

## 5 Conclusion

In this paper, we propose Comet, a communication-efficient and performant approximation framework for private Transformer inference. We specifically unified the hybrid complex protocols into one protocol– inverse square root for non-linear functions. Then, we further carefully design the double approximation method to convert the heavy communication of finding initial approximation to local computation for inverse square root, with our share flooding technique to securely secret sharing under strong assumption satisfaction. Our experimental results show that Comet outperforms prior art with up to 3.9$\times$ less communication and 3.48$\times$ speedups.

## Appendix

### A Preliminary Experiments on Activation Value

We test the activation values distribution of Bert base model in various layers and different dataset, as we trained in Sec. 4.2. The Fig. 4 and Fig. 5 show that the activation values are compact to small absolute magnitude (surround zero) regardless the datasets and layers.

Figure: Figure 4: Activation distribution before GeLU and LayerNorm layer on various dataset (RTE, STS-B,COLA) within feed-forward block 2 of Bert base model
Refer to caption: extracted/5839739/fig/rte-layer2.png

Figure: Figure 5: Activation distribution before GeLU layer on RTE in different block of Bert base model
Refer to caption: extracted/5839739/fig/rte-layer4.png

### B Derivation of SMU function

The SMU smoothed activation function that is proposed in  [] is derived from the equation B.

$$ $\displaystyle max(x_{1},x_{2})$ $\displaystyle=\begin{cases}x_{2}&\text{if $x_{1}\leq x_{2}$}\\ x_{1}&\text{otherwise}\end{cases}$ $\displaystyle=\frac{(x_{1}+x_{2})+|x_{1}+x_{2}|}{2}$ (11) $$

By replacing $|x|$ with $\sqrt{x^{2}+\mu^{2}}$, the smoothed approximation formula can be generated as shown in equation 12.

$$ $\displaystyle f(x_{1},x_{2},\mu)=\frac{(x_{1}+x_{2})+\sqrt{x^{2}+\mu^{2}}}{2}$ (12) $$

To have a smoothed approximation of Parametric Activation function, e.g., Leaky ReLU [], $x_{1}=x$ and $x_{2}=\alpha x$ are substituted and result in equation 13.

$$ $\displaystyle f(x_{1},x_{2},\alpha,\mu)=\frac{(1+\alpha)x+\sqrt{(1+\alpha)x^{2 }+\mu^{2}}}{2}$ (13) $$

For convenience of calculation of inverse square root, we transform it into equation 14.

$$ $\displaystyle f(x,\alpha,\mu)=\frac{1+\alpha}{2}*x+\frac{(1-\alpha)*x^{2}+\mu^ {2}}{2*\sqrt{(1-\alpha)*x^{2}+\mu^{2}}}$ (14) $$

### C Experiments on flexible SMU trainable parameters

We test the model performance with flexible parameters in SMU function to replace the GeLU and ReLU in $Softmax^{*}$ on subset of GLUE benchmark. We follow the training setting and strategy as in Sec. 4.1 and Sec. 4.2. The Table. 5 shows about $1-2\%$ performance boost between unified model (s-GeLU + s-$Softmax^{*}$) and original one (GeLU + Softmax) on the datasets we test when using flexible trainable parameters.

**Table 5: The model performance of a subset of GLUE benchmark with different combination of smoothed maximum unit replacement for GeLU and Softmax function. "$Softmax^{*}$" stands for ReLU replaced Softmax in Sec. 3.1. "s-" stands for the $smu(x)$ smoothed function in Sec. 3.1. Average Pearson and Spearman correlation is reported for STS-B. Matthews correlation is reported for CoLA. Accuracy is reported for other datasets.**
| Bert-base | RTE | MRPC | STS-B | SST-2 | CoLA |
| --- | --- | --- | --- | --- | --- |
| GeLU + Softmax | 70.8 | 88.97 | 88.6 | 92.7 | 58.9 |
| GeLU + s-$Softmax^{*}$ | 68.2 | 87.1 | 86.9 | 91.1 | 56.6 |
| s-GeLU + Softmax | 70.4 | 86.7 | 88.7 | 91.5 | 56.8 |
| s-GeLU + $Softmax^{*}$ | 69.4 | 88.9 | 88.6 | 92.8 | 56.5 |
| s-GeLU + s-$Softmax^{*}$ | 72.1 | 89.5 | 89.2 | 93.8 | 59.4 |
| RoBERTa-base |  |  |  |  |  |
| GeLU + Softmax | 74.2 | 92.3 | 89.1 | 92.1 | 55.7 |
| GeLU + s-$Softmax^{*}$ | 73.4 | 88.9 | 87.9 | 89.1 | 53.9 |
| s-GeLU + Softmax | 72.9 | 89.1 | 89.3 | 88.7 | 55.2 |
| s-GeLU + $Softmax^{*}$ | 73.7 | 91.9 | 89.2 | 88.4 | 55.8 |
| s-GeLU + s-$Softmax^{*}$ | 76.1 | 93.6 | 91.1 | 93.8 | 57.9 |

### D Closeness bound of two shares exponent

We evaluate the closeness bound of exponent parts between two shares to avoid the divergence of Newton method’s iterations and keep the model performance. As the Fig. 6 shows, the upper bound of closeness between shares in exponent is $|5|$, showing as "elbow point", to satisfy the second approximation assumption and result in good initial approximation without divergence.

Figure: (a) performance on RTE
Refer to caption: extracted/5839739/fig/rte-closeness.png

### E Newton Method’s Local Convergence

We give the proof of local convergence of Newton’s Method as following.

###### Theorem 1 .

(local convergence of Newton’s method) Let f be a twice continuously differentiable function defined over $\mathbb{R}^{d}$. Assume that
(1) there exists a neighborhood $N_{\sigma}(x_{*})$ of root of function $x_{*}$ and $M>0$ for which $\|\nabla^{2}f(x)-\nabla^{2}f(y)\|\leq\frac{M}{2}\|x-y\|^{2}$ for any $x,y\in N_{\sigma}(x^{*})$.

###### Proof.

we have

$$ $\displaystyle x_{k+1}-x_{k}$ $\displaystyle=x_{k}-\nabla^{2}f(x_{k})^{-1}\nabla f(x_{k})-x_{*}$ $\displaystyle=x_{k}-x_{*}+\nabla^{2}f(x_{k})^{-1}(\nabla f(x_{k})-\nabla f(x_{ *}))$ $\displaystyle=x_{k}-x_{*}+\nabla^{2}f(x_{k})^{-1}\int_{0}^{1}[\nabla^{2}f(x_{k }+t(x_{*}-x_{k}))](x_{*}-x_{k})dt$ $\displaystyle=[\nabla^{2}f(x_{k})]^{-1}\int_{0}^{1}[\nabla^{2}f(x_{k}+t(x_{*}- x_{k}))-\nabla^{2}f(x_{k})](x_{*}-x_{k})dt$ $$

Then,

$$ $\displaystyle\|x_{k+1}-x_{k}\|$ $\displaystyle\leq\|[\nabla^{2}f(x_{k})]^{-1}\|\|\int_{0}^{1}[\nabla^{2}f(x_{k} +t(x_{*}-x_{k}))-\nabla^{2}f(x_{k})](x_{*}-x_{k})dt\|$ $\displaystyle\leq\|[\nabla^{2}f(x_{k})]^{-1}\|\int_{0}^{1}\|[\nabla^{2}f(x_{k} +t(x_{*}-x_{k}))-\nabla^{2}f(x_{k})](x_{*}-x_{k})\|dt$ $\displaystyle\leq\|[\nabla^{2}f(x_{k})]^{-1}\|\int_{0}^{1}\|\nabla^{2}f(x_{k}+ t(x_{*}-x_{k}))-\nabla^{2}f(x_{k})\|\|x_{*}-x_{k}\|dt$ $\displaystyle\leq\int_{0}^{1}Mt\|x_{k}-x_{*}\|^{2}dt$ $\displaystyle\leq\frac{M}{2}\|x_{k}-x_{*}\|^{2}$ $$

∎

###### Theorem 1 .

###### Proof.

### F SMU Function Explanation

Although our proposed function,$\frac{1+\alpha}{2}x+1/2*\frac{(1-\alpha)x^{2}+\mu^{2}}{\sqrt{(1-\alpha)x^{2}+
\mu^{2}}}$, can be simplified to $\frac{1+\alpha}{2}x+1/2*\sqrt{(1-\alpha)x^{2}+\mu^{2}}$, when switch to such square root version, our efficient initial approximation finding method still works as our method is applicable to $f(x)=x^{1/a}$ ($a\in\mathcal{Z}_{\neq-1,0,1}$) functions, which is of independent interest. However, the Newton-Raphson update formula for square root, $y_{n+1}=1/2*(y_{n}+\frac{x}{y_{n}})$ is inefficient to compute on secret shares mode []. It requires secret share division that needs truncation and extension protocol on given integer ring [].

Therefore, we switch the smu function to the inverse square root version. Since the Newton-Raphson update formula for inverse square root, $y_{n+1}=\frac{1}{2}*y_{n}(3-xy_{n}^{2})$ only requires efficient share multiplications and the inverse square root one is mathematically equal to the square root version. Our method contributes to a more lightweight initial approximation finding for later Newton iterations.

### G Share Conversion

For the conversion between the floating point share in non-linear function and integer shares, a floating-point share ($=M+EL$) is correspond to one input value, according to IEEE 754 protocol. Such input value can be converted to the fixed-point share from decimal value to binary. After multiplied with a pre-defined scale $(=2^{f})$, the $l$-bit fixed-point integer share can be obtained in the given integer ring $\mathcal{Z}_{2^{l}}$ [3]. We give the following share conversion example of 7.25:

$$ $\begin{split}\text{(float-point share)0|10000001|11010000000000000000000 (=108 8946176(=M+EL)decimal)}\longleftrightarrow\\ 7.25\longleftrightarrow\text{(fixed\ point)}0111.01^{(*2^{4})} \longleftrightarrow_{(\div 2^{4})}\text{(fixed\ point\ \ integer\ share)}01110 100(=29\ decimal)\end{split}$ (15) $$

### H Ablation Study

To demonstrate the challenge and show the effectiveness of the proposed share flooding, we conduct the ablation experiment on our method with and without flooding technique. The table 6 shows that the model performance can be significantly harmed without share flooding by 10 loss in corresponding metric. Our technique can successfully avoid such model performance loss while keeping the privacy.

**Table 6: The ablation experiment of Comet method with v.s. without flooding technique.**
|  | STS-B | CoLA | RTE |
| --- | --- | --- | --- |
| Comet without flooding | 72.2 | 49.7 | 62.5 |
| Comet | 87.9 | 57.9 | 71.1 |

We conduct the experiment on the model performance among PUMA, Bumblebee and our method on partial GLUE benchmark datasets, as shown in Table.7. Our method demonstrates competitive model performance when compared to PUMA ([puma]) and BumbleBee ([bumblebee]). For fair comparison, we also compare our method (integrated with Bumblebee) with original Bumblebee in inference time on GeLU protocols (with 14 calls) in Table. 8 in 2-PC mode. We exclude PUMA as it is in 3-PC computation. The Table. 8 shows our method achieve $1.5\times$ speedup when compare to the GeLU function of Bumblebee.

**Table 7: The model performance comparison of Comet, PUMA, and Bumblebee.**
|  | STS-B | CoLA | RTE |
| --- | --- | --- | --- |
| PUMA | 88.4 | 59.2 | 70 |
| Bumblebee | 87.5 | 60.8 | 70.04 |
| ours | 88.8 | 59.4 | 71.3 |

**Table 8: The latency comparison between Bumblebee and Comet with 14 calls of GeLU.**
|  | Bumblebee | Ours | Speedup |
| --- | --- | --- | --- |
| (14 calls) GeLU | 28.9s | 19.2s | 1.5$\times$ |

We also surpass MPCFormer in model performance and avoiding its lengthy knowledge distillation process, with similar inference time performance, as shown in Table. 9.

**Table 9: The accuracy, and inference time (Second) comparison on Bert-base model with MPCFormer ([]). "-" denotes not applicable to the method.**
| Bert-base | Total Time | KD training | STS-B | CoLA |
| --- | --- | --- | --- | --- |
| MPCFormer | 27.7 | 100 | 80.1 | 52.6 |
| Comet | 29.8 | - | 87.6 | 57.9 |

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: $\bullet$ Why the authors have compared their 2PC method to 3PC method CrypTen? Is that the comparison with Crypten's Newton Raphson's method and other approximations, and not exactly with their 3PC settings? 

$\bullet$ Why do the client and server share, which is an integer in the field arithmetic, converted to floating point numbers (Eq 1 and Eq 2)?
soundness: 3
strengths: $\bullet$ Comprehensive evaluation of post-training performance across a variety of downstream tasks, demonstrating robustness, applicability, and generalizability of their proposed approaches. 

$\bullet$ The integration of SMU with trainable parameters as an alternative to standard nonlinearities like GELU and ReLU-approximated Softmax is an interesting approach to simplifying the conventional nonlinearities and reducing communication and latency overheads in private inference. 

$\bullet$ A thorough comparison with previous methods for private inference in transformer-based models.
summary: This paper tackles the communication (and latency) overhead introduced by nonlinear operations (GELU, LayerNorm, and Softmax) in cryptographically secure private inference of transformer-based models. The authors replace GELU and Softmax with a smoothed maximum unit (SMU) function, effectively transforming nonlinearities into inverse square root operations, and developed the methods such as double approximation and share flooding to efficiently compute them in privacy-preserving settings.  

By mitigating the communication overhead associated with these nonlinear operations, this work tackles a major bottleneck for the practical deployment of private LLM services.
weaknesses: $\bullet$ The contributions in this paper are largely engineering tweaks rather than substantive algorithmic advancements. While incorporating SMU (Biswas et al., CVPR'22) as a replacement for traditional nonlinearities like GELU and ReLU-approximated Softmax in a transformer-based model is an interesting approach, it lacks sufficient originality and novelty. Additionally, the use of Network Raphson methods for inverse square root calculations, as previously implemented in CryPTen but with an improved initialization, which reduces the number of iterations to converge, does not offer a particularly compelling approach. These techniques are often highly model- and task-dependent, requiring tuning when applied to other tasks or domains. 

Furthermore, the authors’ proposed share flooding method, which is based on the observation  "that the absolute magnitude of tensor values is closely surrounded around zero (Line#317)," is rather intuitive and lacks deeper insight. Nonetheless, the authors will definitely get the points for the double approximation approach which effectively reduces the communication overhead to constant time complexity for lookup tables. 

Overall, this paper lacks significant research insights or novel observations addressing the challenges of efficient private inference for large language models.

$\bullet$ I also find the comparison with CrypTen unclear, as CrypTen is a 3PC method, whereas the authors use a 2PC method. At one point (Line#482-483), the authors mention excluding PUMA due to its 3PC computation, which raises questions about the consistency of the comparison.

$\bullet$ The authors should consider comparing their approximation methods with the polynomial approximations of LayerNorm, GELU, and Softmax used in *Zimerman et al., Converting Transformers to Polynomial Form for Secure Inference Over Homomorphic Encryption* (**ICML'24**). It’s essential to at least qualitatively contrast and position the merits (efficiency and predictive performance) of their approach with these recent polynomial approximations. 

$\bullet$ I find it quite challenging to follow the draft because the readability suffers from overly long paragraphs. Breaking up the content into shorter, clearer paragraphs would really help make the information easier to digest.

## Correction in the draft

$\bullet$ Figure 1 depicts the 2PC threat model for Post-LN configuration, however, in the caption the authors have mentioned CrypTen (which is 3PC). Also, in the FFN block diagram, one linear layer (after GELU layer) is missing. 

$\bullet$ Line#180: GELU necessitates the Gaussian error function---> BERT and other recent transformer-based models adopt the $Tanh$ approximated GELU implementation, which is significantly faster, even in cryptographic settings (e.g., Bumblebee), than the original GELU   (`torch.nn.GELU()`). See [1] and their implementation in hugging face activation libraries [2]

1. https://paperswithcode.com/method/gelu


2. https://github.com/huggingface/transformers/blob/main/src/transformers/activations.py#L49

---

confidence: 5
contribution: 2
initial_rating: 6
presentation: 2
questions: Please see the weakness section.

One other comment I have is that the experiment testbed is quite outdated. The author should consider more benchmark and larger scale model (This is out-of scope of the paper I think, but would be useful for future research), e.g. Vicuna and MMLU. ChatGPT is not Bert-like models.
soundness: 2
strengths: 1. The methods are well motivated by observations.
2. The approximations are quite novel and interesting.
summary: This paper introduces Comet to tackle to slow private inference for transformer based model. It is a plug-and-play method with an efficient approximation method to reduce the cost in non-linear computation. It achieves 3.9x less communication and 3.5x speed up compared to previous SOTA.
weaknesses: 1. Figures are (in my opinion, not a major weakness) not good enough: (1) Figure 1 seems like hand written, with small fonts and large empty spaces. (2) Figure 2 should be put in appendix or wrapfigure (since they are not the main point of the paper).
2. The 2-relu approximations have been proposed in literature, e.g. MPCFormer. Also, since approximation has been heavily studied, can the author possible make a small table to present what is in literature and what is not (to improve presentation)?

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 2
questions: Do you have code to reproduce your work since the "smu" function is rarely used?
soundness: 3
strengths: 1. Designing an improved approximation for non-linear functions in private Transformer inference is critical, and this paper's adoption of the "smu" function is both innovative and effective, demonstrating strong accuracy after fine-tuning.
2. The paper presents an effective solution for obtaining a good initial approximation in the context of 2PC-based secret sharing, which is a key contribution to reducing communication overhead.
summary: The paper introduces Comet, a framework aimed at efficient, privacy-preserving inference for Transformer models. Comet achieves this by unifying complex non-linear functions, like GeLU and Softmax, under a single protocol based on inverse square root calculations. It introduces a "double approximation" method to eliminate the need for extensive communication when finding initial approximations, further improving communication efficiency with a technique called "share flooding." Experimental results show that Comet reduces communication costs while maintaining good performance on models like BERT and RoBERTa on the GLUE benchmark
weaknesses: 1. The properties of the "smu" function are not sufficiently explained. While experiments indicate that it achieves 1-2% higher accuracy than the original GeLU and Softmax functions, the underlying mechanism remains unclear. Including a graph of the approximated function would help clarify this.
2. The explanation of the double approximation method is confusing. Since MPC typically operates on fixed-point numbers, it is unclear whether the double approximation method is designed for floating-point operations, which are costly in MPC.
3. It is worth noting that other approximation methods like Bolt, Bumblebee, and Puma do not require fine-tuning, and this difference should be mentioned for context.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: - As discussed in the “Weaknesses” section, there is concern about whether the proposed Comet framework can be applied to generative language models. Can the authors provide a feasible approach to extend the proposed framework to generative LLMs like ChatGPT and Llama?
- What if, in a practical application, the server model is not open-source (e.g., users can only call their API, like GPT and Claude)? Is there a way to still conduct private inference?
- The experiments in this paper are conducted on CPUs. However, Crypten supports GPU computing [1]. I wonder how the performance of Comet compares to Crypten on GPUs.
- It seems the major weakness of MPCFormer [2] is that it requires fine-tuning using knowledge distillation (KD). However, the idea makes sense for pure inference. I wonder if Comet could also benefit from KD.
- Encoder models are generally easy to batch. I wonder how different batch sizes affect the computation speed of Comet.

[1] https://github.com/facebookresearch/CrypTen?tab=readme-ov-file#installing-crypten  
[2] https://openreview.net/forum?id=CWmvjOEhgH-
soundness: 3
strengths: - This paper is well-motivated. Studying methods that enable private inference over large language models has great practical potential for privacy-constrained use cases, such as in financial and medical institutions. Private inference methods (e.g., MPC-based approaches) can serve as potential alternatives to pure local computing.
- The paper is well-written, and the techniques proposed in the Comet framework appear sound.
- The experimental results look promising, especially regarding accuracy on the GLUE benchmark.
summary: This paper proposes a secure Multi-Party Computing (MPC)-based private inference framework, called Comet. Comet aims to address the challenge of high communication overhead, particularly for non-linear computations in MPC-based Transformer model inferences. Existing solutions, such as Look-Up Table (LUT) methods and aggressive polynomial approximations, often result in either high communication costs or reduced model performance. To tackle this, Comet harmonizes non-linear functions like GeLU and Softmax into a single approximation function, eliminating the need for heavy communication in finding initial approximations for Newton iterations. A “share flooding” technique is also proposed to further optimize secure two-party computations. Comet was evaluated on encoder-based Transformer models like BERT and RoBERTa (as well as on the evaluation benchmark GLUE) to demonstrate that it is both communication-efficient (i.e., fast in inference) and more accurate over the benchmark compared to existing MPC-based private Transformer inference frameworks.
weaknesses: - The paper was motivated by using ChatGPT as an example, but all experiments were conducted on encoder-based Transformer models like BERT and RoBERTa. I understand that, architecturally, encoder models and decoder models differ mainly in language masking. However, models like GPT are autoregressive, which involves completely different computational intensity. For instance, autoregressive models generate tokens one by one, making it very difficult to batch inputs.
- Although MPC-based frameworks ensure that neither the client nor the server discloses information to each other, in Comet’s setting, the private information includes both the model weights and the users’ data. However, both parties still see the same (non-encrypted) output in these frameworks. While this may be acceptable for classification tasks conducted by BERT, it may not be suitable for generative models like ChatGPT (e.g., generated content could contain private user information).
- The absolute inference speed remains quite slow, e.g., around 30 seconds per batch, as shown in Table 3. Speeds like this make the proposed framework less practical.

Meta-review: This paper tackles the communication (and latency) overhead introduced by nonlinear operations in secure transformer inference. The authors replaced GELU and Softmax with a smoothed maximum unit (SMU) function and developed double approximation and share flooding to efficiently compute them in privacy-preserving settings.

Strength of this paper:
- secure transformer inference is an important problem
- the share flooding technique sounds intuitive

Weakness of this paper:
- several key components of this paper (e.g., approximation of non-linear operations) were developed in previous papers, and this paper seems to simply borrow them into secure inference, hence the ideas are not very refreshing.
- The proposed framework and experiments conducted are mainly for transformer encoders (e.g., BERT), yet the authors used ChatGPT, which is a decoder models with autoregressive generation at inference, as the main driver for this paper. This seems not coherent. 
- The absolute inference time even for BERT remained impractical

The weaknesses outweigh the strengths, hence a rejection.

Decision: Reject