Timezone: »

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

@
in Theory and Practice of Differential Privacy »

A common pain point in differentially private machine learning is the significant runtime overhead incurred when executing Differentially Private Stochastic Gra- dient 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 vector- ization 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.

Author Information

Pranav Subramani (University Of Waterloo)

Undergraduate Student who is interesting in Probabilistic Inference and working with a Professor on a project.

Nicholas Vadivelu (University of Waterloo)
Gautam Kamath (University of Waterloo)

More from the Same Authors