Paper: PriViT: Vision Transformers for Fast Private Inference (ICLR-2024)

Reviewer comment (Novelty):
While the paper introduces an interesting approach, its novelty appears constrained. The core concept of utilizing binary gates to manage the number of nonlinear operations through $\ell_0$ sparsity has been proposed in SNL (Cho et al., 2022b). The proposed method seems like an adaptation of SNL for ViTs. This adaptation is intriguing, but a deeper exploration into its unique contributions compared to SNL would provide more clarity on its novelty.

Author response:
**PriViT's novelty compared to SNL**: It is important to highlight that the SNL algorithm is equipped to linearize only ReLU activations, and more generally is applicable to networks that only exhibit a single type of linearity. The algorithmic innovation in PriViT is to design a method that can jointly linearize two or more types of non-linearities. Additionally, we equip our method to Taylorize the attention mechanism with an appropriate candidate function to make it suitable not only for ViTs, but potentially any future architecture based on attention mechanisms. Finally we propose a GELU replacement strategy to further reduce the latency of ViTs in low latency budget settings.

Reviewer comment (Presentation):
The paper's organization raises some concerns. While critical figures, and experimental results are put in the supplementary material, the main body of the paper exhibits noticeable empty spaces. It would enhance the paper's comprehensibility and impact if these essential elements were incorporated directly into the main content, making the narrative more coherent for readers without necessitating frequent references to supplementary sections.

Author response:
**Paper organization**: We have revised our paper to incorporate your suggestions, we hope the new version is easy to follow for readers. All changes are highlighted in blue, and many figures and tables are pulled from the Appendix into the main paper.

Reviewer comment (Evaluation):
In Figure 1, the computational cost is quantified using the metric of # AND gates. However, the methodology or criteria for determining the number of AND gates associated with a nonlinear operation remains ambiguous in the paper. Clarifying this calculation or providing a concise explanation within the main content would enhance the reader's understanding.

Author response:
**Calculation of AND gates**: The method to determine the number of AND gates is by first breaking down a complex nonlinear function into several basic operations. These operations are already synthesized [TinyGarble] and included in the EMP-Tool library (a commonly-used GC library). Thus we simply add all the AND gates of these basic operations to arrive at the total number of AND gates. We moved relevant figures to the main paper from Appendix to improve readability. Section 2, Figure 2 mentions this explanation in the main paper.

Reviewer comment (Experiments):
The experimental validation is primarily focused on small-scale datasets, which could limit the generalizability of the findings. For a comprehensive evaluation, it would be beneficial for the authors to include results from large-scale datasets like ImageNet. Such experiments would provide a more holistic view of the method's efficacy and scalability.

Author response:
**PriViT on ImageNet**: The suggestion to test PriViT on datasets like ImageNet that have resolution higher than 64 by 64 is essential to highlight the universality of this method. We would like to refer to Section 4.1 (paragraph titled Architecture and data set) where we mention that all images are up sampled to 224 by 224 before being passed in as input to the ViT. Hence all the results presented in the paper are based on the resolution of 224 by 224. To our best knowledge, SENet has not provided publicly available code to train their models, which makes quite difficult to reproduce their method for ViTs.

Reviewer comment (Novelty):
Based on Figure 2, both MPCViT and MPCViT+ seem to outperform the proposed method in achieving a superior Pareto balance between accuracy and latency. This observation raises the question: what distinctive advantages or contributions does the proposed method offer? Is it perhaps more streamlined in terms of training efficiency or some other metric? A clearer delineation of the method's unique strengths would be beneficial for readers.

Author response:
**PriViT's advantage compared to MPCViT**: Based on your suggestions we'd like to present some additional data to highlight the efficient training methodology of PriViT compared to MPCViT. Since training codes are not public for MPCViT, we cannot compare GPU memory usage, so we only compare the number of training epochs as it is reported in their paper. MPCViT undergoes a NAS stage where an architecture is searched for 300 epochs, followed by 600 and 300 epochs of training the searched model on CIFAR 100 and TinyImagenet respectively, whereas PriViT simply trains a ViT in one go. Table 5 in the main paper highlights that for an isoaccuracy of ~64% PriViT requires ~50% less training epochs compared to MPCViT on TinyImagenet. On CIFAR 100 for an isolatency of ~75 M PriViT requires ~50% of the training epochs of MPCViT on TinyImagenet. 

Thus PriViT not only achieves a superior performance on latency, but it does so in a way that is highly efficient, and scalable to larger datasets. This makes PriViT an ideal choice for private inference when a business has personalized models for their individual users, and would like to make them PI friendly.

Reviewer comment (Evaluation):
Observations from Table 4 present some anomalies. Specifically, for PriViT-R, there is a counterintuitive trend where increased latency corresponds to a decline in performance. This is perplexing as one would typically expect a trade-off where longer computation times would yield better results. An explanation or insight into this apparent contradiction would enhance the clarity of the findings.

Author response:
**PriViT-R anomalous behavior**: The unexpected performance of PriViT in Table 4 can be attributed to how PriViT R models are derived. By replacing GELUs in PriViT G models (developed as per our algorithm) with ReLUs, we create PriViT R models. Although individual ReLU and GELU are similar, repeated replacements amplify errors, impacting the model's performance, especially at high GELU budgets. This explains why models with higher latency (82M) underperform compared to those with lower latency. We have revised our paper to recommend PriViT R's use only at low GELU budgets, where it effectively reduces latency without compromising performance.

Reviewer comment (Writing):
The clarity of notational conventions could be enhanced. For instance, the terms PriViT-R and PriViT-G appear in the text without explicit definitions or context. Providing a clear explanation or description of what these notations signify would aid in a more comprehensive understanding of the content for readers.

Author response:
**Notation clarity**: We have updated our main paper to give more clarity on our notations. All updates are in blue.

Reviewer comment (Experiments):
**Experiments:** The authors have conducted sufficient comparative experiments conducted on 3 datasets (CIFAR-10/100 and TinyImageNet). However, the image resolutions are no more than $64\times 64$, which is rather small compared with commonly-used datasets like ImageNet-1k and Caltech-101/256. It will be interesting to see ImageNet results and compare with SENet if possible.

Author response:
**PriViT on Imagenet**: We support the suggestion to test PriViT on datasets such as Imagenet and Caltech-101/256, that have a native resolution higher than 64x64. We would like to point out that  in Section 4.1 (paragraph titled Architecture and data set) where we mention that all images are upsampled to 224x224 before being passed in as input to the ViT, so our PriViT algorithm and implementation natively works for ImageNet resolutions without any changes. The sole reason for focusing on CIFAR-10/100 and TinyImageNet in our experiments is that almost every previous approach on private inference uses these smaller (and lower-resolution) datasets as benchmarks.

Reviewer comment (Experiments):
**Experiments:** The authors' exclusive use of ViT-Tiny for comparison is insufficient to establish the method's universality. It is strongly recommended that the authors broaden their evaluation to include a variety of ViT architectures, such as DeiT, and models with diverse parameters, including ViT-Small and ViT-Base. Such an expanded assessment would be of significant interest to the community, particularly in light of the flourishing development of large Transformer-based vision models.

Author response:
**PriViT on different architectures**: Fig 8 provides a comparison of PriViT over different model variants. We see that PriViT is able to produce a similar latency-accuracy tradeoff across different model sizes, emphasizing its generalizability.

Reviewer comment (Experiments):
**Experiments:** In Table 3, the authors have presented PriViT and MPCViT in distinct hyper-parameter configurations, giving rise to concerns regarding the fairness of the comparisons.

Author response:
**PriViT's comparisons with MPCViT**: We acknowledge the concerns about the fairness of comparisons in Table 3 due to differing hyper-parameter configurations. To re-iterate, PriViT is implemented at 224x224 resolution, which aligns with more challenging “real-world” benchmarks such as ImageNet. This is in comparison with MPCViT, which only supports lower resolutions. This crucial distinction underscores PriViT's broader applicability and superior adaptability. Despite the comparison not being apples-to-apples in hyper-parameters, it demonstrates PriViT's robust performance across more realistic scenarios, taking a step towards private inference for practical transformer architectures.

Reviewer comment (Experiments):
**Experiments:** Furthermore, I strongly encourage the authors to furnish additional results to showcase the superiority of their proposed method. This could encompass metrics like GPU memory usage and training time in comparison to other existing methods.

Author response:
**PriViT's advantage compared to MPCViT**: Based on your suggestions we'd like to present some additional data to highlight the efficient training methodology of PriViT compared to MPCViT. Since training codes for MPCViT have not been made public, we can only compare the number of training epochs as reported in their paper.

MPCViT undergoes a NAS stage where an architecture is searched for 300 epochs, followed by 600 and 300 epochs of training the searched model on CIFAR 100 and TinyImagenet respectively, whereas PriViT simply trains a ViT in a single run. Table 5 highlights that for an isoaccuracy of ~64% PriViT requires ~50% less training epochs compared to MPCViT on TinyImagenet. On CIFAR 100 for an isolatency of ~75 M PriViT requires ~50% of the training epochs of MPCViT on TinyImagenet.

Therefore, PriViT not only achieves a superior performance on latency, but it does so in a way that is efficient and scalable to larger/higher resolution datasets compared to MPCViT.

Reviewer comment (Experiments):
I am curious about the teacher model size and can it boost more performance if we use a larger teacher?

Author response:
Based on your suggestion we are currently running experiments with different teacher models to see its effect on the performance. Due to time constraints, we could not finish our experiments with larger teacher models in this revision, although our preliminary results show promising improvements. We plan to report these findings in a future update to our work.

Reviewer comment (Experiments):
As you states that later layers have a larger number of linearized units, can we first try to linearized the later layers then try to linearized the earlier layers?

Author response:
The suggestion to linearize later layers first is a good one, although it is important to note that we do not know the distribution of nonlinearity beforehand (the PriViT algorithm automatically deduces this by optimizing for the switching patterns for all layers simultaneously).