Posted by Yangsibo Huang, Research Intern, Google Research; Chiyuan Zhang, Research Scientist, Google Research
Large embedding models have become a crucial tool in recommendation systems and natural language processing. These models allow non-numerical data to be integrated into deep learning models by mapping categorical or string-valued input attributes to fixed-length representation vectors using embedding layers. They are widely used in personalized recommendation systems and achieve state-of-the-art performance in language tasks.
Privacy is an important consideration when deploying these models, and various techniques have been proposed to enable private data analysis. One widely adopted method is differential privacy (DP), which limits the exposure of individual user information while still allowing for population-level analysis. DP-SGD is the most commonly used algorithm for training deep neural networks with DP guarantees. However, applying DP-SGD to large embedding models creates scalability challenges because it destroys gradient sparsity, which is essential for efficient training.
To address this issue, we propose a new algorithm called DP-AdaFEST in our paper “Sparsity-Preserving Differentially Private Training of Large Embedding Models” (to be presented at NeurIPS 2023). DP-AdaFEST maintains gradient sparsity by selectively adding noise to a subset of feature rows at each iteration. The key is to make these selections differentially private, achieving a balance between privacy cost, training efficiency, and model utility. Our empirical evaluation shows that DP-AdaFEST significantly reduces gradient size compared to standard DP-SGD while maintaining comparable levels of accuracy. This reduction in gradient size can lead to a 20X improvement in wall-clock time.
To explain the challenges of gradient sparsity and our solution, let’s start with an overview of how DP-SGD works. DP-SGD clips the gradient contribution from each example in a mini-batch and adds Gaussian noise to the average gradient during each iteration of stochastic gradient descent. However, applying DP-SGD to large embedding models is challenging due to non-numerical feature fields and the transformation of words and tokens into dense vectors through an embedding layer. The vocabulary sizes of these features require large embedding tables, but the gradient updates are usually extremely sparse. DP-SGD destroys gradient sparsity by adding noise to all coordinates, which hinders private training of large embedding models.
Our algorithm extends standard DP-SGD with a mechanism to privately select “hot features” that are activated by multiple training examples in the mini-batch. This mechanism computes the contribution count of each feature bucket, clips their counts, adds Gaussian noise, and selects only the features with counts above a given threshold for the gradient update. The mechanism is differentially private, and the privacy cost can be computed by combining it with the standard DP-SGD iterations.
The theoretical motivation behind DP-AdaFEST is that it reduces variance at the cost of slightly increasing bias compared to DP-SGD. DP-AdaFEST adds noise to a smaller set of coordinates, preserving gradient sparsity, while DP-SGD adds noise to all coordinates. We provide more details in Section 3.4 of our paper.
We evaluate the effectiveness of DP-AdaFEST on public datasets, including an ad prediction dataset (Criteo-Kaggle) and a language understanding dataset (SST-2). Our algorithm achieves significantly higher gradient size reduction while maintaining comparable utility compared to the baseline DP-SGD with exponential selection. On the Criteo-Kaggle dataset, DP-AdaFEST reduces gradient computation cost by over 5×105 times while maintaining a comparable AUC. In language tasks, where the vocabulary is smaller and already compact, DP-AdaFEST effectively eliminates dense gradient computation.
In conclusion, our proposed algorithm, DP-AdaFEST, addresses the gradient sparsity problem in private training of large embedding models. It achieves significantly sparser gradients while maintaining comparable levels of accuracy, leading to more efficient and cost-effective training.
Source link
