Scalable and Efficient Training of Large Convolutional Neural Networks with Differential Privacy
Zhiqi Bu, Jialin Mao, Shiyun Xu
Introduction
Deep convolutional neural networks (CNN) are the backbone in vision-related tasks, including image classification , object detection , image generation , video recognition , and audio classification . A closer look at the dominating success of deep CNNs reveals its basis on two factors.
The first factor is the strong capacity of the convolutional neural networks, usually characterized by the enormous model size. Recent state-of-the-art progresses usually result from very large models, with millions to billions of trainable parameters. For example, ImageNet accuracy grows when larger VGGs (increasing 11 to 19 layers69% to 74% accuracy) or ResNets (increasing 18 to 152 layers70% to 78% accuracy) are used . Consequently, the eager to use larger models for better accuracy naturally draws people’s attention to the scalability and efficiency of training.
The second factor is the availability of big data, which oftentimes contain private and sensitive information. The usage of such data demands rigorous protection against potential privacy attacks. In fact, one standard approach to guarantee the protection is by differentially private (DP) training of the models. Since , CNNs have achieved promising results under strong DP guarantee: CIFAR10 achieves 92.4% accuracy in and ImageNet achieves 81.1% accuracy in .
Unifying the two driving factors of CNNs leads to the DP training of large CNNs. However, the following challenges are hindering our application of large and private CNNs in practice.
Challenge 1: Time and space efficiency in DP training. DP training can be extremely inefficient in memory and speed. For example, a straightforward implementation in Tensorflow Privacy library shows that DP training can be slower than the non-DP training, even on a small RNN ; other standard DP libraries, Opacus and JAX , which trade off memory for speed, could not fit a single datapoint into GPU on GPT2-large ; addtionally, slowdown of DP training has been reported in using JAX.
The computational bottleneck comes from the per-sample gradient clipping at each iteration (see (2.1)), a necessary step in DP deep learning. I.e., denoting the loss as , we need to clip the per-sample gradient individually. This computational issue is even more severe when we apply a large batch size, which is necessary to achieve high accuracy of DP neural networks. In , it is shown that the optimal batch size for DP training is significantly larger than for regular training. For instance, DP ResNet18 achieves best performance on ImageNet when batch size is 64*1024 ; and DP ResNet152 and ViT-Large use a batch size in . As a result, an efficient implementation of per-sample gradient clipping is much-needed to fully leverage the benefit of large batch training.
Challenge 2: Do large DP vision models necessarily harm accuracy? An upsetting observation in DP vision models is that, over certain relatively small model size, larger DP CNNs seem to underperform smaller ones. This is observed in models that are either pre-trained or trained from scratch . As an example of the pre-trained cases, previously state-of-the-art CIFAR10 is obtained from a small DP linear model , and the fine-tuned DP ResNet50 underperforms DP ResNet18 on ImageNet . On the contrary, the empirical evidence in DP language models shows that larger models can consistently achieve better accuracy . Interestingly, we empirically demonstrate that this trend can possibly hold in vision models as well.
In this work, we propose new algorithms to efficiently train large-scale CNNs with DP optimizers. To be specific, our contributions are as follows.
We propose a novel implementation, termed as the mixed ghost clipping, of the per-sample gradient clipping for 1D3D convolutional layers. The mixed ghost clipping is the first method that can implement per-sample gradient clipping without per-sample gradients of the convolutional layers. It works with any DP optimizer and any clipping function almost as memory efficiently as in standard training, thus significantly outperforming existing implementation like Opacus .
In some tasks, mixed ghost clipping also claims supremacy on speed using a fixed batch size. The speed can be further boosted (say faster than the fastest alternative DP algorithms and only slower than the non-private training) when the memory saved by our method is used to fit the largest possible batch size.
We provide the first complexity analysis of mixed ghost clipping in comparison to other training algorithms. This analysis clearly indicates the necessity of our layerwise decision principle, without which the existing methods suffer from high memory burden.
Leveraging our algorithms, we can efficiently train large DP models, such as VGG, ResNet, Wide-ResNet, DenseNet, and Vision Transformer (ViT). Using DP ViTs at ImageNet scale, we are the first to train convolutional ViTs under DP and achieve dominating SOTA on CIFAR10/100 datasets, thus bringing new insights that larger vision models can consistently achieve better accuracy under DP.
2 Previous arts
The straightforward yet highly inefficient way of per-sample gradient clipping is to use batch size 1 and compute gradients with respect to each individual loss. Recently, more advanced methods have significantly boosted the efficiency by avoiding such a naive approach. The most widely applied method is implemented in the Opacus library , which is fast but memory costly as per-sample gradients are instantiated to compute the weighted gradient in (2.1). A more efficient method, FastGradClip , is to use a second back-propagation with weighted loss to indirectly derive the weighted gradient.
In all above-mentioned methods and The method in also extends the outer product trick (similar to Opacus, see (2.4)) in to convolution layer, but does not use the ghost clipping trick., the per-sample gradients are instantiated, whereas this can be much inefficient and not necessary according to the ‘ghost clipping’ technique, as to be detailed in Section 3. In other words, ghost clipping proves that the claim ‘DP optimizers require access to the per-sample gradients’ is wrong. Note that ghost clipping is firstly proposed by for linear layers, and then extended by to sequential data and embedding layers for language models. However, the ghost clipping has not been extended to convolutional layers, due to the complication of the convolution operation and the high dimension of data (text data is mostly 2D, yet image data are 3D and videos are 4D). We give more details about the difference between this work and in Appendix F. In fact, we will show that even the ghost clipping alone is not satisfactory for CNNs: e.g. it cannot fit even a single datapoint into the memory on VGGs and ImageNet dataset. Therefore, we propose the mixed ghost clipping, that narrows the efficiency gap between DP training and the regular training.
Preliminaries
Differential privacy (DP) has become the standard approach to provide privacy guarantee for modern machine learning models. The privacy level is characterized through a pair of privacy quantities , where smaller means stronger protection.
A randomized algorithm is -DP if for any neighboring datasets that differ by one arbitrary sample, and for any event , it holds that
In deep learning where the number of parameters are large, the Gaussian mechanism [15, Theorem A.1] is generally applied to achieve DP at each training iteration, i.e. we use regular optimizers on the following privatized gradient:
In words, DP training switches from updating with to updating with the private gradient : SGD with private gradient is known as DP-SGD; Adam with private gradient is known as DP-Adam.
Algorithmically speaking, the Gaussian mechanism can be decomposed into two parts: the per-sample gradient clipping and the Gaussian noise addition. From the viewpoint of computational complexity, the per-sample gradient clipping is the bottleneck, while the noise addition costs negligible overhead.
In this work, our focus is the implementation of per-sample gradient clipping (2.1). We emphasize that our implementation is only on the algorithmic level, not affecting the mathematics and thus not the performance of DP optimizers. That is, our mixed ghost clipping provides exactly the same accuracy results as Opacus, FastGradClip, etc.
2 Per-sample gradient for free during standard back-propagation
In DP training, the per-sample gradient is a key quantity which can be derived for free from the standard back-propagation. We briefly introduce the back-propagation on linear layers, following the analysis from , so as to prepare our new clipping implementation for convolutional layers. Note that the convolutional layers can be viewed as equivalent to the linear layers in Section 2.3.
In the -th layer of a neural network with layers in total, we denote its weight, bias, input and output as respectively, and the activation function as . Consider
Clearly the -th sample’s hidden feature at layer is freely extractable during the forward pass.
Let be the total loss and be the per-sample loss with respect to the -th sample. During a standard back-propagation, the following partial product is maintained,
so as to compute the standard gradient in (2.4). Here is the Hadamard product and is the matrix product. Therefore, is also available for free from (2.3) and extractable by Pytorch hooks, which allows us compute the per-sample gradient by
3 Equivalence between convolutional and linear layer
In a convolutional layerSee a detailed explanation in Appendix B for the operation and the dimension formulae in convolution., the forward pass is
To present concisely, we ignore the layer index and write the per-sample gradient of weight for the convolutional layer, in analogy to the linear layer in (2.4),
Here is the inverse operation of and simply flattens all dimensions except the last one: from to . From (2.6), we derive the per-sample gradient norm for the convolutional layers from the same formula as in [33, Appendix F],
Ghost clipping for Convolutional Layers
Leveraging our derivation in (2.7), we propose the ghost clipping to compute the clipped gradient without ever generating the per-sample gradient . The entire procedure is comprised of the ghost norm computation and the second back-propagation, as demonstrated in Figure 1.
The per-sample gradient norm is required to compute the per-sample in (2.1). While it is natural to instantiate the per-sample gradients and then compute their norms , this is not always optimal nor necessary. Instead, we can leverage (2.7), the ghost norm, to compute the per-sample gradient norm and avoid the possibly expensive per-sample gradient. Put differently, when is small, the multiplication plus is cheap, but the multiplication is expensive. We demonstrate the ghost clipping’s supremacy over complexity empirically in Table 4 and theoretically in Table 2.
2 Second back-propagation: weighted loss leads to weight gradient
We conduct a second back-propagation with the weighted loss to derive the weighted gradient in (2.1), which costs extra time. In contrast, Opacus and JAX generate and store the per-sample gradient for all . Thus the weighted gradient is directly computable from as the memory is traded off for faster computation. However, in some cases like LABEL:tab:imagenet on ImageNet and Table 9 on CIFAR100, we can use larger batch size to compensate the slowdown of the second back-propagation.
Mixed Ghost Clipping: To be a ghost or not, that is the question
While the ghost norm offers the direct computation of gradient norm at the cost of an indirect computation of the weighted gradient, we will show that ghost clipping alone may not be sufficient for efficient DP training, as we demonstrate in Table 4 Figure 3, and LABEL:tab:imagenet. In Table 2, we give the first fine-grained analysis of the space and time complexity for DP training algorithms. Our analysis gives the precise condition when the per-sample gradient instantiation (adopted in Opacus ) is more or less efficient than our ghost norm method. To take the advantage of both methods, we propose the mixed ghost clipping method in Algorithm 1, which applies the ghost clipping or non-ghost clipping by a layerwise decision.
We highlight that the key reason supporting the success of mixed ghost clipping method is its layerwise adaptivity to the dimension parameters, , which vary largely across different layers (see Figure 2). The variance results from the fact that images are non-sequential data, and that the convolution and pooling can change the size () of hidden features drastically.
In the next two sections, we will analyze rigorously the time and memory complexities of the regular training and the DP training, using ghost or non-ghost clippings.
In Algorithm 1, we present the mixed ghost clipping that prioritizes the space complexity by (4.1). We also derive and implement a speed-priority version by comparing the time complexity of ghost norm and gradient instantiation in Table 1. However, the efficiency difference is empirically insignificant and implied by Table 1.
We now break each clipping method into operation modules and analyze their complexities. A similar but coarse analysis from only claims, on sequential layers, space complexity with ghost clipping and without ghost clipping. The time complexity and/or convolutional layers are not analyzed until this work.
Here is the batch size, where is the number of input channels, is the kernel sizes, is the number of output channels, and . We leave the detailed complexity computation in Appendix C. Leveraging Table 1, we give the complexities of different clipping algorithms in Table 2.
2 Layerwise decision in mixed clipping
From the space complexity in Table 2, we derive the layerwise decision that selects the more memory efficient of FastGradClip (gradient instantiation) and ghost clipping (ghost norm):
Therefore, our mixed ghost clipping is a mixup of FastGradClip and the ghost clipping (c.f. Figure 1). We note that the decision by the mixed ghost clipping (4.1) depends on different dimensions: the ghost clipping depends on the size of hidden features (height and width ) which in turn depends on kernel size, stride, dilation and padding (see Appendix B for introduction of convolution), while only the non-ghost clipping depends on the number of channels. In ResNet and VGG, the hidden feature size decreases as layer depth increases, due to the shrinkage from the convolution and pooling operation; on the opposite, the number of channels increases in deeper layers.
As a consequence of decreasing hidden feature size and increasing number of channels, there exists a depth threshold beyond which the ghost clipping is preferred in bottom layers, where the save in complexity is substantial. In Figure 2 and Table 3, as the layer of VGG 11 goes deeper, the hidden feature size shrinks from and the number of channels increases from .
Performance
We compare our ghost clipping and mixed ghost clipping methods to state-of-the-art clipping algorithms, namely Opacus and FastGradClip , which are implemented in Pytorch. We are aware of but will not compare to implementations of these two algorithms in JAX , e.g. , so as to only focus on the algorithms rather than the operation framework. All experiments run on one Tesla V100 GPU (16GB RAM).
We highlight that switching from the regular training to DP training only needs a few lines of code using our privacy engine (see Appendix E). For CNNs, we use models from https://github.com/kuangliu/pytorch-cifar on CIFAR10 () and models from Torchvision on ImageNet () . For ViTs, regardless of datasets, we resize images to and use models from PyTorch Image Models (TIMM) .
We first measure the time and space complexities when the physical batch size is fixed. Here we define the physical batch size (or the virtual batch size) as the number of samples actually fed into the memory, which is different from the logical batch size. For example, if we train with batch size 1000 but can only feed 40 samples to GPU at one time, we back-propagate 25 times before updating the weights for 1 time. This technique is known as the gradient accumulation and is widely applied in large batch training, which particularly benefits the accuracy of DP training .
From Table 4 and the extended Table 6, we see a clear advantage of mixed ghost clipping: our clipping only incurs memory overhead than the regular training, and is the fastest DP training algorithm. In contrast on ResNet18, Opacus uses memory, and FastGradClip uses memory. Even the ghost clipping uses memory of regular training, while being slower than both Opacus and FastGradClip.
Similar phenomenon is observed on ImageNet in LABEL:tab:imagenet: at physical batch size 25, while the mixed ghost clipping works efficiently, we observe that the ghost clipping and Opacus incur heavy memory burden that leads to OOM error on all VGGs and wide ResNets. In fact, the ghost clipping fails in memory on all models except the small AlexNet .
2 Maximum batch size and throughput
Importantly, the speed efficiency in Table 4 can be further boosted, if we use up the saved memory to increase the batch size. To stress test the maximum physical batch size and the throughput of each clipping method, we train ResNet , VGG , MobileNet, ResNeXt ,AlexNet,Wide-ResNet, DenseNet and ViTs on CIFAR10 and ImageNet, as summarized partially in Figure 3 and in LABEL:tab:imagenet, respectively. For example, on VGG19 and CIFAR10, the mixed ghost clipping has a maximum batch size bigger (thus faster) than Opacus, bigger (thus faster) than FastGradClip, and bigger (thus faster) than the ghost clipping. Similarly, on Wide-ResNet50 and ImageNet, the mixed ghost clipping has a maximum batch size bigger than Opacus, bigger than the ghost clipping, and more memory costly than the non-private training.
3 Vision transformers with convolution on ImageNet scale
In addition to training large-scale CNNs such as ResNet152, we apply our mixed ghost clipping to train ViTs, which substantially outperform existing SOTA on CIFAR10 and CIFAR100. Notice that the ViTs are pretrained on ImageNet scale, by which we resize CIFAR images (from pixels to pixels).
It is worth mentioning that ViT is originally proposed as a substitute of CNN. Hence it and many variants do not contain convolutional layers. Here we specifically consider the convolutional ViTs, including ScalableViT, XCiT, Visformer, CrossVit, NesT, CaiT, DeiT, BEiT, PiT, and ConViT. Performance of these ViTs on CIFAR10 and CIFAR100 are listed in Appendix D for a single-epoch DP training and several ViTs already beat previous SOTA, even though we do not apply additional techniques as in (e.g. learning rate schedule or random data augmentation).
By training multiple epochs with best performing ViTs in Table 8 and Table 9, we achieve new SOTA under DP in Table 5, with substantial improvement especially for strong privacy guarantee (e.g. ). Our DP training is at most slower and more memory expensive than the non-private training, even on BEiT large, thus significantly improving the slowdown reported in .
Discussion
We have shown that DP training can be efficient for large CNNs and ViTs with convolutional layers. For example, in comparison to non-private training, we reduce the training time to and the memory overhead to for all vision models examined (up to 303.4 million parameters), including BEiT that achieves SOTA accuracy on CIFAR100 ( absolutely at ). We have observed that for many tasks and large CNNs and ViTs, the memory overhead of DP training can be as low as less than 1%.
We emphasize that our DP training only improves the efficiency, not affecting the accuracy, and therefore is generally applicable, e.g. with SOTA data augmentations in . With efficient training algorithms, we look forward to applying DP CNNs to generation tasks , seq-to-seq learning , text classification , reinforcement learning , and multi-modal learning. Further reducing time complexity and prioritizing speed in DP training is another future direction.
In particular, our layerwise decision principle in (4.1) highlights the advantages of ghost clipping when is small. This advocates the use of large kernel sizes in DP learning, as they shrink the hidden feature aggressively, and have been shown to be highly accurate on non-private tasks .
References
Appendix A Further memory and speed comparison
Appendix B Explaining convolutional layers
In a 2D convolutional layer, the input to the layer has dimension and the folded output has dimension , where is the batch size, is the number of input channels, and is the number of output channels. are the height and width of images (or hidden features in hidden layers). can be calculated by https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html as
Following the above formulae, we recall that in the layerwise decision of mixed ghost clipping (3), the kernel size increases the right hand side and decreases the left hand size. In words, large kernel size always favors the ghost norm over the per-sample gradient instantiation!
To further explain the convolution, we consider the kernel size to (2.5), which establishes the equivalence between linear layer and convolutional layer. See example in https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html.
Appendix C Complexity analysis
In this section, we analyze the time and space complexity of different modules in the DP training pipeline. Our analysis follows a per-layer fashion, as all the dimension constants are layer-specific but ignored only in this section.
We note that the activation is created during the forward pass, and that is created during the random intialization. Since we only initialize and forward pass once, this complexity is the same for all training procedures (DP or non-private), therefore we do not study it. We also omit some trivial operations such as converting from per-sample gradient norm to the clipping factor .
In what follows, we will use the complexity of matrix multiplication repeatedly.
C.2 Back-propagation
Referring to the back-propagation in (2.3), we derive the whole time complexity is the matrix multiplication of , which is and the space complexity is .
The last step (2.6) results in the per-sample gradients, where
which gives time complexity and space complexity (since per-sample gradients are summed in-place).
In total we have time complexity and space complexity for one back-propagation.
For the second round of back-propagation, we add another time complexity but no space complexity as the space is freed by .
C.3 Ghost norm
In this section we study the procedure of computing the ghost norm. That is, from inputs to the output .
As mentioned in (2.7), the clipping norm can be calculated as following:
The final vector-vector product for a batch takes the time complexity is and space complexity .
C.4 Gradient instantiation and the norm
In this section we study the procedure of computing norm via instantiating the per-sample gradients. That is, from inputs , to the intermediate , to the output .
To compute the per-sample gradients, which is not available in the first back-propagation due to the in-place summation, we need to re-compute
For a batch, the time complexity is and the space complexity is .
To calculate the norm of , each with size , the time complexity is and the space complexity is .
C.5 Weighted gradient
To calculate the weighted gradient, , the time complexity is with no space complexity.
C.6 Combining the modules to algorithms
Ghost clipping = Back-propagation + Ghost norm + Second back-propagation
Opacus = Back-propagation + Gradient instantiation + Weighted gradient
FastGradClip = Back-propagation + Gradient instantiation + Second back-propagation
Mixed ghost clipping = Back-propagation + min{Ghost norm, Gradient instantiation} + Second back-propagation
Appendix D ViT details
Our ViTs are imported from PyTorch Image Models . For all ViTs, if they contain the batch normalization, we replace with the group normalization (16 groups). We freeze modules that are not supported by our privacy engine. We do not apply learning rate schedule, random data augmentation, weight standardization, or parameter averaging as in . We describe the models as their configuration argument in .
Appendix E Demo of privacy engine
We demonstrate how to use our privacy engine to train any vision models differentially privately. We term our library as private_vision, which is significantly based on the private_transformers library at https://github.com/lxuechen/private-transformers. We provide two modes through the ‘mode’ argument in the privacy engine: ‘ghost-mixed’ for the mixed ghost clipping, and ‘ghost’ for the ghost clipping.
A special use of our privacy engine is to use the gradient accumulation. This is achieved with virtual step function.
Appendix F Comparison with GhostClip in [33]
We give a thorough comparison between our work and (specifically codebase v0.1.0 which was the public version during the preparation of this paper), which distinguishes our contribution from a simple application of ghost clipping on convolutional layers.
Our contribution is on Conv1d/2d/3d layers, while applies the ghost clipping on linear and embedding layers. To be specific, we show that is layer-dependent (which motivates the layerwise decision in (4.1)), while studies sequential data and is layer-independent. We also precisely quantifies the effect of kernel size/padding/stride on the complexity in DP training in Appendix B.
We provide a fine-grained complexity analysis of the clipping (see Section 4.1), while shows only asymptotic complexity. For example, we show that the space complexity of ghost norm technique is and that of per-sample gradient instantiation is . In contrast, gives and , respectively. We highlight that our mixed ghost clipping, or the layerwise decision (4.1), is only made possible through our complexity analysis.
We additionally analyze the complexity of entire DP algorithms – e.g. Opacus, FastGradClip, and GhostClip, while only focuses on the clipping part of algorithms. Thus their analysis cannot directly help us to compare different DP algorithms, which not only include the clipping but also the back-propagation. Notice that ghost clipping needs two back-propagation but Opacus only needs one back-propagation, so it is insufficient to study the complexity difference between DP algorithms by only looking at the complexity of the clipping part.
Our key contribution is the mixed ghost clipping, which is novel, simple, but extremely important on large image tasks. Our mixed ghost clipping is much more efficient than the vanilla ghost clipping, as visualized in Table 3, Figure 3 and especially LABEL:tab:imagenet (on ImageNet). As a concrete example on ImageNet, ghost clipping incurs huge memory cost on most models (e.g. ResNet18, more than 16GB and thus OOM), while mixed ghost clipping costs only 2.34GB memory for ResNet18 and 7.91GB for ResNet152, almost the same as non-DP training.