[Paper Review] Efficient Pruning for Machine Learning Under Homomorphic Encryption

Ehud Aharoni, Moran Baruch+9 more|arXiv (Cornell University)|Jul 7, 2022

Cryptography and Data Security44 references4 citations

TL;DR

HE-PEx proposes a novel pruning framework for homomorphic encryption (HE)-accelerated machine learning that combines weight pruning, permutation, and expansion to drastically reduce ciphertext volume and inference latency. By reordering weights via a new permutation algorithm and expanding sparse structures, it achieves up to 60% reduction in memory and latency on MNIST autoencoders with only 1.5×10⁻⁵ mean-square reconstruction loss, significantly outperforming standard pruning in HE contexts.

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.

Motivation & Objective

• Address the inefficiency of standard pruning in homomorphic encryption (HE) environments, where sparse weights do not reduce ciphertext count due to packing constraints.
• Overcome the limitation that HE ciphertexts cannot be eliminated even when most elements are zero, due to SIMD packing of multi-dimensional tensors.
• Develop a framework that enables high pruning ratios in HE-encrypted neural networks without sacrificing model accuracy.
• Integrate pruning with HE-friendly packing (tile tensors) and a non-interactive inference model to support real-world privacy-preserving ML workloads.
• Demonstrate that permutation and expansion can recover accuracy loss and unlock the full potential of pruning in HE-protected inference.

Proposed method

• Propose a novel permutation algorithm that reorders network weights to group non-zero values into fewer, denser tiles, increasing the likelihood of eliminating entire ciphertexts during pruning.
• Introduce a weight expansion technique that redistributes sparse weights across expanded tensor dimensions to preserve model accuracy after pruning.
• Integrate pruning with tile tensor packing (from HeLayers), where each tile is packed into a single ciphertext, enabling efficient SIMD operations.
• Design a framework that jointly optimizes pruning ratio, tile size, and accuracy loss using an objective function balancing latency, memory, and reconstruction error.
• Adapt the state-of-the-art pruning method Hunter (Cai et al., 2022) to the non-interactive HE setting, enabling end-to-end pruning and retraining for HE-compatible models.
• Use square activation functions as HE-friendly alternatives to ReLU to maintain compatibility with non-interactive inference and reduce computational overhead.

Experimental results

Research questions

• RQ1 Can standard neural network pruning techniques be effectively applied in homomorphic encryption environments, where ciphertexts cannot be removed even when most elements are zero?
• RQ2 How can permutation of model weights improve the effectiveness of pruning in HE-encrypted neural networks by increasing the number of eliminable ciphertexts?
• RQ3 To what extent can weight expansion mitigate accuracy loss after aggressive pruning in HE-protected models?
• RQ4 What is the achievable trade-off between model compression (ciphertext reduction), latency, memory, and reconstruction accuracy in HE-encrypted autoencoders?
• RQ5 How does the proposed HE-PEx framework compare to existing pruning methods in terms of ciphertext elimination and inference efficiency under HE?

Key findings

• HE-PEx achieves up to 85% reduction in the number of zero tiles for 2×2 tile sizes and 60% reduction for 16×16 tiles on MNIST autoencoders, with only 1.5×10⁻⁵ mean-square reconstruction loss.
• The framework reduces both memory and inference latency by 60% on the largest autoencoder model, demonstrating significant performance gains in end-to-end HE inference.
• Standard pruning techniques eliminate only 17% of ciphertexts when 85% of weights are pruned, whereas HE-PEx increases this to over 85% due to intelligent weight permutation and packing.
• The proposed permutation algorithm enables a 30% reduction in latency even after 95% pruning, indicating that ciphertext elimination is not the sole factor in performance gain.
• The combination of pruning, permutation, and expansion allows HE-PEx to maintain high accuracy while drastically reducing the number of active ciphertexts and HE operations.
• The method is orthogonal to activation function choice, and can be extended to other layers such as convolutional ones, though it is most effective for fully connected layers with tile-based packing.