Efficient Pruning for Machine Learning Under Homomorphic Encryption

Abstract

Privacy-preserving machine learning (PPML) solutions are gaining widespread popularity. Among these, many rely on homomorphic encryption (HE) that offers confidentiality of the model and the data, but at the cost of large latency and memory requirements. Pruning neural network (NN) parameters improves latency and memory in plaintext ML but has little impact if directly applied to HE-based PPML.

We introduce a framework called HE-PEx that comprises new pruning methods, on top of a packing technique called tile tensors, for reducing the latency and memory of PPML inference. HE-PEx uses permutations to prune additional ciphertexts, and expansion to recover inference loss. We demonstrate the effectiveness of our methods for pruning fully-connected and convolutional layers in NNs on PPML tasks, namely, image compression, denoising, and classification, with autoencoders, multilayer perceptrons (MLPs) and convolutional neural networks (CNNs).

We implement and deploy our networks atop a framework called HElayers, which shows a 10–35% improvement in inference speed and a 17–35% decrease in memory requirement over the unpruned network, corresponding to 33–65% fewer ciphertexts, within a 2.5% degradation in inference accuracy over the unpruned network. Compared to the state-of-the-art pruning technique for PPML, our techniques generate networks with 70% fewer ciphertexts, on average, for the same degradation limit.

A Appendix

1.1 A.1 Additional Evaluation

Selection of P4E prune$^\textbf{pack}$ Threshold. We justify our choice of selecting a threshold for the prune$^\textrm{pack}$ step (threshold % of zeros in the tile, above which the tile is pruned) of P4E by running a sweep and reporting the results in Fig. 12, for our AE-Compressor with CIFAR-10. The behavior is similar for other networks and datasets. We observe that large values of the threshold do not introduce sufficient new zero tiles to make any meaningful improvement upon P3E. At the other extreme, too small values of this threshold aggressively remove weights, leading to degradation in test loss. We select 93.8% as our choice here, but posit that our results can be further improved by tuning this as a hyperparameter.

Fig. 12.

Test loss vs. tile sparsity (% of zero tiles) for different prune$^\textrm{pack}$ thresholds of our P4E scheme, considering four tile shapes using the AE-Compressor with CIFAR-10.

Comparison of Permutation Algorithms. We compare different permutation algorithms for the permute step in HE-PEx.

Fig. 13.

Increase in tile sparsity (log scale) with different permutation algorithms for AlexNet on the COVIDx CT-2A dataset, after Lc/L1/Wei pruning.

Figure 13 shows the results of AlexNet with Lc/L1/Wei pruning of 97.5% of the weights, which is the same scenario discussed in Sect. 6.3. As expected, “Random Reordering” for the row-column co-permutation provides little-to-no benefit, due to the sheer number of possible permutations of the entire network. Interestingly, k-Means alone only improves the benefits to <10%. Moreover, the use of a “Grouped Hamming” variant (Sect. 4.2) does not show significance for the smaller tile shapes. However, as seen in the $64\times 64$ case, grouped Hamming distance is critical for grouping zeros together for large tile shapes. Finally, the impact of balancing the clusters in k-Means is paramount, as it brings in the critical information of the first dimension of the tile shape into the optimization algorithm. In summary, our permutation algorithm variant improves the tile sparsity by $\sim$4–1,000$\%$ for this network.

Table 2. Network configurations and hyperparameters. [HL = hidden layer]

Fig. 14.

Flowchart illustrating the P4E strategy with local pruning scope for the prune step and global scope for prune$^\textrm{pack}$. This is executed as a pre-deployment step on the entity that has access to the plaintext network parameters.

1.2 A.2 Experimental Setup

Our experiments involving the methods in Sect. 4.1 are done on a cluster of systems equipped with Tesla A100 GPUs and Xeon Gold 6258R CPUs. We used PyTorch version 1.11.0 accelerated with CUDA 11.6. For end-to-end evaluation using HElayers [2], we used a 44-core Xeon CPU E5-2699 @2.20 GHz machine with 750 GB memory. We observed that for our pruned networks, not all of the cores were fully utilized and performance saturated around 8 cores. Thus, for a fair comparison, we limited the system to use 8 cores for all of our experiments. We use a modified HElayers [2] with the CKKS SEAL [23] implementation targeting 128-bit security, and the reported results are the average of 10 runs.

Network Configurations and Model Hyperparameters. Table 2 shows details of the activation functions used in each network, the network architecture, and training hyperparameters for each of the networks and datasets in Sect. 5.

Enhancing HElayers. We integrate HE-PEx with HeLayers [2] and show a flowchart for P4E in Fig. 14. This can be extended as needed to our other strategies. We illustrate a grid search approach in the flowchart with {thres_all} and {tile_shapes} being a pre-selected set of values. The output of the framework is a deployment-ready model that meets an objective function based on accuracy, latency, throughput, and memory. A local search strategy may be considered to navigate the state space with fewer training and permutation rounds, which are the most time-consuming segments of this flow. The zero-tile–skipping logic is programmed into HeLayers by associating a zero_flag bit with each ciphertext, which is set during the HE Enc procedure (Sect. 3.1).