Fast and Memory Efficient Differentially Private-SGD via JL Projections∗
Year: 2022

Abstract:
Differentially Private-SGD (DP-SGD) of [ACG16] and its variations are the only known algorithms for private training of large scale neural networks. This algorithm requires computation of per-sample gradients norms which is extremely slow and memory intensive in practice. In this paper, we present a new framework to design differentially private optimizers called DP-SGD-JL and DP-Adam-JL. Our approach uses Johnson–Lindenstrauss (JL) projections to quickly approximate the per-sample gradient norms without exactly computing them, thus making the training time and memory requirements of our optimizers closer to that of their non-DP versions. Unlike previous attempts to make DP-SGD faster which work only on a subset of network architectures or use compiler techniques, we propose an algorithmic solution which works for any network in a black-box manner which is the main contribution of this paper. To illustrate this, on IMDb dataset, we train a Recurrent Neural Network (RNN) to achieve good privacy-vs-accuracy tradeoff, while being significantly faster than DP-SGD and with a similar memory footprint as non-private SGD.

heading: 1 Introduction
text: 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.

heading: 1.1 Our Contributions and Techniques
text: 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 cannot 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.

heading: 2 DP-SGD-JL Algorithm
text: In this section we describe our new differentially private optimizers. We will first present DP-SGD-JL, and DP-Adam-JL follows in the same lines and is presented in the Supplementary Material. We begin with an introduction to “Jacobian-vector product” (jvp) which is crucial for our algorithm.
2.1 Jacobian-vector product (jvp)
Given a function f : Rd → R, the gradient of f with respect to θ, denoted by∇θf , is:
∇θf = ( ∂f
∂θ1 , ∂f ∂θ2 , . . . , ∂f ∂θd
) .
Let F : Rd → Rm be some function given by F (θ) = (F1(θ), F2(θ), . . . , Fm(θ)). The Jacobian of F (θ), denoted by∇θF , is the matrix:
∇θF = [ ∂Fi ∂θj ] ij =  ∇θF1 ∇θF2
... ∇θFm  . Most auto-differentiation packages allow for calculating the vector-Jacobian product (vjp) given by uT∇θF = ∑m i=1 ui∇θFi for any u ∈ Rm efficiently using reverse-mode auto-differentiation, which
is the familiar ‘backpropagation’.2 One can also calculate the Jacobian-vector product (jvp) given by
∇θF · v =  〈∇θF1, v〉 〈∇θF2, v〉
... 〈∇θFm, v〉  efficiently using forward-mode auto-differentiation (i.e., the required derivatives are calculated during the forward pass of the network). We refer the reader to the survey on automatic differentiation by [BPRS17] for more details. jvp is implemented using forward-mode auto-differentiation in the recent TensorFlow versions.3 Unfortunately, PyTorch doesn’t have an implementation of forwardmode auto-differentiation. Instead, one can compute jvp using two calls to vjp, this is called the ‘double vjp trick’ (see [Tow]).4 Define G(α) = αT∇θF , which can be calculated using vjp. Note that∇αG = (∇θF )T . Now we can use vjp again on G to calculate jvp as
vT∇αG = vT (∇θF )T = (∇θF · v)T .
In our experiments, we use the efficient implementation of jvp in TensorFlow to compute the Jacobian-vector products as the double vjp trick is a few times slower.
2.2 Algorithm
0.0 0.5 1.0 1.5 2.0 2.5 3.0 0.0
0.5
1.0
1.5
2.0
2.5 3.0 r=1 r=4 r=10 r=30
Figure 2: Distribution of √
1 rχ 2 r for different r.
The main idea of our algorithm is to approximate `2-norms of per-sample gradient norms instead of computing them exactly. And the key tool to achieve this is JL projections. Proposition 2.1 (JL projections). Let y ∈ Rd be any vector. If v1, v2, . . . , vr ∼ N (0, Id) are independent standard Gaussian vectors, then
Mr = √∑r i=1 1 r 〈y, vi〉 2 has the same distri-
bution as ‖y‖2 √ 1 rχ 2 r .
5 In particular E[M2r ] = ‖y‖22 .
Proof. By the properties of the standard Gaussian distribution, 〈y, vi〉 has the distribution of ‖y‖2N (0, 1). And 〈y, vi〉 are independent for i = 1 to r. Therefore ∑r i=1 〈y, vi〉
2 has the same distribution as ‖y‖22 χ2r.
As shown in Figure 2, as r grows larger, the distribution of √
1 rχ 2 r concentrates more around 1
and therefore Mr becomes a better estimate of ‖y‖2 . Using jvp, we can compute projections of per-sample gradients on to standard Gaussian vectors quickly and therefore get good approximations to their norms. This is the main idea of DP-SGD-JL (Algorithm 1).

heading: 2.3 Privacy Analysis
text: Let us first recall the privacy analysis of DP-SGD in Abadi et al. [ACG+16], which consists of three main arguments: 1) First, they show that each step of the DP-SGD algorithm is differentially private,
2In PyTorch, uT∇θF can be calculated as autograd.grad(F,θ,grad_outputs=u). In TensorFlow, this is tf.GradientTape().gradient(F,θ,output_gradients=u).
3Supported in tf-nightly≥2.4.0.dev20200924 as tf.autodiff.ForwardAccumulator(θ,v).jvp(F). JAX also has an implementation of jvp.
4In Pytorch, an implementation of jvp using the double vjp trick exists and can be invoked as torch.autograd.functional.jvp(F,θ,inputs=v).
5χ2r is the chi-square distribution with r degrees of freedom which is the distribution of sum of squares of r standard Gaussians.
Algorithm 1: Differentially private SGD using JL projections (DP-SGD-JL) Input: Examples {x1, x2, . . . , xN}, loss function L(θ) = Ei∈[N ][L(θ;xi)], initialization θ0. Parameters: number of iterations T , learning rates (η1, η2, . . . , ηT ), noise scale σ, batch size B, clipping norms (C1, C2, . . . , CT ), number of JL projections r. for t = 1 to T do
Sample St = {X1, X2, . . . , XB} ⊂ {x1, x2, . . . , xN} uniformly at random.; Define F (θ) = (L(θ;X1),L(θ;X2), . . . ,L(θ;XB)); Sample v1, v2, . . . , vr ← N (0, Id) (where θ ∈ Rd); // JL projections to estimate per-sample gradient norm for j = 1 to r do (P1j , P2j , . . . , PBj)← ∇θF (θt−1) · vj ; // Computed using jvp // Note that Pij = 〈∇θL(θt−1;Xi), vj〉
for i = 1 to B do Set Mi = √ 1 r ∑r j=1 P 2 ij ; // Mi is an estimate for ‖∇θL(θt−1;Xi)‖2 .
// Scale the losses to implicitly clip per-sample gradients Define L̃(θ) = 1B ∑ i∈B min{1, Ct Mi } · L(θ;Xi); // ∇θL̃(θt−1) is the average of clipped gradients, computed by one backprop // Add noise to ∇θL̃(θt−1), the average of clipped gradients g̃t ← ∇θL̃(θt−1) + σ·CtB · N (0, Id); Update θt ← θt−1 − ηtg̃t;
Output: θ0, θ1, θ2, . . . , θT
which follows from the simple observation that norm of the each gradient can be at most C due to clipping and so each step is a Gaussian mechanism. 2) Next, they observe that each mini-batch is a random sample of the dataset, hence one can appeal to the privacy amplification by subsampling. 3) Finally, they do a composition across all the iterations of the DP-SGD algorithm. However, existing works at that time did not give a tight analysis of the privacy loss in steps 2 and 3, hence [ACG+16] invented the Moments Accountant (MA) technique.
Unfortunately, analysis of our algorithm is significantly more involved than that of [ACG+16], as all the above 3 steps pose new technical challenges. First note that unlike DP-SGD, in our algorithm, we only have a rough estimate of the norm of each gradient; hence, the sensitivity of the gradients is a random variable. This means that the each iteration is not a Gaussian mechanism. The Moments Accountant (and the closely related RDP Accountant) cannot be used compute the privacy loss of the composition because the privacy loss random variables for our algorithms do not have exponential moments (for JL dimension r, only moments < r/2 exist).
To overcome these challenges, we use the notion of privacy loss random variables (PRVs) [DR16, MM18, KJH20, GLW21]. In particular, we use the PRV Accountant from [GLW21]. Define δ(X||Y )(ε) = Pr[Y > ε]− eε Pr[X > ε]. For any DP algorithm M , there exists random variables (X,Y ), called its privacy loss random variables (PRVs) such that the privacy curve of M is given by δ(X||Y ), i.e., M satisfies (ε, δ(X||Y )(ε))-DP for every ε. The following theorem shows a lower bound on the privacy curve of each iteration of the DP-SGD-JL algorithm, which is the main technical contribution of this paper. Theorem 2.2. Let Zr = 1√ 1
rχ 2 r
where χ2r is the chi-square distribution with r degrees of freedom.
Each iteration of DP-SGD-JL (before subsampling) satisfies (ε, δ(X||Y )(ε))-DP where X = N ( −Z 2 r
2 , Z2r
) and Y = N ( Z2r 2 , Z2r ) . (1)
From the above theorem, we can prove that the privacy curve of each iteration behaves like δ(ε) ∼ 1 εr/2 for large ε. The proof of Theorem 2.2 can be found in the supplementary material.
If M1,M2, . . . ,Mk are DP algorithms such that Mi has PRVs (Xi, Yi), then the composition of M1,M2, . . . ,Mk has PRVs given by ( ∑k i=1Xi, ∑k i=1 Yi) [GLW21]. If a mechanism M has PRVs
given by (X,Y ), then the subsampled mechanism M ◦ Samplep where Samplep samples p fraction of the data uniformly at random has PRVs given by (Xp, Yp) where:
Xp = log(1 + p(e X − 1)),
Yp = { log(1 + p(eY − 1)) w.p. p log(1 + p(eX − 1)) w.p. 1− p.
(2)
See [GLW21] for a proof of this fact. Combining these facts about PRVs with Theorem 2.2, we can compute the privacy curve of the DP-SGD-JL algorithm. Theorem 2.3. Let p = B/N where B is the batch size and N is the total number of samples. Then DP-SGD-JL run for T iterations satisfies (ε, δ(ε)-DP where
δ ≡ δ ( T∑ i=1 Xip ∥∥∥∥∥ T∑ i=1 Y ip ) where X1p , X 2 p , . . . , X T p are iid copies of Xp and Y 1 p , Y 2 p , . . . , Y T p are iid copies of Yp. Xp and Yp are given by (2) and X,Y are given by (1).
Finally, we numerically compute the distribution of ∑T i=1X i p by using Fast Fourier Transform (FFT)
to convolve the distributions of Xp with itself T times. The distribution of ∑T i=1 Y i p is similarly calculated, see [GLW21] on how to approximate these distributions.
Effect of JL dimension on privacy Figure 3 shows how the JL dimension r effects the privacy curve of DP-SGD-JL. As JL dimension increases, DP-SGD-JL becomes more and more private and converges to the privacy of DP-SGD.

heading: 3 Experiments
text: In this section, we demonstrate experimentally that compared to existing implementations of DP-SGD with exact per-sample gradient clipping, our optimizers have significant advantages in speed and memory cost while achieving comparable accuracy-vs-privacy tradeoff. Moreover our algorithms perform well on a variety of network architectures. The main goal of this section is to give empirical evidences towards the following three strengths of our algorithm alluded in the introduction:
1. Our algorithm is significantly faster compared to per-sample gradient computations and works for any network in a black-box way.
2. Memory footprint of our algorithm is roughly same as non-private SGD. 3. The DP-SGD-JL algorithms with smaller values of JL dimension exhibit similar behavior as
DP-SGD but with orders of magnitude speed up, and hence can be used for hyper-parameter search.
In the following, we write ‘nonDP-SGD’ for the standard non-private SGD and ‘DP-SGD-Vanilla’ for the implementation of DP-SGD in [TP], nonDP-Adam and DP-Adam-Vanilla are similarly defined. We use Tensorflow and [TP] for all our experiments because [Opa] does not support arbitrary network architectures.6 Moreover Tensorflow has an efficient implementation of jvp while PyTorch doesn’t.
Notation: We denote the noise multiplier as σ, clipping norm as C, batch size as B, learning rate as η and the number of epochs as E = BT/N . We fix the privacy parameter δ = 10−5, as done by prior work. We denote the JL dimension used by each optimizer in the parentheses. We use one Tesla P100 16GB GPU for all experiments. In all the experiments, we report the time per epoch by averaging over a large number of epochs. The code for our experiments is available in the supplementary material.

heading: 3.1 Training an LSTM model on IMDb dataset
text: The goal of these experiments is to demonstrate the first strength of our algorithm: it is significantly faster than per-sample gradient computations and works for any network in a black-box way. We
6In Pytorch Opacus github, the LSTM layer is only partially supported, e.g. single directional, single LSTM layer, no dropout layer; other recurrent layers such as GRU are not supported (see Opacus).
train a bidirectional LSTM with an embedding layer on the IMDb dataset for sentiment analysis. We remark that simpler networks not based on RNNs can achieve good accuracy-vs-privacy tradeoff as shown in [BDLS19] and [Pyt]. However, LSTM models based on RNN architectures are widely used in NLP applications, and hence efficient private training of such models remains an important challenge in this area [MRTZ18]. Moreover, as we noted in the introduction, extensions of outer product trick to bidirectional LSTMs as described in [LK21] are significantly more complicated, and require considerable effort to implement. Since authors of [LK21] did not provide the code, we are unable to compare the improvements. Moreover, as also noted in [LK21], the outer product method requires significantly more memory and hence will not scale to large batch sizes, which is very important to achieve good privacy vs utility tradeoff.
We implement DP-Adam-JL using jvp method in TensorFlow. We train the same single-layer bidirectional LSTM7 as in the [Ten] tutorial, using the same IMDb dataset with 8k vocabulary. The dataset has binary labels and is preprocessed by zero-padding the sequences to a length of 150 before being fed into the embedding layer.
Table 2 shows the training time per epoch for different algorithms. As expected, we observe that DPAdam-JL, referred to simply as DP-JL in the table, algorithms with smaller values of JL dimension are significantly faster than DP-Adam-Vanilla. We also notice a linear increase in running-time with incresing JL dimension. However, as we note in the Figure 3 privacy guarantees of DP-Adam-JL algorithm with smaller values of JL dimension are considerably worse than DP-Adam-Vanilla. On the other hand, DP-Adam-JL(30) is 30× faster than DP-Adam-Vanilla while achieving similar privacy guarantees as DP-Adam-Vanilla. When allowed to train for sufficient number of epochs, we observed that all the algorithms achieved same accuracy but with different privacy costs and running times. This three dimensional tradeoff between utility, privacy and speed is depicted in the Figure 1 (see introduction), where we plot the privacy values using a color plot. Remark 3.1. As we can see from Table 2, for a batch size of 256, the slowdown of DP-Adam-Vanilla is more than 256 compared to nonDP-Adam. This is counter intuitive as a naive implementation DP-Adam-Vanilla, by setting the batch size equal to 1 and then doing gradient accumulation across 256 batches should only be 256 times slower. As we show below, this is due to the memory issues of DP-Adam-Vanilla in the implementation in [TP]. Indeed the naive implementation of DP-AdamVanilla in tensorflow using gradient accumulation takes about 4000 seconds per epoch for the same batch size.

heading: 3.2 Memory footprint
text: Another key strength of JL based algorithms is their memory consumption, and the goal of this section is to show this aspect of our new algorithms via experiments. As a proxy for memory consumption,
7We cannot use CuDNNLSTM and instead use LSTM for the following reason. When using CuDNNLSTM (and on GPU), we observe a significant speedup compared to LSTM, but the accuracy is invalid and we further incur a LookupError when computing jvp.
we compare the largest batch size each algorithm can handle without running out of memory. It is known that to achieve good privacy-vs-accuracy tradeoffs for DP-SGD, we need to use large batch sizes [ACG+16]. One way to support large batch sizes is via gradient accumulation; however, this has the disadvantage that one loses parallelism, which in turns leads to slower run times. Hence memory footprint of algorithms also indirectly affects the training time.
We compare our JL algorithms with the implementation of DP-SGD-Vanilla in [TP]. We train a convolutional neural network from [TP] tutorial on MNIST dataset, which has 60,000 training samples.8 As we can see from Table 3 (left), DP-SGD-JL algorithm and nonDP-SGD can both run with the maximum possible batch size of 60,000 whereas DP-SGD-Vanilla can only handle a batch size of at most 500. In general, we believe that the memory footprint of DP-SGD-JL algorithm is very close to that of non-private SGD. To show this, we augment the CNN in [TP] tutorial with dense layers to blowup the model size to 17,146,938 parameters, and repeat the experiment. As we see from Table 3 (right), the largest batch size supported by DPSGD-JL(30) is only a factor 2 away from the largest batch size supported by non-DP SGD. On the other hand, we observe that DP-SGD-Vanilla only supports a batch size of 100.
The above experiments also give a possible explanation of why DP-SGD-Vanilla implementation in TFP has a slowdown that is larger than the batch size, as we observed in the LSTM experiments. Even for MNIST, we observe that DP-SGD-Vanilla running time gets better with batch size in the very beginning but as the batch size becomes larger the running time gets worse, and soon after it runs out of memory.

heading: 3.3 Using DP-SGD-JL for hyper-parameter search
text: As we saw in our experiments summarized in Table 2, our algorithms with small values of JL dimension are orders of magnitude faster than DP-Adam-Vanilla; DP-Adam-JL(1) is about 470x faster and DP-SGD-JL(5) is about 150x faster. However, unfortunately, the privacy bounds we can prove for these algorithms are considerably worse than DP-Adam-Vanilla. Despite this drawback, we observe that behavior of DP-SGD-JL even for small JL dimension is very close to that of DPAdam-Vanilla. Figure 4 plots the accuracy vs epochs for various algorithms training a CNN from [TP] tutorial on MNIST dataset. We observe that as JL dimension increases, the accuracy vs epoch
8We use the implementation and the network from mnist_dpsgd_tutorial_keras.py
curve quickly converges to that of DP-SGD-Vanilla. DP-SGD-JL(3) is already very similar to DPSGD-Vanilla and DP-SGD-JL(20) is nearly indistinguishable. This also lets us hypothesize that the privacy of our algorithms could be much better than what we could prove and that it should converge equally quickly to that of DP-SGD-Vanilla. Thus we believe that DP-SGD-JL(3) or DP-SGD-JL(5) are good candidates for experimentation and hyper-parameter tuning during private training, since their behavior is almost identical to that of DP-SGD-Vanilla while being orders of magnitude faster.
Acknowledgments and Disclosure of Funding
We thank Sergey Yekhanin for his constant support and encouragement during this work. We also thank Lukas Wutschitz and Osman Ramadan for helpful discussions. Finally, we thank the amazing open source community of TensorFlow for their quick response in fixing bugs which was crucial for our experiments ([TFi]).

Review summaries:
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. 

This paper considers the computation and memory cost problem introduced by per-sample clipping in private gradient descent algorithms. To accelerate per-sample clipping, this paper proposed a new framework using JL projections to estimate the per-sample gradient norms. The authors provide rigorous differential privacy analysis of the proposed framework. They also give empirical evaluations of the proposed algorithms in terms of privacy, utility, and computation time.