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

Published: 09 Nov 2021, Last Modified: 04 Jun 2023NeurIPS 2021 PosterReaders: Everyone
Keywords: differential privacy, DPSGD, JAX, vectorization, just-in-time compilation
TL;DR: We show that low-level language features can dramatically speed up differentially private stochastic gradient descent.
Abstract: 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.
Supplementary Material: zip
Code Of Conduct: I certify that all co-authors of this work have read and commit to adhering to the NeurIPS Statement on Ethics, Fairness, Inclusivity, and Code of Conduct.
Community Implementations: [![CatalyzeX](/images/catalyzex_icon.svg) 2 code implementations](https://www.catalyzex.com/paper/arxiv:2010.09063/code)
Code: https://github.com/TheSalon/fast-dpsgd
13 Replies