A Hassle-free Algorithm for Strong Differential Privacy in Federated Learning Systems (2024)

Abstract: Differential privacy (DP) and federated learning (FL) are combined as advanced privacy-preserving methods when training on-device language models in production mobile keyboard applications. DP-Follow-the-Regularized-Leader (DP-FTRL) algorithms, leveraging correlated noise mechanisms such as tree aggregation or matrix factorization, are widely used in practice for their superior privacy-utility trade-off and compatibility with FL systems. This paper presents a novel variant of DP-FTRL by adapting the recent theoretical advancements of the Buffered Linear Toeplitz (BLT) mechanism to multi-participant scenarios. In the FL setting, our BLT mechanism demonstrates enhanced privacy-utility trade-off and improved memory efficiency than the widely used tree aggregation mechanism. Moreover, BLT achieves comparable privacy and utility to the state-of-the-art banded matrix factorization mechanism, while significantly simplifying usage requirements and reducing memory. The flexibility of the BLT mechanism allows seamless integration with existing DP FL implementations in production environments. We evaluate the BLT-DP-FTRL algorithm on the StackOverflow dataset, serving as a research simulation benchmark, and across four on-device language model tasks in a production FL system. Our empirical results highlight the potential of the BLT mechanism to elevate the practicality and effectiveness of DP in real-world scenarios.

## introduction language models (lms) that can predict the next word for input text are a powerful tool for many applications. in mobile keyboard applications, lms are deployed on device to support various features (e.g., auto correction, smart completion and suggestion, and next word prediction) to improve users' typing experience. on-device lms are typically small (less than ten million parameters) due to latency requirement and limited on-device resources. their performance can be significantly improved by training from user data @xcite ; recent work @xcite shows the necessity of training on user data to achieve high utility even when we can access large-scale web data and pre-trained large lms with billions of parameters. as mobile-keyboard user data can be highly privacy sensitive, differential privacy (dp) @xcite and federated learning (fl) @xcite have emerged as best practices for such models. dp provides a mathematical formulation to upperbound the memorization of an individual's information in model training. fl minimizes data exposure by aggregating focused model updates from decentralized data stored only on user devices. dp and fl are combined when training on-device language models in production mobile keyboard applications @xcite . applying dp in a production cross-device fl system is challenging as many dp algorithms require specific pattern of sampling training data to achieve strong privacy-utility tradeoff. however, a cross-device fl system has limited control of sampling as clients can only participate in training when local criteria (e.g., charging, idle, and connected to an unmetered network) are satisfied @xcite . recently, dp-follow-the-regularized-leader (dp-ftrl) algorithms @xcite have achieved superior privacyutility trade-off with simpler client participation requirements, and are used in practice in fl systems @xcite . instead of requiring uniform or poisson sampling of devices as in previous work @xcite , dp-ftrl uses minimum separation (min-sep) to characterizes the participation pattern. min-sep is the smallest number of rounds between the consecutive participation of a client, and smaller min-sep necessitates adding more noise to achieve a desired dp guarantee. minsep is enforced in the fl system by implementing a timer on each device so that a device only becomes eligible for training if a certain period of time (e.g., three days) has passed since their last participation. dp-ftrl algorithms leverage correlated noise mechanisms such as tree aggregation (treeagg) @xcite or matrix factorization (mf) @xcite with the client participation pattern in fl. the banded mf (bandmf) mechanism pre-computes matrices to generate correlated noise from independent noise to achieve stronger dp guarantees than the treeagg mechanism. bandmf is superior when the number of rounds and min-sep can be (accurately) estimated before training to optimize matrices. however, min-sep is only known after training with time-based separation as many system factors may potentially affect training time 1 . furthermore, bandmf consumes more memory for noise generation, and hence is used less often than treeagg in practice. in this work, we focus on the challenges of achieving strong dp guarantees in training production lms in a cross-device fl system. we discuss how we extend the recent theoretical advancements of the buffered linear toeplitz (blt) mechanism from single participation @xcite to multi-participation scenarios, and adapt blt to dp-ftrl. we apply blt-dp-ftrl to fl in practice, demonstrating its advantages in flexibility, ease of use, and privacy-utility tradeoffs. the blt-dp-ftrl algorithm offers flexibility in handling varying numbers of training rounds, and robustness to a wide range of min separation between user participations. furthermore, blt-dp-ftrl simplifies the correlated noise generation process and reduces memory by exploiting the parameterization of the toeplitz matrices. we empirically evaluate blt-dp-ftrl on the stack-overflow benchmark dataset and across four ondevice lm training tasks in a production fl system. our blt-dp-ftrl achieves better privacy-utility trade-off compared to the widely used treeagg mechanism, and comparable results compared to the state-of-the-art bandmf mechanism. blt-dp-ftrl exhibits desirable robustness properties in practice, offering a practical and effective so-1 despite being more complicated in fl systems, it is possible to enforce round-based separation so that a device only becomes eligible for training if min-sep rounds has passed since their last participation. however, it is still challenging to pre-specify min-sep before training due to the dynamics of client availability and population size. if the target min-sep is too large, training might halt because of lacking eligible devices. if the target min-sep is too small, the mf mechanism is not optimal for the correlated noise generation. lution for achieving strong dp in real-world fl systems. ## multi-participation blts we study how to optimize for the blt parameters θ ∈ r d and ω ∈ r d in eq. ( 3 ) for the blt-dp-ftrl algorithm, and account for dp guarantees. particularly, we generalize the blt optimization and dp accounting in @xcite ) from single participation to multiple participations ## simulation experiments we also include treeagg-full using the optimal ("full honaker") decoding (for which a memory-efficient noise generation algorithm is unknown). we observe that all the blts perform competitively with bandmf, and can outperform bandmf when the min-separation differs significantly from the number of bands. for example, with @xmath0,mean) (light blue) blt outperforms bandmf(@xmath1. treeagg-full, and blts. results are summarized in tab. 1. bandmf still achieves the highest performance, as in this scenario we train for a fixed known number of rounds, with an exactly known max participations @xmath2. treeagg-full and blts optimized for only @xmath3, but our multiparticipation-optimized blts are very competitive with bandmf with only 2 or 3 buffers (with a 171× and 114× reduction in runtime memory overhead). in the relatively large signal-to-noise ratio regime ( = 8), blt achieves comparable or even better learning accuracy though the rmsloss (maxloss) is slightly worse. ## production lms for mobile keyboard production setting our blt-dp-ftrl algorithm is flexible in min-sep (shown in sec. 4.2), achieves competitive privacy-utility performance for relatively large signal-to-noise ratio (shown in sec. 4.1), and saves computation and memory cost (shown in tab. 3), which motivates the usage in production fl systems. we follow @xcite for the production setting, and provide additional details including the configuration for baselines in app. g.1. ## concluding remarks this work addresses the critical challenge of achieving strong dp in fl for on-device lms. we have successfully extended the blt mechanism to multiparticipation scenarios and integrated it into the dp-ftrl framework. our blt-dp-ftrl algorithm demonstrates superior privacy-utility tradeoffs compared to the widely-used treeagg mechanism while maintaining its ease of use. furthermore, it rivals the state-of-the-art bandmf mechanism in performance, yet without the associated complexities and high memory costs. through extensive empirical evaluations on both a benchmark dataset and real-world production tasks, we have showcased the practicality and effectiveness of blt-dp-ftrl, paving the way for its broader adoption. the empirical results in this paper primarily focus on the cross-device fl setting where privacy amplification by sampling is challenging in practice. the discussions (e.g. tab. 3) can also be applied to centralized setting for user-level dp or example-level dp. in centralized setting, bandmf @xcite with amplification can achieve better privacy-utility tradeoff measured by rmsloss among the mentioned mechanisms, when number of rounds n and model dimension m is not too large for optimizing and applying the mechanism. when n and m are large, @xcite (similarly, bandfhu @xcite )) can both be applied, where blt has less optimization cost for very large n (shown in fig. 3 ), while band-toep can apply existing amplification by sampling. treeagg @xcite ) that is widely used in industry @xcite , and treeagg-full @xcite that achieves better privacy-utility trade-off but less memory and computation efficient, to represent the tree aggregation mechanisms. bandmf choquette-choo et al. (2023) exploits the banded structure, i.e., @xmath4,@xmath5, ∀|i -j| ≥ b, to simplify the optimization and privacy accounting for mf mechanisms. bandmf successfully applied mf mechanisms to the fl system for the first time. when fixing all the other configurations for training a production lm, bandmf improved the dp guarantee from ρ = 0.52-zcdp by treeagg to ρ = 0.24-zcdp. however, bandmf has to estimate the band size b and total rounds n for optimizing matrices before training, and the performance quickly drops when the actual min-sep b in fl training is smaller than b, or the training round is more than n. bandmf improves memory usage of mf from n × m to b × m for correlated noise, but the typical value of min-sep b in fl is still hundreds to thousands for strong dp guarantees. more recently, bandfhu @xcite and bandtoep (mckenna, 2024) optimize banded toeplitz matrices for larger n and exploit toeplitz structure for computation efficiency, but they have not been shown to outperform bandmf in the fl setting.