PriViT: Vision Transformers for Private Inference
Authors: Naren Dhyani
Jianqiao Cambridge Mo
Patrick Yubeaton
Minsu Cho
Ameya Joshi
Siddharth Garg
Brandon Reagen
Chinmay Hegde
Venue: Accepted by TMLR ()

Abstract: The Vision Transformer (ViT) architecture has emerged as the backbone of choice for state-of-the-art deep models for computer vision applications. However, ViTs are ill-suited for private inference using secure multi-party computation (MPC) protocols, due to the large number of non-polynomial operations (self-attention, feed-forward rectifiers, layer normalization). We develop PriViT, a gradient-based algorithm to selectively Taylorize nonlinearities in ViTs while maintaining their prediction accuracy. Our algorithm is conceptually very simple, easy to implement, and achieves improved performance over existing MPC-friendly transformer architectures in terms of the latency-accuracy Pareto frontier.

Published in Transactions on Machine Learning Research (10/2024)
PriViT: Vision Transformers for Private Inference
Naren Dhyani New York University
whd76820myp..cm edu
Jianqiao Cambridge Mo New York University
jm87822nyy.edd edu
Patrick Yubcaton
wpy2004 Dryy. edu
New York University
Minsu Cho
chomd90@gmaiilcommmm com
New York University
Ameya Joshi New York University
smeyy.j0050@maailccom
Siddharth Garg
sg1 75 75175my...cmm edu
New York University
Brandon Reagen
5rjr.myy..cmmmmmm edu
New York University
Chinmay Hegde
chinmay. Onyu. edu
New York University
Reviewed on OpenReview: https: // openreview. net/ formm? id= 3C=
Abstract
The Vision Transformer (ViT) architect ure has emerged as the backbone of choice for state-of- the-art deep models for computer vision applications. However, VYTs are ill-suited for private inference using secure multi-party computation (MPC) protocols, due to the large nunber of non-polynomial operations (self-attention, feed-forward rectifiers, layer normalization). We develop PriVIT, a gradient-based algorithm to selectively Taylorize nonlinearities in ViTs while maintaining their prediction accuracy. Our algorithm conceptually very simple, easy to implement, and achieves improved performance over existing MPC-Iriendly transformer architectures in terms the latency-accuracy Pareto frontier.
1 Introduction
Deep machine learning models are increasingly being deployed by cloud-based providers, accessible only by API calls. In such cases, user data privacy becomes paramount, motivating the setting of private inference (PI) using secure multiparty computation (MPC). In its simplest form, MPC-based private inference is two-party setup where user (the first party) performs inference of their data model whose weights are owned by the cloud service provider (the second party), with both sides encrypting their inputs using cryptographic techniques prior to inference.
Motivation:
The main technical barrier to widespread deployment MPC-based PI approaches is the lorge number of nonlinear operations present a deep neural network model. Private exccution of linear (or low-dogroe polynomial) operations can be made fast using cryptographic protocols like homomorphic encryption and/or secret sharing. However, private execution of nonlinear operations (such ReLUs softmax operations)

---

Published in Transactions on Machine Learning Research (10/2024)
Table 1: Comparison of various MPC-friendly approaches for deep image classification. NAS stands for neural architecture search; GD stands for gradient descent; CCT stands for Compact Convolutional Transformer. Our approach, PriVIT, adaptively replaces various nonlinearities present transformers with their Taylorized versions in order to reduce PI latency costs without drop in accuracy.
Approach
Arch
Methods
Units removed
Delphi (Mishra et al., 2020)
ConyNets ResNets
NAS
Poly Approx.
ReLU Layers
CryptoNAS (Ghodsi et al., 2020)
NAS
ReLU Layers
Sphynx (Cho et al., 2021)
ResNets
NAS
ReLU Layers
DeepReDuce (Jha et al., 2021)
ResNets
Manual
ReLU layers
SNL (Cho et al., 2022)
ResNets
GD
Individual ReLUs
SENet (Kundu et al., 2023)
ResNets
GD
Individual ReLUs
MPCFormer (Li et al., 2022)
BERT
NAS - Poly Approx.
GELU Layers, Softmaxes
MPCVIT (Zeng et al., 2022)
VIT
NAS - Poly Approx.
GELU layers, Softmaxes
SAL-VVIT (Zhang et al., 2023)
CCT
NAS + Poly Approx.
Self-Attention Layers,
Softmaxes
RNA-VIT (Chen al., 2023)
CCT
Compressed Attention + Poly Approx.
Sele-Attention Layers Softmaxes Individual GELUs, Softmaxes
PriVIT (this paper)
VIT
GD + Poly Approx.
require Yao's Garbled Circuits, incurring high latency and storage overhead. Thus, unlocking fast, accurate, and efficient PI requires rethinking network design.
Consequently, emerging line of work over the last five years has made several advances towards the design of "MPC-friendly" models; see more discussions below in Section 2. These methods approach PI from various angles. Approaches such as Delphi (Mishra et al., 2020) or Circa (Ghodsi et al., 2021) ) propose to replace ReLUs with MPC-friendly approximations, while approaches such as CryptoNAS (Ghodsi et al., 2020) and Sphynx (Cho al., 2021) use neural architecture search (NAS) to search for network backbones with a minimal number of ReLUs. Peng et al. (2023) propose hardware-aware ReLU-reduced networks achieve better latencies. Some of the latest approaches in this direction (SNL. by Cho et al. (2022), and SENet by Kundu et al. (2023)) derive inspiration from network pruning.
However, the above body of work has gaps. The overwhelming majority of Pl-aware model approaches have focused on convolutional architectures, and have largely ignored transformer models. In particular, the proper application of MPC to vision transformer (ViT) architect remains less well-studied; see Table Vision transformers (Dosovitskiy et al., 2020) currently list among the best performing deep models in numerous computer vision tasks, spanning image classification, generation, and understanding, On the other hand, vision transformers very bulky, possessing enormous number of nonlinear operations of different types: GELUs, softmaxes, and layer norms.
We are aware of three published approaches to addressing private inference for vision transformers. MPCVIT (Zeng et al., 2022) proposes a MPC-aware NAS algorithm, along with a polynomial simplifi- cation of the attention mechanism coupled with knowledge distillation. SAL-VIT (Zhang et al., 2023) introduces a Icarnable approximation for the softmax and perform NAS to choose between two types of attention layers. RNA-VIT (Chen al., 2023) proposes an altention architecture with a compressed attention map and Taylor approximation of the softmax. The last two works assume hybrid ViT models, different cryptographic security model, and have not made their models public, making direct apples-to-applies comparisons difficult; in this paper, we will focus on MPCVIT, which (to us) is the gold standard for private inference over vision transformers.
Our contributions and techniques. We introduce PriVIT, very simple algorithm for designing MPC-friendly vision transformers. PriVIT either improves upon (or is competitive with) MPCVIT both in terms of latency and accuracy on Tiny as well CIFAR 10/100. Therefore, PriViT presents
2

---

Published in Transactions on Machine Learning Research (10/2024)
Table Accuracy-lateucy tradeofls between PriVit and MPCVIT. All latencies are calculated with the Secretslow (Ma et al., 2023) framework using the SEMI2K (Cramer et al., 2018) protocol. Detailed methodology is reported in Appendix Left: Comparison of PriVit versus MPCViT on Tinylmagenet. PriVIT achieves 6.77x speedup for an isoaccuracy of approximately 61%. Right: Comparison of PriVit versus MPCVIT on Cifar10. PriVit achieves 1.14x speedup for an isoaccuracy approximately 94%. Bottom: Comparison of PriVit versus MPCVIT on Cifar100.
PriVit achieves 1.05x speedup for isoaccuracy of approximately 78%.
PriViT
MPCVIT
PriVIT
MPCVIT
Acc
Latency (s)
Acc
Latency (8)
Acc
Latency (s)
Acc
Latency (s)
78.88
17.357
62.55
68.095
96.58
7.935
94.3
7.273
78.16
15.303
63.7
53.992
95.88
7.164
94.2
6.791
75.5
12.124
63.36
40.367
95.58
6.789
94.1
6.606
64.46
9.350
62.62
25.996
94.58
6.382
93.6
6.200
PriVIT
MPCVIT
Acc
Latency (s)
Acc
Latency (s)
80.6
7.316
77.8
7.099
79.78
7.12
76.9
6.782
79.27
6.989
76.9
6.590
78.49
6.782
76.4
6.228
algorithmic approach for designing "privacy-friendgy vision transformers; see Table 2 for quantitative comparisons.
At a high level, our approach mirrors the network linearization strategy introduced in the SNL method by Cho et al. (2022). Let us start with a pre-trained ViT model with frozen weights, but now replace nonlinear operations with their switched Taylorized versions, i.e.;:
replace each GELU unit, GELU(x:), by:
GELU(Z)) + (1 c,)x,.
replace each row-wise softmax, Softmax(X,) by:
s,Softmax(X)) + (1 (11s)))+((Auaaaaaa))a SquuaeedAAtt(().)..
where SquaredAttn is just the unnormalized quadratic kernel, and (i, are binary switching variables. These switches decide whether to retain the nonlinear operation, or to replace with its Taylor approximation (linear the case of GELU, quadratic the case of softmax!). Having defined this new network, we initialize all switch variables to 1, make both weights as well switches trainable, and proceed with training using gradient descent.
Some care needs to be taken to make things work. We seek to eventually set most of the switching variables to zero since our goal is to replace most nonlinearities with linear units or low-degree polynomials; the surviving switches should be set to one. We achieve this by augmenting the standard cross-entropy training loss with a L i-ppenalty term that promotes sparsity in the vector of all switch variables, apply homotopy-style approach that gradually increases this penalty sufficient sparsity not reached, and finally binarize the variables via rounding. See Section 3.
Implications. We note that (Cho et al., 2022): judiciously replace both GELUs and softmax operations in vision transformers with their línear (or polynomial) approximations. However, they achieve this via a fairly complex MPC-aware NAS procodure. A major technical contribution their work is the identification of a (combinatorial) search space, along with a differentiable objective to optimize over this space. Our PriVIT algorithm conceptually much
method, MPCVIT, also follows a similar strategy
Via several ablation studies we justify why choose these particular approximations for these functions.
3

---

Published in Transactions on Machine Learning Research (10/2024)
Cost COSLLAAND LDAAD cates in L000
add mul div uep
Sidked
25-
6q
GAUU
Layonomm
Figure Breakdown of latency in ViT-Tiny model of different non-linearittes based Delphi.
Detailed steps of benchmarking the non- linearity cost for softmax.
Figure 2:
simpler, and can be applied out-of-the- box to any pre-trained ViT model. Moreover, we directly obtain a full family of model checkpoints during a single training run, at the cost of storing auxiliary mask variables. We commit to releasing all code and data neoded reproduce our results post-pecr review.
While our focus in this paper is sharply private inference, our results also may hold two broader implications for transformer architectures used in computer vision:
1. First, most nonlinear operations in transformers appear to be superftuous. PriVIT is able to remove nearly 83% of GELUs and 97% softmax operations with less than 0.5% reduction in accuracy over CIFAR-100 KKrizhevsky et al., 2009).
2. Second, given a target overall budget of softmaxes and GELUs, PriViT overwhelmingly chooses retain most of the nonlinearities in earlier layers, while discarding most of the later ones. These suggest that there is considerable room for designing better ViTs than merely stacking identical transformer blocks.
We leave a deeper investigations of these implications to future work.
2 Preliminaries
Private inference. Prior work on private inference (PI) have proposed methods that leverage existing cyyptographic primitives for evaluating the output of deep networks. Cryptographic protocols can be categorized by choice of ciphertext computation used for linear and non-lineas operations. Operations are computed using combination of: (1) secret-sharing (SS) (Shamir, 1979; Micali et al., 1987); (2) partial homomorphic encryptions (PHE) (Gentry & Halevi, 2011), which allow limited ciphertext operations and (3) garbled circuits (GC) (Yao, 1982; 1986).
In this paper, our focus is exclusively on the Delphi protocol (Mishra et al., 2020) for private inference. Wo choose Delphi as a matter of couvenience; the general trends discovered in our work hold regardless of the encryption protocol, and validate this we measure latency of our PriViT-derived models using multiple protocols. Delphi assumes the threat model that both parties are honest-but-curious. Therefore, each party strictly follows the protocol, but may try to learn information about the other party's input based on protocol transcripts (Wang et al., 2022; Peng et al., 2023; Lu et al., 2021; Qin et al., 2022).
Delphi combines cryptographic primitives such as secret sharing (SS) and homomorphic encryptions (HE) for all linear operations, and garbled circuits (GC) for ReLU operations. For convolutional architectures, the authors of Delphi shows empirical evidence that ReLU computation requires 90% of the overall private inference time for typical deep networks. As remody, Delphi and SAFENET (Lou et al., 2021) propose neural architecture search (NAS) to selectively replace ReLUs with polynomial operations. CryptoNAS (Ghodsi et al., 2020), Sphynx (Cho et al., 2021) and DecpReDuce (Jha et al., 2021) design new ReLU efficient architectures by using macro-scarch NAS, micro-search NAS and multi-step optimization respectively.
Only recently has this line of literature turned its focus to transformer-stylle models. Approaches in this vein include MPCVIT (Zeng et al., 2022), SAL-VIT (Zhang et al., 2023), and RNA-VIT Chen et al. (2023). MPC-ViT uses NAS techniques to find cheaper approximations of the softmax attention mochanism, while SAL-VIT and RNA-VIT propose new types of attention mechanisms. The latter two works claim state-of- the art over MPCVIT; however, found that direct. comparison is not possible for two reasons: (i) these
4

---

To begin, we focus on softmax and GeLUs and ignore layernorms; we found that these were far harder Taylorize. For the former, we introduce auxiliary variables to act switches. Given fw, let C and 5 be the total number of GeLUs and softmaxes. Further, let S = [$1, s2.... ss] and C [=[c,, 02......... ca] be collections Published in Transactions on Machine Learning Research (10/2024)
use compact convolutional transformers (CCTs), and (ii) they use the Crypten protocol for benchmarking MPC latency, which assumes the existence of trusted third party (TTP). The lack of open source code or checkpoints for these two works makes it infeasible to port over their results to our comparison framework. Thus, our paper will focus on comparisons to MPCViT.
Protocols for nonlinearities. To standardize across different types of non-linear activations, we compare their Delphi (online) GC computation costs. We use the EMP Toolkit (Wang et al., 2016), widely used GC framework, to generate GC circuits for nonlinear functions.
High-performance GC constructions implement AND and XOR gates, where XOR is implemented using FreeXOR (Kolesnikov & Schneider, 2008) and AND using Half-Gate (Zahur et al., 2015). With FreeXOR, all XOR gates are negligible, therefore we count the number of AND gates as the cost of each nonlinear function (Mo et al., 2023). To be consistent with prior work (Ghodsi et al., 2021), the activation functions also consider value recovery from Socret Sharing. Figure (left) breaks down the GC cost of ViT for different nonlinearities, and (right) shows the # AND gates in Softmax and GeLU. Figure 2 breaks down softmax into fundamental operations, these operations already syuthesized and included in the EMP Toolkit library. Thus simply add all the AND gates of these basic operations to arrive at the total number of AND gates of softmax operations.
Setup. Let fw : IRadd_ + [0. 1| be a vision transformer that takes as input n tokens (each of d dimensions) and outputs a vector probabilities for each of C' classes. Each of these tokens a patch sampled from the original image, X and is indexed by As described, the transformer architecture consists of stacked layers of multi-headed self-attention blocks with nonlinearities like GeLU (Hendrycks & Gimpel, 2016) and Layernorm (Ba et al., 2016). ViTs use dot-product self-attention (sce Equation 1) which additionally consists of row-wise softmax operations.
Softmax (XW, W " XW..
vd
To frame the computational challenges inherent to Vision Transformers (ViTs), we consider the ViT-base (12 layer) model designed for 224 x 224 images. Overall, this architecture consists of(approximately) 726, 000 GeLUs, 28, 000 softmax, and 4000 layer norms. All the non-lincarities, when viewed through the lens of the Delphi protocol, become extremely resource-intensive operations.
3 PriViT: MPC-Friendly ViTs
We propose to learn an architecture that circumvents these computationally heavy operations. Our approach surgically introduces appropriate Taylor approximations of the GeLU and softmax attention operations wherever possible (under the constraint that accuracy drops due to such approximations should be minimal.) The key challenge is to figure out where to perform these approximations.
3.1 Switched Taylorization
of binary switch variables defined for all instances of GeLU and softmax activations. Our goal here is to learn W,S, and C to ensure high accuracy with as few nonlinearit ies as possible. We also use N to denote the number of tokens, H to denote the number of heads and m to denote the size of the token embedding (and consequently the output size the feedforward MLP).
GELUs. In the case of GELU operations, we define a switched version of the GeLU activation:
f(G, c;ELUU()) + (1 c:)x;
(2)
y = [[(,,,,)) .... (Cnn
(3)
5

---

Published in Transactions on Machine Learning Research (10/2024)
where G is the corresponding auxiliary variable for the 7th token, x, is the ith input token embedding of dimension m (m being the ML.P dimension) and y E RR is the output. During training, are initially real-valued, trainable, and are initialized to 1 at the start of training. During inference, we binarize all q using an indicator function, Lase, where e is an appropriately chosen threshold. c; = 1 implies that the GELU is preserved whereas C, = linearizes it.
Softmax attention. The next step is to reparameterize softmax attention. However unlike GELUs, choice of parameterization is not obvious here. As per the Delphi protocol, exponents are expensive in terms of latency; on the other hand, polynomials are comparatively cheaper. Division by a constant is preferable to division by an input-dependent parameter (like in softmax); therefore, we propose a modified 'Squared Attention' operation:
ATUN XW, WX XW.
(4)
wherein we apply pointwise squaring instead of row-wise softmax and divide by the number of tokens. Similar to approach with GFLUs, we further add learnable auxiliary variable, 8; for every row-wise softmax operation in the attention layer.
= s,Softmax (X:) + (1 8; |QQUAREDAATT(X..).
(5)
where X, is the ith row of the attention matrix. As before, 8;8 are initially real-valued, trainable and initialized to 1. The variables are binarized during inference allowing usee of either Softmax or squared attention based on the values of s. Further ablations of different candidate attention functions are presented in the results sections.
3.2 Training PriViT models
To train PriVIT models, we need to train three sets of variables: the weights of the transformer, W, the switch variables for the GELU parameterization, C, and the switch variables for the atlention parametrization, S. Our goal is train a model that minimizes the nunber of nonlinearities satisfy a given nonlinearity budget, that is, |CC|o C, and ||8]|o S, while increasing the overall performance.
This is reminiscent of standard LASSO-style (Tibshirani, 1996) optimization. We therefore propose the following loss function to train the model,
I Iepriott ([ff((X),)) Ag
c
A3
(6)
j=0
where is the standard cross-entropy loss. We then optimize for each of the variables until the required softmax attention and GELU budgets.
Optionally, we can also make use of knowledge distillation during both training and fine-tuning. We introduce KL divergence loss the soft labels generated by the teacher and student ViT model. This loss is added to the Eppioit loss defined in eq. 6. Thus our final minimization objective looks as follows,
I(J[(X).v) + Ag 1
9 lesl + Ao D |8al (xx((w(())
min W.C.S
where T denotes the weights of the teacher model, and Lut is the KL divergence loss.
After every epoch, we count the number of GELUs and softmax attention operations by thresholding the 8; and & values. Once the model satisfies the required budgets, we freeze the chosen GELUs and softmax attention operations by binarizing all s; and e values and fino-tune the model weights for the classification task. See Figure 15 for overview of the method.

---

Published in Transactions on Machine Learning Research (10/2024)
PRIVIT MPCV+
PAVET MPCVET
CTEAR, 10
CIFAR 100
81
Aceurccy (%)
Accuracy (86
96
78
84
83
7.6
7.2
Latency (6)
Latency
Figure 3: Comparison of PriVIT over CIFAR 10/10/ 100 benchmarked against MPCVIT, and MPCVIT
70
Prt-VEE Pot-NT wala RDD
Text Accuracy (86)
Tel. Accuraccy (90)
60
Culkendtii Sedettes Syaord (oe))
50
60
100
180
Latency (M)
0.25
0.5
0.75
Softmax
-10*
Figure 4: We evaluated PriVIT with and without KD. The x axis represents latency measured as per Section 4.1, while the axis shows the accuracy Tinylmagenet. We observe an overall improvement in the latency-accuracy motivaling the of KD.
Figure We evaluated the PriVIT algorithm using three attention operations: Uniform, Linear, and Squared Attention. The x-axis represents the target softmax count, while the y-axis shows the test accuracy CIFAR-100.
BaseuPrite PriVit
Saltaao
25
0AL 500k
1,000
10
Lacos Indis
Figure 7: Comparison softmax distribution in ViT-hase model (Base) versus PriVIT without GeLU linearization. The r-axis denotes the layer index, while the y-axis shows the softmax operations per layer. With a 197 x 197 attention matrix across 12 heads, the ViT-hase model totals 2364 softmax oper- ations per layer. Notably, PriVIT tends to substitute earlier layer softmaxes with linear operations. Top: 1K target softmax; Bottom: 10K target softmax.
Figure 6: Comparison of GELU distribution be- tween ViT-base (Base) and PRiVIT without softmax linearization. The x-axis represents the model's layer index, while the y-axis shows log-scaled GELU op- erations per layer. With input tensor size of 197 x 3072 for the GELU layer, each layer contains 197 x 3072 605184 GELU operations. Top: 150K target GELU, Bottom: 500K target GELU.
7

---

Published in Transactions on Machine Learning Research (10/2024)
Table 3: Lateney comparison between PriVIT and PriVIT w/o Pretrain
PriViT
PriViT w/o pretrain
Latency (M)
Accuracy (%)
Latency (M)
Accuracy (%)
271.59
75.5
234.18
53.57
151.74
69.8
194.78
51.66
128. 23
66.98
167.20
55.59
93.71
60.53
153.31
55.92
4 Results
4.1 Experimental setup
Architecture and datasets. We apply PriViT algorithm to a pretrained checkpoint of ViT-Tiny (Steiner et al., 2021) that is trained on ImageNet-21k (14 million images, 21,843 classes) at resolution 224 x 224, and fine-tuned on ImageNet 2012 (1 million images, 1,000 classes) at resolution 224 x 224. The pretrained ViT Tiny checkpoints are made available by (WinKawaks, 2022). In this research work focus on finetuning existing model checkpoint like ViT Tiny on target standard image classification datasct (CIFAR-10/ (Krizhevsky et al., 2009) and Tiny-lmageNet).. CIFAR-10/ has images of size 32 x 32 while Tiny-ImageNet has 64 x 64. These images were resized to 224 x 224 before being given as an input. CIFAR-IU 10 has 10 classes with 5000 training images and 1000 test images per class. CIFAR-100 has 100 classes with 500 training images and 100 test images per class. Einy-ImageNet has 200 classes with 500 training images and 50 test images per class. We also perform hyperparameter tuning and present more details in Appendix A
VIT teacher pretraining. As the base model, we finetune a pretrained ViT-Tiny on CIPAR-10/100 for 10 epochs. We use AdamW (Loshchilov & Hutter, 2017) as the optimizer with an initial learning rate and weight decay as 0.0001 and 0.0001 respectively, and decay the learning rate after every 30 epochs by factor 0.1. We use the same hyperparameters for the Tinylmagened model.
Joint optimization of student VIT and parametric non linearities. We use Adam (Kingma & Ba, 2014) optimizer with learning rate equal to 0.0001. We use knowledge distillation and use soft labels generated by the teacher model with temperature of 4. The total loss is then, L LeaVIr LKL, where LEAVIr is Equation 6 and LKI. is the KL divergence loss between the logits of teacher and student model. The Lasso coefficient (Tibshirani, 1996) for parametric attention and GELU mask are set to Ag = 0.00003 and X = 0.00003 respectively at the beginning of the scarch. We set warmup epochs to 5 during which we don't change any hyperparameters of the model. Post warmup, we increment >g by a multiplicative factor of 1.1 at the end of each epoch if the number of active GELUUs of current epoch do not decrease by atleast 2 as compared to previous epoch. Note that a GELU/softmax is considered active if it's corresponding auxiliary variable is greater than threshold hyperparameter 0.001. We follow the same approach for Ag with multiplicative factor of 1.1 and an active threshold of 200.
Binarizing parametric nonlinearities, finetuning. When the GELUs and softmax budgets are satisfied, we binarize and freeze the the GELU and softmax auxiliary variables. We subsequently finetune the model for 50 epochs using AdamW with a learning rate 0.0001, weight decay 0.0001 and cosine annealing learning rate scheduler (Loshchilov Hutter, 2016). Our finctuning approach continues to use knowledge distillation as before.
Non-linearity cost comparison. We conduct experiments to assess the computational cost of non-linear functions such as layernorm, softmax, and GeLU in comparison to ReLU within GC. The detailed results are reported in Table 12, brackets considers amortizing to a vector of inpuls, e.g., a Layernorm is an operation over a vector length of 192 is equivalent to 6504x than the cost of a ReLU, It demonstrates that with a vector length of 197, all layernorm and softmax functions incur higher computational costs (i.c., number of ANDs) than ReLU. Specifically, they exhibit costs 6504x, 18586x higher than that of ReLU respectively and for pointwise GELU, we saw cost 270x higher than that of ReLUU The cost. of denominator of layernorm and
8

---

Published in Transactions on Machine Learning Research (10/2024)
Table 4: Comparison of training efficiency between PriVIT and MPCVIT on Vir-Tiny
PriViT
MPCViT
Dataset
Latency
Accuracy
Epochs
Latency
Accuracy
Epochs
TinyImagent
151.75
69.8
293
381.42
62.55
600
128.23
66.98
351
331.35
63.7
600
110.60
61.46
312
307.45
63.36
600
CIFAR 100
88.24
78.5
403
72.21
77.8
900
75.92
77.74
117
71.77
76.9
900
67.54
75.47
498
71.40
76.9
900
softmax can be amortized to the whole vector and thus incur less cost than GELU, We estimate the latency of each model gencrated by PriViT using these conversion factors. To showcase an example, we estimate the non-linearity cost hypothetical model with 1000 softmax operation, 1000 layernorm operations and 1000 GELUs, by taking the weighted sum of each operations with their corresponding latency factor.
GELU replacement post training. Mirroring the MPCVIT+ approach, also report the effect of PriVIT with all GELUs replaced with ReLUs, we call this model PriViT-R, and the original PriViT model PriViT-G. Such an optimization is effective For low GELU budgets as it introduces very minimal errors. For high GELU budgets, the error is quite significant that it affects the overall performance.
4.2 Comparisons standard benchmarks
We benchmark PriVIT against MPCVIT, using the checkpoints publicly shared by the authors of Zeng el al. (2022). We calculate latencies with the Secretflow (Ma et al., 2023) framework using the SEMT2k (Cramer et al., 2018) protocol.
In our evaluation on various datasets, the performance of PriVIT was benchmarked against both MPCVIT and MPCViT+. We measure two metrics of importance - the latency (measured in seconds), and accuracy. An ideal model would achieve high accuracy with low latency.
1. Tiny ImageNet: terms of the Pareto frontier on the Tiny ImageNet dataset, PriViT showcases notable improvement. On Tiny Imagenet, for an isoaccuracy of approximately 63% PriViT achieved 5.77x speedup compared MPCVIT reported in Table 2.
2. CIFAR-10: Similarly, on the CIFAR-10 dataset, PriViT showcases noticeable improvements over MPCViT and MPCViT+. At an isoaccuracy of approximately 94%, PriVIT achieves 1.14x and 1.08x speedup over MPCVIT and MPCVVT+ respectively as reported in Table 2 and Table 9. This represents non-trivial speedup, despite that fact that PriViT's model architecture is significantly larger than MPCViT's model architeclure as noted Table 8. A graph of these values can be found Figure 3.
3. CIFAR-100: Finally, on the CIFAR-100 dataset, PriVIT also showcases noticeable improvements over MPCVIT and MPCViT+. At isoaccuracy of approximately 78%, PriVIT achieves 1.05x and 1.02x speedup over MPCVIT and MPCViT- respectively reported in Table 2 and Table 9. We draw similar conclusion to the analysis in CIFAR-10. PriViT achieves better accuracy with lower latency despite its larger model size. A graph of these values can be found in Figure 3.
4.3 Further comparisons with MPCViT
The majority of our numerical comparisons above are relation MPCViT, which is an important predecessor our work. Our results above show that PriViT achieves better Pareto frontier for private inference than MPCViT by achieving lower latencies iso-accuracies across various datasets.

---

Published in Transactions on Machine Learning Research (10/2024)
However, these comparisons not completely apples-to-apples due to (i) differences in model size, and (ii) differences in pre-training. This is because MPCVIT uses smaller transformer architecture than the standard ViT backbone (which we use).
Therefore, we perform the following additional experiments:
1. We take the pre-trained ViT-Tiny checkpoint used in PriViT and apply NAS described in the MPCVIT paper.
2. We perform 300 epochs of NAS for each dataset in our comparisons (CIFAR-10/100 and Tiny- ImageNet). This is in line with the number of epochs reported in Zeng et al. (2022).
3. We then train this (post-NAS) ViT-Timy on each benchmark dataset for varying values of The μ parameter represents the percentage of attention heads that use ReLU softmax instead of Scaled- Attention.
We train four models per dataset with μ values of 0.1, 0.3, 0.5, and 0.7 to match the training scheme of the MPCViT.
We were not able to replicate the knowledge distillation results present in Zeng et al. (2022). In our experiments found that the loss did not converge using their public code. Therefore, we opted to use standard cross entropy loss function instead, and achieved reasonable accuracies with this method.
6. We trained the models for 300 epochs to match the training time of the MPCVIT work. This leads us to family of 12 models (four for each dataset). We call this "MPC-Tiny".
We present comparisons for PriViT with this new family of models Tables 5, and below.
Table 5: MPC-Tiny (ie. MPCViT with the standard ViT backbone)
Dataset
Accuracy (%)
Latency (seconds)
0.7
CIFAR-10
91.48
42.03
0.5
CIFAR-10
88.48
40.28
0.3
CIFAR-10
93.97
36.80
0.1
CIFAR-10
94.49
34.21
0.7
CIFAR-100
40.22
43.17
0.5
CIFAR-100
79.12
39.37
0.3
CIFAR-100
75.04
36.82
0.1
CIFAR-100
76.21
34.47
0.7
Tiny Imagenet
66.66
42.98
0.5
Tiny Imagenet
65.46
39.30
0.3
Tiny Imagenet
67.57
37.50
0.1 Tiny Imagenet
65.11
34.03
We now analyze our new results, starting with latency comparisons. We first see that all of the models present in MPC-Tiny are significantly slower than their counterparts in PriViT. Notably, the accuracy for the MPC-Tiny models is worse than both PriViT models and MPCVIT models on average. This is likely due to the knowledge distillation procedure not transferring well larger models. We also notice that for each value μ, the latency similar regardless of the datasct, which as expected and provides a good sanity check. In addition, we see that MPC-Tiny models have relatively inflexible lower bound for latency. The p parameter controls how much latency one save by removing ReLU Softmaxes; even after have reached μ = 0.1, the latency is still very significant due to the large model size. This is in contrast to PriViT where we control the nonlinearity Taylorization at a scalar level (instead of attention head level) which enables us to push the latency down to significantly lower levels.
In conclusion, we show that a better comparison of PriViT can be made with MPC-Tiny as opposed to the original published MPCVIT work. Our results show that PriViT has significantly lower latency at similar accuracies compared to MPC-Tiny.
10

---

Published in Transactions on Machine Learning Research (10/2024)
Table The original MPCViT (which uses a small transformer backbone)
μ
Dataset
Accuracy (%)
Latency (scconds)
0.7
CIFAR-10
94.30
7.03
0.5
CIFAR-10
94.20
6.66
0.3
CIFAR-10
94.10
6.38
0.1
CIFAR-10
93.60
6.06
0.7
CIFAR-100
77.80
6.97
0.5
CIFAR-100
76.90
6.62
0.3
CIFAR-100
76.90
6.42
0.1
CIFAR-100
76.40
6.06
0.7
Tiny Imagenet
62.55
69.64
0.5
Tiny Imagenet
63.70
55.45
0.3
Tiny Imagenet
63.36
40.86
0.1
Tiny Imagenct
62.62
26.76
Table 7: PriVIT
(which uses the standard VIT backbone)
Dataset
Accuracy (%)
Latency (Seconds)
CIFAR-10 10
96.58
8.15
CTFAR-10
95.88
7.28
CIFAR-10
95.58
6.88
CIFAR-10
94.58
6.42
CIFAR-10
94.24
6.30
CIFAR-10
93.60
5.91
CI
93.62
5.88
C-
80.60
7.30
CIFAR-100
79.78
7.20
CIFAR-100
79.27
7.04
CIFAR-100
78.49
6.83
Tiny Imagenet
78.88
17.95
Tiny Imagenet
78.16
15.92
Tiny Imagenet
75.50
12.49
Tiny Imagenet
64.46
9.53
11

---

Published in Transactions on Machine Learning Research (10/2024)
Table & Base model architecture of PriVIT and MPCVTT
Model
Layers
Width
MLP
Heads
Image size
Patch size
params (M)
PriVIT
224x221
16x16
5.8
MPCVIT (Tiny Imagenet)
9
192
384
12
64x64
4x4
MPCVIT (Cifar 10/100) (0/1
256
512
32x32
4x4
3.72
Table 9: Accuracy-latency tradeoff MPCViT+. We compare this tradeoff with the PriVit tradeoff values from
Table 2. Left: We first examine Cifar10 at isoaccuracy of approximately 94%. PriVit achieves 1.08x speedup over M.C Right: We then examine Cifar at an isoaccuracy of approximately 78%. PriVit achieves a 1.02x speedup M
Cifar10
Cifar100
Acc
Latency (s)
Ace
Latency (s)
94.3
6.898
77.1
6.903
94.2
6,888
76.8
6.606
93.9
6.606
76.3
6.355
93.3
6.026
76.2
6.021
4.4 Ablation studies
Knowledge Distillation. We incorporate knowledge distillation (KD) alongside supervised learning. To assess the contribution of KD to the overall performance, we trained PriViT on the Tinylmagenet dataset with varying non-linearity budgets. We then compared its performance to a version of PriViT (as outlined in Figure 15) that does not employ a teacher model for knowledge distillation. Our results in Figure 4 indicate that, under identical latency conditions, incorporating KD enhances performance by approximately 5%.
Contribution of pretraining. In PriViT, we utilize pretrained checkpoint, which is subsequently fine- tuned. Post fine-tuning, introduce parametric GeLU and attention mechanisms to decrease non-lincarities in the model. To gauge the impact of using pretrained model on the overall performance, contrast the performance of PriViT with a variant of PriViT that is not built upon pretrained model. Instead, this variant employs weights initialized from scratch and is trained with the same parametric non-linearity mask as used PriViT to minimize non-lineane non-lincaiitics The comparative outcomes these approaches are presented in Table 3. Our findings reveal that, for comparable latencies, PriViT with the pret rained checkpoint. outperforms the alternative by about 14%.
Choice of softmax approximation. We run PriVIT over different softmax budget over CIFAR-100, and report the accuracy of the resulting model versus the number of original softmax attention retained. Lower number of softmax operations implies higher the number of softmax attention replaced with our candidate attention operation. As per Figure 5 we see almost no performance drop for AQQUREDAATN, roughly 5% drop performance for SCALEATTN and 10% drop in performance for UNIFORMATTN in low budgets. Thus SquarodAAtenntion outperformed the others across all softmax budgets, motivating selection replace the standard softmax attention in PriVIT.
Fine- Fine-grained versus layer-wise Taylorization PriVIT employs unique approach where it selectively Taylorizes softmax and GELU operations. To probe the effectiveness of this method, we contrasted it with an alternative PriViT approach that Taylorizes a ViT model progressively, layer by layer. As illustrated in Table 10, our observations underscored the superiority of selective Taylorization.
Visualization of non-linearity distribution. To understand which nonlincaritics are preserved, we investigate the distribution of PriVIT models under different softmax and GELU budgets. From our observations in Figure we conclude that GELUs in earlier encoder layers preferred over the ones in
12

---

Published in Transactions on Machine Learning Research (10/2024)
Table 10: Comparing PriVIT and layerwise linearization of GeLU in VIT model with 200k GeLUs. Six models were generated by replacing two GeLU layers at time with Identity.
Layerwise GELU linearizing
Pri-VIT
Gelu (K)
Acc (%)
Gelu (K)
Acc (%)
96.07
200
95.59
193
95.91
95.34
187
94.28
100
95.58
181
93.33
50
94.98
174
93.04
94.24
164
92.06
93.96
123
82.48
0
56.64
the later layers. From figure we observe a similar trend in softmax distributions. We find this interesting, since the trends reported in earlier work on convolutional networks are the reverse direction: earlier layers tend have larger number of lincarized units. Understanding this discrepancy an interesting question for future work.
Table 4 shows that PriViT, at similar accuracy 64%), requires about half the training epochs compared to MPCVIT on Tinylmagenet. For an isolatency of 75M on CIFAR 100, PriVIT also needs only about 50% of MPCViT's training epochs. This demonstrates PriViT's enhanced efficiency and scalability, making it a promising alternative to MPCViT, particularly situations valuing efficiency and performance.
5 Conclusion
We introduce PriVIT, new algorithm for designing MPC-friendly vision transformers, and showed its competitive performance on several image classification benchmarks. A natural direction of future work is to extend similar techniques for designing other families of transformer architectures, such Swin Transformers, well encoder-decoder transformer architectures. A key limitation of PriViT is the difficulty in properly Taylorizing layer norms without introducing instability in the training; we leave an in-depth study of this question to future work.
Impact Statement
This paper aims to improve the elficiency of private inference for vision transformer architectures. Facilities who wish to perform inference on vision models while keeping user data private (such hospitals) would gain the ability to use third party pretrained models without risking the leakage of confidential data. However, this advancement could also open the door to malicious users who would be able to use models without revealing their intentions. Considering both perspectives, we believe that the pros of advancing this field vastly outweigh the cons and thus believe it will lead to a net-positive impact.
Acknowledgments
This work was supported in part by the US Department of Education's GAANN (Graduate Assistance in Areas of National Need) Fellowship Program. The research was also supported in part by the NSF CAREER award #2340137, NSF award #2154119 and by DARPA under the Data Protection in Virtual Environments (DPRIVE) program, contract HR0011-21-9-00
References
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 2016.
Dake Chen, Yuke Zhang, Souvik Kundu, Chenghao Li, and Peter A. Becrel. Rna-vit: Reduced-dimensson approximate normalized attention vision transformers for latency efficient private inference. Tn 2023
13

---

Published in Transactions on Machine Learning Research (10/2024)
IEEE/ACM International Conference Computer Aided Design (ICCAD), pp. 1-9, 2023. doi: 10.1109/ TCCAD5533002023) 202.00337222
Minsu Cho, Zahra Ghodsi, Brandon Reagen, Siddharth Garg, and Chinmay Hegde. Sphynx: Relu-efficient network design for private inference. arXiv preprint arXiv:2106.11756 2021.
Minsu Cho, Ameya Joshi, Brandon Reagen, Siddharth Garg, and Chinmay Hegde. Selective network linearization for efficient private inference. In International Conference on Machine Learning, pp. 3947- 3961. PMT.R, 2022.
Ronald Cramer, Ivan Damgård, Daniel Escudero, Peter Scholl, and Chaoping Xing. Spd: efficient mpe mod for dishonest majority. In Annual International Cryptology Conference, pp. 769-798. Springer, 2018.
Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoe V Le. Randaugment: Practical automated data augmentation with a reduced search space. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops, pp. 702 703, 2020.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghami, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image worth 16x16 words: Transformers for image recognition at scale. In International Conference Learning Representations, 2020.
Craig Gentry and Shai Halevi. Implementing gentry's fully-homomorphic encryption scheme. In Annual international conference on the theory and applications of cryptographic techniques, pp. 129-148. Springer, 2011.
Zahra Ghodsi, Akshaj Kumar Veldanda, Brandon Reagen, and Siddharth Garg. Cryptonas: Private inference relu budget. In Adv. Neural Inf. Proc. Sys. (NeurIPS), 2020.
Zahra Ghodsi, Nandan Kumar Jha, Brandon Reagen, and Siddharth Garg. Circa: Stochastic relus for private deep learning. In Adv. Neural Inf. Proc. Sys. (NeurIPS), 2021.
A Hassani, S Walton, N Shah, A Abuduweili, J Li, and H Shi. Escaping the big data paradigm with compact transformers. arXiv preprint arXiv:2104.05704, 2021.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08411 2016.
Nandan Kumar Jha, Zahra Ghodsi, Siddharth Garg, and Brandon Reagen. DeepReDuce: Relu reduction for fast private inference. In Proc. Int. Conf. Machine Learning, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Vladimir Kolcsnikov and Thomas Schneider. Improved garbled circuit: Free xor gates and applications. In Automata, Languages and Programming: 35th International Colloquium, ICALP 2008, Reykjavik, Iceland, July 7-11, 2008, Proceedings, Part 11 35, pp. 486 498. Springer, 2008.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. Technical report, U. Toronto, 2009.
Souvik Kundu, Shunlin Lu, Yuke Zhang, Jacqueline Tiffany Liu, and Peter Anthony Becrel. Learning to linearize deep neural networks for secure and efficient private inference. In Proc. Tnt. Conf. Learning Representations, 2023.
Dacheng Li, Hongyi Wang, Rulin Shao, Han Guo, Eric Xing, and Hao Zhang. MPCFormer: Fast, Perfor- mant Private Transformer Inference With MPC. In The Eleventh International Conference on Learning Representations, 2022.
14

---

Published in Transactions on Machine Learning Research (10/2024)
Iya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint
2016.
Nya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint Xiv: 1711::111100011 2017.
Qian Lou, Yilin Shen, Hongxia Jin, and Lei Jiang. Safenet: Asecure, accurate and fast neu-ral network inference. In Proc. Int. Conf. Learning Representations, 2021.
Jiachen Lu, Jinghan Yao, Junge Zhang, Xiatian Zhu, Hang Xu, Weiguo Gao, Chunjing Xu, Tao Xiang, and Li Zhang. Soft: Softmax-free transformer with lincar complexity. In Adv. Neural Inf. Proc. Sys. (NeurIPS), 2021.
Junming Ma, Yancheng Zheng, Jun Feng, Derun Zhao, Haoqi Wu, Wenjing Fang, Jin Tan, Chaofan Yu, Benyu Zhang, and Lei Wang. SecredFlow-SPU: A performant and User-Friendly framework for Privacy-Preserve machine learning. In 2023 USENIX Annual Technical Conference (USENIX ATC 23), pp. 17-33, Boston, MA, July 2023. USENIX Association. ISBN 978-1-939138-35-9 URL https: wwww:usenix.
Silvio Micali, Oded Goldreich, and Avi Wigderson. How to play any mental game. In Proceedings of the Nineteenth ACM Symp. on Theory of Computing, STOC, pp. 218 229. ACM New York, NY, USA, 1987.
Pratyush Mishra, Ryan Lehmkuhl, Akshayaram Srinivasan, Wenting Zheng, and Raluca Ada Popa. Delphi: A cryptographic inference service for neural networks. In 29th USENIX Security Symposium (USENIX Security 20), pp. 2505-2522. USENIX Association, Aug. 2020. ISBN 978-1-939133-17-5 URL https:
//www. www.ix. .eciiii.orr/////e
Jianqiao Mo, Jayanth Gopinath, and Brandon Reagen. Haac: A hardware-software co-design to accelerate garbled circuits. In Procoedings of the 50th Annual International Symposium on Computer Architecture,
pp. 13, 2023.
Hongwu Peng, Shanglin Zhou, Yukui Luo, Nuo Xu, Shijin Duan, Ran Ran, Jiahui Zhao, Shaoyi Huang, Xi Xie, Chenghong Wang, Tong Geng, Wujie Wen, Xiaolin Xu, and Caiwen Ding. Rrnet: Towards relu-reduced neural network for two-party computation based private inference. arxiv preprint: ArXiv:2302.0220222. 2023.
Zhen Qin, Weixuan Sun, Huicai Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjic Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. In Proc. Int. Conf. Learning Representations, 2022.
Adi Shamir. How to share a secret. Commmications of the ACM, (2(1)) 613, 1979.
Andreas Steiner, Alexander Kolesníkov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. arXiv preprint arXiv:2106.10270 2021.
Robert Tibshirani.
Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Socicty. Series B (Methocological), pp. 267 288, 1996.
Hugo Touvron, Matthicu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient. image transformers and distillation through attention, 2021.
Xiao Wang, Alex J. Malozemoff, and Jonathan Katz. EMP-coolkit: Efficient MultiParty computation toolkit. https://github. htt.gggttt..cccmmmttooooooe com/emp-tolkki. 2016.
Yongqin Wang, G Edward Suh, Wenjie Xiong, Benjamin Lefaudeux, Brian Knott, Murali Annavaram, and
Hsien-Hssin $ Loc. Characterization of mac-based private inference for transformer-basee models. In 2022 IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS), pp. 197-197. TEEE, 2022.
15

---

Published in Transactions on Machine Learning Research (10/2024)
WinKawaks.
model, 2022.
https://hhuggnnafcc.ooo hwtp://g...a...o///aaaaaaaaaaaa
-at-tinyypatchh6-222
Andrew Chi-Chih Yao. Protocols for secure computations. In 23rd Annual Symposium on Foundations of Computer Science (sfes 1982), pp. 160-164. TEFE, 1982.
Andrew Chi-Chih Yao. How to generate and exchange secrets. In 27th Annual Symposium on Foundations Computer Science (sfcs 1986), pp. 162 167. IEEE, 1986.
Sangdoo Yun, Dongyoon Han, Scong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 6023 6032, 2019.
Samee Zahur, Mike Rosulek, and David Evans. Two halves make whole: Reducing data transfer in garbled circuits using half gates. In Advances in Cryptolgy-RUROocYYP 2015: 34th Annual International Conference on the Theory and Applications of Cryptographic Techniques, Sofia, Bulgaria, April 26-30, 2015, Proceedings, Part II 34, pp. 220 250. Springer, 2015.
Wenxnan Zeng, Meng Li, Wenjic Xiong, Wenjic Lu, Jin Tan, Runsheng Wang, and Ru Huang. Mpevit: Search- ing for mpe-friendly vision transformer with heterogeneous attention. arXiv preprint 2XXv:2211.139555 2022.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Pas Paz. mixup: Beyond empirical risk minimization. arXiv preprint arXiv:1710.09413 2017.
Yuke Zhang, Dake Chen, Souvik Kundu, Chenghao Li, and Peter A. Beerel. Sal-vit: Towards latency efficient private inference on vit using selective attention search with a learnable softmax approximation. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 5116-5125, October 2023.
Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yi Yang. Random erasing data augmentation. Proceedings of the AAAT conference on artificial intelligence, 2020.
16

---

Published in Transactions on Machine Learning Research (10/2024)
Table 11: Accuracy-lateucy tradeoffs for PriVIT applied to DeiT-Tiny.
Cifar10
Cifar100
Tiny Imagenet
Ace
Latency (s)
Acc
Latency (s)
Acc
Latency (s)
97.12
29.121
85.71
28.753
78.17
28.711
96.92
11.494
85.44
15.036
78.11
25.270
96.81
10.205
84.12
10.260
77.48
17.586
96.28
9.032
82.17
8.339
76.84
13.062
73
5.883
78.92
6.723
75.68
11.481
A
Supplementary results
PriVIT on different size architectures. Figure 8 provides comparison of PriVIT over different VIT variants from 5M (ViT-Tiny) to 80M (ViT-Base) parameters. We see that PriVIT is able to produce a similar latency accuracy trade off across different model sizes emphasizing it's generalizability.
95.5
95
Aecuracy
94.5
94
VIT Base
VIT Tiny
93.5
VIT Small
50
100
150
200
250
Latency (M)
Figure 8: PriVIT produces similar latency accuracy trade off for variety of ViT models over CIFAR 10, highlighting it's scalability to large model sizes
PriVIT different VIT architcctures. We study PriVIT the Data-Efficient Image Transformer (DeiT) (Touvron et al., 2021) see how applicable the algorithm is to other ViT style architectures. We perform latency and accuracy tests for DeiT-Tiny CIFAR-10, CTFAR-100, and Tiny Imagenet. Results are found in Table 11. The accuracy-latency tradeoff is very similar in the DeiT-Tiny model to the ViT-Tiny model. This suggests that PriVit is applicable to ViT style architectures as a whole, and not only to the original VIT architecture.
Analysis of performance degradation. In this analysis, we aim to compare the performance of trained PriViT models with their finetuned versions. Our analysis is based on the class-level accuracy metric from the Tiny ImageNet dataset, which consists of 200 classes. We focus on three specific parameters to understand the performance degradation:
Maximum Difference in Accuracy: We assess the greatest disparity in accuracy across all 200 classes between the PriViT and finetuned models.
Overall Accuracy Difference: We compute the average accuracy difference between the finetuned and the PriVIT models across all 200 classes.
17

---

Published in Transactions on Machine Learning Research (10/2024)
Table 12:
Non- Con-lincarity cost normalized to the cost of one ReluOp which ReLU operation over a scalar value.
PriViT
MPCVIT Tinylmagenet
MPCVIT CIFAR-10/100
Function
# ReluOps
Function
# ReluOps
Function
# ReluOps
(19t
18586
ReLU Softmax((55)
4428
ReLU Softmax(65)
1133
Layernorm(192)
6504
Layernorm
6504
Layernorm(250)
8614
G1LU(1)
270
GcLU(1)
270
GeLU(1)
270
x²
3248
Table 13: Performance degradation of PriViT models compared to finetuned model on tinyimagenet.
Accuracy
Latency (M)
Max Difference
Mean Difference
Variance (10 x 10-3)
69.8
151.75
30.00%
1.85%
6.8
66.98
128.23
34.00%
4.68%
6.9
64.46
110.60
34.00%
7.21%
8
60.53
93.72
40.00%
11.13%
9
59.55
86.51
34.00%
12.12%
9.4
59.58
84.78
36.00%
12.08%
9.8
59.04
84.13
36.00%
12.63%
9.9
58.74
69.42
40.00%
12.92%
9.4
58.2
67.43
36.00%
13.48%
8.7
Variance in Accuracy Difference: We analyze the consistency of the differences in accuracy across the 200 classes by calculating the variance.
Table 13 highlights that average accuracy degradation anywhere between 1-13% for different non-linee non-linearity budgets but certain classes seem to be more adverscly affected even in low budgets the max class level difference in accuracy is consistent around 30%.
CutMix (Yun al., 2019), Mixup (Zhang et al., 2017), Randaugment (Cubuk et al., 2020), and Random Erasing (Zhong et al., 2020) as dala augmentation strategy. We probed multiple hyperparameter strategies for the joint optimization phase of PriViT to ensure consistent good performance over multiple configurations non-linearity budgets of softmax and GELUs. Specifically we describe these strategies as follows:
Hyperparameter Tuning. Following (Hassani et al., 2021)
This strategy involved 10 nost-linearization training epochs. The binarization of auxiliary parameters, and c, occurred late in the process, specifically after the lincarization was complete. The penalty increment condition for this method was checked when the reduction in the softmax and GELU coefficients per epoch was less than 200 and 2, respectively. Both masks began with identical penalties, signifying an 'equal' starting penalty.
Late- BintaiBid Epoch (Strategy 1):
Late- Sate-Binarized Incremental (Strategy 2):
This strategy also encompassed 10 training epochs with late binarization. Here, the penalty increment condition was activated with an increase the softmax and GELU coefficients per epoch. The starting penalty for both masks was 'equal.
Much like Strategy 2, this involved 10 epochs with late binarization and an increment condition based on softmax and GELU coefficient rises. However, the initial penalty was set to 'unequal', making the softmax penalty 20 times higher than the GELU penalty.
Bitte-inarized Divergent Penalty (Strategy 3):
Early-Binee Binarized Incremental (Strategy 4): This strategy shared several similarities with Strategy 2, including training epochs and an increment condition based on coefficient inercases. The difference, however, lay in its early binarization, occurring during the freezing of the auxiliary parameters. The starting penalty was kept. 'equal' for both masks.
18

---

Published in Transactions on Machine Learning Research (10/2024)
Table 14: We test the different finetuning strategies described in A. We run PriViT for different softmax
and GELU budgets and compare the test accuracy of the resulting model over CIFAR-100. We observe that strategy 5 works the best across wide range of target softmax and GELU budgets.
# Softmax
# Gelu
Strategy 1
Strategy 2
Strategy Strategy 4
Strategy (Acc. %)
(K)
(K)
(Acc. %)
(Acc. %)
(Ace. %)
(Ace. %)
10
5
77.68
76.74
77.82
78.83
5
5
76.27
75.99
75.72
77.63
5
1
76.73
75.21
76.24
77.08
10
76.04
75.23
74.65
76.35
1
75.92
74.84
76.45
76.97
5
76.12
74.99
76.32
76.96
Prolonged Early-Binarized Epoch (Strategy 5): Spanning 50 post-linearization training epochs, this strategy adopted an early binarization approach. The penalty increment condition was activated when the reduction softmax and GELU coefficients per epoch was under 200 and 2, respectively. The masks were initialized with 'equal' penalties.
Each of these strategics offered unique configurations in terms of epoch durations, binarization timings, increment conditions, and starting penalties, enabling comprehensive assessment of the PriVIT algorithm's performance under various conditions.
We test the different finetuning strategics described here by taylorizing PriViT for different softmax and GELU budgets and compare the test. accuracy of the resulting model over C1FAR-100. Table highlights the comparative performance of all the strategies that we described. Strategy 5 seems be performing best different configuration of nonlinearity budget which is important as we would want to find the best model peformance for particular non-linearity budget.
Grid search of softmax and GELU configuration. In order to elucidate the muanced tradc-off between softmax and GeLU operations, we executed a systematic grid search across an extensive parameter space encompassing varied softmax and GeLU configurations. Upon analysis models exhibiting iso-latencies, as demarcated by the red lines in figure 9, bocame evident that the tradc-off dynamics are non-trivial. Specifically, configurations with augmented softmax values occasionally demonstrated enhanced performance metrics, whereas in other scenarios, models optimized with increased GeLU counts exhibited superior benchmark results.
Grid search CIFAR-100
Acc.
300
250
79
200
GELU(K)
78
150
100
77
76
0.4
0.6
0.8
Softmaz
-10
Figure 9: PriViT produces Pareto surface mapping the tradeoff between GeLU and softmax budgets over C10AR-100.
19

---

Published in Transactions on Machine Learning Research (10/2024)
Taylorizing only type of non-linearity. The PriVIT algorithm's standout capability is its simultaneous linearization of GELU and softmax operations, enabling a myriad of model configurations. In our focused experiment, we exclusively linearized GELU operations and anchored the auxiliary softmax parameter S, binarizing it to activate only the SoftmaxAttention mechanism. Despite extensive GELU substitutions, as reported in the PriViT model displayed notable resilience on CTFAR-10 and CIFAR-100 dat datasels, with only slight performance drops, underscoring robustness in varied setups.
100
Accuracy (%)
90
80
CIFAR- 100
CIFAR-10
2,000 1,800 1,600 1,400 1,200
1,000
800
600
400
200
# GELU (K)
Figure 10: PriViT's ability to linearize GeLU operations visualized through performance on CIFAR datasets. As GELU operations decrease, CIFAR-100 and CIFAR-10 accuracies are affected, showcasing the trade-off between operation count and accuracy.
Effect of using pre-trained checkpoints. To further investigate why using pretrained checkpoint is improving performance, we report the non-linear distributions searched by PriVIT and compare it with PriViT without pretrain for the nonlinearity budget of 315k and 320k respectively. We observe from our findings in figures 11,12 that the distribution found by the two methods differs across each layer. This supports our theory as how PriVIT operates under a strategic 'top-down' paradigm. Starting with a finc-tuned model, it has the advantage of an architecture that has not just discerned overarching generalization patterns but has also selectively pruned irrelevant information, streamlining its focus for specific downstream task. This reduction of redundancy, undertaken from vantage point of pre-cxisting knowledge basc, gives PriViT an edge.
B Supplementary graphics
The following figure shows graphical representation of the switching operation.
Search granularity. An important characteristic of PriVIT is it's flexibility to search over different granularity of non-linearities. GELU is a pointwise functions, thus PriViT can search either at embedding level or at a token level. On the other hand, softmax is a token level operation, thus cannot be broken into a finer search space. Note that softmax operations can be extended to search over the head space or layer space, and similarly GELU can be searched over the layer space. Figure 14 illustrates the search granularity over token and embedding space.
Parametric mask. Figure 13 is an illustration of the working mechanism of the parametric mask introduced in PriViT. When the parameters are binarized it selects one of the two candidate function in the attention mechanism, and the gelu activation.
PriVIT overview. Figure 15 provides an illustration of the complete PriViT algorithm, there are three distinct phases namely Finetuning the teacher, joint optimization of network and parametric masks, and final finetuning.
20

---

Published in Transactions on Machine Learning Research (10/2024)
200
197
P w/o pretrain PriVic
Gelus
100
1
1
1
2
3
4
5
10
11
12
Layer Index
We compare the distribution of 208 GELU and 200 GELU operations distributed by PriViT w/o pretrain and PriViT respectively over tiny imagenet dataset.
Figure 11:
600
501
I UPPET w/o pretrain PoVV
Softmax
400
352
1033
200
3
2
3
0
3
2
3
3
3
3
3
3
3
2
3
4
5
10
11
12
Layer Index
Figure 12: We compare the distribution of 973 softmax operations and 998 softmax operations operations distributed by PriViT w/o pretrain and PriViT respectively over tiny imagenet dataset.
C
Role of Delphi in our framework
Delphi Mishra et al. (2020) encompasses two primary elements: secure communication protocol and NAS technique, specifically evolutionary method. The secure communication aspect focuses cryptographic protocols for neural networks within a multi-party computation setting. The NAS technique, known as Delphi's planner, is designed primarily to eliminate ReLU operations from neural architectures.
our paper, we utilize Delphi's secure communication protocol for private inference. Furthermore, we took inspiration from Delphi's planner to design an effective method to reduce non-linear operations for attention- based architectures, particularly in Vision Transformers (ViTs). This enhancement was necessary because the original Delphí's planner is not optimally equipped for reducing non-lincarities like softmax, gelu, and layernorms in VIT architectures.
Figure 16 provides an estimate of Delphi's planner lower bound on latency on ViT-Base with 80M parameters. To estimate this lower bound, we run PriViT on ViT-Base but only replace GELU operations with identity operations, till no Gelus were left linearize, ie.. the left. most point the cyan curve indicates a network with no Gelus. All softmax and layernorms are left intact, Delphi's planner lower bound on latency is defined by the red vertical line which represents the latency of model with no GELUs, and all softmax/layernonn intact. We compare this lower bound with the latency accuracy curve produced by PriVIT ViT-Base that Taylorize both softmax and GFLUS.
D Attention Variants
Here we describe formally the different attention variant we ablated. Uniform form attention basically described by the following equation
UNIFORMATIN(X) = (1)
XW
(7)
N
21

---

Published in Transactions on Machine Learning Research (10/2024)
Softmax Attention
SoftmarAtti( + (1 s) SquaredAttn(1)
Rm
Soottma if = 1, (quaredAAtt(())
if 0.
Squared
Attention
1-5
c+ + (1-
Riam
GELU
(GELU(() ifc 1, 1
ifc= 0.
1-c
Figure 13: Parameterized Gelu and Self-Attention operations. Top: Tokens undergo softmax and squared attention in training. Post-training, parameter S is frozen and binarized, selecting only one operation. Bottom: Embeddings pass through GeLU and Identity during training, Afterwards, parameter C' is frozen and binarized, choosing a single operation.
Linear Projection
Projectionn
Parametric
Linear
Gelu
Linear Projectoo
Muli Headed
Linear Projection
Projection
Parametric Gelu
Linear
Figure 14: Left: The green blocks are SQUARED and the grey blocks are Softmax Attention. For parametric attention, tokens emerge from a blend of softmax and square attention (refer to fig 13). Post-training, auxiliary variable S is set to or 1, resulting in Nxll potential combinations per encoder block. Right: The green blocks are Identity function, and the grey blocks are GELU activation. Embeddings combinc GELU and identity operations during training, as seen in fig 13. After training, parameter is frozen and binarized. This yields potential combinations of either 2HNN 2NxHxm xm for each VIT encoder block. Note that GELU being pointwise function, possess the flexibility to expand our search space either to tokens or directly to individual embeddings.
Where N is the number of tokens, so for each token the attention weights are equal hence the name CniformAAtentionn
Attl is the softmax candidate used in the work Zeng et al. (2022) which is essentially described
22

---

Published in Transactions on Machine Learning Research (10/2024)
Finelune Teacher VIT
Joint opómization of network and parametric nonlinearibes
Binarizing parametric nonimearties and
Shudert VIT
Student VIT
Teacher
Teaachr VIT
Allention
Self Attention
1-8
Teacher VIT
1-
Cross Entropy Less
Golu
Coea
i'Pavly
F1
+
Uharperoo
Figure 15: Left: Step Fine-tuning a pretrained VIT over target dataset to produce the 'teacher VIT. Middle: Step Duplicate teacher ViT, introduce parametric GELUs and attention mask to form 'student VIT. Train using cross-entropy loss, KT. divergence, and L, penalty to gradually find sparse mask. Binarize the mask post desired non-linearity budget. Right: Step 3 With frozen, binarized mask, further fine-tune the student model using cross-entropy loss and KL divergence with the teacher.
90
PriVIT
PriVIT (only gelu replaced)
DELPHII lower bound
85
Accuracy
80
75
102
103
Latency (M)
Figure 16: Estimation of Delphi's latency lower bound compared PriViT. We use ViT-Base as our architecture and CIFAR 100 as our dataset.
SCALEATTN(X) XW, W X
XW 03
(8)
23


Reviews:
[{"review_id": "nfRI2GOsST", "reviewer": "Reviewer_fJqe", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "S1 - PriViT utilizes a simple method to replace non-linear functions in ViTs. The methods are well-motivated and simple to adopt in other transformer-based architectures.\n\nS2 - The proposed method shows high performance in both classification and accuracy.\n\nW1 - I wonder why there is no comparison with RNA-ViT in the paper. Since RNA-ViT originates from the approximation of a softmax function, PriViT and RNA-ViT may be similar. It seems that the RNA-ViT method can be applied to the ViT structure as described in the RNA-ViT paper. It would be better for the authors to compare their method to RNA-ViT quantitatively and qualitatively.\n\nW2 - The base models of PriViT and MPCViT are different; furthermore, they have different input image sizes and parameters. For a fair comparison, especially for accuracy, I think using the same architecture is necessary. If it is hard due to any method-related reasons, a deep analysis considering those conditions is needed.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Nowadays, because of user data privacy, the importance of private inference using secure multi-party computation protocols is increasing. However, most researchers have focused on CNN-based architectures because transformer-based architectures are ill-suited for private inference. Unlike CNN-based models, transformer-based models (e.g., ViT) use a large number of non-polynomial operations such as GELUs, softmaxes, and layer norms. The authors replace these nonlinear operations with Taylorization.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I think it would be better for the authors to supplement the comparison parts of the paper. In my opinion, the comparison with RNA-ViT (see W1) and using the same architecture when comparing the methods (see W2) are needed.", "broader_impact_concerns": "It seems that there is no concern about ethical issues.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "lUQKdmZ3ty", "reviewer": "Reviewer_jc19", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The paper exposition is clear, and the method analysis is easy to follow.\n- The proposed method is simple and yet effective in outperforming existing baselines.\n\nWeaknesses:\n- Experiments: The presented comparisons are somewhat limited, and currently not enough to demonstrate the superiority of proposed approach. More specifically, \n  - a) The considered datasets/settings are rather small in terms of image resolution (only up to $64 \\times 64$), which naturally leaves a lot of redundancy in the image when upsampled and does not represent a standard use-case for ViTs. Thus, the authors need to demonstrate the method's effectiveness for datasets with larger and more diverse images (for example, Imagenet-1k) \n  - b) In terms of the comparisons against MPCViT, I understand the reasoning behind comparing latencies with matched accuracies but given that MPCViT employs much smaller models which are likely to have lower accuracy ceilings, the comparison may be somewhat unfair. It would help the paper to provide a comparison to MPCViT used with architectures of a similar scale.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a method for adapting ViT-based model architectures for efficient and secure multi-party computation.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Critical\n  - Please addresses weaknesses i.e. experiments on a dataset with larger dataset as well as fairer comparison to the baseline.\n\n- Recommended:\n  - Exploring the effectiveness of the method when scaling the ViT architecture (for ex. to small/base) would also help strengthen the paper.\n\n- Presentation/Clarity:\n  - The presentation of the main results, split across table 2 and 6 make it rather cumbersome for the reader to follow. It is recommended to consolidate them in one table. \n  - The efficiency comparisons in table 4 are missing details; does the comparison account for the fact that PriViT employs pretrained checkpoints? Does MPCViT employ similar pretraining? Please provide details.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "dxWKuobFee", "reviewer": "Reviewer_S5rT", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- The work features many good ablations and corresponding analysis. I especially liked the following\n  - The comparison between Squared Attention and other replacements for Attention\n  - Comparing PriViT with/without pretraining, and with/without knowledge distillation\n  - Analyzing the distribution of non-linearities replaced over layers\n- The work provides a detailed comparison/contextualization to existing works at the start of the paper\n\n### Weaknesses\n- What are the differences between this work and Cho et al except for replacing Attention with Squared Attention? If I understand correctly, the algorithm in this work is the natural analogue of SNL for transformers. Therefore, the primary contribution comes from the analysis of the algorithm, not the algorithm itself.\n- The baseline considered in this work is very conceptually similar to the method at hand. Does PriViT improve upon other approaches to private inference? Though other papers such as SAL-ViT and RNA-ViT haven't released their models, is it possible to compare against their numbers? \n\n### Questions\n\n- Is the citation for MPCViT on the first line of implication on page 3 meant to be Cho et al or Zeng et al?\n- How much work was it to privatize your algorithm using Delphi? Did it make sense to consider any alternatives? I wonder if practitioners would be interested in reading about these difficulties.\n- Is it possible to provide a plot demonstrating the Pareto frontier of accuracy vs inference latency for prior/current work? I understand the table already contains the information in numeric form.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The work aims to develop private inference for ViT's in an efficient, accuracy-preserving way. Unfortunately, ViT traditionally involves a lot of operations that are not amenable to standard private inference protocols such as Delphi. This work addresses this by taking a pretrained ViT and training it to switch out as many difficult operations (GeLU, Attention) for simpler alternatives (Identity, SquaredAttention) while still maintaining low error and distance to the original model. This method appears to achieve SOTA tradeoff between accuracy and inference latency. Other interesting findings include that many non-linearities are superfluous (since they can be replaced with simpler counterparts) and more important at the start (since these tend to be replaced less).", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "All in weaknesses/questions above.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors are recommended to consider the reviewer comments and thoroughly address them in the camera-ready version. Having said that, I do believe the paper passes the bar of the TMLR acceptance.
Decision: Accept (poster)