Source-linked AI summary
Opacus: User-Friendly Differential Privacy Library in PyTorch
Ashkan Yousefpour, Igor Shilov, Alexandre Sablayrolles, Davide Testuggine, Karthik Prasad, Mani Malek, John Nguyen, Sayan Ghosh, Akash Bharadwaj, Jessica Zhao, Graham Cormode, Ilya Mironov
TL;DR
DP-SGD requires per-sample gradients, but micro-batching is slow and conventional frameworks do not expose the needed intermediate computations. Opacus provides a PyTorch library using vectorized per-sample-gradient computation, a simple and flexible API, and broad layer support. It outperforms existing DP frameworks overall while remaining competitive with custom JIT-compiled implementations, with a speed–memory trade-off addressed by virtual steps.
Problem
DP-SGD needs per-sample gradients, but standard frameworks expose only batch-averaged gradients and micro-batching can underutilize hardware accelerators.
Method
Opacus uses vectorized batched per-sample-gradient computation within a simple, flexible PyTorch library supporting private training across diverse model layers.
Results
Opacus outperforms existing differential-privacy frameworks overall and performs competitively with custom JIT-compiled DP-SGD implementations across models and datasets.
Takeaways & Limitations
Opacus supports practical private PyTorch workflows with two-line integration, broad predefined and custom-layer support, and improved efficiency over micro-batching.
Takeaways & Limitations
Opacus’s vectorized per-sample-gradient computation creates a speed–memory trade-off because stored gradient tensors grow with batch_size.
Abstract
from arXiv · showhide
We introduce Opacus, a free, open-source PyTorch library for training deep learning models with differential privacy (hosted at opacus.ai). Opacus is designed for simplicity, flexibility, and speed. It provides a simple and user-friendly API, and enables machine learning practitioners to make a training pipeline private by adding as little as two lines to their code. It supports a wide variety of layers, including multi-head attention, convolution, LSTM, GRU (and generic RNN), and embedding, right out of the box and provides the means for supporting other user-defined layers. Opacus computes batched per-sample gradients, providing higher efficiency compared to the traditional "micro batch" approach. In this paper we present Opacus, detail the principles that drove its implementation and unique features, and benchmark it against other frameworks for training models with differential privacy as well as standard PyTorch.
1 Background and Introduction
DP-SGD protects individual samples by computing, clipping, aggregating, and noising per-sample gradients, but standard frameworks expose only batch-averaged gradients. Opacus addresses this implementation gap with vectorized computation designed to improve speed while preserving usability and flexibility.
- Motivation: DP-SGD limits the influence of any individual sample by privatizing every model-parameter update.Its privacy guarantee is quantified by ϵ and δ.
- Motivation: Standard PyTorch and TensorFlow interfaces expose batch-averaged gradients rather than the per-sample gradients required by DP-SGD.This lack of access complicates direct implementation of private training.
- Motivation: Micro-batching computes correct per-sample gradients but can be very slow because it underutilizes GPUs and TPUs optimized for batched computation.The naïve procedure separates each batch into micro-batches of size one, then clips and noises the resulting gradients.
- Contribution: Opacus replaces micro-batching with performance-improving vectorized computation while emphasizing simplicity and flexibility.The paper evaluates Opacus against other DP-SGD frameworks and standard PyTorch.
2 Design Principles and Features
Opacus combines a compact PyTorch-facing API with vectorized private-gradient computation and broad extensibility. Its features cover privacy accounting, model validation, sampling, hardware utilization, distributed training, and support for predefined and custom layers, while virtual steps address its speed–memory trade-off.
- Design principles: Opacus’s design principles are simplicity, flexibility, and speed through a compact API, rapid prototyping support, and vectorized computation.The library is intended for researchers and engineers already using PyTorch and Python.
- API: PrivacyEngine.make_private() wraps a model, optimizer, and data loader into differentially private counterparts.The wrapped components compute per-sample gradients, clip and noise gradients, and use Poisson sampling.
- API: Users can add Opacus to an existing training script by changing as few as two lines of code, while advanced users can modify private components without changing their interface.This balances a simple default workflow with flexibility for customized behavior.
- Features: Opacus provides privacy-budget tracking, direct targeting of an (ϵ, δ) budget, model validation, Poisson sampling, accelerator-efficient execution, and distributed training.Its accountant supports monitoring and early stopping, while validation rejects layers that mix information across batch samples.
- Trade-offs: Opacus faces a speed–memory trade-off because storing per-sample gradients increases gradient-tensor size by the batch_size factor.Virtual steps decouple memory-limited physical batch sizes from convergence- and privacy-driven logical batch sizes.
- Features: Predefined support covers convolution, multi-head attention, LSTM, GRU, generic RNN, normalization, and embedding layers, with decorator-based registration for custom layers.Custom layers can supply their own per-sample-gradient calculation method.
3 Benchmarks
Opacus is benchmarked against other DP-SGD frameworks and standard PyTorch across end-to-end tasks, while layer-level experiments quantify runtime and memory overhead. Results show strong dependence on framework, batch size, layer type, and parameter-to-feature ratio.
- End-to-end benchmarks: Opacus is evaluated against JAX (DP), Custom TFP (XLA), BackPACK, PyVacy, and standard PyTorch across four end-to-end training tasks.The tasks use CNNs on MNIST and CIFAR-10 and embedding and LSTM networks on IMDb.
- End-to-end benchmarks: At batch size 2048, Opacus achieves the lowest runtime on CIFAR-10 and remains within 1.4× to 3× of JAX (DP) on the other listed tasks.On IMDb with the LSTM network, Opacus is 7× slower than Custom TFP (XLA) and 2.4× slower than JAX (DP).
- End-to-end benchmarks: 40× mean per-epoch runtime reduction results for Opacus when increasing batch size from 16 to 2048, compared with 12.8× for JAX (DP) and 6.3× for Custom TFP (XLA).PyTorch without DP shows a 37.8× reduction, BackPACK 10.5×, and PyVacy 1×.
- End-to-end benchmarks: Enabling DP with Opacus incurs 2× to 2.9× runtime overhead on MNIST, CIFAR-10, and IMDb with the embedding network, versus 25× to 30× on IMDb with the LSTM network.The paper attributes the larger LSTM overhead to the custom RNN-based modules and their interaction with GradSampleModule.
- Layer-level overhead: 1.2× to 2.9× runtime overhead occurs for convolutional, normalization, and multi-head attention layers, while linear and embedding layers reach 5.5× and 49× respectively.RNN-based layers reach up to 18× runtime overhead, which decreases as batch size increases.
- Layer-level overhead: 129× and 334× peak allocated memory overhead occurs for linear and embedding layers, whereas normalization, multi-head attention, and RNN-based layers reach up to 1.5×.The larger overheads arise when the parameter-to-feature ratio L/C is substantial relative to batch size.
4 Related Work
Prior work develops alternative gradient-clipping methods and privacy-learning frameworks, while Opacus is situated among approaches with differing layer support and privacy mechanisms.
- Goodfellow’s gradient-clipping method computes per-sample gradient norms but is restricted to fully connected layers, while Rochette et al. extend it to CNNs.
- Lee and Kifer compute per-sample gradient norms directly using two backpropagation passes, whereas Opacus obtains per-sample gradients in one pass.
- Ghost clipping extends the Goodfellow method to sequential inputs and is reported as memory-efficient with good throughput for large-Transformer fine-tuning.
- Existing frameworks include TensorFlow Privacy, PyVacy, BackPACK, and Objax, with BackPACK supporting fully connected and convolutional layers but not recurrent or residual layers.
- Secure hardware and secure multi-party computation provide alternative approaches to machine-learning privacy and security beyond DP-SGD.
5 Conclusions
Opacus targets simple, flexible, and fast differentially private deep-learning training in PyTorch. The paper reports performance advantages over existing DP frameworks and competitiveness with custom JIT-compiled implementations across models and datasets.
- Opacus is a PyTorch library designed for simplicity, flexibility, speed, and compatibility with existing machine-learning pipelines.
- Opacus outperforms existing differentially private training frameworks and performs competitively with custom JIT-compiled DP-SGD implementations across models and datasets.
- Opacus is actively maintained as an open-source project, with planned extensions for custom components, efficiency, and PyTorch ecosystem integration.
A Micro-Batching
Micro-batching computes per-sample gradients by processing each example separately, but Opacus instead derives vectorized layer formulas and implements them with batched tensor operations.
- A Micro-Batching: Micro-batching processes each training example separately to obtain per-sample gradients, then resets the model gradients between examples.
- B Vectorized Computation: Opacus computes per-sample gradients by deriving a formula for each layer and implementing it with a single einsum application.
- B Vectorized Computation: For a linear layer, the forward pass maps input matrix X to output matrix Y through Y = WX, with columns representing batch elements.
- B Vectorized Computation: Opacus requires per-sample derivatives of the loss with respect to weights rather than gradients averaged over the batch.
- B Vectorized Computation: The highway gradient is the derivative of the loss with respect to a hidden-layer activation for one batch element.
- B Vectorized Computation: Opacus replaces summed matrix multiplication with a batched outer product, using einsum for vectorized multi-dimensional linear algebra.
C Detection of DP Violations
Opacus checks model compatibility with differentially private training by validating supported layers and detecting batch-level operations or unprotected tracked statistics. These checks can be inspected alongside intermediate gradient states during training.
- C Detection of DP Violations: Opacus checks whether all model layers are supported and whether operations violate DP compatibility through batch-level computation or unprotected extra information.
- C Detection of DP Violations: Opacus disallows batch normalization because it shares information across samples and disallows tracked instance-normalization statistics outside DP guarantees.
- C Detection of DP Violations: The compatibility checks are not exhaustive because Opacus cannot verify sample independence or detect every extraneous statistic maintained by a model.
- C Detection of DP Violations: Opacus exposes normal, per-sample, clipped-and-aggregated, and final noisy gradients through fields such as grad_sample and summed_grad.
- C Detection of DP Violations: Calling optimizer.zero_grad() clears the gradient fields after the training step.
E.1 End-to-end benchmarks
The end-to-end benchmarks compare cumulative framework runtime over 20 epochs at batch size 512, highlighting the first-epoch cost of JIT compilation.
- JIT compilation results in a slower first epoch for frameworks that use it.
- 101× and 625×: JAX (DP) and Custom TFP (XLA) incur these first-epoch overheads relative to subsequent epochs because of JIT compilation.The comparison covers cumulative runtime over 20 epochs across end-to-end model training tasks.
E.2 Microbenchmarks
The microbenchmarks compare runtime and peak memory for supported recurrent and attention layers across standard PyTorch, Opacus modules, and differentially private Opacus modules.
- Figure 5 reports mean runtime in milliseconds and peak allocated memory in megabytes for four attention and recurrent layer types.
- Peak memory overhead is small where applicable and arises from GradSampleModule wrapping and collecting per-sample gradients.Opacus custom modules generally use slightly less memory than corresponding torch.nn modules.
- Tables 2 and 3 provide the raw runtime and memory data underlying the microbenchmark figures.Table 4 additionally breaks down CUDA memory usage and related quantities by layer and batch size.
E.3 Experiment Setup
The experiments use cloud-based NVIDIA A100 systems and report standardized runtime, memory, software-version, and layer-size measurements for the benchmarks.
- End-to-end benchmarks use an Intel Xeon CPU, NVIDIA A100 SXM4 with 40GB VRAM, and 83GB RAM.They run in separate Docker containers based on nvidia/cuda:11.4.2-cudnn8-devel-ubuntu20.04.
- Microbenchmarks use an Intel Xeon Platinum CPU, NVIDIA A100 SXM4 with 40GB VRAM, and 1.1TB RAM.CUDA memory is allocated in blocks of 512.
- Table 2 measures mean milliseconds for one forward and one backward pass, comparing torch.nn, Opacus, and GradSampleModule configurations.Its factor is DP-enabled layer runtime divided by torch.nn runtime.
- Table 3 measures peak allocated CUDA memory in megabytes for one forward and one backward pass across the same module configurations.Its factor is DP-enabled peak memory divided by torch.nn peak memory.
- Table 4 reports input, labels/output, and torch.nn module sizes by layer and batch size.The corresponding Opacus configurations are almost identical within ±1%, and C divides input plus twice label size by batch size.
- Table 5 lists software versions used in the end-to-end and microbenchmarks.The microbenchmarks use only Python, PyTorch, and Opacus among the listed software components.