Source-linked AI summary
Large-Scale Differentially Private BERT
Rohan Anil, Badih Ghazi, Vineet Gupta, Ravi Kumar, Pasin Manurangsi
TL;DR
The paper addresses the accuracy and efficiency challenges of training BERT-Large with differentially private SGD. It uses mega-batches and an increasing batch size schedule, achieving 60.5% masked language modeling accuracy at a 2M batch size while reducing training examples by up to 14%.
Problem
DP-SGD training faces substantial accuracy gaps and per-example gradient-clipping overhead, especially for complex tasks.
Method
The paper implements DP-SGD for BERT-Large using million-example mega-batches, gradient-SNR insights, and an increasing batch size schedule.
Results
60.5% masked language modeling accuracy is achieved with DP at a batch size of 2M.
Takeaways & Limitations
Mega-batches improve DP-SGD utility, while increasing batch sizes reduce examples seen by up to 14% while matching fixed-batch accuracy.
Takeaways & Limitations
The experiments use a fixed step budget, while fixed-epoch studies at very large batch sizes require hyperparameter retuning that remains an open challenge.
Abstract
from arXiv · showhide
In this work, we study the large-scale pretraining of BERT-Large with differentially private SGD (DP-SGD). We show that combined with a careful implementation, scaling up the batch size to millions (i.e., mega-batches) improves the utility of the DP-SGD step for BERT; we also enhance its efficiency by using an increasing batch size schedule. Our implementation builds on the recent work of [SVK20], who demonstrated that the overhead of a DP-SGD step is minimized with effective use of JAX [BFH+18, FJL18] primitives in conjunction with the XLA compiler [XLA17]. Our implementation achieves a masked language model accuracy of 60.5% at a batch size of 2M, for $ε= 5.36$. To put this number in perspective, non-private BERT models achieve an accuracy of $\sim$70%.
1 Introduction
The paper addresses privacy, utility, and efficiency challenges in private machine learning by establishing a DP BERT-Large pretraining baseline. It reports that mega-batches and increasing batch sizes improve accuracy or efficiency while retaining strong masked-language-model performance.
- Motivation: Differential privacy addresses concerns that training data may leak through memorization in deployed machine-learning models.DP controls how much an algorithm’s output distribution changes when one training example is added or removed.
- Motivation: DP-SGD applies broadly but can incur substantial accuracy gaps and training overhead, especially on complex tasks.The cited example reports an accuracy gap of approximately 25% for CIFAR-10 under reasonable privacy settings.
- Contributions: The paper establishes a high-accuracy baseline for differentially private BERT-Large pretraining using a carefully implemented DP-SGD variant.The implementation is designed to train relatively quickly on state-of-the-art hardware.
- Contributions: Scaling batch sizes to approximately 2M improves the utility of each DP-SGD step and yields 60.5% masked-language-model accuracy.The approximately 2M batch size is 32× larger than previously used for non-private BERT training.
- Contributions: An increasing batch-size schedule improves training efficiency while matching the accuracy of fixed-batch training.Proof-of-concept experiments reduce the total number of examples visited by up to 14% to reach the same accuracy.
- Related work: The work builds on prior research showing that software advances can reduce DP-SGD overhead and supports scaling private BERT pretraining.The paper positions prior JAX/XLA-based DP-SGD speedups as a foundation for this scaling effort.
2 Preliminaries
The preliminaries define neighboring datasets and differential privacy through a probability bound on algorithm outputs. They also note standard mathematical properties that motivate DP’s use in machine learning.
- Definitions: Neighboring datasets differ by adding or removing a single training example.This relation defines the datasets compared by the privacy guarantee.
- Definitions: An algorithm is (ε, δ)-differentially private when output probabilities on neighboring datasets satisfy a specified multiplicative-plus-additive bound.The bound is Pr[A(X) ∈ S] ≤ e^ε · Pr[A(X′) ∈ S] + δ for every output subset S.
- Properties: Differential privacy supports composition, post-processing, and group privacy.These properties are cited as mathematical reasons for DP’s significant interest in the literature.
3 Algorithm
The algorithm combines DP-SGD with Adam and weight decay: it clips per-example gradients, adds calibrated noise, and updates model parameters using noisy averages. Privacy across varying batch sizes is accounted for through composed Rényi differential privacy.
- Algorithm: The paper uses DP-SGD with Adam and leaves higher-order optimization methods for future work.Adam follows prior work showing effective tuning up to batch sizes of 65K.
- Algorithm: Each training step samples examples, clips their gradients, adds noise to the average gradient, and updates parameters with Adam and weight decay.The noise multiplier is computed using the moment accountant, also called Rényi differential privacy accounting.
- Algorithm: Algorithm 1 specifies DP-SGD with Adam, including batch sizes, clipping norm, noise multiplier, momentum, weight decay, and learning rate parameters.The procedure runs for a predetermined number of steps with potentially varying batch sizes.
- Privacy accounting: For varying batch sizes, each step receives an Rényi-DP cost that depends on its batch size and fixed noise multiplier.The per-step costs are composed and translated into an (ε, δ)-DP guarantee.
4 Experimental Setup and Tuning
The experiments pretrain BERT-Large privately on a large Wikipedia-and-Books corpus using JAX/FLAX and TPU hardware, with hyperparameters tuned at 32K batch size and transferred to other batch sizes. The setup specifically addresses scale-invariant layers, whose interaction with DP noise motivates larger weight decay.
- Experimental setup: BERT-Large uses 24 transformer blocks, 1024 hidden dimensions, 16 attention heads, and 340M parameters in a JAX/FLAX reimplementation run on TPUv3-1024.JAX was selected for JIT compilation and compiler fusion that reduce DP-SGD runtime overhead.
- Pretraining dataset: The pretraining data combines Wikipedia and BooksCorpus, containing 346M two-sentence examples with 128 tokens each and 20 masked tokens per example.The objective predicts masked tokens and sentence order; typical non-private pretraining reaches 70% MLM accuracy.
- Experimental setup and tuning: Hyperparameters are tuned with 30 Bayesian-optimization trials at batch size 32K, then transferred to other batch sizes using warmup followed by quadratic learning-rate decay.The tuning objective maximizes MLM accuracy over 10K examples, and the tuning itself is non-private.
- Scale-invariance and large weight decay: Layer normalization makes some layers scale-invariant, so increasing weight norms leaves outputs unchanged while shrinking their gradients inversely with weight norm.For a fully connected layer, the gradient is expressed through the loss gradient with respect to the preactivation and the input vector.
- Scale-invariance and large weight decay: DP noise tends to increase scale-invariant weights’ Frobenius norms and shrink their gradients, motivating substantially larger Adam weight decay than in non-private training.This interaction is relevant to normalization layers including layer normalization, batch normalization, and weight normalization.
- Scale-invariance and large weight decay: The straightforward gradient-projection remedy is difficult because scale invariance must be inferred beforehand, and BERT-Large embeddings are jointly but not individually scale-invariant.The embedding interaction arises because layer normalization is applied after aggregating wordpiece, positional, and token-type embeddings.
5 Experimental Results
The experiments show that privacy level and batch size strongly affect private BERT-Large training: larger batches improve gradient-SNR and accuracy, while an increasing schedule improves efficiency without reducing accuracy.
- 5.1 Varying privacy parameter ϵ: At ϵ = 1.08, MLM accuracy is 33.2%, rising to 42.85% at ϵ = 10.6; the experiments use ϵ = 5.36 thereafter.The privacy parameter is varied with batch size 65,536, 20K training steps, and 7,500 warmup steps.
- 5.2 Varying batch size under a fixed step budget: The batch-size study uses a fixed 20K-step budget because high-batch private training still requires efficient hyper-parameter retuning for fixed-epoch comparisons.Hyper-parameters are transferred from 32K tuning while ϵ remains fixed at 5.36.
- 5.2 Varying batch size under a fixed step budget: Larger batch sizes yield more favorable gradient-SNR during training and higher overall MLM accuracy.Gradient-SNR is the ratio between the aggregated-gradient norm and noise-vector norm.
- 5.2.2 Batch size schedule improves efficiency: An increasing schedule raises batch size from 262K to 1M over 7.5K steps, improves training efficiency, and matches fixed-batch accuracy while reducing examples seen by 14%.The schedule increases the batch size by 196K every quarter of the 7.5K-step interval.
- 5.3 On scaling up to mega batch sizes: Gradient accumulation enables a 2M batch size, producing 60.5% MLM accuracy with DP.The 2M batch is 32× larger than previously reported for non-private BERT training.
6 Conclusions
The work establishes a high-accuracy DP BERT-Large baseline by combining mega-batches, gradient-SNR analysis, normalization-aware training insights, and an increasing batch-size schedule.
- 6 Conclusions: The implementation establishes a DP BERT-Large baseline using mega-batches and additional training insights.The approach builds on recent software and hardware advances.
- 6 Conclusions: Mega-batches improve DP-SGD utility, while an increasing batch-size schedule improves efficiency and matches fixed-batch accuracy.The schedule yields up to a 14% reduction in examples visited to achieve the same accuracy.