Differentially Private Optimization on Large Model at Small Cost

Zhiqi Bu, Yu-Xiang Wang, Sheng Zha, George Karypis

Introduction

Deep learning with differential privacy (DP; (Dwork et al., 2006)) has shown strong performance while guaranteeing rigorous protection against privacy risks, especially on large models that tend to memorize and leak the training data (Carlini et al., 2021; Haim et al., 2022; Shokri et al., 2017). For example, recent advances have shed light on the success of DP GPT2 (Li et al., 2021; Bu et al., 2022b; Yu et al., 2021), which achieves $64.6$ BLEU scoreBLEU (BiLingual Evaluation Understudy) is a metric (0-100) for automatically evaluating translated text. BLEU $>60$ is considered as ”very high quality, adequate, and fluent translations, often better than human”. at strong privacy guarantee ( $\epsilon=3$ ), on the text generation task using E2E restaurant review dataset. This is only marginally below the standard non-private GPT2 (BLEU score 66.8). Similarly, on computer vision tasks ( $\epsilon=2$ ), DP vision transformers and ResNets have obtained $97.1\%/86.2\%$ accuracy on CIFAR10/100 by (Bu et al., 2022a) and over $81\%$ accuracy on ImageNet by (De et al., 2022; Mehta et al., 2022).

However, DP training of large neural networks is well-known to be computationally burdensome in comparison to the standard training, in terms of both the training time and the memory cost. For instance, training a small recurrent neural network (0.598M parameters) experiences a $1000\times$ slowdown using DP optimizers in Tensorflow-Privacy (TF-Privacy) library in (Bu et al., 2021a), and training a small convolutional neural network (CNN, 0.605M parameters) on CIFAR10 has a $24\times$ slowdown with Tensorflow 2 and the XLA compiler (Subramani et al., 2021). Even with SOTA efficient implementations, large models such as RoBERTa (Liu et al., 2019), GPT2 (Radford et al., 2019), ResNet (He et al., 2016), VGG (Simonyan & Zisserman, 2014), ViT (Dosovitskiy et al., 2020) and its variants, experience about $2\sim 3\times$ slowdown in Pytorch (Li et al., 2021; Bu et al., 2022a) and $2\sim 9\times$ slowdown in JAX (Kurakin et al., 2022; De et al., 2022), with possibly $4\sim 20\times$ memory overhead (Bu et al., 2022a; Li et al., 2021; Subramani et al., 2021) if not running out of memory.

The efficiency bottleneck in DP deep learning lies in the per-sample gradient clipping, which restricts the magnitude of each per-sample gradient in the mini-batch. Applying the clipping jointly with the Gaussian noise addition, one can privately release the gradient to arbitrary optimizers like SGD and Adam, and thus guarantee the privacy of the training as described in Section 1.3:

At high level, the DP training is a system effort consisting of multiple parts:

parameter efficiency: last layer (linear probing), LoRA, Adapter, BiTFiT;

implementation: Opacus, GhostClip, Book-Keeping;

Previous works have tackled the efficiency bottleneck with various approaches. One approach (part II) focuses on the parameter efficiency by partially training a neural network, in contrast to fully fine-tuning all model parameters, e.g. only the last output layer (Tramer & Boneh, 2020), the adapter layers (Houlsby et al., 2019; Mahabadi et al., 2021), or the Low-Rank Adaptation (LoRA) (Hu et al., 2021; Yu et al., 2021). For example, (Mehta et al., 2022) accelerate the DP training on ImageNet (Deng et al., 2009) up to $30\times$ by only training the last layer of ResNet152. Noticeably, parameter efficient fine-tuning does not improve on the efficiency in terms of complexity per parameter, rather than reducing the number of parameters. Furthermore, this approach oftentimes leads to some accuracy degradation compared to DP full fine-tuning (Bu et al., 2020; Mehta et al., 2022; Li et al., 2021; Yu et al., 2021).

An orthogonal approach, including this work, focuses on the computation efficiency (part III), i.e. reducing the time and space complexity through efficient implementations, without modifying the DP optimizers (part I) and thus not affecting their performance. We will elaborate on existing methods in Section 1.2. Additionally, these methods can be compiled on different platforms (part IV) such as Tensorflow 2(XLA), JAX and Pytorch (Li et al., 2021; Subramani et al., 2021; De et al., 2022; Kurakin et al., 2022), where remarkable speed difference has been observed in some cases, even with the same implementation. For example, (Subramani et al., 2021) implemented DP-SGD using JAX and claimed its efficiency advantage over the same algorithm using Tensorflow or Pytorch.

[Algorithm] We propose the book-keeping (BK) algorithm that makes existing DP optimizers fast and memory efficient, especially comparable to non-private optimizers. We demonstrate BK via the computation graph in Figure 1. The highlight is that BK only uses one back-propagation and never instantiates per-sample gradients $\{\frac{\partial\mathcal{L}_{i}}{\partial\mathbf{W}}\}_{i=1}^{B}$ .

[Analysis] We analyze the complexity to show that BK has almost the same time and space complexity as non-DP training, especially when the feature dimension is small (see Table 5).

[Extension] We strengthen BK using a layerwise decision to mix with Opacus (see Section 3.2), which proves to be efficient when the feature dimension is large (and difficult for GhostClip). We also extend BK to the parameter efficient fine-tuning such as DP LoRA and Adapter.

[Codebase] We develop a Pytorch (Paszke et al., 2019) codebase for our BK algorithm, leveraging the auto-differentiation technique on the computation graph and a new trick in LABEL:app:auto-differentiation. We highlight that our codebase can automatically switch the standard training of any model to its DP version, by adding a single piece of codes.

[Experiments] We demonstrate the amazing efficiency of BK on training large models, saving the memory up to $10\times$ and boosting the speed by $30\%\sim 5\times$ than previous DP implementations.

2 Related works

Previous arts have developed different implementations of the same DP optimizer in Equation 1. Among these implementations, the tradeoff between the time and space complexity has been constantly maneuvered. TF-Privacy (Tensorflow, ) back-propagates a vectorized loss $[\mathcal{L}_{1},\cdots,\mathcal{L}_{B}]$ to compute the per-sample gradients, each from one back-propagation, which is memory-efficient but slow. Opacus (Yousefpour et al., 2021) and (Rochette et al., 2019) accelerate the training significantly using the outer product trick in (Goodfellow et al., 2014), though incurring heavy memory burden so as to store the per-sample gradients. This memory burden is partially alleviated in FastGradClip (Lee & Kifer, 2020) by sharing the space complexity in two rounds of back-propagation, hence almost doubling the time complexity. In ghost clipping (Goodfellow, 2015; Li et al., 2021; Bu et al., 2022a), the per-sample gradients can be clipped without being instantiated, thus both time and space complexity can be further improved if the feature dimension is small. We refer interested readers to Figure 3 and Appendix C for algorithmic details of these implementations.

We now compare BK to different implementations in Table 2 and Figure 2. In what follows, $B$ is the batch sizeIn this work, we report the physical batch size, which affects the efficiency but not the accuracy; the accuracy is only affected by the logical batch size, which can be implemented through the gradient accumulation of physical batch size., $T_{(l)}$ is the feature dimensionFor non-sequential data, $T=1$ ; for texts, $T$ is the sequence length, which is layer-independent; for images (or videos), $T_{(l)}$ is the height $\times$ width( $\times$ time) of hidden feature representation, which is layer-dependent., $d_{(l)},p_{(l)}$ are the input or output dimension of a layer.

3 Preliminaries

We work with the $(\epsilon,\delta)$ -DP by (Dwork et al., 2006), defined in Appendix A, which makes it difficult for any privacy attacker to distinguish or detect an arbitrary training sample, even with full access to the model. In deep learning, DP is achieved by training on the private gradient in Equation 1 with any optimizer such as SGD, Adam, FedAvg, etc. Essentially, the private gradient is the addition of Gaussian noise to the sum of clipped per-sample gradients, which guarantees the DP protection through the privacy accounting theorems (Abadi et al., 2016; Mironov, 2017; Dong et al., 2019; Zhu et al., 2021; Gopi et al., 2021; Koskela et al., 2020).

Book-keeping: Efficient DP training in low dimension

The main computational bottleneck of DP training comes from the per-sample gradient clipping, or from the computation of per-sample gradient norms, to be exact. One widely used approach in Opacus, TF-privacy, and FastGradClip, is to instantiate the per-sample gradients and then deriving their norms. Straight-forward implementation of this approach on a mini-batch of per-sample losses requires $B$ rounds of back-propagation (unacceptable slowdown) or $B\times$ gradient storage (unacceptable memory burden; see Opacus in Figure 2). Consequently, these implementations are not suitable for large model training. For instance, (Li et al., 2021) shows that, when training GPT2-large (774M parameters), Opacus (Yousefpour et al., 2021) and JAX (Subramani et al., 2021) cannot fit even one single sample into a 24GB GPU.

An alternative approach, termed as the ghost clipping (GhostClip), directly computes the per-sample gradient norms without computing the gradients themselves. This is made possible, unfortunately, through two rounds of back-propagation. During the first back-propagation, one uses the regular loss $\sum_{i}\mathcal{L}_{i}$ and extracts the activation tensor and the output gradient $(\bm{a},\frac{\partial\mathcal{L}}{\partial\bm{s}})$ . One can use an algebraic trick in Equation 2 to compute the per-sample gradient norms $\{\|\frac{\partial\mathcal{L}_{i}}{\partial\mathbf{W}}\|\}_{i}$ and the clipping factors $\{C_{i}\}_{i}$ in Equation 1. During the second back-propagation, one uses the reweighted loss $\sum_{i}C_{i}\mathcal{L}_{i}$ whose gradient is directly the weighted gradient $\sum_{i}C_{i}\bm{g}_{i}$ , which constitutes the private gradient we need. Note that this double back-propagation roughly doubles the training time (or to be more precise, $10/6\approx 1.667\times$ when $T$ is small; but this approach loses its advantage when $T$ is large), as shown in Table 2).

To make the DP training as efficient as the standard training, we propose the book-keeping technique (BK) that $\langle 1\rangle$ only requires a single round of back-propagation, like Opacus and the standard training; $\langle 2\rangle$ does not instantiate the per-sample gradients, like GhostClip.

BK algorithms in their base forms are built on GhostClip and especially the ghost norm trick, so as to avoid instantiating the memory costly per-sample gradients: as can be seen in Algorithm 1 and Figure 3, $\frac{\partial\mathcal{L}_{i}}{\partial\mathbf{W}}=\bm{a}_{i}^{\top}\frac{\partial\mathcal{L}}{\partial\bm{s}_{i}}$ is not computed throughout the training. In comparison to GhostClip, our significant improvement is solely on the speed (see Table 2) through two novel tricks: the book-keeping and the ghost differentiation. The entire BK algorithm is built on the understanding of computation graph in Appendix A. Note that these tricks also offer improved efficiency for existing implementations, to be presented in Section 2.4. We now elaborate on these tricks.

$\text{BK (base)}=\underbrace{\text{ghost norm}}_{\text{from GhostClip}}+\underbrace{\text{book-keeping}}_{\text{ours}}+\underbrace{\text{ghost differentiation}}_{\text{ours}}$

The ghost norm trick (Goodfellow, 2015) computes the gradient norm without the gradient: while the gradient is instantiated by the multiplication in Equation 2, the gradient norm can be computed without $\bm{a}_{i}$ meeting $\frac{\partial\mathcal{L}}{\partial\bm{s}_{i}}$ . This trick is applicable to generalized linear layers including the linear, the embedding (Li et al., 2021), and the convolution layers (Bu et al., 2022a). We emphasize that these generalized linear layers represent 99.9% of the trainable parameters in modern neural networks.

without actually computing $\frac{\partial\mathcal{L}_{i}}{\partial\mathbf{W}}=\bm{a}_{i}^{\top}\frac{\partial\mathcal{L}}{\partial\bm{s}_{i}}$ . Here ‘vec’ means flattening the $T\times T$ matrix to a vector. This trick is particularly efficient when $T$ is small, reducing the space complexity from $O(Bpd)$ to $O(BT^{2})$ by Table 3.

This trick improves the time complexity by removing the second back-propagation from GhostClip. Our idea is to book-keep and re-use the output gradient $\frac{\partial\mathcal{L}}{\partial\bm{s}_{(l)}}$ , which is deleted after the first back-propagation of GhostClip and must be re-computed during the second back-propagation. The difference between GhostClip and BK is clearly illustrated via a line-by-line comparison in Algorithm 2. In fact, denoting the total number of model parameters as $M=\sum_{l}p_{(l)}d_{(l)}$ , our trick reduces the time complexity from $10BTM+O(BT^{2})$ by GhostClip to $8BTM+O(BT^{2})$ according to Table 3. In contrast to Opacus, which book-keeps the per-sample gradients $\bm{g}^{(l)}_{i}$ using $O(BM)=O(B\sum_{l}p_{(l)}d_{(l)})$ memory, we instead book-keep the output gradient with substantially cheaper $O(BT\sum_{l}p_{(l)})$ memory when the feature dimension $T$ is small.

This trick improves the time complexity on the first back-propagation in GhostClip, further reducing from $8BTM+O(BT^{2})$ to $6BTM+O(BT^{2})$ in Table 2. Our idea is to only compute the output gradient $\frac{\partial\mathcal{L}}{\partial\bm{s}_{(l)}}$ but not the (non-private) parameter gradient $\frac{\partial\mathcal{L}}{\partial\mathbf{W}}$ . That is, we break the $4BTM$ time complexity of the full back-propagation into two sub-processes, each of $2BTM$ complexity, and remove the unnecessary one.

To be more specific, during the back-propagation of Opacus and GhostClip, the output gradient $\frac{\partial\mathcal{L}}{\partial\bm{s}}$ and then the parameter gradient $\frac{\partial\mathcal{L}}{\partial\mathbf{W}}=\bm{a}^{\top}\frac{\partial\mathcal{L}}{\partial\bm{s}}$ are computed. However, we can stop after we obtain $\frac{\partial\mathcal{L}}{\partial\bm{s}}$ : we only need the output gradient to compute the clipped parameter gradient $\frac{\partial\sum_{i}C_{i}\mathcal{L}_{i}}{\partial\mathbf{W}}$ in Line 9 of Algorithm 1. Therefore, the ghost differentiation trick sets all parameters to not require gradients. See the technical details in LABEL:app:auto-differentiation, including the origin parameter trick that propagates on a computation graph even when no parameters require gradients.

2 Complexity of DP implementations: a modular analysis

In this section, we analyze the complexity of DP implementations from their opearation modules. We summarize the time and space complexity in Table 3 and give the derivation in Appendix B. We will refer to these modules by indices, e.g. 2a for the computation of output gradient.

Now we are ready to decompose each implementation, following the flowcharts in Figure 3. Consequently, we can easily write down the complexity of different implementations in Table 2. Such a modular analysis displays the clear effects of the tricks in BK algorithm: the ghost norm trick removes the memory costly 4 from Opacus and FastGradClip; the book-keeping trick removes the 2b in the second back-propagation of FastGradClip and GhostClip; the ghost differentiation trick removes the 2b in the first back-propagation of Opacus and GhostClip.

3 BK is optimally efficient in low dimension

When the feature dimension $T$ is small, we claim that BK is almost as efficient as the standard non-private training, with a negligible $O(BT^{2})$ time and memory overhead by Table 2:

Now, we discuss the cases where the data has low dimension and thus $T$ is small. Generally speaking, the feature dimension $T_{(l)}$ depends on both the data and the model.

For non-sequential input and 1D audio data, $T=1$ . For sequential data such as texts ( $T$ being sentence length) or time series ( $T$ being time duration), $T_{(l)}$ is fixed across layers. In this case, BK is efficient on short-sequence datasets including GLUE (Wang et al., 2019) (e.g. SST2/QNLI/MNLI/QQP) and natural language generation datasets (e.g. E2E/DART), since $T^{2}\ll p_{(l)}d_{(l)}$ . For instance, (Yu et al., 2021; Li et al., 2021; Bu et al., 2022b) applied GPT2 on E2E dataset, which has a sequence length $T\approx 100$ and the number of parameters $p_{(l)}d_{(l)}$ per layer is $2-4$ M; (Yu et al., 2021; Li et al., 2021) applied RoBERTa-large on GLUE datasets, which has a sequence length $T=256$ and the number of parameters per layer is $1-4$ M. As illustrated in Figure 5 and Table 1 (extended in LABEL:tab:GPT_max_throughput), BK improves the throughput of existing implementations by $25-388\%$ on multiple language tasks in (Li et al., 2021; Bu et al., 2022b), with minor memory overhead compared to GhostClip and non-private training.

However, on the convolution layers with image data, $T_{(l)}$ is the product of hidden feature sizes (c.f. Section 3 in (Bu et al., 2022a)), thus $T_{(l)}$ depends on the original image size and network architecture. For example, larger kernel size/dilation/stride in convolution layer reduces $T_{(l)}$ , while larger images have larger $T_{(l)}$ at each layer. Therefore, BK (and GhostClip) may suffer on when training ResNet on ImageNet $(224\times 224)$ , as we show in Figure 6 (see also Table 7 in (Bu et al., 2022a)), although training the same network efficiently on CIFAR10/100 $(32\times 32)$ .

4 Applying our tricks to existing implementations

Our tricks in Section 2.1 can also improve other existing implementations, reducing the time complexity of GhostClip from $10BTpd+2BT^{2}(p+d)$ to $6BTpd+2BT^{2}(p+d)$ , that of Opacus and FastGradClip from $8BTpd$ to $6BTpd$ . We highlight that these improved implementations are leveraged to design hybrid implementation in Section 3.2. In addition to DP full fine-tuning, BK is demonstrated in LABEL:app:param_eff_BK to also apply to the parameter efficient fine-tuning like Adapters (Houlsby et al., 2019) and LoRA (Hu et al., 2021).

Hybrid Book-keeping: Efficient DP training in high dimension

In previous section, we have analyzed DP implementations in the small $T$ regime, where the ghost norm-based GhostClip and BK are efficient. Nevertheless, in the large $T$ and large model regime, none of the base implementations may be efficient (see Figure 6) and we turn to hybrid methods.

A closer look at the space complexity in Table 3 shows that, the ghost norm trick is favored over the per-sample gradient instantiation if and only if $2T^{2}_{(l)}<p_{(l)}d_{(l)}$ , where $p_{(l)}d_{(l)}$ is the number of parameters at one layer. When this criterion is violated for large $T$ , GhostClip/BK (base) can significantly under-perform Opacus/FastGradClip, as shown in Figure 6, Figure 7 and LABEL:tab:layerwise_complexity.

Similar to Section 2.3, we discuss two cases where $T$ is large. For paragraph or document-level language tasks like WikiHop (Welbl et al., 2018) and TriviaQA (Joshi et al., 2017), $T$ can range from $2000\sim 20000$ to train large language models, which makes $2T^{2}=8\sim 800$ M. For example, LLAMA (Touvron et al., 2023) is trained with token length $4096\leq T\leq 8192$ and GPT-3 (Brown et al., 2020) is trained with token length $T=2048$ .

For image tasks, particularly on CNN, $T_{(l)}$ varies at each layer with large values on top layers, as the features are less compressed by convolution and pooling. Taking ImageNet and the first convolution layer of VGG11 as an example (see Table 3 of (Bu et al., 2022a)), $2T_{(1)}^{2}=5\times 10^{9}\gg p_{(1)}d_{(1)}=1.7\times 10^{3}$ . Consequently, ghost norm-based implementations (i.e. GhostClip and BK) costs more than 40GB memory on ResNet18, under $B=32$ , while Opacus only costs 2.5GB. This curse of dimension grows from a difficult issue on ImageNet to an impossible challenge on videos or high-resolution images, e.g. GhostClip cannot train ResNet18 with even one single CelebA-HQ image ( $1024\times 1024$ ) using a 40GB GPU.

In short, the ghost norm trick is inefficient for large $T$ and the per-sample gradient instantiation is inefficient for large model. Hence, we must hybridize the base implementations.

2 Hybrid implementations via layerwise decision

We adopt the same layerwise decision as (Bu et al., 2022a), known as the mixed ghost norm technique: we use the ghost norm trick on a layer if $2T_{(l)}^{2}<p_{(l)}d_{(l)}$ , and instantiate per-sample gradients otherwise. Therefore, the space complexity of computing the per-sample gradient norm reduces to $\min\{2T_{(l)}^{2},p_{(l)}d_{(l)}\}$ , which is significantly cheaper than either the ghost norm or the per-sample gradient instantiation in high dimension, as depicted in Table 4 and Figure 7. Consequently, over all layers, the space complexity is lower than both constituting methods, e.g. saving more than $10\times$ memory for the per-sample gradient clipping on ResNet18 (see more models in LABEL:tab:layerwise_complexity).

In contrast to the mixed ghost clipping (MixGhostClip) in (Bu et al., 2022a), which hybridizes FastGradClip and GhostClip, we boost the efficiency by hybridizing our BK with the improved FastGradClip/Opacus in Section 2.4. We propose BK-MixOpt (and BK-MixGhostClip as an intermediate product only for comparison) and use MixGhostClip as a reference point,

The hybrid BK algorithms are presented in LABEL:alg:BKAL-mixed. We summarize the layerwise complexity in Table 5, from which we derive the overall complexity in LABEL:tab:complexity_overhead and observe that BK has almost the same complexity as non-DP training. Note that in low dimension, the mixed ghost norm is equivalent to the ghost norm, hence MixGhostClip/BK-MixOpt is equivalent to GhostClip/BK, respectively.

3 Effect of model architecture & feature dimension on hybridization

We dive deeper to understand when the hybridization favors the ghost or non-ghost norm tricks.

From a model architecture viewpoint, transformers such as ViT, RoBERTa, GPT tend to prefer the ghost norm: for moderate-sequence text data and moderate-dimension image data, hybrid BK algorithms are close or equivalent to the base BK algorithm (see right-most plot in Figure 7). However, CNN prefers the per-sample gradient instantiation at top layers, and there exists a depth threshold below which the ghost norm is more efficient. Hence the hybridization is necessary to take advantages of both worlds.

From the feature dimension viewpoint, larger input enlarges this depth threshold, e.g. from the 9-th layer of ResNet18 to the 17-th layer in Figure 7, when the image size increases from $224\times 224$ to $512\times 512$ . We visualize this pattern on various models in LABEL:app:effect_of_hybrid. In particular, we observe in LABEL:tab:complexity_overhead that when $T$ is large, both per-sample gradient instantiation (Opacus) and ghost norm trick (GhostClip) are significantly dominated by our BK algorithms.

Instructions to use the codebase

In this section, we demonstrate how to modify a standard training script to its DP variantsThat is, our codebase can easily adapt to any per-sample gradient clipping function and privacy accouting methods. by one piece of code.

from fastDP import PrivacyEngine import torch.functional as F optimizer = torch.optim.Adam(model.parameters()) privacy_engine = PrivacyEngine( model,epochs, batch_size,sample_size, target_epsilon,target_delta) privacy_engine.attach(optimizer) logits = model(data) loss = F.cross_entropy(logits, labels) loss.backward() optimizer.step() optimizer.zero_grad()

We highlight that our codebase automatically modifies the training for any network architecture and any optimizer. Additionally, it is designed to work compatibly with large-scale training techniques, such as the gradient accumulation, the parameter-efficient fine-tuning (e.g. LoRA and BiTFiT (Bu et al., b)), and the parallel distributed learning (e.g. ZeRO (Bu et al., a)).

Discussion

In this work, we propose the Book-Keeping (BK) algorithms to effciently implement DP optimizers using three tricks: ghost norm, book-keeping, and ghost differentiation. Our BK reduces the time and space complexity of DP training to the similar level of the standard training. Specially, we develop hybrid BK to overcome the computational challenge of training large models with high-dimensional data, and we extend BK to parameter efficient fine-tuning such as LoRA and Adapter.

One limitation of this work is that BK (and GhostClip) only applies to the weights, not the biases, and only on the generalized linear layers, i.e. the embedding, the linear, and the convolution layers. However, this is a minor concern because the weights in the generalized linear layers constitute 99.9% of the trainable parameters (see LABEL:tab:applicable_to_linear_layer).

Implementation-wise, our codebase is automatic (allowing any model to be DP optimized) and future-proof (allowing any training setting, including the user-level DP and the distributed learning). However, although BK is theoretically as fast as the standard training for small $T$ , we observe some gap between the theoretical complexity and the hardware throughput in practice. This gap is mainly due to the mechanism of Pytorch hooks which can be possibly optimized by customizing the CUDA kernel or using the symbolic programming. We expect this gap to be closed by future research.

References

Appendix A Background

We formally introduce the differential privacy (DP).

A randomized algorithm $M$ is $(\varepsilon,\delta)$ -differentially private (DP) if for any two neighboring $S^{\prime}$ is a neighbor of $S$ if one can obtain $S^{\prime}$ by adding or removing one data point from $S$ . datasets $S,S^{\prime}$ , and for any event $E$ ,

Clearly, stronger DP (smaller $\epsilon,\delta$ ) indicates the higher difficulty for privacy attackers to extract information from the training data.

DP can be achieved by adding Gaussian noise to a bounded-sensitivity function (see Theorem A.1 of (Dwork et al., 2014)). In deep learning, this function is the sum of per-sample gradients $\sum\bm{g}_{i}$ and the bounded sensitivity is $R$ (that is guaranteed through the gradient clipping after which the per-sample gradient norm is at most $R$ ). Note that the Gaussian noise magnitude is proportional to the sensitivity: $\sigma_{\text{DP}}=\sigma R$ in Equation 1, and $\epsilon(\delta)$ only depends on $\sigma$ , not on $R$ . The derivation from $(\sigma,T,p)$ in Algorithm 1 to $\epsilon$ can be done through various methods in Section 1.3.

A.2 Computation graph

We elaborate on the computation graph presented in Figure 1. For DP and the standard training, the forward pass is the same: we pass through the layers

For the backward propagation, there are two sub-processes:

the computation of output gradient for all layers,

i.e. the output gradient meets with the weight $\mathbf{W}$ ;

the computation of parameter gradient only for trainable parameters,

i.e. the output gradient meets with the activation tensor $\bm{a}$ .

Note that foward pass, output gradient, and parameter gradient have the same time complexity of $2BTM$ ( $B$ being the batch size, $T$ being the feature dimension, e.g. the sequence length in texts, and $M$ being the model size).

For example, GhostClip (Li et al., 2021) and MixGhostClip (Bu et al., 2022a), which use one forward pass and double backward propagation, have a time complexity of $10BTM+O(BT^{2})$ , while the standard training which uses one forward pass and a single backward propagation has a time complexity of $6BTM$ .

Appendix B Complexity analysis for one layer

Let us consider a layer without bias term for simplicity:

We now break down the time and space complexities for each operation in the training. Notice that we focus on major complexities, e.g. ignoring cubic terms like $BTp$ when higher order terms like $BTpd$ or $BT^{2}p$ exist.

B.2 Back-propagation: output gradient

The complexity to compute the output gradient is incurred by the chain rule: for a single sample,

where $\phi$ is the non-linear activation function. We compute the matrix multiplication $\frac{\partial\mathcal{L}}{\partial\bm{s}_{(l),i}}\mathbf{W}_{(l)}$ first, with time complexity $2BTpd$ and space complexity $pd+BTd+BTp$ . Then the elementwise product uses time complexity $2BTd$ and space complexity $BTd$ .

B.3 Back-propagation: parameter gradient

This module could represent different operations in different DP implementations. In the first back-propagation of GhostClip and the only back-propagation of Opacus, it computes $\frac{\partial\mathcal{L}}{\partial\mathbf{W}}=\frac{\partial\sum_{i}\mathcal{L}_{i}}{\partial\mathbf{W}}$ ; in the second back-propagation of Ghost/FastGradClip/BK, it computes the clipped gradient $\frac{\partial\sum_{i}C_{i}\mathcal{L}_{i}}{\partial\mathbf{W}}$ . Regardless of the cases, the operation always takes the same format as

This tensor multiplication has time complexity $2BTpd$ and space complexity $pd$ unless all per-sample gradients are stored.

B.4 Ghost norm

Ghost norm is the operation taking $\bm{a}_{i}$ and $\frac{\partial\mathcal{L}}{\partial\bm{s}_{i}}$ as the input and outputs the per-sample gradient norm. According to Equation 2 and Appendix C.3 of (Bu et al., 2022b), this operation computes $\bm{a}_{i}\bm{a}_{i}^{\top}$ and $\frac{\partial\mathcal{L}}{\partial\bm{s}_{i}}\frac{\partial\mathcal{L}}{\partial\bm{s}_{i}}^{\top}$ , taking the time complexity $2BT^{2}d$ and $2BT^{2}p$ respectively, and the space complexity $BT^{2}$ for each variable. Hence total time complexity is $2BT^{2}(p+d)$ and total space complexity is $2BT^{2}$ .

Alternatively, one can instantiate the per-sample gradients and then compute their norms. This is different than the computation of parameter gradient in the back-propagation: that computation is an efficient tensor multiplication while this operation consists of $B$ matrix multiplication.

This operation has time complexity $2BTpd$ and space complexity $Bpd$ to store all per-sample gradients. Computing the norms is cheap enough to be neglected.

B.5 Weighted sum of per-sample gradient

In contrast to double back-propagation, which indirectly derives the clipped gradient by differentiating the reweighted loss $\sum_{i}C_{i}\mathcal{L}_{i}$ at a cost of $O(BTpd)$ , this operation directly computes the clipped gradient under almost no time complexity. Noticeably, this is only possible if per-sample gradients are readily instantiated and stored.