Fast and Memory Efficient Differentially Private-SGD via JL Projections

1 Introduction . Over the past decade , machine learning algorithms based on ( deep ) neural architectures have lead to a revolution in applications such as computer vision , speech recognition and natural language processing ( NLP ) . An important factor contributing to this success is the abundance of data . For most of these applications , however , the training data comes from individuals , often containing personal and sensitive information about them . For example , natural language models for applications such as suggested replies for e-mails and dialog systems rely on the training of neural networks on email data of users [ CLB+19 , DBS19 ] , who may be left vulnerable if personal information is revealed . This could happen , for example , when a model generates a sentence or predicts a word that can potentially reveal private information of users in the training set . Many studies have shown successful membership inference attacks on deep learning models [ SSSS17 , CLE+19 ] . Indeed , in a recent work , [ CLE+19 ] show that “ unintended memorization '' in neural networks is both commonplace and hard to prevent . Such memorization is not due to overtraining [ TLL95 , CLE+19 ] , and ad hoc techniques such as early-stopping , dropout etc. , do not prevent the risk of privacy violations . Moreover , [ Fel20 ] shows that memorization is in fact necessary , provably , for some learning tasks . Thus , to prevent unintended privacy breaches one needs a principled approach for private training of deep learning ∗Author ordering is alphabetical . Work done when the first and the last two authors were interns at Algorithms group , Microsoft Research Redmond . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . models . In this paper we study training neural networks with differential privacy , a mathematically rigorous notion of privacy introduced in the seminal work of [ DMNS06 ] , and focus on user level privacy . Definition 1.1 ( ( ε , δ ) -DP ) . We say that an algorithm M is ( ε , δ ) -DP if for any two neighboring databases D , D′ and any subset S of outputs , we have Pr [ M ( D ) ∈ S ] ≤ eε Pr [ M ( D′ ) ∈ S ] + δ . Besides being a provable privacy notion , it has been shown that deep learning models trained with DP protect against leakage of sensitive information ; we refer the readers to [ CLE+19 , ACG+16 ] for more details . In a highly influential line of work , [ SCS13 , BST14 , ACG+16 ] introduced a differentially private version of stochastic gradient descent ( DP-SGD ) for training deep learning models , and showed that it is possible to achieve reasonable accuracy-vs-privacy tradeoff on common benchmarks such as MNIST and CIFAR10 . Since then , there has been a vast body of work building on and extending the DP-SGD algorithm ; we refer the readers to [ MRTZ18 , BDLS19 , CLE+19 , TAM19 , AMR+20 , ZWB20 , CWH20 , BKM+20 ] . The DP-SGD and its variations such as DP-Adam differ from their non-private counter parts in two crucial ways : • Gradient Clipping : In each iteration of DP-SGD , we clip each per-sample gradient to have ` 2-norm at most some fixed parameter C. This step ensures that the sensitivity of the average gradient is bounded , which is crucial for privacy analysis . Computing the norms of per-sample gradients is the most expensive step in the algorithm . Efficient implementations of backpropagation such as in TensorFlow and PyTorch only maintain the average of persample gradients across a batch by default . Therefore , getting the norm of each per-sample gradient requires significantly more time and memory . • Adding Noise : Once clipped gradients are averaged across a batch , DP-SGD algorithm adds carefully calibrated noise , typically sampled from the Gaussian distribution , to ensure privacy . The analysis of DP-SGD in [ ACG+16 ] then follows from a careful tracking of privacy budget lost in each iteration , for which they introduced a novel technique called moments accountant , which was later generalized as Renyi differential privacy by [ Mir17 ] . Tighter and accurate privacy accountants based on the notion of privacy loss random variables were given in [ KJH20 , KJPH21 ] . While DP-SGD has been shown to achieve reasonable accuracy-vs-privacy tradeoff [ BDLS19 , ACG+16 ] , and arguably is the only known algorithm for training deep neural networks , its use in real-world deep learning has been rather limited . One of the primary reasons for this is the training time of DP-SGD compared to SGD . In DP-SGD , per-sample gradients are computed at a heavy cost in runtime , especially for large deep learning models . The naive approach of setting the batch size to 1 is too slow to be practical as we completely lose the benefits of parallelization . This problem has attracted significant attention in the community and has been noted in popular implementations of DP-SGD including [ TP ] and [ Opa ] ( Pytorch DP ) . Many strategies have been proposed to circumvent this issue , and they fall broadly into the following categories : • Microbatching : DP-SGD implementation in Tensorflow Privacy allows dividing a batch into several microbatches and clipping the gradient at the microbatch level ; the per-sample gradients in a microbatch are first averaged and then clipped , and finally these clipped gradients are averaged across microbatches . Thus , if each microbatch is of size L , then it gives a speedup of L over the usual DP-SGD . Unfortunately , the sensitivity goes up by a factor of L , so we need to add L times more noise . In our experiments , we observe that this often leads to poor accuracy even for moderate values of L. • Multiple method : In this approach , proposed by [ Goo ] and implemented as the vectorized DP-SGD in [ TP ] , one copies the model as many times as there are samples in a batch . As each copy of the model is used on only one example , we can compute the per-sample gradients in parallel . This approach improves speed at the cost of memory and is impractical for large models . • Outer product method : This strategy was proposed by [ Goo15 ] for fully-connected networks and later generalized by [ RMT19 ] to include convolutional layers . The norms of per-sample gradients are computed exactly using outer products between the activations and backpropagated gradients across adjacent layers . In a very recent work , [ LK21 ] showed how to extend this approach to recurrent layers . A drawback of this approach is that it does not work for all network architectures in a black-box manner , and needs careful implementation for each network architecture . Furthermore , the current implementations of these methods require significantly more memory than non-private SGD as shown in [ SVK20 ] . • Compiler Optimization : A completely different approach towards mitigating the problem was suggested by [ SVK20 ] . They showed that by exploiting language primitives , such as vectorization , just-in-time compilation , and static graph optimization , one can implement DP-SGD algorithm significantly faster . They demonstrated these ideas in two frameworks : JAX and TensorFlow . While we believe this is exciting progress , the ideas in [ SVK20 ] are specific to these JAX and TensorFlow implementations ( as of today ) and present a non-algorithmic approach to this problem . As summarized in Table 1 , none of these approaches for speeding up DP-SGD completely solve the problem and fall short in at least one dimension . In this work , we propose a new algorithmic framework based on JL-projections for fast and memory efficient differentially private training of deep neural networks , which bypasses the expensive step of exactly computing per-sample gradient norms . 1.1 Our Contributions and Techniques . The main idea of our algorithm is to approximate the per-sample gradient norms instead of computing them exactly . Johnson–Lindenstrauss ( JL ) projections provide a convenient way to approximate the ` 2 norm of a vector ; simply project the vector onto a uniformly random direction , the length of the projection ( scaled appropriately ) is a good approximation to the ` 2-norm of the vector . By doing more such projections and averaging , we get even better approximation to the true ` 2-norm . Moreover , there is an efficient way to obtain such projections using forward-mode auto-differentiation or Jacobianvector product ( jvp ) ( see Section 2.1 for details ) . jvp can be calculated during the forward pass making it very efficient . Since this makes the sensitivity itself a random variable , the privacy analysis is significantly harder than the traditional DP-SGD . Moreover the Moments Accountant and related Renyi DP Accountant are not applicable to analyze the privacy of our algorithm since the associated privacy loss random variables have unbounded moments . Instead , we convolve the privacy loss distributions using Fast Fourier Transform ( FFT ) as done in [ KJH20 ] . To summarize , the key contributions of this paper are : • Our algorithms DP-SGD-JL and DP-Adam-JL are considerably faster than previously known differentially private training algorithms that require exact per sample gradient norms , and work for all network architectures . The privacy-vs-accuracy tradeoff achieved by our algorithms is comparable to the existing state-of-the-art DP-algorithms . • Memory footprint of our algorithms is nearly the same as that of non-private training algorithms . This allows us to work with larger batch sizes ( which is crucial for achieving good privacy-vs-accuracy tradeoffs ) without resorting to gradient accumulation . This also improves the running time of training . This is because we never need to calculate ( and store ) all the per-sample gradients simultaneously , whereas DP-SGD will need to store all the per-sample gradients simultaneously . • Compared to DP-SGD , our analysis of privacy is more involved . Since we only approximate the per-sample gradient norms , we can not precisely bound sensitivity . Therefore the analysis requires significantly new ideas , which could be of independent interest . 100 101 102 103 104 105 Time ( sec ) 0.5 0.6 0.7 0.8 0.9 Te st A cc ur ac y nonDP-Adam DP-Adam-Vanilla ( TFP ) DP-Adam-JL10 DP-Adam-JL20 DP-Adam-JL30 5 10 15 20 25 30 35 40 45 Figure 1 : Performance of various algorithms training an RNN on IMDb dataset using a batch size of 256 . The privacy parameter ε is color coded for δ = 10−5 . In DP-Adam-JLr , r refers to the JL dimension , i.e. , the number of JL projections used to approximate per-sample gradient norms . DPAdam-Vanilla ( TFP ) is the implementation of DP-Adam in [ TP ] library . • We demonstrate these improvements by training an RNN using layers such as bidirectional LSTM , embedding layer , fully connected etc . on the IMDb dataset . As can be seen from Figure 1 , our algorithms are significantly faster than current implementations of DP-SGD while achieving similar privacy-vs-accuracy tradeoff . • Our algorithms introduce a new knob , the dimension of JL-projection , which allows us to do a tradeoff between training time and privacy , which was not possible in earlier algorithms . All hyperparameters being the same , smaller JL dimension will give much better running time with an increase in privacy budget . • Our experiments show that although the privacy bounds we could prove for DP-SGD-JL are not so great for very small JL dimensions ( see Figure 3 ) , their behavior ( accuracyvs-epoch curve ) converges very quickly to that of DP-SGD-Vanilla . Figure 4 shows that DP-SGD-JL ( 3 ) is already very close to DP-SGD-Vanilla and DP-SGD-JL ( 20 ) is almost indistinguishable . We find these properties of DP-SGD-JL algorithms to be very useful during initial stages of experimentation and hyper-parameter tuning for private training .

Summary: Differentially private stochastic gradient descent requires clipping the gradient computed on each example before aggregation. That is, we compute the 2-norm of the gradient vector and, if it is too large, we scale it down. Surprisingly, this step entails a significant computational overhead relative to non-private stochastic gradient descent. The reasons for this slowdown are not entirely clear to me. It seems to be a software issue deep in the bowels of TensorFlow and PyTorch -- somewhere between the Python and the Hardware this operation is poorly optimized.  This paper proposes to alleviate the problem by performing approximate, rather than exact clipping. Specifically, a Johnson-Lindenstrauss sketch is used to approximate the 2-norm of the vector instead of exactly computing it. This JL approximation can be computed more efficiently than the exact norm. However, there is a tradeoff between speed and approximation quality as we can alter the dimension of the JL sketch.  Technically, the most involved aspect of this work is the privacy analysis. Since the gradients are only approximately bounded, we cannot apply the usual analysis. This algorithm does not satisfy Renyi differential privacy as there is heavy-tailed privacy loss due to the possibility of drastically underestimating the norm. Instead an entirely new analysis is required based on the recent Fourier accounting approach.  Experimental results demonstrate that there is a significant improvement in runtime. However, the privacy parameters also degrade significantly.