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

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