Federated Learning (FL) (McMahan et al., 2017(McMahan et al., , 2018;;Kairouz et al., 2019) is designed to collaboratively train a global model on decentralized data across user clients while protecting data privacy. FL emerged as an effective privacy-preserving solution of training (language) models, as rich text data are generated by users, which may contain sensitive and personal information. After McMahan et al. (2017) proposed to train on-device recurrent neural networks, FL has been widely used in various natural language processing applications and products, including next-word prediction (Hard et al., 2018), keyword spotting (Hard et al., 2020), and out-ofvocabulary word discovery (Chen et al., 2019).
To further protect user privacy, Differential Privacy (DP) (Dwork et al., 2006;Dwork, 2011;Dwork and Roth, 2014;McMahan et al., 2018) is introduced to provide formal privacy guarantees of models trained by federated learning. DP for deep learning explicitly adds random noise with bounded sensitivity to a training process (e.g., DP-SGD (Abadi et al., 2016)), ensuring a quantifiable similarity in output model distributions when the training dataset changes. When combining DP with FL, a variant of DP-SGD called DP-FedAvg (McMahan et al., 2018)) is applied to guarantee user-level DP (Dwork, 2010). Current research primarily focuses on applying user-level DP to small on-device models with fewer than 10 million parameters (McMahan et al., 2018;Kairouz et al., 2021;Ramaswamy et al., 2020). The model size is limited due to challenges such as significant DP noise required to preserve privacy (Li et al., 2021) and the communication costs in cross-device FL.
Recent advances in large language models (LLMs) (Thoppilan et al., 2022;Radford et al., 2019;Brown et al., 2020;Devlin et al., 2019;Raffel et al., 2020) have revolutionized natural language processing (NLP) and achieved unprecedented performance on various tasks such as text generation, machine translation, and sentiment analysis. However, their success comes at a cost of requiring massive amounts of computational resources, making them difficult to deploy on resource-constrained devices such as smartphones, tablets, or other edge devices. Additionally, there are concerns regarding the user privacy in various aspects such as memorizing personal information in training, and exposing private query in inference.
Recent work explore incorporating public information to improve privacy-utility trade-off in applying DP for (large) LMs (Yu et al., 2022;Li et al., 2021). Public data (Amid et al., 2021) or other side information (Li et al., 2022) are also studied for (DP) FL. In non-DP FL settings, Nguyen et al. (2022) studies the effect of initializing from a pretrained model. However, it is an open question on how to leverage the power of pre-trained LLMs to facilitate private FL for on-device LMs.
In this work, we answer the question through systematic study aimed at enhancing private federated learning for on-device LMs with public pretrained LMs. Specifically, Our approach involves leveraging both public data and pre-trained LLMs to improve differentially private federated learning for on-device models by techniques of public pre-training and distillation. Additionally, we propose a novel distribution matching algorithm, which is backed by theoretical analysis, to sample public data closely resembling the private data distribution, which significantly increases sample efficiency in public training. Moreover, our extensive empirical results align with our theoretical predictions, further substantiating our approach. Our work complements existing research by utilizing LLMs to improve public training through knowledge distillation for private cross-device federated learning, and achieve a strong privacy-utility trade-off with substantial improvements on sampling efficiency for public data. Our method points to a novel direction of efficiently enhancing private FL with public pretraining data and LLMs.
We summarize our contributions as follows: • We focus on improving private federated learning for language modeling tasks and explore ways to leverage public data and pre-trained LLMs for tokenizers, training protocols, and data (sub)sampling.
• We conduct comprehensive studies and compare the use of Sentence Piece tokenizers from public LLM and unigram tokenizers from private corpus. We find that adopting public tokenizers from LLMs can not only prevent the potential privacy leakage from the private tokenizer vocabulary, but also lead to better learning utility with DP guarantees.
• For training protocol, we propose to leverage public LLM to teach private on-device LMs by knowledge distillation. We demonstrate that distilling public LLM to pre-train on-device LM can lead to more than 7% accuracy improvement with tight privacy bound (ε = 1.77). Moreover, it can achieve high data efficiency of using only 1% of the public data compared to that in public pre-training without LLM, and attain better accuracy.
• We further propose a novel distribution matching method that leverages both private on-device LMs and public LLMs to select public records close to private data distribution. We show that using 0.08% of carefully sampled public data to train ondevice LM can lead to comparable performance as public pre-training on-device LMs with the whole pre-training corpus. Moreover, it reduces the public training time from more than one week to a few hours. Our method is grounded in theoretical analysis, which is corroborated by our extensive empirical results.
2 Differentially Private Federated Learning for On-device LMs
In this section, we walk through the preliminaries of differentially private federated learning of language models following the cross-device federated learning literature (McMahan et al., 2018;Kairouz et al., 2019Kairouz et al., , 2021)). We also introduce the experimental setup used throughout this paper.
Cross-device Federated Learning. McMahan et al. (2017) introduce federated learning to collaboratively train LMs for next-word prediction from decentralized user data on a large number of mobile devices without directly sharing the private data. A common training algorithm of federated learning is FedAvg (McMahan et al., 2017), where each client downloads the current model from the centralized server, computes an update by performing local computation on their dataset (e.g., running SGD) and sends the update back to the server. The server aggregates the updates across clients to update the global model and send the updated model back to local clients to achieve the goal of collaborative learning without directly accessing the training data on each user's mobile device.
In our experiments, we follow previous work (Kairouz et al., 2021;Amid et al., 2021;Wu et al., 2022) and sample 100 clients in each training round. Each client uses a batch size of 16 for local training. We set the training rounds T = 1600 in total. User-level Differential Privacy. To further protect user privacy, Differential Privacy (DP) (Dwork et al., 2006;Dwork, 2011;Dwork and Roth, 2014) was introduced to provide a formal privacy guarantee for federated learning. Definition 2.1 provides a formal definition of (ε, δ)-DP by bounding the change in output distribution caused by a small input difference (or, adjacent datasets) for a randomized algorithm. In the FL setting, it is preferable to bound the output distribution caused by different users in order to protect the privacy of each client's whole dataset. Specifically, adjacent datasets of D and D for userlevel differential privacy (Dwork, 2010) are defined as: D can be obtained from D by adding or subtracting all the records of a single user/client, which determines the unit of privacy guarantees.
In our experiments, we use DP-FTRL (Kairouz et al., 2021) for privacy accounting and private federated training, which can achieve strong privacy guarantee in practical FL scenarios (Xu et al., 2023). We use δ = 10 -6 and consider two ε bounds: a tight privacy bound with ε = 1.77 by using a large noise multiplier m = 8.83, and a slightly loose privacy bound with ε = 18.71 and noise multiplier m = 1.13. We present more hyperparameter tuning details in Appendix §C. On-device LMs. Due to the limited memory constraints of mobile devices, on-device LMs are relatively small (usually less than 10M parameters). In our work, we focus on two types of on-device autoregressive LMs: LSTM (Hochreiter and Schmidhuber, 1997) and transformers (Vaswani et al., 2017). More model details can be found in Appendix §B.2. Pre-trained LLMs. In addition to the on-device LMs trained on private datasets, this work also assumes that we have access to LLMs pre-trained on a large public corpus to aid private learning. Specifically, we use LaMDA (Thoppilan et al., 2022) 2B throughout this work as an example, and conduct a systematic study of leveraging LLMs to help private training of on-device LMs. Datasets. We focus on next word prediction task on the StackOverflow benchmark dataset (2019) for private federated learning. Since StackOverflow is naturally keyed by users, each client in FL is a user in the Stack Overflow online forum. The examples of a client are sentences of questions and answers posted by a specific user. We follow (Reddi et al., 2021; Kairouz et al., 2021) to construct a validation set of 10K samples, and a test set of 16.5M samples. Our evaluation metric is in-vocabulary next word (token) prediction accuracy, which is computed 18.73 18.42 21.44 21.19 17.39 17.57 20.48 20.48 20.48 24.45 16 17 18 19 20 21 22 23 24 25 Next Token Prediction Next Word Prediction Next Token Prediction Next Word Prediction Unigram Tokenizer, vocab=10000 Unigram Tokenizer, vocab=32000 SentencePiece Tokenizer, vocab=32000 𝜀 = 1.77 𝜀 = 18.71 Figure 1: Next word (token) prediction accuracy for ondevice LSTM with different tokenizers in the private FL.
as the ratio of accurately predicted in-vocabulary words to the total number of words in the sequence (excluding OOV tokens).
In addition to StackOverflow as the (private) dataset, we use the realnews variant c4/realnewslike of C4 dataset (Raffel et al., 2020), as the public dataset. We analyzed the sources of the public C4 dataset and the Stackoverflow dataset for private training, and verified that there is no explicit overlap between public C4 dataset and the private StackOverflow dataset. More details can be found in Appendix §B.1.
The success of publicly pre-trained LLMs motivate us to have retrospective views on further improving private on-device LMs. In this section, we explore inpiration from LLMs: the use of subword tokenizers and a large public corpus for pre-training. We apply them to on-device LMs, and observe that both techniques bring significant performance improvement for private FL.
Tokenizer is an important module of LMs, which transforms natural languages into a sequence of predefined symbol sets (vocabulary). Prior work in the literature of private FL of LMs (McMahan et al., 2018;Kairouz et al., 2021;Amid et al., 2021) use word-level unigram tokenizers potentially directly built from user data, which may need additional privacy budget (Ponomareva et al., 2022;Bagdasaryan et al., 2022).
Recent LLMs adopt sub-word tokenizers (Kudo and Richardson, 2018;Sennrich et al., 2016;Schuster and Nakajima, 2012), which mitigate most outof-vocabulary (OOV) problems and yield state-ofthe-art performance across different downstream tasks. This motivate us to replace the prior wordlevel unigram tokenizers with public sub-word tokenizers. Specifically, we use SentencePiece tokenizer (Kudo and Richardson, 2018) from LaMDA.
To conduct comparison between unigram tokenizers and subword tokenizers for next word (token) prediction task, we convert the next word prediction accuracy into next token prediction accuracy. This conversion is achieved through splitting each word using the SentencePiece tokenizer. We consider all tokens within a word as accurate if the predicted word is correct. We compare standard SentencePiece models (vocabulary size = 32K) with unigram tokenizers that selects the top-k frequent words from user data with k = 10K or 32K as vocabulary.
We present the private FL accuracy on the Stack-Overflow dataset in Figure 1. For the unigram tokenizer, using a larger vocabulary size in the DP setting can result in a slight performance drop, which can be different from the observation in non-DP settings (Charles et al., 2022;Xu et al., 2022a). It is possible that the parameter increase of the embedding layer enlarges the effect of DP noise and hurts the final accuracy. However, for next token prediction accuracy, although the public Sentence-Piece tokenizer from LaMDA also consists of 32K tokens, it can significantly improve the private FL accuracy upon the unigram tokenizers, especially with smaller DP noise and ε = 18.71. We also observe that SentencePiece tokenizer finds no OOV tokens in the StackOverflow dataset, thus yielding the same high prediction accuracy with or without the OOV token. Therefore, we use SentencePiece tokenizer in the rest of this paper.
In addition to the use of subword tokenizers, LLMs benefit from pre-training on a large public corpus (Li et al., 2022;Yu et al., 2022). In this section, we explore pre-training on-device LMs on public corpus to improve private federated learning. Pre-training Details. We use the standard autoregressive language modeling loss L LM to pre-train on-device LMs on the public C4 dataset, which takes around 1, 400K steps (over a week of single GPU time) to process the entire dataset with the batch size of 512. We then use the publicly pre-trained checkpoint as the start point for private federated learning. We leave more details in §B.2. Results. We present the next token prediction accuracy on the private StackOverflow dev set in Table 1. We observe that the accuracy on the private dataset significantly improves after pre-training for different different privacy budgets, shedding light on an effective way to boost private FL perfor- mance. We also observe that after pre-training, it gives reasonable zero-shot accuracy on the private dataset even without private training (round=0).
On one hand, the cost of public pre-training for on-device LMs is still expensive on a large public corpus (around a week of GPU time). On the other hand, existing LLMs are well pre-trained and demonstrate promising performance across a variety of downstream tasks. This motivates us to explore on whether we can leverage existing LLMs to improve the sample efficiency of pre-training ondevice LMs. In this section, we answer the question above with systematic studies and show that we can improve the sample efficiency by using only 1% of pre-training data and distillation from LLMs, achieving similar or even better performance than using 100% of pretrianing data without distillation.
Inspired by the literature of model compression (Sun et al., 2020;Jiao et al., 2019), we use knowledge distillation to transfer the knowledge from trained LLMs into on-device LMs during pretraining. The distillation pipeline contains the following two steps: Building a distillation corpus. Given an input sequence from the public pre-training corpus, the LLM outputs the probability distribution over the vocabulary for next token prediction at each decoding step. To construct a distillation corpus, we save the top-k logits with k nonzero entries z T from the teacher LLM as a silver-label dataset. In this way, the distillation corpus is model-agnostic, and thus can be applied to different variants of on-device LMs for pre-training. Moreover, selecting a reasonable top-k for the logits can both help compress the distillation corpus to a moderate size and filter out noisy signals from tokens with low output probabilities. Public pre-training with distillation loss. Since we align the tokenizer of the on-device LM with the LLM to share the same vocabulary, we can 0 5 10 15 20 25 Distillation Coverage of C4 dataset (%) 25 26 27 28 29 30 Next Token Prediction Acc Accuracy v.s. Distillation Coverage eps=18.71, w/ public pre-training, covering 100% c4 eps=18.71, w/o public pre-training eps=18.71, w/ public pre-training eps=18.71, w/ public pre-training + distillation, topk=10 (a) Acc. v.s. distillation steps (ε = 18.71) 0 5 10 15 20 25 Distillation Coverage of C4 dataset (%) 21 22 23 24 25 26 27 Next Token Prediction Acc Accuracy v.s. Distillation Coverage eps=1.77, w/ public pre-training, covering 100% c4 eps=1.77, w/o public pre-training eps=1.77, w/ public pre-training eps=1.77, w/ public pre-training + distillation, topk=10 (b) Acc. v.s. distillation steps (ε = 1.77) 2.0 2.5 3.0 3.5 4.0 4.5 5.0 Distillation Coverage of C4 dataset (%) 26.0 26.2 26.4 26.6 26.8 27.0 27.2 27.4 27.6 Next Token Prediction Acc Accuracy v.s. Distillation Coverage eps=1.77, w/ public pre-training + distillation, topk=10 eps=1.77, w/ public pre-training + distillation, topk=5 eps=1.77, w/ public pre-training + distillation, topk=3 eps=1.77, w/ public pre-training + distillation, topk=1 (c) Acc. v.s. top-k logits (ε = 1.77) align the output distribution of on-device LMs and LLMs by the cross-entropy loss. Formally, for next token prediction task, given the output logits from student on-device LMs z S , the gold label from the pre-training corpus y, and the logits from the distillation corpus of LLMs z T , we add an additional knowledge distillation loss L KD = CE(z S /t, z T /t) to the pre-training language modeling loss L LM = CE(z S , y) as our public pretraining loss L pub = L LM + βL KD where t is the temperature. More distillation details are in §B.3.
After public pre-training with knowledge distillation, We use the checkpoints at different pretraining steps as the start point for private federated learning. Our main results can be found in Table 2. We show that by using 1% C4 dataset for pre-training with knowlegde distillation, we can significantly improve the sample efficiency without hurting but even improving the private FL accuracy for both LSTM and transformers, when compared with public pre-training on the whole C4 dataset.
The sample efficiency improvement thus reduces the pre-training cost from one week to around one day, shedding light on a promising direction to improve the efficiency and utility of private FL.
Ablation studies on distillation steps. To understand whether distillation for more epochs can help with private FL, we conduct a set of ablation studies on distillation steps given different privacy budgets as shown in Figure 2b and 2a. Specifically, we use the checkpoints at different distillation steps to initialize on-device LSTM and report the next word prediction accuracy after private FL at round 1600. We observe a consistent performance improvement when the distillation covers less than 5% of the C4 dataset. But when we pre-train the LM for more epochs, the improvement becomes marginal. This suggests that teaching on-device LMs via LLMs can converge quickly within a few iterations.
Abaltion studies on top-k logits. We take the topk logits of the LLM to construct our distillation datasets and pre-train the on-device LMs. Here, we conduct an ablation study by pre-training different on-device LMs with different k and evaluate how top-k logits in distillation can impact the accuracy of private FL. We present our empirical results in Figure 2c and Appendix Figure 4. We observe that pre-training with a larger k is more helpful to achieve better downstream accuracy on private data.
To have a reasonable trade-off between dataset size and pre-training performance, we use top-k = 10 in all the following experiments.
In the previous section, we achieve compelling performance by employing LLM distillation using only 1% of the randomly sampled pre-training corpus. Now we further investigate the possibility of improving sample efficiency by selectively identifying public samples that align with the distribution of private samples. To this end, we propose a novel distribution matching method to sample public records for pre-training with a novel theoretical analysis jointly considering public-private distribution shift and DP mechanism. We demonstrate that by carefully selected 0.08% of public samples, we can pre-train on-device LMs that perform as well as using 1% of public samples with distillation. This approach significantly improves sample efficiency, providing an additional knob of using public pre-training for private on-device models.
We hypothesize two principles to sample public records to match the private distribution: (i) the probability of the public sample x on the private data distribution p priv (x) is high, which can be approximated by the prediction of the on-device LMs trained on the private dataset; (ii) the probability of a public sample x on the public data distribution p pub (x) is also high, as we expect those samples are easy-to-learn (Swayamdipta et al., 2020) and of high data quality in the public corpus. The probability p pub (x) can be approximated by the public pre-trained LLMs.
To verify our hypothesis, we visualize the perplexity (PPL) distribution of public samples and private samples evaluated by both a privately finetuned on-device LM and a public pre-trained LLM in Figure 3. To have an "oracle" on-device LM that well captures the private data distribution, we fine-tune it on the private data without DP noise to overfit the private data distribution. We randomly sample 10k records from the public dataset and private dataset, respectively. We observe that the private dataset mostly concentrates on the regime with low PPL evaluated by the public and private LMs, whereas the public dataset is more diverse and distributed across a broader range of PPL values. The distribution visualization confirms our hypothesis to select public samples from the lower left corner, which correspond to samples with high probabilities p pub (x) and p priv (x) on public and private data distribution (i.e., low perplexity evaluated by public and pirvate LMs).
In practice, we do not have an "oracle" on-device LM trained on private data for distribution match. Instead, we propose to fine-tune an on-device LM with DP for certain rounds T < T before consuming all the privacy budgets, and then use the checkpoint at round T with DP guarantee to ap- proximate p priv (x) and perform distribution matching to sample public records. This post-processing based on a DP checkpoint will not incur any additional privacy cost. Thereafter, we can use the sampled public records to further train the private checkpoint at round T , as a way for efficient public (pre-)training. Following the strategy in §4, we also employ the distillation loss to better train the on-device LM with carefully sampled public records to further enhance the sample efficiency. Lastly, we use the remaining privacy budgets to fine-tune the on-device LM until reaching round T , and evaluate its next token prediction accuracy at the dev and test sets. We term the paradigm of two-stage private learning combined with public training as "public mid-training". This approach differs from "public pre-training", which involves public pre-training prior to private FL. We present the distribution matching protocol in Algorithm 1.
In this section, we provide the theoretical analysis of our distribution matching protocol to present the intuition behind our selection hypothesis. In essence, the goal of our distribution matching algorithm is to have a good estimator for the private distribution. However, characterizing the distribution shift in the context of differential privacy is a challenging problem, in that the private models are trained with DP noise, which can yield an inac- curate estimation of private data distribution, and thus add the complexity to our analysis.
Problem Setup. Define the text data domain as X . Denote ℓ pub : X → R as the log-density function of the public data distribution (i.e., ℓ pub (x) = log p pub (x) where p pub (x) is the public data density estimated by public LLMs), and ℓ priv as the accurate log-density function of the private data distribution (i.e., ℓ priv (x) = log p true priv (x) where p true priv (x) is the true private data density). However, due to limited private data sampled from the true private data distribution and DP noise injected in the private FL, we can only obtain an inaccurate estimation lpriv = log p priv (x) of the true private log-density ℓ priv , where p priv (x) is the private data density estimated by private on-device LMs. Note that we use the hat notation lpriv to denote that it is an estimation of the true private log-density ℓ priv .
We can view the estimation lpriv is a random variable where the randomness comes from: (i) that the private dataset we have is sampled from the private data distribution; and (ii) the randomness in the algorithm of obtaining lpriv based on the private dataset, e.g., differential privacy. Following previous work (Jiang et al., 2023), we make a standard assumption. We assume the estimated private data log-density function is an unbiased estimator, i.e., E[ lpriv ] = ℓ priv . Since ℓ pub may not be ideal because of public-private domain shift, and lpriv may mot be ideal because of its DP noise, ℓ pub and lpriv are neither good estimators for ℓ priv . Can we leverage both of the information and form a function ĥ : X → R that combines ℓ pub and lpriv such that ĥ is a good estimator for ℓ priv ? In the following analysis, we choose ĥ = 1 2 ℓ pub + 1 2 lpriv and analyze when and why it can be a better estimator to the true private log-density ℓ priv than ℓ pub and lpriv .
We need some mathematical tools to define what does it mean to be "better". Concretely, we need a metric to measure the distance between functions. This can be done by having an inner product •, • in the function space of H = {f : X → R}, and hence the norm in the function space H is f = f, f for ∀f ∈ H. Our analysis holds with any choice of the inner product as long as it does not make the log-densities norm infinite. We discuss a concrete choice of the inner product and its relation to the KL divergence in Appendix §D.
With the norm as a "ruler", we are able to define the following key quantities that formally characterize the setting. 1. Public-Private Domain Distance.
Let d pub, priv = ℓ pubℓ priv denote the distance between the public data log-density ℓ pub and the true private log-density ℓ priv . 2. Private Domain Randomness. Let σ 2 priv = E[ lprivℓ priv 2 ] denote the randomness of the estimated private log-density, i.e., the quality of the estimated private log-density lpriv The above definitions are important because the quality of a private log-density estimator would depend on the public-private domain shift and the private domain randomness as we show next.
] characterise how good f is as an estimator of the true private data log-density ℓ priv for any random function f ∈ H. Consider the following three quantities: 1. (ℓ pub ) characterizing the error of the public log-density function ℓ pub to approximate ℓ priv 2. ( lpriv ) depicting the error of the noisy private log-density function lpriv to approximate ℓ priv 3. ( ĥ) characterizing the error of ĥ = 1 2 ℓ pub + 1 2 lpriv to approximate ℓ priv . Then,
Interpretation Theorem 5.1 implies that:
• ( ĥ) ≤ • ( ĥ) ≤ min{ (ℓ pub ), ( lpriv
Combining the above, we have the following conclusion: recall ĥ = 1 2 ℓ pub + 1 2 lpriv = 1 2 log(p pub (x)p priv (x)). We can expect that ĥ is better than either ℓ pub or lpriv for any settings. Moreover, we can expect ĥ to be better than both ℓ pub and lpriv if (i) there is a domain shift between the public-private domain; and (ii) our estimated private log-density lpriv is noisy in an extent comparable to the domain shift. We leave the full proof and additional discussion in Appendix D.
Experimental Setup. We set T = T /2 = 800 rounds for the first-stage private federated learning.
We use q = 0.08% of the whole pre-training corpus for public training, which reduces the public training time from more than 1 weeks to a few hours with a single GPU. For the public mid-training setting, we also evaluate how LLM distillation and distribution matching can impact the FL accuracy, respectively. We run all the experimental settings for three times and report the average and standard deviation of test accuracy on the private StackOverflow dataset.
We present the results of on-device LSTM and transformers in Table 2. In the pre-training setting (T = 0), we show that we cannot further improve the sample efficiency from 1% to 0.08% with LLM distillation improves the sample efficiency, as the final accuracy after private FL significantly decreases. In comparison, in the mid-training setting (T = T /2), using LLM distillation on the 0.08% of randomly sampled pre-training corpus already gives better performance than pre-training. Moreover, with distribution matching to carefully sample public data, we further improve the private FL accuracy, attaining comparable performance to the setting using the whole public corpus for pre-training. Ablation studies on p pub (x). Our distribution matching algorithm leverages both on-device LM and LLM to sample data close to the private distribution. To understand how the use of LLM (p pub (x)) impact the sampling quality, we con- duct an ablation study to sample a subset of D based on top log p priv (x) values alone instead of log p priv (x)+log p pub (x). We use the p priv -sampled D for public mid-training and report the test accuracy of three runs for both on-device LSTM and transformers given different privacy budgets in Table 3. The experimental findings corroborate our theoretical analysis. Specifically, when on-device language models (LMs) are trained with high noise levels (ε = 1.77), we find that a combined utilization of both on-device LMs and LLMs consistently yields superior performance. This is because the estimated private log-density lpriv is noisy to a degree comparable to the domain shift, making ĥ a more reliable estimator than lpriv . Conversely, when ondevice LMs are trained with low noise (ε = 18.71), the performance difference between models with and without p pub is negligible. This indicates that the noise introduced by differentially private (DP) training is not as significant as the distribution shift, allowing lpriv to serve as a good estimator.
Ablation studies on T . T separates two-stage private federated learning and determines the timing for distribution matching and public training. In this ablation study, we evaluate the dev set accuracy of on-device LSTM given different T and privacy budgets, as shown in Table 4 and Appendix Table 5. From the table, we can see that the ondevice LSTM achieves the best private FL accuracy given T = T /2 = 800. We think the reasons are as follows: when T = 0, we cannot perform distribution matching as the on-device LM is not trained on the private dataset yet, and thus we can only use the randomly sampled data for pre-training; when T = 400, the on-device LM could not be well trained on the private data distribution, thus yielding worse distribution matching quality; when T = 1200 and T = 1600, the private on-device LM is biased towards the public data distribution due to public training, thus giving worse private FL accuracy. As a result, we use T = 800 in our main experiments, as it balances the private federated training and public training to have satisfactory distribution matching capabilities without biasing too much towards the public data distribution.
In this work, we propose to improve private federated learning by using LLMs in public training. We leverage LLMs to aid public training of on-device LMs via distribution matching to sample public data close to private data distribution, which further improves the effectiveness and efficiency of public training, demonstrating strong private learning accuracy while minimizing the need for large amounts of public training data. Our work sheds light on a promising direction to improve private federated learning with public LLMs.
This work has paved the way for enhancing the utility of differentially private on-device FL models, using large-scale public data and LLMs, but we also acknowledge the following limitations:
• Data Distribution Matching: The proposed distribution matching algorithm aims to sample public data close to the private data distribution. The choice of ĥ can be data dependent and a weighted combination of ℓ pub and lpriv , i.e., ĥ = (1β)ℓ pub + β lpriv where β ∈ [0, 1], as mentioned in Appendix §D.3. In practice, the optimal β can be an important hyper-parameter to tune the distribution matching algorithm. Our work mainly leverages ĥ = 1 2 ℓ pub + 1 2 lpriv to analyze when and why a better estimator to the true private log-density ℓ priv than ℓ pub and lpriv . We leave it as important future direction to get the optimal β theoretically and empirically.
• Computational Resources: The use of largescale public data and LLMs can improve the privacy-utility trade-off in DP FL models, but this often comes at the cost of computational resources. Our work mainly focuses on LaMDA 2B as an example of LLM due to the lack of computational resources. While our main focus does not lie in the knowledge distillation, we leave it as future work to extend the size of LLMs in public pre-training.
Private Federated Learning in On-device NLP Federated learning is designed to collaboratively training NLP models without sharing sensitive user data to protect user privacy. Given relatively small model sizes, state-of-the-art differentially private (DP) learning algorithms (McMahan et al., 2018;Kairouz et al., 2021) have enabled on-device LMs to achieve strong downstream task utility with reasonable userlevel differentially privacy guarantee (Dwork, 2010). The success of private FL has also led to real-world applications such as GBoard, which uses on-device LMs for next word prediction (Hard et al., 2018;Ramaswamy et al., 2020). Recent advances in DP optimization (Kairouz et al., 2021) further improves upon the state-of-the-art DP-SGD algorithm (Abadi et al., 2016), providing a practical tool to analyze privacy bound for federated learning.
Privacy-preserving Large NLP Models Scaling up LMs with more data and parameters has significantly improved performance and achieved great success in a variety of NLP tasks. Moreover, recent studies show that LLM has great potential in private learning. For example, Kerrigan et al. (2020) show that public pre-training is helpful for downstream DP fine-tuning. Follow-up studies argue that large pre-trained LMs can be strong differentially private learners with parameter-efficient fine-tuning (Yu et al., 2022;Bu et al., 2022) or full model fine-tuning (Li et al., 2021), narrowing the gap between non-private training and private training. Ganesh et al. (2023) also provide theoretical groundings on the necessity of involving public training into private learning. Motivated by the recent success of LLMs, our work performs comprehensive studies on how to use public data and existing LLMs to help private training of cross-device FL models.
Model Compression for Pre-trained LMs One promising approach to address the resource limitations of LLMs is to compress them into smaller models through various techniques such as knowledge distillation (Jiao et al., 2019;Sun et al., 2020;Wang et al., 2020), or pruning (Elbayad et al., 2020;Gordon et al., 2020). While these techniques have demonstrated success in reducing the size of pre-trained LMs, most resulting models are still too large (with over 10 million parameters) to be effectively deployed on resource-constrained devices. In our work, we also explore the use of knowledge distillation in public training, but with a primary focus on leveraging LLMs to improve sample efficiency in pre-training on-device LMs. We aim to improve the private FL performance of on-device LMs while minimizing the need for large amounts of training data. We recognize that private federated learning can further benefit from advanced model compression techniques, and we leave this as a promising and orthogonal future direction for research in this area.
StackOverflow contains 342K clients for training with 135.8M examples. In this section, we detail the method used to verify that there is no explicit overlap between the public C4 dataset and the private StackOverflow dataset utilized in our study. We explored C4 which has multiple variantsfoot_0 : c4/en, c4/realnewslike, and c4/webtextlike.
To verify this hypothesis, we conducted a rigorous comparison of these two datasets and its variants. Specifically, we compared the unique identifiers (e.g., URL for webpages in the C4 dataset, and post ID for StackOverflow posts) between the two datasets.
No matching identifiers were found between the c4/realnewslike and the StackOverflow dataset. Thus we use the c4/realnewslike variant as our public pretraining corpus throughout the experiment.
Through this comprehensive comparison, we have confirmed that there is no explicit overlap between the public C4 dataset and the private StackOverflow dataset. This conclusion is critical to our study as it ensures that the integrity and privacy-preserving conditions of our experiment are maintained. Ablation studies on the timing T for mid-training T separates two-stage private federated learning and determines the timing for distribution matching and public training. In this ablation study, we evaluate the dev set accuracy of on-device LSTM given different T and privacy budgets, as shown in Table 4 and Appendix Table 5. From the table, we can see that the on-device LSTM achieves the best private FL accuracy given T = T /2 = 800. We think the reasons are as follows: when T = 0, we cannot perform distribution matching as the on-device LM is not trained on the private dataset yet, and thus we can only use the randomly sampled data for pre-training; when T = 400, the on-device LM could not be well trained on the private data distribution, thus yielding worse distribution matching quality; when T = 1200 and T = 1600, the private on-device LM is biased towards the public data distribution due to public training, thus giving worse private FL accuracy. As a result, we use T = 800 in our main experiments, as it balances the private federated training and public training to have satisfactory distribution matching capabilities without biasing too much towards the public data distribution.
We need to define a meaningful distance metric in order to define the closeness of two log-density functions. To do this, we can choose any inner product •, • in the function space of H = {f : X → R}. Note that the log-density functions ℓ pub , ℓ priv , lpriv ∈ H. Accordingly, the norm in the function space H is denoted as • and by definition ∀f ∈ H : f = f, f . We note that our analysis works for any choice of the inner product as long as they don't make the log-densities norm infinite. For a concrete example, we discuss a generalization of the L 2 inner product, i.e., the L π inner product where π is a distribution on X . Formally, for this example of H = L π we define f, g π = E x∼π [f (x)g(x)] and
The L π is a rather general definition that is common in the literature of Bayesian coresets (Zhang et al., 2021;Campbell and Broderick, 2019) and kernel machine (Rahimi and Recht, 2007). For example, it recovers L 2 if π is chosen to be the uniform distribution on X .
Moreover, if we choose π = p priv as the private data density, we can show that for any probability
https://www.tensorflow.org/datasets/catalog/c4