Paper: xMLP: Revolutionizing Private Inference with Exclusive Square Activation (ICLR-2024)

Reviewer comment (Evaluation):
The author mentioned that they use Delphi's protocol to evaluate the proposed architecture inference. However, adding too many linear operations here would also increase the cost in offline phase. As the HE-based computation on lots of linear computation is quite costly.

Author response:
Regarding the concern that "adding too many linear operations here would also increase the cost in the offline phase," it's important to note that we use the combination of  3x3 depth-wise convolution and linear layers (matrix multiplication which is equivalent to 1x convolution). The MACs of linear operations on one xMLP block in our xMLP-M16 is 4.34M (4.2M for the fully connected layer, and 147K for the depth wise convolution layer), when compares to the standard convolution layer in ResNet18, the MACs are 18.8M, 38M, 9.4M, 2.4M in the first, second, third and forth block of ResNet. Our approach actually reduces computational requirements compared to standard convolutional. Moreover, convolution operations are not as efficient as fully connected layers in Homomorphic Encryption (HE).

Reviewer comment (Evaluation):
Besides, the multiple layers here would also increase the communication cost during offline phase. As you have to do ciphertexts exchange for preparing the secret sharing masks.

Author response:
In response to the concern that "the multiple layers here would also increase the communication cost during the offline phase," we emphasize that xMLP achieves superior results with fewer layers compared to ResNet models. Additionally, the number of layers requiring the exchange of latent features is only correlated with the number of activation layers. Consecutive linear layers do not necessitate repeated HE exchanges and can be combined in one exchange.

Reviewer comment (Experiments):
This paper lacks comparison with the latest methods. Authors should compare with the following papers:

a. Jha, Nandan Kumar, and Brandon Reagen. "DeepReShape: Redesigning Neural Networks for Efficient Private Inference." arXiv preprint arXiv:2304.10593 (2023).

b. Souvik Kundu, et al. Learning to linearize deep neural networks for secure and efficient private inference. International Conference on Learning Representation, 2023.

c. Kundu, Souvik, et al. "Making Models Shallow Again: Jointly Learning to Reduce Non-Linearity and Depth for Latency-Efficient Private Inference." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

d. Zeng, Wenxuan, et al. "MPCViT: Searching for Accurate and Efficient MPC-Friendly Vision Transformer with Heterogeneous Attention." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

e. Zhang, Yuke, et al. "SAL-ViT: Towards Latency Efficient Private Inference on ViT using Selective Attention Search with a Learnable Softmax Approximation." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

f. Dhyani, Naren, et al. "PriViT: Vision Transformers for Fast Private Inference." arXiv preprint arXiv:2310.04604 (2023).

Author response:
Dear reviewer,

Thank you for your insightful comments and the opportunity to clarify aspects of our research. We appreciate the opportunity to address your concerns and elaborate on our work.

We compare our results with the recent published works SENet (b. Souvik Kundu, et al. Learning to linearize deep neural networks for secure and efficient private inference. International Conference on Learning Representation, 2023).
While SENet's comparison with SNL shows only slightly better ReLU efficiency, it's important to note that SENet still requires 50K ReLUs to achieve 75.25% accuracy, and this is close to the smallest SNL model compared in our paper which has a latency of 1.07 seconds. Our xMLP-M16 only has 0.27 seconds online latency with 75.52% accuracy. This shows our xMLP remains the advantage when compared to recent works. We will include the results that compare with more recent works in our paper.

Paper [c]  uses knowledge distillation while we do not, thus we cannot directly compare with.

Paper [d][e] uses 3pc protocol for latency measurements, that we cannot directly compare with since our evaluations are under 2pc setting.

We hope this response addresses your concerns and provides further clarity on our work. We are grateful for your feedback and look forward to any additional insights you may have.

Reviewer comment (Experiments):
**Comparison with SOTA in PI:** One of the core issues with this paper is the lack of comparison with existing SOTA in PI, SENets [2]. Hence, the claim of the authors for pushing the SOTA in PI is not valid.

Author response:
While SENet's comparison with SNL shows only slightly better ReLU efficiency, it's important to note that SENet still requires 50K ReLUs to achieve 75.25% accuracy, and this is close to the smallest SNL model compared in our paper which has a latency of 1.07 seconds. Our xMLP-M16 only has 0.27 seconds online latency with 75.52% accuracy. This shows our xMLP remains the advantage when compares to the recent works. We will include the results that compare with more recent works in our paper.

Reviewer comment (Evaluation):
**FLOPs comparison are not included when comparing CNN-based PI:** Another major concern with this paper is excluding the FLOPs count when comparing with CNN-based PI, in Table 2. The authors have included the params counts, which are inconclusive as the parameter counts do not signify anything in the context of PI overheads. Since this paper, which used MLP-mixer as a backbone network, compares with the prior PI methods, which used CNN as backbone architecture, comparing only the non-linearity overheads does not tell the entire story (when comparing two distinct architecture). Moreover, the FLOPs cost cannot be ignored in PI, as they carry significant latency overhead for end-to-end latency [1]. The assumption of processing all the FLOPs offline is valid only when optimization is performed for a single inference request in isolation. However, in real-world scenarios, private inference requests arrive at non-zero rates, and processing the entire FLOPs offline becomes impractical due to limited resources. Consequently, offline costs are no longer truly offline, and FLOPs start affecting real-time performance, as illustrated in Figure 7 of the paper [1]. This effect can be exacerbated by networks with higher FLOP counts, as proposed by the authors.

Author response:
The ResNet18 in our experiments have 3.4 GFLOPS and xMLP-M16 has 0.4 GFLOPs for the same input size (32x32 rgb image). We will include more comparisons of FLOPS in our paper. We also want to emphasize that during the offline phase that convolution can be slower than it is on plaintext as they require more computations due to its complexity, and they are not as optimized in HE as they are on plaintext. Therefore, we focused on comparing the end-to-end online latency.

Reviewer comment (Experiments):
**Ambiguity in timing analysis:** The micro-benchmarking for timing analysis performed in Section 4.3.2 is full of ambiguity. First, the paper said that it used a 2pc protocol, and then it included CrypTen for benchmarking. Note that Delphi and CryptFlow2 have 2pc implementation; however, CryptTen has 3pc implementation as they used TTP for generating beaver triples in the offline phase. It seems that authors have mixed 2pc and 3pc implementation for their benchmarking.

Author response:
We focus on online latency, CrypTen framework use Torch for computations and can provide good performance for online inference and GPU accelerations. We have modified the CrypTen framework to perform the DELPHI's online inference. The trusted third party (TTP) in CrypTen can be seen as an alternative to offline processing, they both only need to perform once for each model and have no impact on the online inference.

Reviewer comment (Novelty):
**Lack of any new insight:** The discussion presented in Section 3.1 for the sparsity-inducing property of ReLU, redundancy in global features, and "information compounding" is not novel. See [9].

Author response:
Our novel contribution is not centered on ReLU's sparsity-inducing property. Instead, we emphasize that this property leads to suboptimal results with square activations in CNNs. Using MLP-like structures, we effectively mitigate the issues associated with square activations.

Reviewer comment (Novelty):
**Authors are ill-informed about the relevant literature:** Authors need to tone down the claim for substituting all the ReLUs with polynomials for the first time. Prior work [6, 7, 8] replaced all the ReLUs with polynomials, and some of them demonstrated their efficacy on the ImageNet dataset too. Nonetheless, the prior work required to have higher-degree polynomials [6] or needed to use LayerNorm to mitigate the accuracy drop stemming from substituting all the ReLUs with polynomials [7]. Also, on page 9, the authors claim that implementation for offloading private computation to GPUs is not available. See the relevant work in [3,4,5].

Author response:
While these works utilized polynomial activation, they also simultaneously used ReLU functions [8],(https://github.com/zarekxu/QuadraLib/blob/main/image_classification/models/resnet.py#L83C18-L83C18). We also highlight that high-degree polynomials lead to impractical latency in PI [6]. Additionally, LayerNorm, requiring normalization operations, is unsupported in existing PI frameworks[7], even with Newton-Raphson's approximation of $1/x$, it requires at least quadratic computations. Thus, our use of simple quadratic functions and the absence of BatchNorm are significant contributions that should not be ignored. The comparison with [6][7][8] are unfair in the context of PI.  

The references [3][4][5] doesn't support for Garbled Circuits (GC) computations on GPU. Instead, they ORCA[3] and CryptGPU[5] are based on secret sharing, and Piranha[4] uses edaBits. We want to emphasize, even with recent ReLU protocols, square activations (multiplications) are still significantly faster than ReLU on both CPU and GPU. For example, as reported in CryptGPU[5], the VGG-16 online inference latency for ReLU with a batch size of 32 using CryptGPU on CIFAR dataset (32x32 image size) still takes 4.74 seconds, while the online inference of square activations in xMLP-M16 with the same batch size and image size only takes 0.29 seconds on CPU and 0.12 seconds on GPU.

Reviewer comment (Experiments):
In this paper, Table 4 reports that the performance of the square function is very poor in CNNs, e.g., ResNets. This is due to the lack of sufficient surveys.

Author response:
We also emphasize that our contribution lies in the fact that instead of simply assuming quadratic activation functions are insufficient, we believe that the primary reason quadratic functions do not perform well in ResNets is due to limitations in the network architecture.

Reviewer comment (Novelty):
The results of this paper have very marginal novelty, such that they show that good results can be obtained with a square function in MLPmixer-type models. Also, they do not provide any information on the coefficient of used square functions.

Author response:
Our research uniquely demonstrates that a simple square function, $y=x^2$ (the coefficient for the quadratic term is 1, for the linear term is 0, and the constant term also has a coefficient of 0.), and it achieves results close to those with ReLU.