Federated Learning of Gboard Language Models with Differential Privacy

Introduction

FL and Gboard LMs. In cross-device federated learning (FL), client devices collaboratively train a model without directly exchanging their local data [Kairouz et al. 2019]. Google Keyboard (Gboard) was an early adopter of FL to train models that improve the user experience, following data minimization principles [Bonawitz et al. 2021] to protect users’ privacy from some risks. Language models (LMs) are trained with FL to support various features in Gboard, including Next Word Prediction (NWP), Smart Compose (SC), and On-The-Fly rescoring (OTF). As illustrated in Fig. 1, NWP [Hard et al. 2018] uses an LM to suggest a word, which is triggered after a previous word is committed; SC provides longer inline suggestions to accelerate typing, which can be triggered per character when the confidence is high; OTF is used to re-rank the candidate words generated during typing before a word is committed.

We train LMs with the same neural network (NN) architecture described in [Hard et al. 2018]: a one-layer LSTM/CIFG of 670 hidden neurons, with input and output word-based embeddings of dimension 96. OTF LMs use a larger vocabulary ($\sim$ 30K words) compared to NWP LMs ($\sim$ 10–20K words); the number of parameters for models with a 10K/20K/30K vocabulary is 2.4M/4.4M/6.4M, respectively. SC is a downstream task that reuses NWP LMs without any retraining from data. We train NWP LMs and OTF LMs from populations of devices categorized by language and location. For example, en-US NWP denotes the task of training NWP model on data generated by devices using English in the United States.

Federated Averaging (FedAvg) [McMahan et al. 2017] and variants [Wang et al. 2021] are popular FL training algorithms in practice. In each communication round, the server will orchestrate a small subset of client devices for training and aggregate the resulting model deltas to update the global model. In a successful round, the system guarantees the number of clients participating in training is at least as large as the configured report goal [Bonawitz et al. 2019]. A model is typically tested and deployed after training for several thousands of rounds. Top-1 in-vocab accuracy is used to track the utility during training and additional metrics for A/B testing are introduced in Section 3.

DP and DP-FTRL.

Differential privacy (DP) can be combined with FL to provide a formal guarantee that the trained model will not memorize specific users’ data, which provides stronger privacy protection by executing data anonymization principles [Bonawitz et al. 2021; Wang et al. 2021]. Ramaswamy et al. 2020 applied DP-FedAvg [McMahan et al. 2018; Geyer et al. 2017], a variant of DP-SGD [Abadi et al. 2016] for user/client-level DP, to train production LMs in FL. Ramaswamy et al. 2020 demonstrated anonymization via empirical auditing techniques by Carlini et al. 2019 but did not provide a formal DP guarantee. Achieving a strong formal DP guarantee for DP-FedAvg would require privacy amplification-by-sampling, which necessitates sampling clients uniformly at random on each round. However, a cross-device FL system has limited control over client sampling as devices have to satisfy local criteria such as being charging and connected to an unmetered network to be eligible for participation [Bonawitz et al. 2019; Balle et al. 2020]. In contrast, we deploy a recent algorithm, DP-FTRL [Kairouz et al. 2021b], allowing us to achieve strong privacy and utility for production models without uniform sampling assumptions.

Contributions.

We discuss our strategy and experience of training Gboard LMs with FL and DP. We introduce an algorithm that enables adaptive clipping [Andrew et al. 2021] in DP-FTRL [Kairouz et al. 2021b] (Section 2.1), which can reliably estimate the clip norm to reduce hyperparameter tuning. We discuss the impact of scaling up computation and limiting client participation (Section 2.2), and identify the algorithm and system configurations for the regime of strong privacy and utility. We also successfully apply pre-training (Section 2.3) to improve privacy and utility, which is (to the best of our knowledge) the first time pretraining is applied to training a DP model directly from users’ data.

We combine DP-FTRL with secure aggregation (SecAgg) to further strengthen the data minimization properties of our approach (Section 2.4). Fig. 2 provides a system overview of the techniques for training Gboard language models with federated learning and differential privacy. Finally, we summarize concrete suggestions for practitioners training differentially private models to deploy in production in (Section 2.5), and present and analyze twenty Gboard LMs trained with formal DP guarantees (Section 3). We are happy to announce that all the next word prediction neural network LMs in Gboard now have DP guarantees, and all future launches of Gboard neural network LMs will require DP guarantees.

DP FL in Practice

As described in Algorithm 1, we apply DP-FTRL in FL by modifying the FedAvg algorithm: clip the model update $\Delta$, and add noise when updating the global model. Two additional hyperparameters are introduced for DP: the clip norm $C$, which bounds the norm of $\Delta$, and the noise multiplier $z$, which determines the standard deviation $zC$ for the added Gaussian noise. We discuss clip norm in this section and defer the discussion of noise multiplier and other privacy related hyperparameters to Section 2.2.

One step of DP-FTRL with adaptive clipping using $\sigma_{b}$ noise standard deviation on the clipped counts $\sum b_{i}^{t}$ and $z_{\Delta}$ noise multiplier on the vector sums $\sum\Delta_{i}^{t}$ is equivalent to one step of non-adaptive DP-FTRL with noise multiplier $z$ if we set $z_{\Delta}=\left(z^{-2}-(2\sigma_{b})^{-2}\right)^{-\nicefrac{{1}}{{2}}}$.

In practice, Algorithm 1 slightly inflates the noise for the model from $zC$ to $z_{\Delta}C$ and requires restarts that complicate the privacy accounting for DP-FTRL. Moreover, we find that a fixed clip norm can achieve comparable or slightly better model utility, and is more robust in experiments with large report goal. For example, adaptive clipping for the de-DE NWP model experiences catastrophic failure and makes no progress in the first 1000 rounds.

Nevertheless, adaptive clipping can reduce hyperparameter tuning for many tasks when privacy budget allows. Fig. 3 shows the evaluation accuracy and corresponding clip norm for DP training the en-GB NWP model with report goal 6500 and noise multiplier 7. The adaptive clip curve starts from a small initial clip norm to avoid catastrophic failure due to large initial noise and eventually catches up on accuracy. The estimated clip norm (quantile $\gamma=0.5$) stabilizes and we can fix the clip norm to $5$ based on the estimated value. The clip norm is relatively insensitive, especially when tuning together with the server learning rate. However, clip norm can have a wide tuning range across tasks and models, and quantile-based estimation is still useful for estimating a clip norm to be fixed.

2 DP parameters and system configuration

The privacy guarantees of DP-FTRL [Kairouz et al. 2021b] are affected by several factors: noise multiplier $z$, number of total rounds $T$, max participation (MaxP) of a client, and min separation (MinS) of rounds between the participation of the same client. The noise multiplier is a conventional parameter for controlling privacy-utility trade-off: large noise achieves strong privacy guarantees but can potentially hurt the utility. Achieving the same utility with smaller rounds $T$ can significantly improve the privacy guarantees. Next, we discuss the effect of MaxP and MinS, and the privacy-utility-computation trade-off for system configuration.

DP-FTRL achieves strong privacy if each client only participates once during training, or the number of client participation is limited when a client can participate multiple times. Two parameters are introduced to characterize client participation for DP-FTRL: the maximum participations (MaxP) of a client in all training rounds and the minimum round separation (MinS) between any single client’s two participations. MaxP and MinS are correlated as MaxP is upper bounded by rounds $T$ divided by MinS. In general, for fixed rounds $T$, decreasing MaxP and increasing MinS can lead to stronger privacy guarantees without changing utility. In addition, Cho et al. 2023 suggests potential advantage of increasing MinS for utility.

When using the worst-case MaxP estimated by rounds $T$ divided by MinS, Fig. 4(c) shows increasing MinS can achieve stronger privacy measured by smaller zCDP values. However, the maximum MinS is limited by the population size divided by the number of clients per round lower bounded by the report goal. For example, when the report goal is 6500 for small population of around $10^{6}$, MinS has to be smaller than $153$ rounds, so strong privacy guarantees are difficult to achieve when training for 3000 rounds. While we cannot measure the precise population size in the FL system due to client dynamics, we estimate the population size of various Gboard tasks as ranging from $0.8$ million to $16.6$ million in Table 1.

Report goal.

We study report goal for privacy-computation trade-off based on a hypothesis used in [McMahan et al. 2018; Kairouz et al. 2021b; Xu et al. 2022]: for sufficiently large data, the utility is approximately non-decreasing if the noise multiplier and clients per round (lower bounded by report goal) proportionally increase. We provide empirical justification to this hypothesis by comparing the evaluation accuracy of two training runs: one with a report goal of 500 and noise multiplier of 0.54, versus another of report goal 6500 and noise multiplier 7. On more than three Gboard language tasks, we observed that the final utility remains similar, or slightly better for larger report goals. Moreover, using a larger report goal speeds up learning at the beginning of training. Based on the hypothesis, we plot Figs. 4(a) and 4(b) by linearly increasing report goal and noise multiplier, and assuming the MinS is set to the maximum possible value (population divided by report goal) for strong privacy. Though a large report goal can limit the MinS, it generally leads to stronger privacy guarantees for reasonable population size and total rounds.

System configuration.

According to Figs. 4(a) and 4(b), we choose a large report goal of 6500 supported by the large scale FL systems and aim for maximum MinS for DP-FTRL. To control MinS in practice, a timer is introduced on clients in the FL system so that a client will only become eligible to participate in training (again) after a certain period of time has passed. McMahan and Thakurta 2022 used a timer period of 24 hours to train the es-ES NWP model, which led to an observed MinS of 313. The MinS of es-ES is upper bounded by $4.21M/6500\sim 647$ and can be potentially improved by increasing the timer period. We increase the timer period in the unit of 24 hours due to the uneven diurnal participation pattern [Yang et al. 2018; Zhu et al. 2022], and generally observe that MinS can proportionally increase with the timer period before reaching the possible maximum. However, there are many factors in the FL system that may affect the wall clock training speed, which makes it challenging to optimize the timer period to maximize MinS.

3 Public pretraining

We explore pretraining on public data for production models, which were shown to substantially improve model utility in DP simulations [Li et al. 2021; De et al. 2022; Yu et al. 2021; Xu et al. 2022; Wang et al. 2023]. We pretrain a model for each Gboard language task using the multi-lingual C4 dataset [Raffel et al. 2019; Xue et al. 2020] collected from public web pages. Fig. 3(a) shows that pretraining can reduce $\sim 1000$ rounds to reach a given utility threshold under the same noise multiplier, which can significantly improve the privacy guarantees as shown in Fig. 4.

We additionally observe that: (1) it is challenging to fine-tune from a pretrained model when the word embeddings are shared for input and output to reduce the parameter size of LMs for on-device deployment; (2) the accuracy may decrease in the first a few rounds of fine-tuning; (3) pretraining helps with diminishing marginal returns: at some point further pretraining does not necessarily improve the final performance. Therefore, we use models with separate input and output embeddings and pretrain with half of the C4 dataset for Gboard LMs.

4 Combining with secure aggregation

Secure aggregation (SecAgg) Bonawitz et al. 2017 is a cryptographic multiparty computation protocol ensures that the central server can only access the aggregated update from a large set of clients, preventing inspection of individual client updates. We combine SecAgg and DP-FTRL to provide strong data minimization and anonymization protection [Bonawitz et al. 2021]. This work considers central DP and honest-but-curious server, and not the setting where the DP mechanism is applied distributively (i.e. on the client) as in Kairouz et al. 2021a; Agarwal et al. 2021. We describe the algorithm in Algorithm 2, and provide detailed discussion tackling the main challenge: how we can properly calibrate the sensitivity for DP when using SecAgg.

and applies the DP-FTRL protocol to $\Delta$. SecAgg’s finite group, $M$, dictates the number of bits needed per parameter, and $d\log_{2}M$ bits are needed per model update vector.

Sensitivity analysis.

Challenges and future work.

The large report goal requirement for strong DP guarantees is challenging for SecAgg in practice, which requires a slightly different system configuration. The SecAgg training speeds we observe are still notably slower, and we leave for future work potential improvements such as compression for communication efficiency [Chen et al. 2022], new DP methods to reduce report goal [Choquette-Choo et al. 2022], and embedding compression to reduce round time [Shu and Nakayama 2017].

5 Recommended strategies and practices

We summarize our strategy for training Gboard LMs with DP. (1) Pre-train the model on public datasets if possible. (2) Choose the maximum noise multiplier that meets the utility target based on small report goal simulation experiments on public datasets that is similar to the production task. (3) Based on the target number of rounds and estimated population, linearly increase the report goal and noise multiplier to meet the privacy target, and choose a large report goal supported by the system. If the privacy target is unachievable, fix the report goal to maximum, and increase the noise multiplier to target on a model with suboptimal utility. (4) Estimate the possible maximum MinS based on chosen report goal and estimated population, and configure the timer period to approach the MinS; use previous experience of model training speed if applicable. (5) If the hyperparameters (e.g., learning rates) are known from previous experiments or simulation on public datasets, apply DP-FTRL with adaptive clipping (Algorithm 1) without manual tuning to try meet the privacy and utility goals. Note that Algorithm 1 needs to account the noise inflation and restart for privacy guarantees. (6) If Algorithm 1 fails or stronger privacy and utility are desirable, we can run a few small report goal experiments with Algorithm 1 that tune quantile $\gamma$ and server learning rate $\eta_{s}$, select the best learning rate, and fix the clip norm based on the estimation; and run DP-FTRL with large report goals. (7) SecAgg can be used for all experiments, and precise MaxP and MinS are computed by post-processing for privacy accounting.

Deploying DP LMs

We introduce metrics in A/B test to measure the utility of Gboard LMs. (1) Picked Rate (PRate): the ratio of picked candidates among the NWP predictions; or SC predictions when it is triggered. (2) Accuracy (Acc): the ratio of candidates matching the final committed words among the NWP model predictions. (3) Trigger Rate: the ratio of words with SC triggered among all committed words, which is an important metric when PRate is fixed. (4) Word Modified Ratio (WMR): the ratio of words being modified during typing or after committed; improvement is shown by reduction. (5) Word Per Minute (WPM): the number of committed words per minute.

Privacy guarantees.

Same as [McMahan and Thakurta 2022], the zero-one device neighboring relationship ([Kairouz et al. 2021b, definition 1.1]) is adopted for DP. For user’s with a single device, device-level DP corresponds directly to user-level DP. Our privacy guarantee holds for all well-behaved clients during training, and we do not account for privacy cost of modest amount of hyperparameter tuning. DP is measured by the zero-Concentrated DP (zCDP) [Bun and Steinke 2016] guarantee that has been used by US census bureau [US Census Bureau 2021], and can be easily converted to $(\epsilon,\delta)$-DP. We use the privacy accounting in [Kairouz et al. 2021b, appendix D] implemented in Tensorflow Privacy [TFP Authors 2022], and follow the guidelines outlined in [Ponomareva et al. 2023, Sec. 5.3] to report detailed narratives of privacy guarantees in Appendix A.

Implementation.

Experimental setup.

We apply the strategy in Section 2.5 to train Gboard LMs with DP. We present NWP results in Table 1, and OTF results in Table 2. As Smart Compose (SC) reuses NWP LMs, SC has the same DP guarantees as NWP models by the post-processing property [Dwork et al. 2014]. Following es-ES NWP model in [McMahan and Thakurta 2022], we choose noise multiplier 7 and report goal 6500 based on simulation in [Kairouz et al. 2021b] on public StackOverflow dataset [TFF Authors 2022a]. We pretrain the models on public datasets and configure the timer period to control client participation, separately for different tasks. We use DP-FTRL with adaptive clipping and small report goal 500 to tune server learning rate and estimate the clip norm. Interestingly, we observe the learning rate and clip norm to be consistent for various Gboard LMs, and tuning seems to be unnecessary. DP-FTRL with fixed clip and large report goal is used to run the final model for deployment.

Result analysis.

All NWP and OTF models in Tables 1 and 2 are trained with stronger guarantees (smaller zCDP) compared to zCDP $>2.6$ used by US Census Bureau [US Census Bureau 2021]. For five NWP models in Europe (DE, GB, FR, IT, PT), the DP NN models significantly improve the utility compared to previous N-gram models. On en-US, pt-BR and en-IN, DP NN models also achieve comparable, or slightly better utility compared to their non-private versions as the strong models. SecAgg is successfully applied to en-US and es-ES, and can achieve good privacy-utility trade-off with a smaller number of rounds, likely due to the system configuration that results in more clients per round. However, SecAgg is also notably slower. There is a general positive correlation between the estimated population size and privacy guarantees.

However, only a few tasks approach the possible maximum MinS for strong privacy guarantees, which highlights the challenge of both estimating population and controlling client participation. Longer training rounds are often used for NWP (compared to OTF) as the non-private NN baselines are strong, and to improve the downstream SC performance. As an example, we train es-ES NWP for 1900 rounds with a pretrained model, while the previous models [McMahan and Thakurta 2022] is trained for 2000 rounds without pretraining. Our es-ES NWP model slightly improves the utility measured by PRate and Acc, and improves the zCDP bound from 0.81 to 0.35 due to the larger MinS by timer configuration. We highlight that our es-ES model at round 1240 already achieves similar NWP utility and a strong privacy guarantee, but the utility of SC keeps improving with training. Compared to the previous model in [McMahan and Thakurta 2022], our model improves the SC trigger rate by 4.23% at round 1240, and 9.51% at round 1900.

Concluding remarks

We discuss our experience and summarize our strategy for training production Gboard LMs with FL and DP. We propose an algorithm applying adaptive clipping [Andrew et al. 2021] in DP-FTRL [Kairouz et al. 2021b] to reduce the hyperparamter tuning. We discuss the impact on privacy and utility of several important factors: the clip norm, report goal, client participation, and pre-training. Our study highlights the importance of system and algorithm co-design for differential privacy in practice, the challenges of tuning in FL systems, and opportunities to improve the scalability and stability of FL with DP and/or SecAgg. More than twenty LMs with formal DP guarantees are trained and launched to support Gboard NWP, SC, and OTF features, including en-US and es-ES NWP models additionally trained with SecAgg. Our experience demonstrates the possibility of training DP models for practical applications when a large scale system is available for large scale data. Therefore, Gboard is introducing and enforcing a new policy: DP has to be applied in all future training and launching of Gboard LMs.

The authors would like to thank Stanislav Chiknavaryan, Adria Gascon, Zachary Garrett, and Timon Van Overveldt for infrastructure configuration support; Swaroop Ramaswamy, Om Thakkar, Abhradeep Thakurta for early discussion on models and algorithms; Jeremy Gillula for internal review process; Xu Liu, Shumin Zhai, and Daniel Ramage for leadership support.

References

Appendix A Reporting privacy guarantees

This section clarifies the nuances of the reported DP guarantees following the guidelines outlined in [Ponomareva et al. 2023, Sec. 5.3]

DP setting. This a central DP guarantee where the service provider is trusted to correctly implement the mechanism.

Data accesses covered: The DP guarantee applies to all well-behaved clients Clients that faithfully follow the algorithm including participation limits. Due to the design of the algorithm, a mis-behaved client does not adversely affect the DP guarantee of any well-behaved clients. in a single training run. We do not account for hyperparameter tuning, or the selection of the final model checkpoint using evaluation metrics or A/B testing in our guarantees. Public multilingual C4 data [Raffel et al. 2019; Xue et al. 2020] is used for pre-training.

Final mechanism output: Only the final model checkpoint is released for production launches, however the mechanism’s output is technically the full sequence of privatized gradients, and so the guarantee also applies at this level, and hence all intermediate models are protected (including those sent to devices participating in federated learning).

Unit of privacy. Device-level DP is considered, i.e., the notion of adjacency is with respect to arbitrary training datasets on each client device, and the device might have an arbitrarily large local dataset containing arbitrary training examples. For user’s with a single device, this corresponds directly to user-level DP; for devices shared with multiple users, this provides a stronger notion of DP than user-level; for a user with multiple devices that happen to both participate in training the model, the notion is weaker, but group privacy can be used to obtain a user-level guarantee.

Adjacency definition for “neigbouring” datasets: We use the zero-out definition [Kairouz et al. 2021b]. This is a a special form of the add-or-remove definition, where neighboring data sets differ by addition/removal of a single client. In the absence of a client at any training step, we assume that the client’s model update gets replaced with the all zeros vector. This assumption enforces a subtle modification to the traditional definition of the add/remove notion of DP which allows neighboring data sets to have the same number of records.

Type of accounting used: Both $\rho-$zCDP [Bun and Steinke 2016] accounting, and PLD accounting [DP Team 2022] for $(\epsilon,\delta)-$DP are used.

Accounting assumptions : Each client only participates limited times during the training, and there are at least a min-separation number of rounds between two consecutive participation of a client, i.e., MaxP and MinS as discussed in Section 2.2. Client participation is enforced by a timer on clients in the cross-device FL system.

The formal DP statement: The launched Gboard LMs have $\rho-$zCDP range in (0.2, 2). We also transform zCDP to $(\epsilon,\delta)-$DP by PLD accounting [DP Team 2022]: given $\delta=10^{-10}$, the smallest zCDP $\rho=0.25$ corresponds to DP $\epsilon=4.49$; the largest zCDP $\rho=1.86$ corresponds to DP $\epsilon=13.69$.

Transparency and verifiability: We open sourced our core implementation code in TensorFlow Federated and Tensorflow Privacy. Key portions of the cross-device FL system are also open sourced.