🔬 Research: Fast differentially private SGD with JAX - paper/code

[8bitmp3]

Just wanted to highlight an awesome paper by @pranavsubramani, @n2cholas and @hoonose (University of Waterloo). It's called Enabling Fast Differentially Private SGD via Just-in-Time Compilation and Vectorization 👏.

Repo: https://github.com/TheSalon/fast-dpsgd
JAX code: https://github.com/TheSalon/fast-dpsgd/blob/main/jaxdp.py
arXiv: https://arxiv.org/pdf/2010.09063.pdf

TL;DR (paper)

"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 ex- ploiting 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..."

Results (paper)

"...We demonstrate that one can mostly eliminate the significant running time overhead of differentially private SGD by exploiting language primitives such as vectorization, just-in-time (JIT) compilation, and static graph optimization. These features are core primitives within JAX and TensorFlow 2 (TF2), both tensor-processing libraries by researchers at Google. These frameworks combine JIT compilation backed by the Accelerated Linear Algebra (XLA) [Google] compiler with auto differentiation for high-performance machine learning. As we will see, JAX is consistently the fastest method for running DPSGD, with running times comparable to the non-private case. Our custom TensorFlow Privacy (TFP) implementation (referred to as Custom TFP), which takes advantage of new features in TensorFlow 2 as well as XLA, demonstrates similar performance to JAX and significantly outperforms the existing TFP library..."
...
"3. We demonstrate that methods which use vectorization, JIT compilation, and static graph optimization are consistently the fastest and more memory-efficient: specifically, JAX and our modification of TFP.
"4. We find that, despite similarities in the compilation pipeline, JAX is generally faster than our customized TFP. We examine and discuss compiled XLA assembly to explain the discrepancy."

"...For the fully-connected and MNIST convolutional networks, JAX or Custom TFP almost entirely remove the overhead due to privacy. In fact, the running times are significantly better than some alternatives without privacy."

Some background on DP and challenges with compute (paper)

"...Lack of support for fast computation of per-example gradients has been noted and lamented numerous times in discussion forums and GitHub issues for both TensorFlow and PyTorch..."

"...Numerous attempts to avoid these computational roadblocks have been proposed... As mentioned in the literature (and thoroughly explored in this paper), all existing approaches seem to incur moderate to severe running time overhead versus non-private SGD, with slowdowns as large as two orders of magnitude... Additionally, Thomas et al. (working in an unspecified framework) document a slowdown from 12 minutes to 14 hours due to the introduction of differential privacy, a 70x slowdown. The effect of these slowdowns can range from an inconvenience when it comes to rapid prototyping of smaller models, to prohibitively expensive for a single training run of a larger model. Overcoming this obstacle is an important step in helping differentially private machine learning transition from its present nascent state to widespread adoption."

(Hi @iamtrask (OpenMined), I thought you'd also be interested 👋)

[8bitmp3]

@n2cholas did your research get a mention in https://ai.googleblog.com/2021/01/google-research-looking-back-at-2020.html 👀

[n2cholas]

yes!! 🎉🎉 @pranavsubramani @hoonose

[8bitmp3]

Congrats on the award @n2cholas 🏆 https://cs.uwaterloo.ca/news/nicholas-vadivelu-receives-2021-jessie-w-h-zou-memorial-award

"I don't really understand how, but despite being an undergrad, he already has a professional-level understanding of machine learning tools and techniques. E.g., he knew about JAX far earlier than most, which was a key part of this paper." - https://twitter.com/thegautamkamath/status/1397206194230632452

cc @mattjj @jakevdp @skye @jekbradbury @saeta

[n2cholas]

Thank you!! 😀😀

[8bitmp3]

Congrats @pranavsubramani, @n2cholas and @hoonose

Your research was mentioned in the Google AI blog post: Applying Differential Privacy to Large Scale Image Classification

Scaling Differential Privacy with JAX
Exploring multiple architectures and training configurations to research what works for DP can be debilitatingly slow. To streamline our efforts, we used JAX, a high-performance computational library based on XLA that can do efficient auto-vectorization and just-in-time compilation of the mathematical computations. Using these JAX features was previously recommended as a good way to speed up DP-SGD in the context of smaller datasets such as CIFAR-10.

New research by @AlexeyKurakin et al:

• GitHub: https://github.com/google-research/dp-imagenet
• arXiv: Toward Training at ImageNet Scale with Differential Privacy