Timezone: »

Enabling Fast Differentially Private SGD via Just-in-Time Compilation and Vectorization
Pranav Subramani · Nicholas Vadivelu · Gautam Kamath

Thu Dec 09 08:30 AM -- 10:00 AM (PST) @

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: one is JAX, which provides rich support for these primitives through the XLA compiler. We also rebuild core parts of TensorFlow Privacy, integrating more effective vectorization as well as XLA compilation, granting significant memory and runtime improvements over previous release versions. Our proposed approaches allow us to achieve up to 50x speedups compared to the best alternatives. Our code is available at https://github.com/TheSalon/fast-dpsgd.

Author Information

Pranav Subramani (University Of Waterloo)
Nicholas Vadivelu (University of Waterloo)

I'm a student interested in **machine learning** and **programming languages**. Currently, I'm making machine learning inference faster at **NVIDIA** ([TensorRT](https://developer.nvidia.com/tensorrt)) and studying Computer Science & Statistics at the **University of Waterloo**. Previously, I did deep learning research on collaborative self-driving at **Uber ATG**, helped improve the neural network optimization algorithm [K-FAC](https://github.com/tensorflow/kfac) at **Google Brain**, and developed data-driven models to identify fraud at **John Hancock Financial**. I've also worked as a Research Assistant for **Prof. Lin Tan** in deep learning-driven software analysis and **Prof. Pascal Poupart** on neural network parameter learning. If you love tidbits of technical info, follow me on [twitter](https://twitter.com/nicvadivelu). Feel free to reach out through email: nicholas.vadivelu@gmail.com or book some time with me through [calendly](https://calendly.com/nicv)--I love to chat!

Gautam Kamath (University of Waterloo)

More from the Same Authors