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

1 Introduction . Machine learning has recently experienced tremedous growth , being used to solve problems with unprecedented accuracy in a myriad of domains . However , not all domains are alike—while many datasets are freely available , we often would like to train ML models on sensitive data . Troublingly , it has been demonstrated that disregarding these concerns , or even using heuristic and best-effort privacy approaches , can result in significant leakage of private information [ 9 , 12 ] . Differential privacy ( DP ) [ 16 ] has emerged as a strong and rigorous notion of data privacy , capable of protecting against privacy violations in a variety of settings . One of the workhorse algorithms in machine learning is stochastic gradient descent ( SGD ) which has a differentially private analogue , DPSGD [ 43 , 8 , 2 ] which was introduced as a drop-in replacement for SGD . The primary differences include a per-example gradient clipping step and a batch-level addition of Gaussian noise . While these modifications seem relatively innocuous , they have so far led to non-trivial costs in terms of running time and final accuracy . In this paper , we address and mitigate the running time overhead of DPSGD . Most modern machine learning frameworks allow efficient access to average minibatch gradients , and not at a per-example level . Access to these objects is critical in DPSGD , as well as other applications ∗Equal contribution . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . beyond privacy [ 52 ] . Lack of support for fast computation of per-example gradients has been noted and lamented numerous times for both TensorFlow [ 35 , 42 , 3 ] and PyTorch [ 32 , 6 ] . Numerous attempts to avoid these computational roadblocks have been proposed . Goodfellow [ 20 ] proposed an algorithmic solution for computing per-example ` 2-norms of the gradients for fullyconnected networks . Other proposed solutions work by exploiting Jacobians [ 14 ] or parallelizing over the batch dimension [ 4 ] . Several of these approaches are are restricted to specific types of architectures—for example , [ 20 ] is restricted to fully-connected layers , though [ 41 ] extends this to convolutional layers , and a very recent work [ 33 ] further considers layers including recurrent networks , attention , and more . BackPACK [ 14 ] currently supports only fully connected and convolutional layers , and while the paper states that it can be extended to recurrent and residual layers , GitHub issues related to implementation of these features have been open since November 2019 [ 13 ] . Facebook ’ s Opacus [ 51 ] emphasizes speed and scalability as the main selling points . We briefly mention microbatching , in which comparatively small subsets of the minibatch called “ microbatches ” of points are averaged before clipping , reducing the number of clipping operations ( and thus the running time ) , at the cost of requiring additional noise to achieve the same privacy guarantee . Since this generally results in significantly worse accuracy , we do not investigate it further in our work . A more thorough description of approaches is provided in Section 2.1 . As mentioned in the literature ( and thoroughly explored later 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 . For instance , Carlini et al . [ 12 ] comment “ Training a differentially private algorithm is known to be slower than standard training ; our implementation of this algorithm is 10-100x slower than standard training , ” where their implementation is based on TensorFlow Privacy . Additionally , Thomas et al . [ 45 ] 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 . 1.1 Results . 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 [ 19 , 10 ] and TensorFlow 2 ( TF2 ) [ 1 ] , both tensor-processing libraries from Google . These frameworks combine JIT compilation backed by the Accelerated Linear Algebra ( XLA ) [ 22 ] just-in-time compiler ( JIT ) 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 leverages vectorization and XLA compilation in TensorFlow 2 , demonstrates similar performance to JAX and significantly outperforms the existing TFP library . These changes have since been merged into TensorFlow Privacy . Our primary contributions are as follows : 1 . We thoroughly benchmark several frameworks and libraries for DPSGD . 2 . We extend TensorFlow Privacy to support TF2 , more efficient vectorization , and XLA compilation , significantly improving its running time in most cases ( referred to as Custom TFP in this paper ) . We also contribute a variant of our implementation to TensorFlow Privacy , which is now the fastest DPSGD algorithm the library provides . 3 . We demonstrate that methods which use vectorization , JIT compilation , and static graph optimization are consistently the fastest and most memory-efficient : specifically , JAX and Custom TFP . 4 . We find that , despite similarities in the compilation pipeline , JAX is generally faster than Custom TFP . We examine and discuss compiled XLA assembly to explain the discrepancy . 5 . Finally , our supplement contains code to reproduce these experiments , as well as guide researchers and engineers in developing fast code for private ML . Table 1 summarizes some of our experimental results , with median running time per epoch for a variety of settings . JAX and Custom TFP are consistently the fastest . We observe dramatic improvements for LSTMs [ 28 ] , potentially significant enough to bring LSTMs from impractical into the realm of feasibility . JAX is able to privately train these models 17x and 50x faster than the best alternative.2 Examining the overhead due to privacy : JAX ’ s running time increases by roughly 2x , compared to factors closer to 10x for alternatives . For fully-connected and 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 . Recall that these are per-epoch times : while an improvement of 0.5 seconds might seem insignificant , this can add up when training for many epochs . We perform an ablation study ( Table 2 ) for all models to pinpoint the source of all improvements . While our investigations show the consistent and substantial superiority of JAX for fast private machine learning , these benefits remain relatively unknown . Though a small number of experts are aware [ 44 ] , and the official JAX repo contains a toy demonstration [ 25 ] , before the initial posting of our paper a Google Scholar search revealed only two papers which use JAX for differential privacy [ 49 , 38 ] , and neither emphasizes or even comments on the computational advantages of JAX . Similarly , while efficient per-example gradients have been studied in TensorFlow [ 5 ] , efficient application of these techniques is not readily available to privacy researchers . We hope that our investigation will document this phenomenon and encourage others to adopt it for their private machine learning needs . 1.2 Simultaneous and Subsequent Work . Simultaneous to our work , [ 11 ] employ Johnson-Lindenstrauss projections to quickly approximate per-example gradient norms . This is an algorithmic modification , and will not be functionally equivalent to DPSGD – similar to microbatching , there is a time-accuracy tradeoff ( though not as severe in this case ) . Subsequent to the initial posting of this paper , we worked with Google engineers to implement our improvements into TensorFlow Privacy . Vadivelu contributed an JAX implementation of DPSGD to the Optax library [ 27 ] . [ 7 ] employed our findings to efficiently privately train BERT-Large . In a recent Opacus whitepaper [ 51 ] , the authors repeat some of our experiments on more recent versions of these frameworks ; we defer to their work for discussion of these results . 2 Description of Approaches . 2.1 Libraries Enabling DPSGD . JAX [ 19 , 10 ] . JAX is a recently introduced framework for machine learning , defined by its automatic differentiation capabilities and JIT compilation via the XLA compiler [ 22 ] . Programs written in pure Python and JAX ’ s NumPy [ 26 ] API can be translated to an intermediate language ( XLA-HLO ) to be JIT compiled , i.e. , to generate custom assembly instructions for the hardware . This enables optimizations such as kernel-fusions , buffer reuse , improved memory layout , and more . Additionally , one of the core functions present in JAX is VMAP , a vectorized map , which enables easy-to-write and efficient batch level parallelism that is fundamental to DPSGD . As we will demonstrate , these enable the fastest approach for DPSGD that we are aware of . 2Note that a confirmed bug in TF2 currently prevents us from running Custom TFP in these cases . Custom TFP . Vectorization and XLA-driven JIT compilation is also available in TensorFlow 2 [ 1 ] , which we leverage in our implementation , Custom TFP . With these primitives , we achieve performance comparable to JAX and surpassing existing DPSGD implementations in TensorFlow . We augment TensorFlow Privacy to better utilize tf.vectorized_map and follow TensorFlow 2 best practices while retaining the existing functionality . Chain-Rule-Based Per-Example Gradients [ 20 , 41 ] . This suite of techniques is implemented on top of PyTorch [ 39 ] . They support efficient GPU-accelerated per-example gradients for fullyconnected layers via [ 20 ] , as well as convolutional layers via [ 41 ] , which we describe in the detail in the following paragraphs . Let C , D , T , and B refer to the number of input channels , output channels , the spatial dimension , and the batch size . The shape of the input x is ( B , C , T ) . The conventional formula for the discrete convolution can be written as : C−1∑ c=0 K−1∑ k=0 x [ b , c , t+K ] h [ d , c , k ] . The gradient of this expression can be efficiently computed via automatic differentiation [ 40 ] . PyTorch ’ s automatic differentiation can not be parallelized across the batch dimension b [ 41 ] , which is required to backpropagate through the above expression . Instead , they rewrite the convolution as follows : C/G−1∑ c=0 K−1∑ k=0 x [ b , c , g C G , t+K ] h [ d , g , c , k ] , where G is the number of groups and the shape of x is ( 1 , B , C , T ) . The initial convolution is 1-dimensional , while the above expression includes an added dimension . Similarly , to allow backpropagation through a k-dimensional convolutional layer , a ( k + 1 ) -dimensional convolutional layer is required . This can be achieved by utilizing the group attribute in the convolution function in PyTorch , since splitting it into groups implies that the same convolution is applied to each individual group . BackPACK [ 14 ] . The chain rule gives the following expression for the gradient of a loss function : ∇θ ( i ) ` ( θ ) = ( Jθ ( i ) z ( i ) n ) T L−1∏ j=i ( J z ( j ) n z ( j+1 ) n ) T ( ∇ z ( L ) n ` n ( θ ) ) . In order to compute this quantity , one requires the ability to multiply the Jacobian by a vector and by a matrix , which is not currently supported in PyTorch ’ s automatic differentiation framework . In BackPACK , Dangel et al . [ 14 ] extend several layers within PyTorch to support fast Jacobian-vector and Jacobian-matrix products in order to extract quantities like individual gradients , variance , ` 2- norm of the gradients , and second-order quantities . In particular , to extract first-order gradients , their method multiplies the transposed Jacobian with the outputs of the layer : 1 N ∇θ ( i ) ` ( θ ) = 1 N ( Jθ ( i ) z ( i ) n ) T ( ∇ z ( i ) n ` ( θ ) ) , where i = 1 , . . . , N and each θ ( i ) has a gradient which is of shape ( N , d ( i ) ) . BackPACK provides efficient computation for the transpose of the Jacobian as well as the Jacobian . Opacus [ 51 ] . Opacus is a library for training PyTorch models with differential privacy , recently released by Facebook . It supports per-example gradients , using PyTorch ’ s forward and backward hooks to propagate gradients . They provide support for several PyTorch layers including LSTM layers , which are not supported in either of the previous two frameworks . Note that Opacus does not support PyTorch ’ s nn.LSTM but instead implements a separate opacus.layers.DPLSTM , with adjustments that allow individual gradients to propagate through it . PyVacy [ 48 ] . Before the release of Opacus ( and its predecessor PyTorch-DP ) , PyVacy was the most popular library for DP machine learning in PyTorch . PyVacy has no custom support for parallelization across the batch dimension for any layer since it processes each sample individually ( by way of a for-loop ) . This generally leads to a large increase in runtime for models trained using PyVacy . TensorFlow Privacy [ 37 ] TensorFlow Privacy is a library for differentially private machine learning , built on top of TensorFlow . TensorFlow Privacy has general support for a vectorized implementation of DPSGD via vectorized_map which allows it to parallelize across the batch dimension , used to extract per-example gradients . The library recently introduced a TensorFlow 2 compatible API that leverages GradientTape.jacobian to compute per-example gradients , which we compare seperately to the TensorFlow 1 API in our experiments .

Summary: This paper investigates three language primitives, namely vectorization, just-in-time (JIT) compilation and static graph optimization, that can be leveraged to speed up the process of stochastic gradient descent (DPSGD) in training ML models with differential privacy. It demonstrates by experiments that these techniques can lead to significant reductions in per epoch training time of neural networks compared to naïve DPSGD. It also performs ablation studies to further investigate the impact of each technique and find that JIT alone provides the most improvements.