Optimizing private AI inference with new protocols and hardware design
OptiPrime: Optimizing Private Inference through Protocol-Hardware Co-design
Hardware ArchitectureCryptography and SecurityMachine Learning
Summary
Private AI inference ensures that your data stays secret when AI models make predictions, but this can be very slow because of complex encryption steps. The authors introduce OptiPrime, which combines new ways of handling encrypted data with smart hardware tricks to reduce the amount of data moving around and speed up computation. They also compress certain pieces of data to reduce memory use and improve performance. Their method runs several times faster than previous approaches on both regular computers and specialized chips.
What this means in practice
- •For cloud service providers: Run private AI inference on client data with improved speed by integrating the new protocol and hardware optimizations to reduce communication and memory bottlenecks.$Commercial implications: Enables offering faster privacy-preserving AI services to customers by accelerating encrypted AI computations with optimized protocols and hardware.
- •For data center operators: Deploy optimized hardware and protocols to decrease latency and resource use when handling confidential AI inference workloads.
Authors
Jiangrui Yu, Ye Yu, Si Chen, Chenqi Lin, Wenxuan Zeng, Junfeng Fan, Mingyu Gao, Meng Li
Abstract
Private deep neural network (DNN) inference based on hybrid homomorphic encryption (HE) and multi-party computation (MPC) can protect user data with a formal guarantee, but at the cost of significant latency overhead due to HE. Customized HE accelerators have been proposed and have achieved orders-of-magnitude speedup for individual HE operations. However, when directly applying a commercial HE accelerator to state-of-the-art HE-MPC frameworks, we observe only limited end-to-end performance gain. This is because HE-MPC frameworks often require wireless transmission of input and output ciphertexts for each HE operation, leading to a severe network communication bottleneck. To overcome this challenge, we introduce OptiPrime, a protocol-hardware co-optimization framework for efficient private DNN inference. OptiPrime features a novel HE protocol for convolutions that substantially reduces the number of transmitted output ciphertexts and mitigates the network communication bottleneck. Meanwhile, as the new protocol introduces complex computation for fewer output ciphertext, we observe new memory access challenges due to a high volume of weight plaintexts and intermediate ciphertexts. Hence, we further propose a lightweight compression system for the weight plaintexts, reducing memory traffic by 10 times, as well as a specialized dataflow to maximize on-chip data reuse of intermediate ciphertexts. Extensive experiments show that our framework outperforms the Cheetah baseline by at most 5.7 times on CPUs and 4.2 times with an accelerator.