Enabling Fast Differentially Private SGD via Just-in-Time Compilation and Vectorization

Abstract

A common pain point in differentially private machine learning is the significant runtime overhead incurred when executing Differentially Private Stochastic Gradient Descent (DPSGD), which may be as large as two orders of magnitude. We thoroughly demonstrate that by exploiting powerful language primitives, including vectorization, just-in-time compilation, and static graph optimization, one can dramatically reduce these overheads, in many cases nearly matching the best non-private running times. These gains are realized in two frameworks: JAX and TensorFlow. JAX provides rich support for these primitives as core features of the language through the XLA compiler. We also rebuild core parts of TensorFlow Privacy, integrating features from TensorFlow 2 as well as XLA compilation, granting significant memory and runtime improvements over the current release version. These approaches allow us to achieve up to 50x speedups in comparison to the best alternatives. Our code is available at https://github.com/TheSalon/fast-dpsgd.

arXiv commentTo appear in NeurIPS 2021

References in corpus (10)

• Deep Learning with Differential Privacy
• The Secret Sharer: Evaluating and Testing Unintended Memorization in Neural Networks
• Protection Against Reconstruction and Its Applications in Private Federated Learning
• Private Mean Estimation of Heavy-Tailed Distributions
• Efficient Per-Example Gradient Computations
• Tempered Sigmoid Activations for Deep Learning with Differential Privacy
• CoinPress: Practical Private Mean and Covariance Estimation
• Opacus: User-Friendly Differential Privacy Library in PyTorch
• Efficient Per-Example Gradient Computations in Convolutional Neural Networks
• Auto-Vectorizing TensorFlow Graphs: Jacobians, Auto-Batching And Beyond

Cited by in corpus (10)

• Adversarial Attacks Against Deep Generative Models on Data: A Survey
• Large Language Models Can Be Strong Differentially Private Learners
• PPFL: Privacy-preserving Federated Learning with Trusted Execution Environments
• Opacus: User-Friendly Differential Privacy Library in PyTorch
• Fast and Memory Efficient Differentially Private-SGD via JL Projections
• Deep Learning with Label Differential Privacy
• Privacy Regularization: Joint Privacy-Utility Optimization in Language Models
• Information-constrained optimization: can adaptive processing of gradients help?
• D3p -- A Python Package for Differentially-Private Probabilistic Programming
• Towards Prior-Free Approximately Truthful One-Shot Auction Learning via Differential Privacy