CryptGPU: Fast Privacy-Preserving Machine Learning on the GPU

### Overview CryptGPU is an end-to-end framework enabling privacy-preserving machine learning (PPML) inference and training directly on GPUs. As machine learning models scale from small networks (like LeNet on MNIST) to deeper architectures (such as ResNet and VGG-16 on ImageNet), the primary performance bottleneck shifts from secure multi-party computation (MPC) communication overhead to raw computation, with linear layers accounting for up to 98% of training time in prior work. To bridge the massive performance gap between plain text and private machine learning—which historically reached 23,000x for inference (CryptFlow) and 42,000x for training (Falcon)—CryptGPU keeps all computations on the GPU, yielding substantial speedups across both linear and non-linear layers.

### Key Topics - **Privacy-Preserving Machine Learning (PPML):** Addressing data privacy challenges in collaborative healthcare and diagnostic scenarios where patient data and proprietary model weights must remain confidential during inference and training. - **Threat Model & Security Assumptions:** Operating under a three-party computation (3PC) framework with semi-honest security and an honest majority, permitting at most one semi-honest corrupt party who follows protocol execution but attempts to learn additional information. - **Bridging MPC Fixed-Point Arithmetic and GPU CUDA Kernels:** Overcoming the lack of native CUDA kernels for 32-bit or 64-bit integer fixed-point arithmetic by embedding integer operations into floating-point hardware. - **GPU-Friendly MPC Protocols:** Designing parallel, branch-free non-linear operations (such as approximation-based division and limit-based exponentiation) to maximize GPU parallel execution rather than relying on control-flow-heavy approaches like garbled circuits.

### Technical Details CryptGPU is built on top of PyTorch and CrypTen, structuring its architecture around three primary pillars: 1. **Replicated Secret Sharing:** Utilizing 3-party replicated additive secret sharing as the core primitive, where each party holds two out of three shares, efficiently supporting addition and multiplication. 2. **Embedding Fixed-Point Arithmetic in Floating-Point Computations:** Standard MPC fixed-point representations require 32-bit or 64-bit integer operations, but standard GPU CUDA acceleration is built around floating-point formats. CryptGPU exploits the fact that 64-bit double-precision floats have 52 bits of precision, enabling exact integer representations between -2^52 and 2^52. Using bilinearity, large integers (32-bit or 64-bit) are decomposed into sums of products of smaller 16-bit integers (e.g., (a1 + a2)(b1 + b2)), ensuring intermediate products never exceed 52 bits and execute without precision loss. This functionality is exposed as a reusable `CudaLongTensor` abstraction. 3. **GPU-Friendly Non-Linear Protocols:** Recognizing that GPUs struggle with conditional branching, CryptGPU replaces branch-heavy garbled circuits with approximation-based methods for private division, and uses limit-based approximations to turn exponentiation into sequences of parallel multiplications.

Empirical evaluations indicate that while CryptGPU shows little benefit on small workloads where communication dominates, it significantly outperforms prior state of the art on larger workloads: - **Inference:** Delivers a 2.5x to 3.7x improvement over Falcon and CryptFlow on networks such as ResNet and VGG-16 on Tiny ImageNet. - **Training:** Outperforms Falcon by 7x on small models and up to 36x on large models (e.g., AlexNet on Tiny ImageNet). - **Microbenchmarks:** Microbenchmark tests reveal that running private convolution on GPU is 100x faster than on CPU, and private ReLU is 10x faster.

### Takeaways CryptGPU demonstrates that GPU acceleration is critical for practical privacy-preserving machine learning at scale, effectively mitigating the severe computation bottlenecks in linear layer evaluation. While the current system is limited to single-GPU computation, future work includes expanding to multi-GPU clusters and designing dedicated MPC protocols optimized specifically around GPU parallel capabilities.