Towards Efficient and Scalable Implementation of Differentially Private Deep Learning

paper starts from the issue of weaken privacy guarantees in current improper implementations of Poisson subsampling. However, the rest of the paper discuss the computation cost of DP-SGD, which is unrelated.

Very recent work [3,4] (ICML 2024/NeurIPS 2024) has shown that basing the accounting on Poisson subsampling while implementing a different sampling strategy can lead to a significant difference in the actual privacy loss compared to the accounted one. Especially noteworthy is Figure 3 (middle) of [4] which illustrates the required $noisemultiplier$ for the different sampling mechanisms with the Poisson subsampling $noisemultiplier$ being constantly lower than the others. Furthermore, [4] shows that Poisson subsampling is crucial in achieving the optimal utility under DP (Figure 3, left) while other sampling mechanisms can be beneficial in the non-DP setting (Figure 3, right). Another very recent paper [5] (preprint available on 15 November) also showcases this problem in their Figure 1. We already cited the ICML paper [3] but we added a citation to the NeurIPS paper [4] and the arXiv pre-print [5] in the revision.

Implementing Poisson subsampling is not trivial as Poisson subsampling results in varying (logical) batch sizes at every step of DP-SGD. This is especially challenging for the JAX implementations as changing tensor sizes (can) trigger recompilations as also mentioned by [4]. We believe that this is the reason why prior JAX implementations do not correctly implement the Poisson subsampling but instead use other sampling methods while using the privacy accounting for Poisson subsampling. We would like to clarify that none of the original versions of the compared implementations apart from Opacus correctly implements Poisson subsampling and that we add the Poisson subsampling to these to allow for a fair comparison of throughput and maximum possible physical batch size. Furthermore, we propose the masked DP-SGD algorithm as an alternative to the naive implementation of Poisson subsampling in JAX that is prone to recompilations.

In Sec. 2.1, the paper quickly jumps into Virtual batching and discuss some detailed batch sizes (25000, 300, etc), which seems to be more appropriate in Sec. 3. How are paragraphs of Sec. 2.1 relate to each others?

Thanks for the feedback. We restructured the Sections 2 and 3 to improve the readability of them and hope that these changes address your criticism. The mentioned batch sizes are for illustrative purposes and we added a “For example” to the sentence to make that clearer.

First, the author acknowledged that their current implementation is incomplete and does not support all network types yet.

We believe there is a misunderstanding as some of the efficient algorithms like GhostClipping do not support all layers, but this is an algorithmic limitation not something that can be trivially fixed. We do not propose these algorithms but benchmark all of them under the requirement of using Poisson subsampling. We would like to highlight that our own contribution (masked DP-SGD implementation in JAX) should support all layers and we are not aware of any limitations regarding that.

Second, using JAX or lower precision for speed up training is very common and too naive, even in other generic non-private training

The implementation using Poisson subsampling is not straightforward as the naive implementation in JAX suffers from recompilations (see also [4] and that can be overcome using our proposed method masked DP-SGD).

The authors need to present their proposed implementation and explain how, whether it insures correct privacy guarantee of Poisson subsampling,

Based on your feedback and the feedback of other reviewers we improved the writing of Section 2 and 3 and added a short paragraph about our proposed method. We agree with you that this was not written clearly enough in the first version of the paper.

How algorithms are optimized in Sec. 5.2 to decrease the cost? From the writing, it seems that the way to "optimize" is just using a better GPU (A100 instead of V100).

Thanks to the reviewer for the comments. We understand that there is some confusion about it. We changed the Section 5.1 title from Optimized gradient clipping algorithms, to Efficient gradient clipping algorithms, to highlight that the algorithms (Ghost Clipping, Mix Ghost Clipping, Opt Mix) by themselves are an efficient way of calculating the gradients and present an alternative to the Opacus implementation, but that we are not optimizing them.

[3] Chua, L., Ghazi, B., Kamath, P., Kumar, R., Manurangsi, P., Sinha, A., & Zhang, C. How Private are DP-SGD Implementations?. ICML 2024

[4] Chua, L., Ghazi, B., Kamath, P., Kumar, R., Manurangsi, P., Sinha, A., & Zhang, C. Scalable DP-SGD: Shuffling vs. Poisson Subsampling. NeurIPS 2024

[5] Annamalai, MSMS, Balle, B, De Cristofaro, E, Hayes, J. To Shuffle or not to Shuffle: Auditing DP-SGD with Shuffling. arXiv:2411.10614

In the caption for figure 2, is there a typo when saying lower throughput is better?

Thanks for pointing out this unclarity! We clarified this in the paper and renamed the y-axis-label to “Relative Slowdown” and changed the caption of Figure 2 to: “Relative slowdown in mean throughputs between Opacus per-example clipping and the non-private baseline using one A100 GPU. It is defined as private-throughput/non-private-throughput, this means that lower is better. It shows how many times private training is more expensive with a relative slowdown of 1 indicating that Opacus is as fast as non-private training.”

In lines 399 to 400, could you elaborate on why larger models are too expensive (to have benefit from lower precision)?

Thanks for the question! We clarified this point in the respective discussion of Figure 5: ”Models that are too small do not gain much from TF32, and the larger ones have too little maximum physical batch size to benefit (See detailed discussion of this in Appendix C).”

In Appendix C, we write: ’The speedup observed in Figure 5 peaks at the ”base” model. We believe that the reasons are the following: Speed-ups resulting from TF32 can significantly vary on per case basis as “all storage in memory and other operations remain completely in FP32, only convolutions and matrix-multiplications convert their inputs to TF32 right before multiplication.” (Stosic & Micikevicius, 2021). Until now, TF32 precision benchmarks have been limited to non-DP applications which was one of the reasons we wanted to discuss our observations in DP context. It appears the effectiveness of TF32 arithmetic peaks at “base” configuration. This due to a mix of reasons which are difficult to quantify exactly. Firstly, it is likely that matrix multiplication kernel dominance peaks at this configuration i.e. we have the most parameters whilst the batch size dimension also remains sufficiently large. With large and huge model variants the parameter count still increases but at the cost having very small batch dimension of 10 and 3, respectively. Secondly, we observe similar trend in Figure 2(a) where the discrepancy between dp and non-dp grows as model size gets bigger. This suggests that the dominance of DP operations also grows with the model size. None of the DP-operations are cast as matrix-multiplications and hence won’t benefit from TF32.’

In the JAX paragraph of section 2.2, could you elaborate on how masking achieves poisson sampling?

Based on your feedback and the feedback of other reviewers we improved the writing of Section 2 and 3 and added a short paragraph about our proposed method. We agree with you that this was not written clearly enough in the first version of the paper.

Could the authors provide some explanations on the novelty or difficulty in implementing the Poisson sampling?

Poisson subsampling results in varying (logical) batch sizes at every step of DP-SGD. This is especially challenging for the JAX implementations as changing tensor sizes (can) trigger time-consuming recompilations. We believe this is the reason why prior JAX implementations do not correctly implement the Poisson subsampling but instead use other sampling methods while using the privacy accounting for Poisson subsampling. Concurrent work that will be published at NeurIPS [1] has also acknowledged issues with JAX recompiling due to varying batch sizes and proposed a different but non-trivial solution.

Opacus backend does have optimized implementation of gradient clipping and ghost clipping is now supported as well, and torch compile should be able to improve speed and memory consumption in some cases.

We would like to thank the reviewer for the question. In our original submission we briefly (but not prominently enough) mention that we tried

• We tried ghost clipping but we ran into issues (Section 3). We made this now more prominent by mentioning it in Table 1 with an “?” and in the caption: ”Benchmarked DP-SGD frameworks and libraries. Note that Opacus implements Ghost Clipping but in our experiments the loss does not decrease, thus indicating a problem”. During the rebuttal period we installed the most recent version of Opacus from the git repository but still observe that the model is not learning when training with Ghost Clipping. In fact, while Ghost Clipping was released in version 1.5 of opacus [2] it seems that the maintainers are still resolving issues regarding it [3]. We believe that the findings regarding the throughput of Ghost Clipping will transfer between the implementation of the algorithm in the libraries that we used and Opacus and as the underlying algorithm is the same.
• we also tried the new torch compile but did not find any benefits with opacus (Section 6) as the speedups are minimal and the compilation time “eats” up these benefits while changing tensor sizes (e.g., through changes in the logical batch sizes) are also retriggering recompilations as in the naive JAX implementation. We will post throughput numbers ASAP.

Also when comparing with Opacus, the backend used by Opacus is not mentioned, Opacus have a few backend methods for efficient clipping, like functorch based implementation and backward hook used approach

Thanks for raising this point, we looked into this more closely and compared the throughput of all available grad_sample models (“hooks”, “functorch”, “ExpandedWeights” and “GhostClipping”). Functorch yields the same throughput as hooks (our original mode) while expanded weights crashes and GhostClipping does learn anything (see above). We added these details to the paper in Section 3 and Appendix A.3:

”We benchmark a native PyTorch (Ansel et al., 2024) implementation with PyTorch based libraries Opacus (Yousefpour et al., 2021) (details on gradsampling in Appendix A.3), PrivateVision (Bu et al., 2022), and FastDP (Bu et al., 2023), see Table 1.”

In Appendix A.3: *”Opacus supports multiple different gradient sampling methods as indicated in the documentation [4]. In our original experiments we used the grad sample mode hooks that is the default. This will use custom opacus modules when they are defined for that layer and functorch as a fallback. Based on the feedback by a reviewer we tried out different methods listed in the documentation for both ResNet and ViT models:

• $functorch$: We forced opacus to use functorch but did not observe any significant speed differences to using $hooks$. This is in line with the opacus documentation which writes that the speed is $0-50$% slower than $hooks$
• $ExpandedWeigths$: We tried using this approach but ran into runtime errors. Interestingly, when looking through the issues others have reported issues[5,6] but it seems to be more a PyTorch problem and has not been addressed for years. According to the documentation $ExpandedWeights$ is still in beta status.
• $GhostClipping$: Note that this method only works for ViT as described in Section 5.1. We did not manage to decrease the loss with this implementation due to the implementation in opacus being unstable but think that the speedups should be similar as observed in our experiments in Section 5.1 as the underlying algorithm is the same.”*

[1] Chua, L., Ghazi, B., Kamath, P., Kumar, R., Manurangsi, P., Sinha, A., & Zhang, C. Scalable DP-SGD: Shuffling vs. Poisson Subsampling. NeurIPS 2024

[2] Opacus v1.5 https://github.com/pytorch/opacus/releases/tag/v1.5

[3] Some bug fix in the implementation https://github.com/pytorch/opacus/pull/664

[4] https://github.com/pytorch/opacus/tree/61ae0ea4fb37a835e93040b5de19e8dfcd465a07/opacus/grad_sample

[5] https://github.com/pytorch/opacus/issues/584

[6] https://github.com/pytorch/opacus/issues/464

Summary of "lessons learnt"

Thanks for this suggestion! We added Table 4 to the Section 8 (Conclusion) with a brief summary of the lessons learnt.

whether this refers to, e.g., a single epoch

Thanks for the suggestion, we changed this in the contributions: “We find that non-optimized training with DP-SGD costs per-epoch between 2.6 and 3.2 times more than non-private training for ViT and 4 to 8 times for ResNets (See Section 4).”

In your experiments you use a privacy budget of eps=8, which is arguably on the upper end of what is usually understood to be a "private model". How does this choice impact the outcome of your experiments?

In our experiments, we only measure the throughput which we measure as processed examples per second and only use the test accuracy as a metric for validating that all approaches are implemented correctly. We agree that decreasing the epsilon will most likely have an effect on the optimal hyperparameters in terms of test accuracy but in terms of throughput only the (logical) batch size will have an effect on the throughput as it changes the number of steps per epoch but it is likely that the change is small [1,2].

The other hyperparameters like $learning rate$, $clipping bound$ and $noise multiplier$ will not have an effect as the computational cost does not change when modifying the values used for them. The $number of epochs$ won’t change the throughput as we divide the time by the number of processed samples.

In Table A3, you compare non-private models with optimal DP hyperparameters against DP models. Why? How is this a useful or fair comparison?

We agree with the reviewer that the optimal hyperparameters for non-DP training would differ from the optimal hyperparameters with DP but we focus on the computational cost of DP deep learning and want to showcase the cost of the additional operations done under DP regardless of the optimal hyperparameters that could be very problem specific.

One illustration of the difference between the optimality of hyperparameters can be found in concurrent work that will be published at NeurIPS [4] in their Figure 3 that showcases that large batch sizes with Poisson subsampling is not optimal for non-DP training in comparison to Poisson subsampling with large batch sizes being optimal for DP training.

What makes it difficult to implement Poisson subsampling?

Poisson subsampling results in varying (logical) batch sizes at every step of DP-SGD. This is especially challenging for the JAX implementations as changing tensor sizes (can) trigger recompilations. We believe this is the reason why prior JAX implementations do not correctly implement the Poisson subsampling but instead use other sampling methods while using the privacy accounting for Poisson subsampling. Concurrent work that will be published at NeurIPS [4] has also acknowledged issues with JAX recompiling due to varying batch sizes and proposed a different but less efficient solution for our setting.

While you say that Poisson subsampling is often not properly implemented, I understand that all the methods you compare against do in fact implement it. Is this the case? And if yes, why do you not compare against methods which do not implement Poisson subsampling?

We would like to clarify that none of the implementations but Opacus correctly implements Poisson subsampling and that we add the Poisson subsampling to these to allow for a fair comparison of the throughput and maximum possible physical batch size. Furthermore, we propose the masked DP-SGD algorithm as an alternative to the naive implementation of Poisson subsampling in JAX that is prone to recompilations. We do not compare to implementations that don’t use Poisson subsampling as very recent work (NeurIPS 2024) [4] shows that Poisson subsampling is crucial in achieving the optimal utility under DP (Figure 3, left) when all sampling mechanism are applied correctly while other sampling mechanisms can be beneficial in the non-DP setting (Figure 3, right). Furthermore, [3,4] ([3] being from ICML 2024) have shown that basing the accounting on Poisson subsampling while implementing a different sampling strategy can lead to a significant difference in the actual privacy loss compared to the accounted one. Especially noteworthy is Figure 3 (middle) of [4] which illustrates the required $noise_multiplier$ for the different sampling mechanisms with the Poisson subsampling $noisemultiplier$ being constantly lower than the others. Another very recent paper [5] (preprint got published on 15 November) also showcases this problem in their Figure 1. We already cited the ICML paper [3] but we will add a citation to the NeurIPS paper [4] and the arXiv preprint [5] in the revision of the paper.

These are the main changes in the revision. We added the reviewers that requested the changes.

• Changed Figure 2 axis label (L1Bo)
• Restructured Sections 2 and 3 and adding Table 1(gRw1, ZyeT)
• Added more details on our method masked DP-SGD (ZyeT, L1Bo)
• Added Table 4 with lessons learnt to Conclusion (reviewer gRw1)
• Cited very recent works on Poisson subsampling and privacy implications [1,2] (Support general discussions with most reviewers.)
• Appendix C and main text discussion of benefit of TF-32 for models of different sizes (L1Bo)
• Appendix A.3 and main text discussion of other clipping modes in Opacus (srbv)

[1] Chua, L., Ghazi, B., Kamath, P., Kumar, R., Manurangsi, P., Sinha, A., & Zhang, C. Scalable DP-SGD: Shuffling vs. Poisson Subsampling. NeurIPS 2024

[2] Annamalai, MSMS, Balle, B, De Cristofaro, E, Hayes, J. To Shuffle or not to Shuffle: Auditing DP-SGD with Shuffling. arXiv:2411.10614

Reviewer ZyeT brought to our attention that there might be some unclarity about the privacy/utility trade-off of the benchmarked methods and we want to clarify that the optimized algorithms do not have any impact on the trade-off unless the precision changes.

We would like to point the reviewer(s) to Table A2 in the Appendix where we have updated the utility numbers for fine-tuning a ViT Base on CIFAR-100 using all benchmarked approaches and showcase that there is no significant difference between the benchmarked methods. We additionally added pointers to the Table A2 in multiple places in the paper (Section 5.2 and 6) to make this point more clear.

All of the methods achieve a test accuracy of ~82% under $\epsilon=8$ and $\delta=1e-5$ when using FP32. This is expected as the optimizations do not change anything theoretical about the operations that could alter the utility/privacy trade-off but make the operations more computationally efficient. (The tiny differences might stem from non-deterministic changes due to the change of operations or CUDA optimizations).

In our experiments changing to TF32 did not change the accuracy in comparison to using FP32 which has also been observed in more comprehensive benchmarks in the non-private application [15].

Details on the change of numbers

All approaches now use the exactly same weights vit-base-patch16-224-in21k (PyTorch previously used a slightly better pre-training as the weights were vit_base_patch16_224.augreg_in21k but these seem to be unavailable in JAX) and furthermore all approaches zero the weights of the classification layer (We missed this earlier). Note that all of these changes do not affect the throughput or maximum physical batch size at all but only the utility.

Why is no change in utility/privacy expected

This is expected as the optimizations do not change anything theoretical about the operations that could alter the utility/privacy trade-off but make the operations more computationally efficient. One of the papers that we benchmark [14] also highlights this in their abstract where they write: “We propose an efficient and scalable implementation of this clipping on convolutional layers, termed as the mixed ghost clipping, that significantly eases the private training in terms of both time and space complexities, without affecting the accuracy.”

Our own method masked DP-SGD builds on-top of the same operations as typical DP-SGD but avoids recompilations. There is no change in terms of the clipping operation, noise addition or sampling but our proposed method enables the use of JAX for DP-SGD using Poisson subsampling and is more efficient than concurrent work that will be published at NeurIPS 2024 [4] as discussed with reviewer ZyeT and showcase in Appendix D.

What is the difference between using Poisson subsampling and not Poisson subsampling in terms of utility/privacy trade-off

We cited multiple papers [3,4,5] in the revision that discuss the importance of Poisson subsampling for achieving the optimal privacy/utility trade-off.

[3] Chua, L., Ghazi, B., Kamath, P., Kumar, R., Manurangsi, P., Sinha, A., & Zhang, C. How Private are DP-SGD Implementations?. ICML 2024

[4] Chua, L., Ghazi, B., Kamath, P., Kumar, R., Manurangsi, P., Sinha, A., & Zhang, C. Scalable DP-SGD: Shuffling vs. Poisson Subsampling. NeurIPS 2024

[5] Annamalai, MSMS, Balle, B, De Cristofaro, E, Hayes, J. To Shuffle or not to Shuffle: Auditing DP-SGD with Shuffling. arXiv:2411.10614

[14] Bu, Z., Mao, J., & Xu, S. Scalable and efficient training of large convolutional neural networks with differential privacy. NeurIPS 2022.

[15] Dusan Stosic and Paulius Micikevicius. Accelerating AI training with NVIDIA TF32 tensor cores. https://developer.nvidia.com/blog/accelerating-ai-training-with-tf32-tensor-cores/, 2021