Paper: SELECTFORMER: PRIVATE AND PRACTICAL DATA SELECTION FOR TRANSFORMERS (ICLR-2024)

Reviewer comment (Theory):
While a semi-honest setting is convenient and efficient for the research purpose, a setting against malicious adversary is more practical. The assumption based on faithful protocol execution by both parties is not quite strong in the real-world application.

Author response:
Response: We agree with the limitation of a semi-honest setting. Yet, it is still unsolved and sees active research. We hope the results from semi-honest settings could pave the way to more challenging settings with malicious parties. 
References:
[1] P. Mohassel and Y. Zhang. SecureML: A system for scalable privacy-preserving machine learning. In IEEE Symposium on Security and Privacy, 2017.
[2] S. Tan, B. Knott, Y. Tian, and D. J. Wu. CryptGPU: Fast privacy-preserving machine learning on the GPU. In arXiv 2104.10949, 2021.
[3] P. Mohassel and P. Rindal. ABY3: A mixed protocol framework for machine learning. In ACM Conference on Computer and Communications Security (CCS), 2018.

Reviewer comment (Novelty):
Since the setting of model and data owner is based on two parties. Instead of MPC, have you considered to solve this problem in the zero-knowledge proof setting? ZK approach allows model owner to select and appraise the data from data owner by privately committing parameters and data, and providing verifiable results without data leakage. It may be more efficient than MPC (not a conclusion and probably interesting for the future investigation).

Author response:
**A1**: That’s an interesting direction that we are looking forward to. Besides, we will cite and discuss ZKP in the Related Work section.

Reviewer comment (Theory):
In the threat model of Section 2.1, can you explain more about what can be revealed (like the purchase budget B mentioned in the Section 4.1) and what should be private (like model parameters and data points)? Or also merge the privacy guarantee in the Section 4.1 into the threat model.

Author response:
**A2**: Thanks for the suggestion! We have added a new Figure 1 and a “Workflow” paragraph in Section 1 to provide more encryption/reveal information. We intentionally keep “Privacy guarantee” in Section 4.1 to better introduce the multipass selection workflow. Therefore readers will not be confused by encountering these details too early.

Reviewer comment (Reproducibility):
What is the threshold of phase i? How should we determine when the early or later phases are?

Author response:
**A3**: We interpret “threshold” and “early or later phases” as how to decide the number of phases and the dimension of MLPs.
The number of phases depends on the trade-off between selection performance and computing resources. More phases may have better selection performance but run longer.
The MLP dimensions in each phase are determined by the schedule via offline grid search(Section 5.4 MLP hidden dimensions, Table 5, and a complete Table 6 in the Appendix).
For complex datasets, adding one larger dimension phase may increase the selection accuracy (Table 6 in the Appendix).

Reviewer comment (Novelty):
The first contribution—emulates high-dimensional nonlinearity with low-dimensional MLPs—is not a novel technique. Besides, training MLPs to approximate every non-linear operator is very costly. Does author evaluate the time cost here?

Author response:
**A1**: Thank you for pointing out the paper, but this is not applied to our problem because of the goal and the cost. (1) DeepONet wants to solve a different problem from us. It wants to solve the differential equations which will find solutions in input X’s domain. But we want to find an approximation of a function that gives us an output in a different range. (2) Our models have different inputs. Because DeepONet wants to solve differential equations, their inputs consist of both sides of the equation. They require both the left and right-hand side: u(x) and y. However, the only input of our MLPs is just data value x without the right-hand side y. (3) They need more than one big NN which is expensive on MPC. The Trunk net and Branch net have width==40 and depth==3 and 2 respectively. By contrast, our MLP approximation has as low as width==2 and depth==2. Therefore, DeepONet is much more expensive than ours. (Figure 2 has the MPC cost breakdown of each layer). The cost will be even more prohibitive on MPC with the Stacked DeepONet, which has much more Branch net inside.

There is another paper [The-X, ACL2022] mentioned in Section 1 we got inspiration from. We agree that approximation is a well-known trick but the particular design is important, we have revised our draft to clarify (introduction and contribution1). 

We will highlight our key discovery: MLP approximation, while cannot be used for generic transformer inference (because of low accuracy), is uniquely suitable for data selection. Such a novel use of MLP approximation is discovered for the first time. 

To show that our discovery is non-trivial, we will add the following supporting experiments, which show that MLP approximation, when used in generic inference, results in poor accuracy. Following MPCFormer, we replace three nonlinear modules in a BERT-base model with our approximations: Attention Softmax, Attention LayerNorm, and Feedforward Network Layernorm. There are 3*12 = 36 MLPs in the BERT model. In the experiments, two LayerNorm approximations are always data-aware, while there are both data-aware and fixed approximations of Softmax because of varying attention masks. The BERT with MLP approximation will be fine-tuned on the
benchmark before inference. Three versions of attention softmax approximations achieve no better than random guess performances: 
(1) Using fixed MLP for attention softmax. Our BERT with approximation achieves only 52.92% accuracy. 
(2) Using data-aware attention softmax approximation, which zeros out masked values on MLP outputs, has 49.08% accuracy. If we normalize the remaining values to [0, 1], the accuracy remains 49.08%. 
To better understand the influence of each nonlinear approximation, we did an ablation study that keeps just one kind of MLP (12 of them) in BERT. However, none of them achieve better than 50.92% accuracy. We further notice that adding just one attention softmax MLP to the model has no impact on the inference accuracy. However, adding one LayerNorm MLP will degrade the accuracy by 0.85% on average. These results show that having just one MLP will hurt the inference performance a lot; having 12 or even 36 MLPs will degrade model performance drastically. 
(3) Removing the attention mask to make MLP data-driven. It always has 50.92% accuracy, no matter with all three approximations or just one approximation. 
The poor inference performance and good data selection performance show that the MLP approximation is specifically suitable for data selection while impractical for model inference directly.

**A2**: MLP training is very cheap actually. They are simple models and we generate data in large batch size==128. The total training time for softmax approximation takes 20.88 minutes for 20K epochs, 5.25 minutes for LayerNorm approximation, and 0.68 minutes for the fuse approximation with the same number of epochs on a Nvidia RTX 2080Ti GPU. It takes only 26.81 minutes of MLP training time in total for one phase, which is cheap compared with tens of hours of selection time.
Additionally, MLP training should be done before private data selection. Therefore the training cost is even less because of training just once.

Reviewer comment (Novelty):
The MPC protocol hiding the computation behind the communication data exchange is not novel. Any MPC protocol would try to minimize the total latency in this way. It is not unique here. Thus, this contribution is a weak statement.

Author response:
**A3**: To clarify, MPC protocols do NOT specify the order of computation and communication (unless they are dependent); it is the role of MPC frameworks.  To our knowledge, common MPC frameworks such as Crypten do NOT hide computation behind communication automatically. 

We will highlight that: our scheduling does not just simply hide computation delays. It has to create such hiding opportunities (by co-executing tasks across the batch boundaries, see Sec 4.4)

We hope these additional explanations can address your questions above. We have updated the manuscript for better clarity (new Figure 1, Section 1 Workflow, and our first contribution).

Reviewer comment (Novelty):
This work seems to simply combine the techniques of data selection and secure inference on LLMs.

Author response:
We need to clarify that we don’t just "simply combine" some techniques. the approximation and proxy model design part is non-trivial. They are efficient in our data selection case, while just borrowing current methods will bring prohibitive private selection costs as mentioned in Section 1.

Reviewer comment (Theory):
Replacing high-dimensional nonlinearity with low-dimensional MLPs seems less general.

Author response:
There are prior works of approximating nonlinearity with MLP. As mentioned in Section 1, we are inspired by THE-X[ACL2022]. They use MLP approximation in homomorphic encryption. But we improve its efficiency by lowering the dimension, approximating the nonlinear operation as a whole, and data-driven parameters training.

As mentioned by other reviewers, Lu, Lu, et al. "Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators." Nature machine intelligence 3.3 (2021): 218-229. proposes using fully connected network to precisely approximate nonlinear operations for identifying differential equations.

We will highlight our key discovery: MLP approximation, while cannot be used for generic transformer inference (because of low accuracy), is uniquely suitable for data selection. Such a novel use of MLP approximation is discovered for the first time. 

We acknowledge that low-dimensional MLPs are NOT for generic transformer inference (New experiment results are in the Appendix). However, We customize it and make it work for an important problem (secure data selection). In other words, we contribute a specialty technique for an important, pervasive problem.

Reviewer comment (Novelty):
The batch evaluation is a widely used method in PPML and lacks novelty.

Author response:
We will highlight that: our system goes beyond standard batch evaluation in PPML, which will result in IO starvation (Sec 4.4). Our system co-executes tasks across batch boundaries in order to create parallelism. At a high level, it brings the job-stealing principle from parallel computing into the context of PPML.

Reviewer comment (Experiments):
Is this non-linear evaluation method only applicable in the proposed data selection setting? Can it be extended to general MPC-based LLM?

Author response:
**A1**: We interpret “LLM” in the question as Transformer models. However, we believe our method can also used on LLMs since they are also transformer-based.
Yes, it is for selection, not for inference.

Reviewer comment (Novelty):
What are the differences between MPC-based data selection on LLMs and secure inference on LLMs?

Author response:
**A2**: Selection has a lower requirement of the proxy model capacity. It only requires the output entropy’s rank to be consistent with that of a full model, but inference needs the proxy model to have the same output logits as those of the full model.

Reviewer comment (Experiments):
Why do the authors focus on LLMs? Can this method be applied to CNNs?

Author response:
**A3**: We focus on Transformer models because they have many softmax and LayerNorm modules in each layer, which are very costly for MPC. We didn’t try CNN because they have just a few softmax and LayerNorm modules.
We believe our method can be extended to CNN models, because our MLP approximation doesn’t rely on model architectures.

Reviewer comment (Reproducibility):
Does this work use a third party for generating correlated randomness for MPC similar to MPCFormer?

Author response:
**A4**: In principle, our design needs no third-party assistance. In particular, the Beavor Triples needed by us can be generated by oblivious transfers offline, without any third party.

Reviewer comment (Writing):
The paper does not clearly explain why each specific technique is necessary and what problems they aim to solve, making it difficult to grasp the reasons behind their use.

Author response:
Response: In Section 1, we talk about the challenges of private selection over MPC and explain the motivation of each contribution. In our Goal & techniques on page 2, we explicitly give the problems we want to tackle for each technique and the benefit it can bring.

Reviewer comment (Reproducibility):
The paper does not offer a clear protocol or source code to show exactly how each operation is performed; the technical aspects remain ambiguous without a protocol or algorithmic explanation.

Author response:
Response: We interpret “protocol” in the question as a generic description, not specific to the MPC protocol. Section 1 provides a clear overview of our system operation. (i.e. “protocol”). We interpret “source code” in the question as pseudo-code. Our innovation is a new proxy model architecture and selection workflow. Hence. figures are more suitable in this case.

Reviewer comment (Writing):
Overall, the paper is not yet ready for conference presentation and requires a complete rewrite to clarify each algorithm and protocol and make it more reader-friendly.

Author response:
Response: In Section 1, we provide a clear overview of our system operations. In terms of MPC protocol, we mentioned it in Section 2.1. In Section 2.2 and Section 4, we explain our machine learning algorithm in detail.

Reviewer comment (Reproducibility):
It is unclear how each operation is computed using MPC.

Author response:
**A1**: In Section 2.1 Threat model, we mention that all operations are supported by Crypten [Crypten, NeurIPS2021]. Before MPC operations, all data must be first encrypted. This includes both the model parameters as well as the inference data. On a high level, they can be seen as matrices. After encryption, the encrypted matrices still have the same shape as before the encryption, but all the entries now contain randomized values. Furthermore, during the MPC operations (addition and multiplication), the data entries remain randomized, but their shapes are still publicly revealed. It’s until both parties decide to share their respective data share to each other and add them up, that the data is decrypted.

Reviewer comment (Reproducibility):
It is unclear how entropy is calculated if MPC is used.

Author response:
**A2**: In Section 4.3, we mention that entropy is predicted by an MLP. Besides, entropy can also be calculated by the formula, they are logarithms and multiplications supported by Crypten.

Reviewer comment (Reproducibility):
It is unclear how the index with the highest entropy is jointly found.

Author response:
**A3**: We use QuickSort (Section 4.1 Multipass selection) to find the index with the highest entropy. But since MPC doesn’t support control flow, we reveal the comparison results between two data entropy. Since our goal is to find the index with the highest entropy, this revealing doesn’t hurt the privacy assumption.

Reviewer comment (Reproducibility):
The definition and details of the QuickSelect algorithm are not provided.

Author response:
**A4**: It is a textbook algorithm, c.f. Hoare, C. A. R. (1961). "Algorithm 65: Find". Comm. ACM. 4 (7): 321–322. doi:10.1145/366622.366647.

Reviewer comment (Presentation):
The interpretation of Figures 1 and 2 is unclear.

Author response:
**A5**: In Section 4.2, we have a detailed explanation of our algorithm and these two figures (They are Figures 3 and 4 now).

Reviewer comment (Writing):
The role and precise definition of the proxy model are not provided.

Author response:
**A6**: In Section 2.2(2) Proxy models, we cite prior works of creating lightweight proxy models for data selection and mention we tailor new proxy model architectures for our goal. In Section 4.2, we explain how to generate proxy models in detail, including their architectures and training process.

Reviewer comment (Reproducibility):
The dimensions and hyperparameters used for the MLP that replace the nonlinear function are not specified.

Author response:
**A7**: In Figure 4 and Section 5.4 MLP hidden dimensions, we mention that dimensions are 2/8/16 for a three-phase selection.

Reviewer comment (Experiments):
The data used for proxy model training is not described.

Author response:
**A8**: In Section 4.1 Pre-selection bootstrap, we explain that the model owner will use the bootstrap data S_boot to generate proxy models $M^^_{1..N}$.