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

Understand

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.