Large-Scale Differentially Private BERT (2022)

Abstract: In this work, we study the large-scale pretraining of BERT-Large (Devlin et al., 2019) 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 the training efficiency by using an increasing batch size schedule. Our implementation builds on the recent work of Subramani et al (2020), who demonstrated that the overhead of a DP-SGD step is minimized with effective use of JAX (Bradbury et al., 2018; Frostig et al., 2018) primitives in conjunction with the XLA compiler (XLA team and collaborators, 2017). Our implementation achieves a masked language model accuracy of 60.5% at a batch size of 2M, for epsilon=5, which is a reasonable privacy setting. To put this number in perspective, non-private BERT models achieve an accuracy of ∼70%.

## introduction the widespread deployment of machine learning in recent years has raised serious concerns about the privacy of users whose data is used during the training process (see, e.g., @xcite ). these concerns are exacerbated by the well-known memorization behavior exhibited by deep neural networks @xcite and in particular large language models @xcite . to mitigate these concerns, the framework and properties of differential privacy (dp) @xcite provide a compelling approach for rigorously controlling and preventing the leakage of sensitive user information present in the training dataset. loosely speaking, dp guarantees that the output distribution of a (randomized) algorithm does not noticeably change if a single training example is added or removed; this change is parameterized by two numbers ( , δ)-the smaller these parameter values, the more private the algorithm. we refer the reader to section 2 for the formal definition of dp, and to @xcite for a thorough overview. motivated by these concerns, there has been a significant body of work on training private ml models. notably, @xcite presented a generic recipe for training ml models with dp. while their dp-sgd framework is quite robust as it applies to arbitrary neural network architectures, it faces two challenges that have significantly limited its practical deployment: note that on simple tasks such as mnist digit classification, the accuracy of dp models is not too far from that of non-private models. however, for more complex tasks such as cifar-10, the accuracy gap is very large, ∼25% for reasonable privacy parameter settings. for even more complicated tasks such as cifar-100, the inefficiency of dp-sgd has for several years precluded the training of dp neural networks. these limitations have made the dp training of a complex language model such as the bidirectional encoder representation (bert) @xcite a daunting task. very recently @xcite tackled the challenge of fine-tuning bert with dp and relied on non-private pretraining on the combined wikipedia and bookscorpus @xcite datasets. in this work, we take a step further and consider the task of pretraining a bert-large model with dp. obtaining a pre-trained model with dp guarantee allows us to employ it for multiple downstream tasks without violating the privacy of the data used in pre-training 1 . pretraining bert-large, however, is a computationally intensive task even without privacy; with privacy, using dp-sgd to pretrain bert-large poses significantly more computational challenges. our work shows how to overcome these barriers. we present an implementation of a variant of dp-sgd that, surprisingly, can train (relatively) quickly on state-of-theart hardware and achieve good accuracy. ## preliminaries similar to most previous works on dp ml, we say that two datasets x and x are neighboring, if x results from adding or removing a single training example from x. definition 1 (differential privacy @xcite ). for any real @xmath0, 1], a randomized algorithm a is ( , δ)differentially private (dp) if for every pair x, x of neighboring datasets and every subset s of outputs of a, it is the case that dp has seen significant interest in the literature due to its nice mathematical properties such as composition, post-processing, and group privacy (see, e.g., the book of @xcite ). ## algorithm at a high-level, we use the dp-sgd algorithm @xcite with the adam optimizer @xcite . our choice of adam follows the work of @xcite , who showed that tuning adam works well up to large batch sizes of 65k. as our goal is to establish a baseline for bert pretraining with dp, we leave the investigation of higher-order methods @xcite to future work. at each step of training, we randomly select a prespecified number of examples, compute and clip their gradients and add appropriate noise to the average gradient to ensure privacy. to compute the noise multiplier, we use the privacy loss distribution (pld) method (e.g., @xcite ) implemented in google dp accounting library foot_1 . ## experimental setup and tuning we present results on pretraining a bert architecture @xcite , focusing on its larger variant, aka "bert-large", which is a transformer model @xcite containing 24 transformer blocks with 1024 hidden dimensions and 16 self attention heads. it has 340m parameters (1.297 gib). the training setup is a reimplementation of the official bert codebase 3 in the jax framework @xcite with the flax library @xcite . our choice of jax was motivated by the recent results of @xcite , where they demonstrate that jax features such as just-in-time (jit) compilation and compiler fusion result in low runtime overheads for the dp-sgd step and match or commonly exceeds the performance of various other frameworks. all experiments were carried out on google tpus @xcite , using the tpuv3-1024 configuration. ## experimental results in this section, we study the effect of the privacy parameter and fix a particular value of to aim for the highest attainable mlm accuracy, by changing the batch size. ## conclusions we build off the recent advances in software @xcite @xcite and hardware @xcite and establish a baseline for bert-large pretraining with dp. we achieve high accuracy for the model by scaling up the batch size to millions of examples (mega-batches) and using additional insights such as improving the trainability of networks under normalization layers and measuring the gradient-snr metric. we proposed a proof of concept batch size increasing schedule and demonstrate an efficiency improvement. an interesting direction is to improve the efficiency of dp training further by leveraging recent advances such as higher-order methods @xcite , local loss optimization @xcite , scaling techniques that exploit increased parallelism: such as online distillation @xcite and ones that use lower memory @xcite , automatic tuning of hyper-parameters meta optimization @xcite , more efficient architectures such as @xcite for longer sequences, multistep training at mega-batch sizes @xcite , loss functions that are robust to label noise @xcite .