[论文解读] Efficient Pruning for Machine Learning Under Homomorphic Encryption

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

Cryptography and Data Security参考文献 44被引用 4

一句话总结

HE-PEx 提出了一种用于同态加密（HE）加速机器学习的新型剪枝框架，结合了权重剪枝、排列和扩展技术，显著降低了密文体积和推理延迟。通过一种新型排列算法重新排序权重，并扩展稀疏结构，其在 MNIST 自编码器上实现了高达 60% 的内存和延迟减少，且平均平方重建损失仅为 1.5×10⁻⁵，显著优于标准剪枝在 HE 环境下的表现。

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.

研究动机与目标

• 解决标准剪枝在同态加密（HE）环境中效率低下的问题，因为在打包约束下，稀疏权重无法减少密文数量。
• 克服 HE 密文无法被消除的限制，即使大多数元素为零，也由于多维张量的 SIMD 打包机制而无法移除。
• 开发一种框架，使 HE 加密神经网络能够实现高剪枝率，同时不牺牲模型精度。
• 将剪枝与 HE 友好的打包方式（分块张量）及非交互式推理模型相结合，以支持现实世界中的隐私保护机器学习工作负载。
• 证明排列与扩展技术能够恢复精度损失，并充分发挥剪枝在 HE 保护推理中的全部潜力。

提出的方法

• 提出一种新型排列算法，通过重新排序网络权重，将非零值聚集到更少、更密集的块中，从而提高剪枝过程中可消除整个密文的可能性。
• 引入一种权重扩展技术，将稀疏权重重新分布到扩展后的张量维度中，以在剪枝后保持模型精度。
• 将剪枝与分块张量打包（来自 HeLayers）集成，其中每个块打包为单个密文，从而支持高效的 SIMD 操作。
• 设计一种联合优化剪枝率、块大小与精度损失的框架，利用平衡延迟、内存与重建误差的目标函数。
• 将最先进的剪枝方法 Hunter（Cai et al., 2022）适配到非交互式 HE 环境中，实现 HE 兼容模型的端到端剪枝与微调。
• 采用平方激活函数作为 ReLU 的 HE 友好替代方案，以保持与非交互式推理的兼容性并降低计算开销。

实验结果

研究问题

• RQ1 标准神经网络剪枝技术能否在同态加密环境中有效应用，其中即使大多数元素为零，密文也无法被移除？
• RQ2 通过重新排列模型权重，如何提升剪枝在 HE 加密神经网络中的有效性，从而增加可消除密文的数量？
• RQ3 在 HE 保护模型中，权重扩展在激进剪枝后能多大程度上缓解精度损失？
• RQ4 在 HE 加密自编码器中，模型压缩（密文减少）、延迟、内存与重建精度之间可实现怎样的权衡？
• RQ5 与现有剪枝方法相比，所提出的 HE-PEx 框架在 HE 环境下的密文消除能力与推理效率如何？

主要发现

• 在 MNIST 自编码器上，HE-PEx 实现了 2×2 块大小下最多 85% 的零块减少，以及 16×16 块大小下 60% 的减少，平均平方重建损失仅为 1.5×10⁻⁵。
• 该框架在最大自编码器模型上将内存和推理延迟均减少了 60%，显著提升了端到端 HE 推理的性能。
• 当剪枝 85% 的权重时，标准剪枝技术仅能消除 17% 的密文，而 HE-PEx 由于智能权重排列与打包，将这一比例提升至 85% 以上。
• 所提出的排列算法即使在 95% 剪枝后仍能实现 30% 的延迟降低，表明密文消除并非性能提升的唯一因素。
• 剪枝、排列与扩展的结合使 HE-PEx 在大幅减少活跃密文数量与 HE 操作数的同时，仍能保持高精度。
• 该方法与激活函数的选择正交，可扩展至其他层（如卷积层），但对基于分块打包的全连接层效果最佳。