Large-Scale Differentially Private BERT
Rohan Anil, Badih Ghazi, Vineet Gupta, Ravi Kumar, Pasin Manurangsi · 2022
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 = 5, which is a reasonable privacy setting.To put this number in perspective, non-private BERT models achieve an accuracy of ∼70%.