Enabling Fast Differentially Private SGD via Just-in-Time Compilation and Vectorization
Authors: Pranav Shankar Subramani
Nicholas Vadivelu
Gautam Kamath
Venue: NeurIPS 2021 Poster (2021)
Keywords: differential privacy
DPSGD
JAX
vectorization
just-in-time compilation

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.


Enabling Fast Differentially Private SGD via bust-in-Time Compilation and Vectorization
Pranav Subramani"
Nicholas Vadivelu"
Cheriton School of Computer Science University of Waterloo
Cheriton School of Computer Science
University Waterloo
sutiivvssbbaman'' @uwaterloo... ca
abvadive@@uwteerloo.c
Gautam Kamath
Cheriton School of Computer Science
University of Waterloo
gocsail. mim.mit.... edu
Abstract
A common pain point in differentially private machine learning the significant runtime overhead incurred when executing Differentially Private Stochastic Gra- dient Descent (DPSGD), which may be large 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 JAX, which provides rich support for these primitives through the XLA compiler. We also rebuild core parts of TensorFlow Privacy, integrating more effective vector- ization well as XLA compilation, granting significant memory and runtime improvements over previous release versions. Our proposed approaches allow us to achieve up 50x speedups compared to the best alternatives. Our code is available at lhttps://githuu: http://iithuu..cmmThhteeee wwww/o/ewwlooo/eeooooooooooooeee fast-aaat-taee dpsgd
Introduction
Machine learning has recently experienced tremedous growth, being used solve problems with unprecedented accuracy in a myriad of domains. However, not all domains are alike-while many datasets freely available, it has been demonstrated that disregarding these concerns, or even using heuristic and best-effort
often would like train MI models sensitive data. Troublingly,
privacy approaches, can result in significant leakage of private information [9] [2]. Differential
privacy (DP) [66] 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 differentially private analogue, DPSGD 43 which introduced as drop-in replacement for SGD. The primary differences include per-example gradient clipping step and a batch-level relatively innocuous, they have so far led
addition of Gaussian noise. While these modifications 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 per-example level. Access to these objects is critical in DPSGD, as well other applications
Equal contribution.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

beyond privacy [52]. Lack of support for fast computation of per-example gradients has been noted and lamented numerous times for both TensorFlow 335 42 and PyTorch [32 6].
Numerous attempts to avoid these computational roadblocks have been proposed. Goodfellow [20] proposed an algorithmic solution for computing per-example fg-norms of the gradients for fully- connected networks. Other proposed solutions work by exploiting Jacobians [14] or parallelizing over the batch dimension 4 Several of these approaches are are restricted specific types of architectures-for example, 20 is restricted to fully-connected layers, though AI 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 [T3]. Facebook's 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 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 rumning time overhead versus non-private SGD, with slowdowns as large as two orders of magnitude. For instance, Carlini et private algorithm is known to be slower than standard training; our implementation of this algorithm 10-100x slower than standard training," " where their implementation based on TensorFlow Privacy. Additionally, Thomas et al. 45 document 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 a larger model. Overcoming this obstacle an important step in helping differentially private machine learning transition from its present nascent state to widespread adoption.
comment "Training differentially
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-tmme (JIT) compilation, and static graph optimization. These features are core primitives within JAX
19 10 and Tensor- Flow 2 (TF2) [Л] both lensor-processing libraries from Google. These frameworks combine JIT compilation backed by the Accelcrated Linear Algebra (XLA) (22) just-in-time compiler (JIT) with auto differentiation for high-performant machine learning. As 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 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 XLAA compilation, significantly improving its running time in most cases (referred to as Custom TFP in this paper). We also contribute 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- efficien:: 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: Median running time (s) per epoch of training various models at batch size 128. FCNN stands for Gally-Connected Neural Network; CNN stands for Convolutinal Neural Network.
Private Training
Non-Private Training
Architecture JAX/Custom TFP
Best Alternative
Best Time
FCNN
0.21
0.77
0.55
CNN
7.3
12
1.7
LSTM
8.2
407
4.8
Table summarizes some of experimental results, with median running time per epoch for 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 Examining the overhead due to privacy: JAX's 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 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 Shanson-LLindestraant 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 time-accuracy tradeoff (though not as severe this case).
Subsequent to the initial posting this paper, we worked with Google engineers to implement our improvements into TensorFlow Privacy. Vadivelu contributed JAX implementation of DPSGD to the Optax library (27). 7 employed our findings to efficiently privately train BERT-Large. In recent Opacus whitepaper [5] [1,, the authors repeat some of our experiments more recent versions of these frameworks; we defer to their work for discussion of these results.
Description of Approaches
2.1 Libraries Enabling DPSGD
JAX a recently introduced framework for machine learning, defined by its auto- matic differentiation capabilities and JIT compilation via the XL.A compiler (22). Programs written in pure Python and JAX's NumPy (26) API can be translated to an intermediate language (XLA-HI.O) to be JIT compiled, i.c., to generate custom assembly instructions for the hardware. This enables optimizations such as fernel-fusions buffer reuse, improved memory layout, and more. Additionally, one of the core functions present in JAX is VMAP, vectorized map, which enables easy-to-we write and efficient batch level parallelism that fundamental to DPSGD. As we will demonstrate, these enable the fastest approach for DPSGD that are aware of.
JAX 19 10
Note that confirmed bug in TF2 currently prevents us from running Custom TFP in these cases.
3

---

Custom TEP. Vectorization and XLA-driven JIT compilation is also available in TensorFlow 2 M. which 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 vectorized map and follow TensorFlow best practices while retaining the existing functionality.
Chain-Ruus Baiee-Rulee-Baaaae 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 fully-
connected layers via (20), well as convolutional layers via [4]] 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 (B,CC C, C) The conventional formula for the discrete convolution be written
C-1K-1 ΣΣ Ex CIOO R=U
The gradient of this expression can be efficiently computed via automatic differentiation [40]. Py- Torch's ss automatic differentiation cannot be parallelized across the batch dimension & [41] which is required to backpropagate through the above expression. Instead, they rewrite the convolution as follows:
C/G-1K-1
C + K Σ Dxbocottt
g.ckk,
where G is the number of groups and the shape x is C,T). The initial convolution is 1 1-dimensional, while the above expression includes an added dimension. Similarly, allow back- propagation through a j-dimennional convolutional layer, a (k + 1) 1)-dimensional convolutional layer 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:
L-1
Vcco ((0) (Jacn (1)
(V 2p (z))
In order to compute this quantity, one requires the ability to multiply the Jacobian by vector and by a matrix, which not currently supported in PyTorch's automatic differentiation framework. In BackPACK, Dangel et al. [14] extend several layers within PyTorch to support fast Jacobian-vecton and Jacobian-matri products order to extract quantities like individual gradients, variance, l2- 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 (Jocs 2 ())) _ ((O)),
o00 (())
where = 1,..., N and each ((1) has a gradient which is shape (N, d()). BackPACK provides efficient computation for the transpose of the Jacobian well the Jacobian.
Opacus [51]. Opacus is library for training PyTorch models with differential privacy, recently released by Facebook. It supports per-example gradients, using PyTorch's 's forward and backward hooks to propagate gradients. They provide support for several Py PoTcch layers including LSTM layers, which are not supported in either of the previous two frameworks. Note that Opacus does not support PyTorch's 'ss nn. LnTM but instead implements a separate opacus 1ayers.DPLSTM, adjustments that allow individual gradients to propagate through it.
DPLSTM, with
Vacy [88. Before the release of Opacus (and predecessor PyTorch-DP), Vyyacy was the most popular library for DP machine learning in PyTorch. Vaca has no custom support for parallelization across the batch dimension for any layer since it processes each sample individually (by way of for-loop). This generally leads to large increase in runtime for models trained using Py Vacy.

---

TensorFlow Privacy [37 TensorFlow Privacy is a library for differentially private machine learn- ing, built on top of TensorFlow. TensorFlow Privacy has general support for vectorized implemen- tation of DPSGD via vectorized_ maa which allows it to parallelize across the batch dimension, used extract per-example gradients. The library recently introduced TensorFlow 2 compatible API that leverages GradientTape jacobian to compute per-example gradients, which we compare seperately to the TensorFlow API in our experiments.
2.2 Notable Framework Features
TensorFlow and JAX use a static graph to track computation in order optimize execution and compute gradients. This means the sequence of operations is traced and a large proportion of shapes are determined during the first invocation of the function, allowing for kernel fusion, buffer reuse, and other optimizations on subsequent calls. PyTorch uses a dynamic graph to track computation flow in order to compute gradients, but does not optimize execution. This enables increased dynamism in the shapes and types of computations, the cost of losing all the aforementioned optimizations.
Static versus Dynamic Graph.
Grappler versus XLA. TensorFlow has two optimization engines: Grappler [2] and XLA. Grap- pler, TensorFlow's ' original graph optimizer, takes as input the computation graph and is able to prune dead nodes, remove redundant computation, improve memory layouts, and more. XLA, TensorFlow's new optimizing compiler, can perform the same optimizations as Grappler, in addition to generating code for fused kernels. For this reason, XLA has the potential to extract more performance out of TensorFlow graphs than Grappler, but does not always accomplish this due to Grappler's 's maturity.
JAX and XLA. JAX was built from the ground up leverage XLA, and so many of its operations map directly to XLA primitives. We often observe that JAX is able to extract better performance out of XLA than TensorFlow.
Pytorch and Static Graphs. Recently, PyTorch has released the capability to JIT compile its code through torch. jit or PyTorch XLA (18]. Due to the early nature of these two efforts, they were not successful JIT compiling the methods we tried, thus we do not consider them further.
Just-In- Lnnt-ln-Tiee Compilation JIT compilation is a method of compilation that happens at runtime, as opposed before program execution. JAX and TensorFlow perform JIT compilation by recording operations that are executed on tensors/arrays (i.e. "tracing"), generating the low-level instructions to perform these operations, optimizing these instructions, then producing fast low-level kernels. The XLA compiler requires all array shapes to be known at trace-time so it can statically determine how much memory should be allocated for each operation. While compilation can be relatively slow compared to execution (~10x the time), you only pay this price once at the first training iteration, provided the input shapes do not change.
3 Empirical Findings
We evaluate the aforementioned implementations of DPSGD in runtime and memory consumption on three datasets: CIFARIO (II, a dataset of small colour images with 60,000 training examples of
size 32 x3 each, IMDb (36, movie review sentiment classification dataset with 25,000 training examples padded to a sequence length of 256 each, and Adult (15), containing 45,220 examples with 104 features, which was preprocessed via methods from 229.. These datasets are available for open use and do not contain personally identifiable information or offensive content.
We perform our evaluations on three different architectures. We start with the smallest dataset, Adult, training 5.532-paramedt fully-connected neural network (FCNN). Then, we train a CIFARIO convolutional neural network classifier architecture with 605,228 parameters used by Papernot et al. 388. For IMDb, we use an LSTM network with 1,081,002 parameters, demonstrating the method on relatively large model. This selection covers the common data and architecture types at realistic sizes for differentially private learning. In particular, we did not consider the exceptionally large models which are not prevalent in non-private machine learning. Additional experiments can be found in the supplement that are omitted due to space constraints. These experiments cover wider range of parameters and architectures to elucidate the benefits of using XLA JIT for DPSGD.
5

---

Median Runtime Private Connected Network (FCNN)
20-
10006
10.
JAX
Oustom TFP PILA) Custom TFP
TFP (XLA)
TFP
TFP GILA)
TFP
Opocus
BackPACK
Che
Pylacy
Ulbrary
Modian Runtima for One Non-Private Epoch Fully Connected Neural Network (FCVN)
13
10-
Tensortoon (RLA)
Tenserfoww
Terrofiow (XLA)
Twnsorf ow
PyTecch
Ubrrar
Figure 1: Runtimes for the fully connected network on the Adult dataset. We observe that JAX and Custom TFP the fastest by large margin in both settings, with DPSGD having low overhead over the non-private setting. The y-axis is truncated for clarity.
Median Runtime Epoch CIFARIO Convolutional Network
150-
Modian hurtma
200
Jax
Custom TFP (LA) Custom
TFP2
(XLA)
BECKPACK
Median Runtime Non-Private CIFAR10 Convolutional Netwark
Murdan Rurtime
350
202-
JAX
TensorFiww (XLA)
Temsoflow 2
Tensrflow 1 (R0A)
Tensorllow
PyTorch
Library
Figure 2: Runtimes for the CNN on the CIFAR10 dataset. We observe that JAX has the fastest runtime private case while TensorFlow 2 with XLA and JAX are the fastest in the non-private case at most batch sizes. Similar to the MNIST case, TFP struggles at the largest batch size due to an inability properly parallelize per-gradient computation with this level of memory consumption. The y-axis is truncated for clarity.
We compare number of different methods for fast DPSGD, including JAX 199 10, BackPACK (4). the chain-rule based method (CRB) 20 41), Opacus (31) Py Vacy [48], TensorFlow Privacy (TFP) (37) and our modification of TFP, dubbed Custom TFP. For TensorFlow based frameworks, we evaluate performance both with and without XLA JIT compilation both TF and TF 2. We refer TensorFlow and JAX modern XLA-compiled libraries.
These architectures and datasets are evaluated in terms runtime at batch sizes 16, 32, 64, 128, and 256. This showcases comparison of runtimes across variety of batch sizes to present a holistic picture the rumning times, as well as demonstrating the impact memory utilization on runtime. Each experiment was for 20 epochs and the median epoch run-time is reported. The variance of these experiments was low enough that the errors bars are negligible (full results provided in the supplement). Outside of the initial compilation time required for the static graph frameworks, the runtime showed little variance between epochs, resulting in narrow confidence intervals for the epoch runtime (data available upon request). We preprocess all the data in advance and use identical generator- based dataloader for all frameworks, to ensure consistency.
All experiments were run on Ubuntu 18.04 with Intel Core i7-7800X CPU (3.50 GHz, 6 cores), NVIDIA GTX Titan V GPU (12GB VRAM), and 32GB of RAM. The code is provided in the supplement and is publicly released.

---

Median Runtime
Epoch
313
200-
132
100-
JAx
Custom TFP DLA) Custom
TFP
(XLA)
Opacus
BackPACK
Median Runtime Non-Private LSTM Network
Rundim txed
250
283
JaX
TesADrFion: (XUA))
Temooflow? 2
Tenorfiow 1 (R0A)
Tensvflow
PyTorch
Library
Figure 3: Runtimes for the LSTM network on the IMDb dataset. JAX is by far the fastest option, resulting in roughly 50x speedup for batch size 256. The quadratic memory cost of Opacus prevents us from evaluating this implementation at these hatch sizes, however, we observe a median runtime of 1024. 16s at batch size 10. Excessive memory consumption prevent from evaluating Custom TFP and TFP at larger batch sizes. An open TensorFlow bug prevents us from evaluating Custom TFP (XLA) in this setting [46). BackPACK and CRB do not support embedding layers. The y-axis is truncated for clarity.
First, we evaluate the FCNN model (Figure D where we observe that JAX and Custom TFP are significantly faster than the other options in both the private and non-private setting. With such simple architecture, the compiler can perform significant optimizations. Notably, JAX and Custom TFP show little overhead over their non-private counterparts. The non-statically compiled frameworks (apart from Py Vacy) remain competitive this setting due to the shallow network size and low parameter count. TFP without XLA performs unexpectedly poorly, due to the Jacobian computation which does not account for per-example independence of gradients.
We then evaluate the CIFARIO CNN model (Figure 2). where JAX is by far the most performant implementation the private case, and the modern XLA-compiled frameworks are the fastest non-private methods. Custom TFP is noticeably slower than JAX- we conjecture that this is due to different utilization of the JIT compiler, see Section Also, at the largest batch size, TFP's performance deteriorates due to the memory consumption preventing it from effectively parallelizing per-example gradient computation.
For the final runtime experiment we evaluate the LSTM network (Figure 3p. A TensorFlow 2 bug 46 and lack of support from BackPACK and CRB prevent from evaluating this setting those implementations. We see similar phenomenon to the CIFARI0 CNN case for TFP: at larger batch sizes, the library fails parallelize effectively pessimizing the runtime. TensorFlow and Py PyTorch benefit from fast cuDNN LSTM implementation in the non- con-private case which they fail to leverage the private case, explaining the significant difference performance. JAX, on the other hand, uses an LSTM implementation based on primitive operations, which allows it to retain similar performance in both the private and non- private settings.
To understand the importance of vectorization (via VMAP) and JIT compilation (JIT), we ablate JAX's performance these tasks with and without these components (Table 2}. We observe that JIT alone provides up to a 435x improvement, and VMAP alone provides up to a 64x improvement, When used in tandem, both complement each other, providing up to a 5160x improvement.
Similarly, we ablate the components Custom TFP (Table 2}. TensorFlow, vectorized, map automatically compiles the code, preventing from ablating it alone. We observe that XLA in TensorFlow able to reclaim performance lost by not vectorizing, seeing that the non VMAP XLA performance comes close to VMAP Graph performance in many settings. We further the non-compiled runtimes are not as extreme as seen in JAX: TensorFlow is better optimized to run reasonably fast in all settings.
Finally, we explore the memory consumption behaviour these implementations (Table observing that running time has strong negative correlation to memory consumption. The modern XLA-

---

Table 2:
Ablation of JIT compilation and vectorization in Custom TFP (left) and JAX (right). Median runtime per epoch for run of 20 epochs batch size 128 for DPSGD. In TensorFlow, one can not use vectorization without graph compilation, which is why there is no standalone vectorization. Empty entries are due confirmed active bugs TensorFlow [47 466. We exclude LSTMs, as JAX runs out of memory without JIT compilation for this model.
Graph
JIT
VMAP
FCNN
CNN
LSTM
JIT
VMAP
FCNN
CNN
147
561
717
1240
3840
12.7
85.5
361
19.2
64.8
/
1.83
71.5
2.85
84.0
/
0.770
33.8
68.6
/
/
0.239
7.28
/
0.209
23.3
Table 3: Maximum batch size supported by each library before encountering out of memory errors. Missing entries represent missing functionality or bugs in the frameworks. PyVacy handles examples sequentially, giving constant memory consumption with respect batch size.
Library
CNN
LSTM
Library
CNN
LSTM
JAX
10,448
11,984
JAX (DP)
4,264
2,487
TensorFlow 2 (XLA)
15,040
Custom TFP (XLA)
3,144
TFP 2 (XLA)
168
TensorFlow 2
11,328
9,221
Custom TFP
1,944
137
TFP 2
104
105
TensorFlow 1 (XLA)
10,880
5,070
TFP (XLA)
168
88
TensorFlow PyTorch
11,480
5,264
TFP
104
105
10,752
9,943
Opacus
1,920
10
BackPACK
1,216
CRB
2,184
Vacy
00
00
compiled libraries provide impressive batch size capability. Also, the frameworks except JAX and Pyyacy struggle with batch sizes on the LSTM. Since Py Vacy processes examples sequentially, it has a constant memory consumption with respect batch size, effectively trading off running time for optimal memory use. The other frameworks are crippled without access to their fused cuDNN LSTM implementation, while JAX has no issues as its L.TT composed of primitives. Finally, due the specialized per-example gradient computation afforded by CRB for convolutions, it shows the best memory utilization among the PyTorch frameworks, even beating Custom TFP (without XLA) the CIFAR10 CNN case.
4
Discussion
JAX and Custom TFP's runtime advancements can be primarily attributed to the advancement of the compiler present both of these languages. The XLA compiler performs variety operations ranging from memory scheduling to kernel fusion. The memory optimizations are vital for larger models where DPSGD becomes memory-bound algorithm. One of the core features XLA buffer reutilization which has significant impact on the maximum memory used [22]. Furthermore, the memory scheduler can mitigate peak memory usage to prevent a runtime exception for overusing available memory.
The effectiveness of XLA demonstrated through the peak batch size experiment (which serves as proxy for memory efficiency): in both the private and non-private settings, XLA far exceeds alternatives in the peak batch size it supports. Through this experiment, we also see the benefits of using small operation primitives as opposed to large fused kernels: TensorFlow and Py PyTorch both leverage the optimized cuDNN kernel the non-private setting for performance (23 17,, but cannot in the private setting, leading to significantly worse performance. Modifying all existing cuDNN kernels to enable use in the private case would require non-trivial engineering investment. JAX instead focuses optimizing operation primitives,
in foreign computational circumstances,

---

its performance is comparatively strong, Succinetly, JAX sacrifices the ability to use highly optimized fused kernels for generalizability.
Our implementation of Custom TFP leverages the vectorized map for both the forward and backward pass, unlike existing implementations which only use vectorized map for the backward pass (identical to the JAX version). This gives the compiler explicit information about the independence of batches in the computation, enabling significant optimizations. explaining the improvement in Custom TFP compared to the existing TensorFlow implementations. Details about the implementation are provided in the supplementary code.
We notice that the runtimes and memory consumption are different between Custom TFP and JAX despite having the same backend compiler. We investigate differences in the XLA-assembly for simple code segments in both frameworks in the supplemental material. TensorFlow 1.0 does not have the same capability integrate with XLA as TF does. The primary optimization available is autoclustering, which we observe does not optimize operations nearly as much as the full JIT.
Through the ablation study of VMAP and JIT in JAX shown in Table we observe that both components complement each other. In general, we observe that JIT provides the larger performance gain, shown with the FCNN. For the CNN, the large matrix operations coupled with JAX's s asynchronous execution [4] allow reasonable utilization of the GPU even without VMAP, which is why we observe less of an improvement from JIT alone in this experiment. The runtimes without these two primitives is significantly slower than the other frameworks this is because JAX was built ground-up to leverage these primitives [0.
In the ablation for Custom TFP in Table 2 we see some key differences with JAX. First, the mechanism for VMAP in TensorFlow 2 is different from that in JAX: JAX performs op-by-op batching without compilation, while TensorFlow always compiles the vectorized map [24]. We observe that the non-compiled code still runs in reasonable amount of time since TensorFlow is optimized to have a competitive eager execution, while JAX is not. Also, in TensorFlow, XLA's JIT compilation without vectorization is often able to bring the runtime performance close to that of the graph mode with vectorization, implying that XLA is able to recognize opportunities for parallelization even when the user does not explicitly request it. Finally, we observe the benefit of having fast, fused implementation for LSTMs: while JAX ran out of memory outside of a compiled and vectorized context, Custom TFP is able to achieve reasonable runtimes by leveraging the fused kernel.
4.1 Drawbacks
While integrating XLA into DPSGD presents a massive runtime and memory advantage, there is a cost. XLA is subset of all permissable operations in JAX and TensorFlow, requiring users to be cognizant of functionality they use. For example, j jax. numpy. unique (x) produces a result whose shape not known compile time, preventing it from being used with JIT compilation. More subtle errors involving unintended recompilation of code are also possible which can lead to enormous slowdowns.
Another potential downside is that these benefits are not observed the batch size is large enough and the bottleneck becomes the actual network evaluation. In situation like this, all of the frameworks would perform at approximately the same speeds. However, in practice, this is rarely true and moreso in differentially private machine learning where models not particularly large.
5 Conclusion and Future Work
We have demonstrated that language primitives like vectorization, JIT compilation, and static graph optimization can dramatically improve the running time of DPSGD, realized by JAX and our Custom TFP. In particular, we find that using JAX can almost entirely remove the computational overhead introduced by DPSGD, thus alleviating a major pain point of private machine learning practitioners.
In our work, we focus on conventional set-ups for academic researchers; for future work, it would be insightful to explore the performance of distributed DPSGD, as distributed set-ups are becoming increasingly commonplace. Furthermore, implementing a PyToorch JIT compatible version of DPSGD could provide an alternative to TensorFlow and JAX, particularly if said implementation compatible with Py Torch XLA. Though these two compilation systems are immature compared to TensorFlow

---

and JAX, they are rapidly improving and should not be ignored. Outside of Python, there are powerful autooiiferentiation methods in other more perfomant languages such as Julia [3] and Swift [50] which are worthy of study.
Acknowledgements
We would like to thank Roy Frostig for helpful discussions on JAX, Steve Chien and Shuang Song for their work in implementing our improvements TensorFlow Privacy, Xi He and Om Thakkar for valuable feedback on drafts of this work, and several JAX, TensorFlow, and Opacus developers who helped answer our issues, including James Bradbury, Peter Buchlovsky, Peter Hawkins, Matthew Johnson, Karthik Prasad, Github handle ravikyram, and Qianli Scott Zhu.
Funding Transparency Statement
Funding in direct support of this work: an NSERC Discovery Grant, Compute Canada RRG grant, and University of Waterloo startup grant.
References
[1] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp. Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems. fittpa://wwt //wwww/www... tensorflow orgg 2015. URL /ttpp:: http:////wwwwwwtttllll.ooom www.srffooooooooooo org/ Software available from tensorflow.org.
[2] Martin Abadi, Andy Chu, Ian Goodfellow, H Brendan McMahan, Ilya Mironov, Kunal Talwar, and Li Zhang. Deep learning with differential privacy. In Proceedings of the 2016 ACM Conference Computer and Communications Security, CCS 16, pages 308-318, New York,
NY, USA, 2016. ACM.
[3] act65. Custom gradient aggregation methods. tensorfloo//isues//Sti December 2017. URL
http://ggihub..cm/eeesrtrt
kensorffoox/isaues//57//
[4] Ashish Agarwal and Igor Ganichev. Auto-vectorizing tensorflow graphs: Jacobians, auto- batching and beyond. arXiv preprint arXiv: 2019.
[5] Ashish Agarwal and Igor Ganichev. Auto-vectorizing tensorflow graphs: Jacobians, auto- batching and beyond, 2019.
[6] alexdepremia. [feature request] expanding gradient function with variances. httpss//github com comm/pytorcch/pyoocch/iiie hwwp:orrrrrrriisss//////// June 2018. URL https:///ithub. http://ittuu..coooyttoo/t com/pytorch/ sytorch/issues/B88
[7] Rohan Anil, Badih Ghazi, Vineet Gupta, Ravi Kumar, and Pasin Manurangsi. Large-scale differentially private BERT. arXiv preprint arXiv:2108.0162. 2021.
[8] Raef Bassily, Adam Smith, and Abhradeep Thakurta. Private empirical risk minimization: Efficient algorithms and tight error bounds. In Proceedings of the 55th Annual IEEE Symposium on Foundations of Computer Science, FOCS '1, pages 463 Washington, DC, USA, 2014, IEEE Computer Society.
[9] Abhishek Bhowmick, John Duchi, Julien Freudiger, Gaurav Kapoor, and Ryan Rogers. Pro- tection against reconstruction and its applications in private federated learning. arXiv preprint arXiv: 1rXiv:1812.0098. 2018.
10

---

[10] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, and Skye Wanderman- Milne. JAX: composable transformations of Python+NumPy programs, 2018. WtRLhttp://// thub. cww coom/oogle//
[1]] Zhiqi Bu, Sivakanth Gopi, Janardhan Kulkarni, Yin Tat Lee, Judy Hanwen Shen, and Uthaipon Tantipongpipat. Fast and memory efficient differentially private-sgd via JL projections. In Advances in Neural Information Processing Systems 34, NeurIPS '21. Curran Associates, Inc.,
2021.
[12] Nicholas Carlini, Chang Liu, Ülfar Erlingsson, Jemnjj Kos, and Dawn Song. The secret sharer: Evaluating and testing unintended memorization neural networks. 28th USENIX Security Symposium, USENIX Security 11, pages 267-284. USENIX Association, 2019.
[13] Felix Dangel. Support for recurrent units. httsss//gttoooo.o///ddaggg///// issues/ 16 2019. URL httpt:://tgtiiiiiie htt.tchwb.ooommmmm/ haap:///aaaaagaaaaaaaaaaaa//a//a
[14] Felix Dangel, Frederik Kunstner, and Philipp Hennig. BackPACK: Packing more into backprop. In Proceedings of the 8th International Conference on Learning Representations, ICLR '20, 2020.
[15] Dheeru Dua and Casey Graff, UCI machine learning repository, http:://archii iic.eeei...... uci edu/ml 2017.
ics wuu.uciiuuuu. edduml License: Open Data Commons.
[16] Cynthia Dwork, Frank McSherry, Kobbi Nissim, and Adam Smith. Calibrating noise sensi- tivity in private data analysis. In Proceedings of the 3rd Conference on Theory of Cryptography, TCC '06, pages 265-284, Berlin, Heidelberg, 2006. Springer.
[17] Facebook. Lstm. ht http://sww..........////////////
rnn NLST
URL
modules/run rnn. #LSTM
[18] FaceBook. Pytorch xla, 2020.
[19] Roy Frostig, Matthew James Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing. In The 1st Conference on Systems and Machine Learning, SysML '18 2018.
[20] Ian Goodfellow. Efficient per-example gradient computations. arXiv preprint arX1v:1510.01799 2015.
[21] Google. Tensorflow graph optimization with grappler.
org/ URL https: //op::://www... weaasorfloo..orgguuieeeeeeeeeeee
guide/ guide/graph.oo gwraph_optimizaaiie.
optimization License: Apache 2.0 License.
[22] Google. XLA: Optimizing compiler for machine learning. htttps://w/i http///wweessorfoo.. org/
xla
URL htt.ttss//ww..esssoollloooo
License: Apache 2,0 License.
[23] Google.
keras/ Barras/aayrs/LSTX http:///////////oooo.oo///ooo/on
lf.keras.lyerrs.stmm
keras/layers/LsTT http:://///.....ooooooooo///////
[24]
Google. ff.vectorieed_magt map. vectorizze__ag
http://www.tnnsrrf..og
vectorized_map URL
[25] Google. differentially py.
http:://///gggiiti.ooooooo/////a oxa.gl//aaffeearrlllyyrrraaegggr 2019. htttsp////
com/
License:
Apache License 2.0.
[26] Charles R Harris, K Jarrod Millman, Stéfan I van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J Smith, et al. Array programming with numpy. Nature, 585(7825)3357-333 2020.

---

[27] Matteo Hessel, David Budden, Fabio Viola, Mihaela Rosca, Eren Sezener. and Tom Hennigan. Optax: composable gradient transformation and optimisation, JAX!, 2020. URL http 7/g thub. com coom/deeemind///ppeeeeeeeeeeeeee optax
[28] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1780, 1997.
[29] Roger Iyengar, Joseph P Near, Dawn Song, Om Thakkar, Abhradeep Thakurta, and Lun Wang. Towards practical differentially private convex optimization. In Proceedings of the 40th IEEE Symposium on Security and Privacy, SP 19, pages 299-316, Washington, DC, USA, 2019, IEEE Computer Society.
[30] Julia. Juliadiff. https://wwwn http///wwwwiliiifff.f...... .rgg URL
org/
[31] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
[32] Frederik Kunstner. [feature requestj simple and efficient way to get gradients of each element of sum. http:////.gggta....ooooooooooo/a BBgthh/..cc//poooo//////////////
May 2018. URL https
[33] Jaewoo Lee and Daniel Kifer. Scaling up differentially private deep learning with fast per-
example gradient clipping. arXiv preprint arXiv:2009.03 2020.
[34] Anselm Levskaya and Matthew James Johnson. Asynchronous dispatch. https://
jax. rex.thaethedsss.c http///aoo..attooaaaaaaatttaaaea
dispatchhdsphhhcccccacccc html 2019. URL https://jat
readthedocs, Saaattoooss./////aaatttte///ssep
[35] Zachary C. Lipton. Gradients of non-scalars (higher rank jacobians). https://github
January 2016. URL https:///ithubb com/
[36] Andrew I.. Maas, Raymond F. Daly, Peter T. Pham, Dan Huang. Andrew Y. Ng. and Christopher Potts.
Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pages 142-150, Portland, Oregon, USA, June 2011 Association for Computational Linguistics. URL http::/www aclweb. org http/aathhooogy/11-111155 License: No explicit license, but explicitly is made available for research.
[37] Nicolas Papernot, Andrew Galen, and Steven Chien. Tensorflow privacy, https://github com//eessorflo//piivvee December 2018. URL https://github.cr wtw.:ttub.cmmm/eeesnffiiif.....e License: Apache 2.0 License.
privacy
[38] Nicolas Papernot, Abhradeep Thakurta, Shuang Song, Steve Chien, and Úlfar Erlingsson. Tempered sigmoid activations for deep learning with differential privacy. arXiv preprint arXiv:2007. /4/9/, 2020.
[39] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vitito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In Advances in Neural Information Processing Systems 32, pages 8024-8035. 8035. Curran Associates, Inc., 2019. URLhttp::/papert
wwwpuriis.cccpppppape 9015 pytorch- an- ann-imprrative-sty-l. high www.oh-ppoooomoneeeeeeeeee deep wttpap/wwe.etiea.aoooooooooroooa wwwp:aaaeneeeeeeeeeeeeeee-----ee
library
License: BSD. pdf
[40] Louis B Rall. Automatic Differentiation: Techniques and Applications, volume 120 of Lecture Notes in Computer Science. Springer, 1981.
[41] Gaspar Rochette, Andre Manoel, and Erie W Tramel. Efficient per-example gradient computa- tions in convolutional neural networks. arXiv preprint 1rX0011112.00015, 2019.
12

---

[42] seerdecker. Provide unaggregated gradients tensors.
httee://gthubb..o///osorro1/d Censorffloo/issuee/48377 4897 October 2016. URL https///gttuub..omm/eeossttttt
kensort\lox/issues/4000
[43] Shuang Song, Kamalika Chaudhuri. and Anand D Sarwate. Stochastic gradient descent with differentially private updates. In Proceedings of the 2013 IEEE Global Conference on Signal and Information Processing, GlobalSIP '13, pages 245-248, Washington, DC, USA, 2013. IEEE Computer Society.
[44] Kunal Talwar. Personal communication, July 2020.
[45] Aleena Thomas, David Adelani, Ali Davody, Aditya Mogadala, and Dietrich Klakow. Inves- tigating the impact of pre-trained word embeddings memorization in neural networks. In Proceedings of the 23rd International Conference on Text, Speech and Dialogue, TSD '20, 2020.
[46] Nicholas Vadivelu. Xla compilation does not work with embeddings layer. https //github. httpshhwwwoo/...eooooo////oooo/a
October 2020. URL https
[47] Nicholas Vadivelu. Xla compilation bug. http:///s.gggta.....oo//o//ooo/p
issues/43723 October 2020. URL http:://gttabbccao/aaaooroooooo0
issues/43723
[48]
Chris Waites. URL
March 2019. License: Apache 2.0 License.
[49] Chris Waites and Rachel Cummings. Differentially private normalizing flows for privacy- preserving density estimation. ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models, 2020.
[50] Richard Wei, Dan Zheng, Marc Rasi, and Bart Chrzaszez. Differentiable pro- gramming manifesto. URL https://github. httpshhu...oomaeaceooemmemoeeeee http///p//sstitt//oo////a/////ea Differentiablerrgre awg.ing.mnnn....mmm. mdl
[51] Ashkan Yousefpour, Igor Shilov, Alexandre Sablayrolles, Davide Testuggine, Karthik Prasad, Mani Malek, John Nguyen, Sayan Gosh, Akash Bharadwaj. Jessica Zhao, Graham Cormode, and Ilya Mironov. Opacus: User-friendly differential privacy library in PyTorch arXiv preprint a0Xiv:2109. 12298, 2021.
[52] Peilin Zhao and Tong Zhang. Stochastic optimization with importance sampling for regularized loss minimization. In Proceedings of the 32nd International Conference on Machine Learning, ICML '15, pages 1-9. JMLR, Inc., 2015.
13


Reviews:
[{"review_id": "tYPaDmNtGIq", "reviewer": "Reviewer_BcVc", "summary": "The paper is presenting a survey of the different privacy preserving libraries available to researchers to date. It compares their runtime and memory usage on FCNN, CNN and LSTM.\n\nThe main contribution of this paper is providing a current state of the art of the different privacy preserving libraries.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The main high level concerns about the paper are the following:\n- What is the difference between Section 3 and 4? In particular, Section 3 provide both the experimental setup as well as interpretation of the results and Section 4 seems to be rephrasing these interpretations. I think that separating Section 3 to be only the experimental setup and Section 4 being only result interpretation would make the overall paper clearer and less repetitive.\n- Some of the comparisons, especially in Table 1 are a bit misleading.\n  - The LSTM results discussed in particular seem to be a special case because most methods do not actually work in that context\n  - While the description discusses the overhead due to moving from non-private -> private, the table put next to each other the private numbers on one hand and the non-private ones on the other hand. Making it hard to relate to the description.\n  - If the goal is to show the absolute runtime while doing privacy preserving training, then the table layout is fine but the non-private section is not useful and the description should focus on that.\n- The equations for CRB and BackPACK are not very useful. While the dimension size are introduced, none of the other symbols are which makes it hard to make sense of what these equations are representation and why they are important. Maybe an illustration for CRB would make the point come across more effectively, without the need for introducing all the terms. \n- The evolution of the runtimes as a function of the batch size is shown in all the graphs but never discussed in details.\n- l312-313, the claim than \"using JAX can almost entirely remove the overhead of DPSGD\" is a bit strong, especially because for LSTM, it is still significantly slower based on Table 1: 3.6 -> 8.2. Other parts of the paper like l87-88 does mention with that statement that this is only true for FCNN and CNN.\n\nSome smaller things:\n- The CNN line in Table 1 doesn't match the runtimes from Figure 2 (while the two other lines match the ones in Fig 1 and Fig 3). Could you clarify that?\n- In Table 1, it is a bit surprising to put jax/TFP in the same category. My understanding of this table is to recommend a given framework that will give consistently better result. But recommending one of the two sounds much less useful as the user would have to write all his code twice.\n- The sentence l88-89 that it is faster than other frameworks without privacy seems like it is comparing two things that are not comparable? why is it relevant to compare these two numbers? In particular, how do you take into account the big difference between framework in the non-private case when making that comparison?\n- Some of the sentence feel a bit superlative. For example the \"substancial superiority\" l92. Sentences like the one l110 \"the fastest approach [...] that we are aware of\" seems more appropriate.\n- You mention the jit compilation time l199-200. Could you give an order of magnitude for this? I have memories of some complicated architecture taking hours to compile in some older frameworks. Is that the order of magnitude we should expect here?\n- Could you clarify the claim l256-257 \"XLA far exceeds alternatives in the peak batch size\" From the numbers in the table 3, it seems that Tensorflow 1 with XLA actually has a smaller peak batch size than Tensorflow 1 without XLA.\n\nSome typos, small fixes\n- l35: there is an extra \"are\"\n- Table 1 caption defines FCNN and CNN but not LSTM.\n- l274-275: the sentence is not english\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The limitations section is interesting and accurate as far as I can see.\nOne question I would have is about the claim that current privacy preserving models are small. Is that due to their increased cost and is that claim going to still be true given the improvements presented here that show that the extra cost is actually minimal.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "qHFhs04RoRo", "reviewer": "Reviewer_hnfP", "summary": "The paper explores how to accelerate the Private SGD algorithm via JIT compilation and vectorization. Specifically, the author benchmarks several frameworks and libraries for DPSGD, extends TFP with vectorization and XLA compilation, and compare the runtime speed with existing libraries.", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "*Originality: Are the tasks or methods new? Is the work a novel combination of well-known techniques? (This can be valuable!) Is it clear how this work differs from previous contributions? Is related work adequately cited?*\n\nI would say the methods is not new. Just-in-time compilation is a well-known techniques which is adopted by deep learning frameworks in these years, like XLA, TVM, Glow, etc. Static graph optimizations and HPC techniques like vectorizations is also well adopted in frameworks.\n\n*Quality: Is the submission technically sound? Are claims well supported (e.g., by theoretical analysis or experimental results)? Are the methods used appropriate? Is this a complete piece of work or work in progress? Are the authors careful and honest about evaluating both the strengths and weaknesses of their work?*\n\nThe submission is technically sound and benchmark result is convinced.\n\n*Clarity: Is the submission clearly written? Is it well organized? (If not, please make constructive suggestions for improving its clarity.) Does it adequately inform the reader? (Note that a superbly written paper provides enough information for an expert reader to reproduce its results.)*\n\n\nThe paper can be improved by stressing the novel points and technique/methods. Too much content is focus on previous works and benchmark.\n\n*Significance: Are the results important? Are others (researchers or practitioners) likely to use the ideas or build on them? Does the submission address a difficult task in a better way than previous work? Does it advance the state of the art in a demonstrable way? Does it provide unique data, unique conclusions about existing data, or a unique theoretical or experimental approach?*\n\nI appreciate the work, that it can be beneficial for the community to have a fast DPSGD library. But it is more like an engineering implementation with mature techniques and is lack of novelty for me.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The limitations and potential negative societal impact is well addressed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1.5"}, "extra_text": {}}, {"review_id": "LBJbFfuuGES", "reviewer": "Reviewer_27Gt", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper aims to address a very important problem: to alleviate the significant and sometimes prohibitive overhead of training ML models with DPSGD. Currently, DPSGD is the most widely adopted general method to train neural network models with privacy protections for user data in the training set. However, its use is limited partly due to the large overhead current implementation of DPSGD will incur. Making DPSGD faster will increase its adoption in academia and industry.\n\nThe authors benchmark several frameworks and libraries for DPSGD and demonstrate that vectorization, JIT compilation and static graph optimization leads to faster and more memory efficient DPSGD implementations. They extend TensorFlow Privacy to support TensorFlow 2 and contribute a variant of their implementation to TensorFlow Privacy.\n\n Strength: \n+ The problem this paper aim to address is an important one. \n+ The experiments are clearly designed and demonstrate that JIP and vectorization lead to significant improvement in per epoch training time.\n+ The paper is overall well written. The presentation is clear.\n\nWeakness:\n- The biggest weakness is that the contribution of this paper is unclear. While it is good to have experiments to show that vectorization, JIT and graph optimization can improve the performance of DPSGD, these findings are not surprising. JIT and vectorization are known to provide big improvements in training efficiency. They are supported in existing compilers like JAX and XLA, and in TensorFlow and PyTorch for training (non-private) models. Thus, it seems that the biggest contribution of this work is on the engineering side: the authors contribute a variant of their implementation to TensorFlow Privacy and extend TensorFlow Privacy to support TensorFlow 2. This is not saying that the contribution is not important. Such implementation, if made publicly available, is a valuable addition to the DP and ML community. However, the paper did not provide details for the reviewers to assess those contributions. What are the technical difficulties? How much engineering efforts are required to solve those difficulties? I suggest that the authors revise the manuscript to clearly state the novelty of their work.\n\n- The author did experiments on three neural networks: a fully connect network with ~5.5K parameters, a CNN with 605K parameters and a LSTM with ~1M parameters. While these experiments are nice, I would also like to see if their method scales to very large models frequently used in vision and NLP, like BERT or GPT2. Those large models have hundreds of millions or even billions of parameters. It is important for privacy practitioners to know the feasibility to train those big ML models, or just to fine tune them with a downstream task using fast DPSDG.\n\nOther Comments\n- The author mentioned in line 258 that “…TensorFlow and PyTorch both leverage the optimized cuDNN kernel in the non-private setting for performance, but cannot in the private setting, leading to significantly worse performance.” However, the authors did not explain why cuDNN cannot be used in the private setting. Is it just an engineering issue or not feasible at all?\n- The author mentioned in their contributions (line 77): ”…JAX is generally faster than Custom TFP. We examine and discuss compiled XLA assembly to explain the discrepancy.” However, such discussion is not present in the paper.\n\nTypos\nline 35: are are -> are", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed the potential negative societal impact of their work by pointing out that improved efficiency of DPSGD could increase the accessibility to differential privacy which, when not used properly, could lead to a misleading or false sense of privacy, and suggesting that these issues can be avoided via training and consultation.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "5GqqCFsLhWI", "reviewer": "Reviewer_7jTZ", "summary": "The authors proposed methods to improve the runtime of Differential Privacy SGD using language primitives. DPSGD suffers from performance degradation compared to SGD as most of the libraries are optimized for batch computation. The authors provided implementation of a high performance DPSGD in JAX and Tensorflow. They also compared existing libraries for DPSGD and show that JAX has a superior performance. ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality:\n- The authors provide new implementations for DPSGD in JAX and tf2. \n- They have listed and reviewed all the previous work on DPSGD.\n\nQuality:\n- The experiment are done on small models that are used for differential privacy research on well known dataset.\n- The authors also compare their implementation with other existing methods in terms of performance and memory footprint for FNN, CNN, and LSTM models. \n- They also provided an ablation study on their implementations. \n\nClarity:\n- The paper is well written and well organized and was easy to follow. \n\nSignificance: \n- DP researchers and practitioners might find the work and its results interesting since it provides them with better performing libraries.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "-The authors mentioned the limitations and possible societal impact of their work. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes an approach to speed-up practical implementations of DP-SGD, a popular approach to privacy-preserving machine learning.

The reviewers recognized the significance of the work and its impact for practitioners. The author response helped to further clarify some points. One reviewer (who did not participate in the discussion) was not convinced by the novelty, as the contribution could be seen as mainly engineering. But in the end the previous arguments prevailed.

Therefore, the paper is accepted. The authors are encouraged to incorporate the additional details and discussion from the response.
Decision: Accept (Poster)