SlowGuess/NeurIPS_2023_OCR huggingface dataset:
badihghazi@gmail.com
Yangsibo Huang*
yangsibo@princeton.edu
pritishk@google.com
ravi.k53@gmail.com
pasin@google.com
Amer Sinha
amersinha@google.com
Chiyuan Zhang
chiyuan@google.com
As the use of large embedding models in recommendation systems and language applications increases, concerns over user data privacy have also risen. DP-SGD, a training algorithm that combines differential privacy with stochastic gradient descent, has been the workhorse in protecting user privacy without compromising model accuracy by much. However, applying DP-SGD naively to embedding models can destroy gradient sparsity, leading to reduced training efficiency. To address this issue, we present two new algorithms, DP-FEST and DP-AdaFEST, that preserve gradient sparsity during private training of large embedding models. Our algorithms achieve substantial reductions  $(10^{6} \times)$  in gradient size, while maintaining comparable levels of accuracy, on benchmark real-world datasets.
Large embedding models have emerged as a fundamental tool for various applications in recommendation systems [CKH+16, WFFW17, GTY+17, ZZS+18, NMS+19, WSC+21] and natural language processing [MSC+13, PSM14, DCLT19]. Those models map categorical or string-valued input attributes with large vocabularies to fixed-length vector representations using embedding layers, which enable integrating non-numerical data into deep learning models. These models are widely deployed in personalized recommendation systems and achieve state-of-the-art performance in language tasks such as language modeling, sentiment analysis, and question-answering.
In the meantime, the use of large embedding models raises significant concerns about privacy violations, as it often entails the processing of sensitive individual information such as personally identifiable information, browsing habits, and dialogues [CTW+21, HSC22, GHZ+22]. Various techniques have been proposed to enable private data analysis. Among those, differential privacy (DP) [DMNS06, DKM+06] is a widely adopted notion that could formally bound the privacy leakage of individual user information while still allowing for the analysis of population-level patterns. For training deep neural networks with DP guarantees, the most widely used algorithm is DP-SGD [ACG+16]. It works by clipping the per-example gradient contribution, and noising the average gradient updates during each iteration of stochastic gradient descent (SGD). DP-SGD has
*This work was partially done while the author was interning at Google Research.
demonstrated its effectiveness in protecting user privacy while maintaining model utility in a variety of applications [PTS+21, LTLH21, DFAP21, TB21, KCS+22, YNB+22, DBH+22, DGK+23].
However, applying DP-SGD to large embedding models presents unique technical challenges. These models typically contain non-numerical feature fields like user/product IDs and categories, as well as words/ tokens that are transformed into dense vectors through an embedding layer. Due to the large vocabulary sizes of those features, the process requires embedding tables with a substantial number of parameters. In contrast to the number of parameters, the gradient updates are usually extremely sparse because each mini-batch of examples only activates a tiny fraction of embedding rows. This sparsity is heavily leveraged for industrial applications [GGDS19,  $\mathrm{WWL}^{+}22$ ] that efficiently handle the training of large-scale embeddings. For example, Google TPUs [AF22], a family of accelerators designed specifically for training large-scale deep networks, have dedicated hardware such as the SparseCore  $\left[\mathrm{JKL}^{+}23\right]$  to handle large embeddings with sparse updates. It is shown [Sna22] that this leads to significantly improved training throughput compared to training on GPUs, which did not have specialized optimization for sparse embedding lookups at the time. On the other hand, DP-SGD completely destroys the gradient sparsity as it requires adding independent Gaussian noise to all the coordinates. This creates a road block for private training of large embedding models as the training efficiency would be significantly reduced compared to non-private training.
Our contributions. In this paper, we present new algorithms that target the issue of diminished gradient sparsity encountered when employing DP-SGD with large embedding models due to the addition of dense noise; this issue represents a significant practical obstacle for leveraging, during DP training, the hardware accelerators commonly used in non-private training of large embedding models. While there are previous studies of the abstract problems of sparse vectors in DP, we study the practical problem of privately training large embedding models, and through systematic evaluations, demonstrate that significant gradient size reduction is achievable in DP training of real-world large-scale embedding models (e.g., recommender systems, language models with  $>30\mathrm{M}$  embedding parameters) with minimal loss in utility. Our contributions can be summarized as follows:
$\triangleright$  We propose two new algorithms to preserve gradient sparsity in DP-training of large embedding models while maintaining their utility (Section 3). The first algorithm, Filtering-Enabled Sparse Training (DP-FEST), selectively adds noise to pre-selected embedding rows during training. To improve the adaptivity to varying training dynamics across mini-batches and to accommodate online training, our second algorithm, Adaptive Filtering-Enabled Sparse Training (DP-AdaFEST), adaptively computes a gradient contribution map for each mini-batch, and only injects noise into the gradients of top contributors (in a manner that complies with the DP requirements).
$\triangleright$  We demonstrate the effectiveness of our algorithms on four benchmark datasets for online advertising and natural language understanding. Our results indicate that these algorithms can meet a wider range of application-specific demands for utility and efficiency trade-offs at a more granular level (Section 4.2). Notably, they achieve a substantially sparser gradient, with a reduction in gradient size of over  $10^{6} \times$  compared to the dense gradient produced by vanilla DP-SGD, while maintaining comparable levels of accuracy.
$\triangleright$  We further showcase the potential applicability and effectiveness of our algorithms to the setting of online learning with streaming data, a common setup in practice, which highlights the versatility and broad impact of our algorithms (Section 4.3).
2.1 Embedding Layers and Gradient Sparsity
Models with large embedding layers are widely used in application scenarios that need to process high-dimensional sparse inputs such as words/tokens in language models [MSC+13, PSM14, DCLT19] and categorical features associated with users or items in recommendation systems [CKH+16, WFFW17, GTY+17, ZZS+18, NMS+19, WSC+21].
Let  $c$  be the number of possible values of an input feature (aka vocabulary size in text models); the feature values are also called feature buckets (string-valued categorical features are commonly preprocessed via hashmap bucketization). An input  $\mathbf{x}$  with the  $i$ th feature value can be represented as a one-hot vector  $\mathbf{e}_i \in \mathbb{R}^c$  that is 1 on the  $i$ th coordinate and 0 everywhere else. Let  $\mathbf{W} \in \mathbb{R}^{c \times d}$  be the parameters of the embedding layer (i.e., the embedding table), where the output dimension  $d$  is the
(a) An illustration of embedding lookup.
(b) Embedding gradient sparsity in Ads model.
Figure 1: Illustration of embedding lookup (a) and gradient sparsity of the Criteo pCTR model (b). We analyze the gradient sparsity, averaged over 50 update steps, of the top five categorical features (out of a total of 26) with the highest number of buckets, as well as the sparsity of all categorical features. For more information about this dataset and model, see Section 4.1.1.
embedding size. The output of the embedding layer is computed as  $\mathbf{z} = \mathbf{W}^{\top}\mathbf{x}$ , which is a linear map. Due to efficiency considerations, in practice, embedding layers are rarely implemented via matrix multiplication (matmul) on one-hot encoded inputs.
For a one-hot input  $\mathbf{x} = \mathbf{e}_i$ , the embedding output is the  $i$ th row of the embedding table:  $\mathbf{z} = \mathbf{W}[i,:]$ . Figure 1a illustrates the embedding table lookup operation. Let  $\partial \mathcal{L} / \partial \mathbf{z} \in \mathbb{R}^d$  be the partial derivative of the loss with respect to the embedding output. The gradient of the embedding table is  $\nabla W = \mathbf{x} \otimes \partial \mathcal{L} / \partial \mathbf{z}$ , where  $\otimes$  is the outer product. For a one-hot input  $\mathbf{x} = \mathbf{e}_i$ , the result of this outer product is a sparse matrix whose  $i$ th row equals  $\partial \mathcal{L} / \partial \mathbf{z}$ . For mini-batch SGD training, the number of non-zero rows in the batch averaged gradient of the embedding table is upper bounded by the batch size, which is typically orders of magnitude smaller than  $c$ . Consequently, the gradient sparsity (the fraction of zero gradient coordinates) for large embedding models is very high (Figure 1b).
Due to this structured sparsity, in practice, both forward and backward computations of an embedding layer are implemented efficiently with gathers and scatters<sup>1</sup>, without expensive matmul. This difference from vanilla linear layers is crucial in real-world applications as the embedding tables are usually very large, with the vocabulary size  $c$  ranging from tens of thousands (language models) to millions (recommendation systems). Moreover, in some large recommendation system models, there are hundreds of different categorical features, each with a different embedding table. Therefore, maintaining the gradient sparsity is critical for any training process.
For notational simplicity, the description above focuses on single-variate features, which only activate one feature value at a time. In practice, multi-variate features activating multiple values are also used. In this case, it is common for the embedding layer to output the average or sum vectors corresponding to each activated value. We also note that in recommendation systems, each categorical feature generally has a separate embedding table; But for text models, all the words/tokens in an input document share the same embedding table.
2.2 Differentially Private Stochastic Gradient Descent
Differential privacy  $(DP)$  [DMNS06,  $\mathrm{DKM}^{+}06]$  is a mathematical framework for ensuring the privacy of individuals in datasets. It can provide a strong guarantee of privacy by allowing data to be analyzed without revealing sensitive information about any individual in the dataset. Formally, a randomized algorithm  $\mathcal{A}$  satisfies  $(\varepsilon, \delta)$ - $DP$  if for any two neighboring datasets  $\mathcal{D}$  and  $\mathcal{D}'$  (i.e., datasets such that one can be obtained from the other by adding/ removing one example), and any subset  $\mathcal{S}$  of outputs, it holds for privacy parameters  $\varepsilon \in \mathbb{R}_{>0}$  and  $\delta \in [0,1)$  that
$$
\Pr [ \mathcal {A} (\mathcal {D}) \in \mathcal {S} ] \leq e ^ {\varepsilon} \cdot \Pr [ \mathcal {A} (\mathcal {D} ^ {\prime}) \in \mathcal {S} ] + \delta .
$$
DP Stochastic Gradient Descent (DP-SGD)  $\left[\mathrm{ACG}^{+}16\right]$  is a recipe for training a deep learning model with DP by modifying the mini-batch stochastic optimization process through the use of per-example gradient clipping and Gaussian noise injection. When training an ML model  $f$  parameterized by  $\theta$  with the per-example loss function  $\ell (\cdot ,\cdot)^2$  on dataset  $\mathcal{D}$ , each optimization step  $t$  involves randomly
${}^{1}$  Gather/scatter refers to a memory addressing technique that enables simultaneous collection (gathering) from or storage (scattering) of data to multiple arbitrary indices. Refer to this wiki page for more details.
2The specific loss depends on the particular task and model (e.g. cross-entropy loss for classification).
sampling a mini-batch  $\mathcal{B}_t$ . Given  $\mathcal{B}_t$ , DP-SGD starts by computing the per-example gradient for each  $(x_i, y_i) \in \mathcal{B}_t$ , where  $x_i$  is the feature vector and  $y_i$  is the corresponding label, as follows:
$$
\mathbf {g} _ {t} (x _ {i}, y _ {i}) \gets \nabla_ {\theta_ {t}} \ell (f _ {\theta_ {t}} (x _ {i}), y _ {i}).
$$
It then clips the gradient  $\ell_2$ -norm to a maximum  $\ell_2$ -norm of  $C$  as:
$$
\left[ \mathbf {g} _ {t} \left(x _ {i}, y _ {i}\right) \right] _ {C} := \mathbf {g} _ {t} \left(x _ {i}, y _ {i}\right) / \max  \left(1, \frac {\| \mathbf {g} _ {t} \left(x _ {i} , y _ {i}\right) \| _ {2}}{C}\right).
$$
Finally, it produces the private gradient  $\hat{\mathbf{g}}_t$  by injecting Gaussian noise into the sum of the clipped per-example gradients as:
$$
\hat {\mathbf {g}} _ {t} \leftarrow \frac {1}{\| \mathcal {B} _ {t} \|} \left(\sum_ {i} \left[ \mathbf {g} _ {t} \left(x _ {i}, y _ {i}\right) \right] _ {C} + \mathcal {N} \left(0, \sigma^ {2} C ^ {2} \mathbf {I}\right)\right), \tag {1}
$$
where  $\mathcal{N}(0,\sigma^2 C^2\mathbf{I})$  is a Gaussian distribution with mean 0 and covariance  $\sigma^2 C^2\mathbf{I}$ , and the noise multiplier  $\sigma$  is computed from  $(\varepsilon, \delta)$  by inverse privacy accounting (e.g., [ACG+16]).
Although DP-SGD has been demonstrated to be an effective algorithm for training ML models with DP, (1) requires adding noise to all the coordinates of the gradient, which would completely destroy any existing sparsity structure in the original gradients. This densification of sparse gradients is especially problematic for large industrial-scale embedding models, where the sparsity is heavily leveraged to improve efficiency, e.g., using dedicated APIs [GGDS19, AF22,  $\mathrm{WWL^{+}22}$ , Sna22].
3 Sparsity-Preserving DP-SGD
In this section, we introduce two algorithms for preserving the gradient sparsity of DP-SGD when training large embedding models. Our main intuition is that in real-world datasets, some buckets of categorical features may be much more frequent, and therefore contain more significant or relevant information, than others. However, vanilla DP-SGD adds noise to the model's parameter gradients, even to those with minimal impact on the model's accuracy, which is not only wasteful but also can reduce the training efficiency by disrupting gradient sparsity (as discussed earlier in Section 2.2).
3.1 Filtering-Enabled Sparse Training (DP-FEST)
A simple solution to this problem is frequency filtering, where we pre-select the top-  $k$  most informative (i.e., frequent) buckets of each categorical feature before training, and add Gaussian noise only to the gradients of these selected buckets during training; this is akin to training a smaller embedding model using a subset of the buckets. Frequency filtering significantly reduces the noise added to the gradients by restricting the noise to only the most impactful subset of features.
There are two possible ways of selecting the top- $k$  buckets. First, if there is prior information on bucket frequency available publicly (e.g., token frequency in the pre-training set of the model for language tasks), then it can be used to select the top- $k$  buckets. If such prior knowledge is unavailable, we can run DP top- $k$  selection on the training dataset by adapting the algorithm of [DR21]; see Appendix B.1 for details.
3.2 Adaptive Filtering-Enabled Sparse Training (DP-AdaFEST)
The limitation of frequency filtering in DP-FEST is its inability to adapt to varying activation patterns across different mini-batches during training. This issue becomes particularly problematic when dealing with streaming data, such as those used in recommendation systems. In such scenarios, the application of frequency filtering poses a challenge as it impedes training until a substantial amount of data is gathered to accurately estimate the frequency information.
To address this issue, we propose DP-AdaFEST, an adaptive method for selecting the most informative buckets using mini-batch information, as outlined in Figure 2 and Algorithm 1. For each mini-batch, DP-AdaFEST first computes the per-example gradient and a gradient contribution map, which is a binary vector indicating which categorical feature buckets have been activated for each example in the mini-batch (Line 5). These per-example gradient contributions are then aggregated (with per-example clipping applied) across the mini-batch in order to construct the batch-wise gradient contribution,
Step 1: Compute contribution statistics for all buckets in the feature vocabulary.
Step 2: Clip the contribution and add Gaussian noise.
Step 3: Keep only the buckets that have a noisy contribution that exceeds  $\tau$  for gradient update (highlighted in purple).
Figure 2: Illustration of the process of DP-AdaFEST on a synthetic categorical feature which has 20 buckets. We compute the number of examples contributing to each bucket, adjust the value based on per-example total contributions (including those to other features), add Gaussian noise, and retain only those buckets with a noisy contribution exceeding the threshold for (noisy) gradient update.
Algorithm 1: DP-AdaFEST: Adaptive Filtering Applied Sparse Training.
Input: Learning rate  $\eta$ , batch size  $B$ , number of training steps  $T$ , dataset  $\mathcal{D}$ , clipping norms  $C_1, C_2$ , threshold parameter  $\tau$ , noise multipliers  $\sigma_1, \sigma_2$ . Output: DP-trained model parameters  $\theta$ . Initialize model parameters  $\theta$  randomly;  
for  $t = 1$  to  $T$  do  
Sample a mini-batch  $\mathcal{B}_t$  of size  $B$  from  $\mathcal{D}$ ;  
for  $i = 1$  to  $B$  do  
 $\begin{array}{r}\lfloor \end{array}$  Compute per-example gradient  $g_{i}$  and gradient contribution map  $v_{i}$ , where  $v_{i}[j] := 1[g_{i}[j,:]\neq 0]$ ;  
Compute the private contribution map  $V_{t} = \sum_{i\in \mathcal{B}_{t}}[v_{i}]_{C_{1}} + C_{1}\cdot \mathcal{N}(0,\sigma_{1}^{2}\cdot I_{c})$ ;  
for  $i = 1$  to  $B$  do  
 $\begin{array}{r}\lfloor g_i[j,:]\leftarrow 0,\end{array}$  for all  $j$  such that  $V_{t}[j] < \tau$ ;  
 $G_{t} = \sum_{i\in \mathcal{B}_{t}}[g_{i}]_{C_{2}} + C_{2}\cdot \mathcal{N}\left(0,\sigma_{2}^{2}\cdot \mathrm{diag}\left((1_{\{V_{t}[j]\geq \tau\}})_{j\in [c]}\otimes 1_{d}\right)\right);$   
Update the model parameters:  $\theta \gets \theta -\eta G_t$
and Gaussian noise with scale  $\sigma_{1}$  is added to the resulting vector to ensure DP (Line 6). The noisy batch-wise gradient contribution is then thresholded using a parameter  $\tau$  to exclude non-significant gradient entries to which only a few examples in the mini-batch contribute (Line 8); This thresholding mechanism helps focus on the most relevant and informative features while reducing the impact of noisy or insignificant contributions. Finally, the algorithm updates the model parameters by adding gradient noise with scale  $\sigma_{2}$  only to the "surviving" gradient entries (Lines 9 and 10). By dynamically adapting feature selection based on bucket contributions in each mini-batch during training, our algorithm achieves adaptive feature selection and maintains the desired DP guarantees. Note that we apply the standard DP-SGD with noise multiplier  $\sigma_{2}$  to update parameters in non-embedding layers.
3.3 Privacy Accounting
The privacy accounting of DP-AdaFEST can be done in a similar way to that of DP-SGD. The main observation is that the privacy cost of Algorithm 1 is dictated by the Gaussian noise mechanism steps in Lines 6 and 9. In particular, the privacy cost of a single iteration is equivalent to that of the composition of two Gaussian mechanisms with noise scale  $\sigma_{1}$  and  $\sigma_{2}$  respectively, which in turn is equivalent to the privacy cost of a single Gaussian mechanism of noise scale  $\sigma = (\sigma_{1}^{-2} + \sigma_{2}^{-2})^{-1/2}$ . Thus, overall, the privacy cost of the entire algorithm is no more than that of DP-SGD with noise scale  $\sigma$ , and identical other parameters of batch size and number of steps. Finally, we adopt the standard approach of analyzing the Poisson subsampled Gaussian mechanism [ACG+16], using numerical algorithms based on privacy loss distributions [KJH20, GLW21, GKKM22, DGK+22]. In particular, we use the open-source implementation available in Google's DP library [Goo20]. We provide more details in Appendix C for completeness.
3When implemented naively, this step has a memory consumption that scales with the number of embedding rows. In Appendix B.2, we discuss an alternative method for sampling the surviving coordinates of the mask in a memory-efficient manner.
An important caveat to the privacy analysis is that we follow a common practice of analyzing all the algorithms where the mini-batches are Poisson subsampled, wherein each example is included independently with a certain probability, even though, the actual implementation involves selecting mini-batches of fixed sizes after shuffling. See Section 4.3 of  $\left[\mathrm{PHK}^{+}23\right]$  for more context. But since this caveat applies to all algorithms similarly, we consider their relative comparison to be fair.
3.4 Bias-Variance Trade-offs
Apart from the computational advantage, to understand the question of when DP-AdaFEST would lead to more accurate models as compared to DP-SGD, we look through the lens of bias-variance trade-offs in stochastic convex optimization. Consider a convex objective  $\mathcal{L}(\cdot)$  over  $\mathbb{R}^D$ , for which we have a gradient oracle, that given  $\theta$ , returns a stochastic estimate  $g(\theta)$  of  $\nabla \mathcal{L}(\theta)$ . We say that the gradient oracle has bias  $\alpha$  and variance  $\sigma^2$  if  $g(\theta) = \nabla \mathcal{L}(\theta) + \zeta(\theta)$  such that  $\| \mathbb{E}\zeta(\theta)\|_2 \leq \alpha$  and  $\mathbb{E}\|\zeta(\theta) - \mathbb{E}\zeta(\theta)\|_2^2 \leq \sigma^2$  holds for all  $\theta$ . When optimizing over a convex set  $\mathcal{K} \subseteq \mathbb{R}^D$ , projected gradient descent with step size  $\eta$  is defined as iteratively performing  $\theta_{t+1} \gets \Pi_{\mathcal{K}}(\theta_t - \eta g(\theta_t))$ , where  $\Pi_{\mathcal{K}}(\cdot)$  is the projection onto  $\mathcal{K}$ . We recall the following guarantee on the expected excess loss obtained using standard analysis of projected stochastic gradient descent (see e.g. [Haz22]).
Lemma 3.1. For an  $L$ -Lipschitz loss function, and a gradient oracle with bias  $\alpha$  and variance  $\sigma^2$ , projected gradient descent over a set  $\mathcal{K}$  with diameter  $R$ , with step size  $\eta = \frac{R}{\sqrt{((L + \alpha)^2 + \sigma^2)T}}$  achieves
$$
\mathbb {E} \left[ \mathcal {L} \left(\frac {1}{T} \sum_ {i = 1} ^ {T} \theta_ {i}\right) \right] - \mathcal {L} (\boldsymbol {\theta} ^ {*}) \leq \frac {R}{\sqrt {T}} \cdot \sqrt {(L + \alpha) ^ {2} + \sigma^ {2}} + \alpha R.
$$
Suppose the loss is  $L$ -Lipschitz. Ignoring the potential bias introduced due to clipping, DP-SGD of noise scale  $\sigma$  uses a gradient oracle with zero bias  $(\alpha = 0)$  and variance  $D\sigma^2$ . On the other hand, consider a hypothetical setting where DP-AdaFEST truncates  $\gamma$  fraction of the gradient due to masking in Line 8, resulting in the final gradient that is supported on  $h$  coordinates  $(h \ll D)$ . In this case, the bias introduced is  $\gamma L$ , but the variance introduced is only  $h\sigma^2$ . Plugging into the above excess risk bound, we would expect DP-AdaFEST to achieve a smaller excess loss when
$$
\sqrt {L ^ {2} (1 + \gamma) ^ {2} + h \sigma^ {2}} + \gamma L \sqrt {T} <   \sqrt {L ^ {2} + D \sigma^ {2}}. \tag {2}
$$
This implies that a smaller truncated fraction within DP-AdaFEST could potentially yield superior utility compared to vanilla DP-SGD. Even though the discussion above is for the average iterate, and in the convex setting, our experiments corroborate this observation in training deep learning models.