Training Mixed-Domain Translation Models via Federated Learning (2022)

Abstract: Training mixed-domain translation models is a complex task that demands tailored architec- tures and costly data preparation techniques. In this work, we leverage federated learning (FL) in order to tackle the problem. Our investiga- tion demonstrates that with slight modifications in the training process, neural machine trans- lation (NMT) engines can be easily adapted when an FL-based aggregation is applied to fuse different domains. Experimental results also show that engines built via FL are able to perform on par with state-of-the-art baselines that rely on centralized training techniques. We evaluate our hypothesis in the presence of five datasets with different sizes, from different domains, to translate from German into English and discuss how FL and NMT can mutually benefit from each other. In addition to provid- ing benchmarking results on the union of FL and NMT, we also propose a novel technique to dynamically control the communication band- width by selecting impactful parameters during FL updates. This is a significant achievement considering the large size of NMT engines that need to be exchanged between FL parties.

## introduction federated learning (fl) is a rapidly growing field in the machine learning community. the reason for this popularity is because of its decentralized and private nature. model training in fl is distributed over multiple nodes where each node could be an independent piece of hardware with its own isolated data. this unique feature enables building high-quality models that benefit from external resources without requiring access to them. although fl is a relatively new field @xcite , it has drawn researchers' attention and the community has witnessed a rapid growth. fields such as computer vision have adapted quickly to the fl framework @xcite @xcite @xcite , but others, such as natural language processing (nlp), have not been as quick and only recently have begun to explore @xcite . we believe that the reason for this slow(er) integration of nlp and fl is because representation learning in nlp is complicated and downstream tasks require large and data-hungry models. these requirements are heavy-handed for any fl model and slow down the unification. however, real-world nlp problems necessitate the use of distributed solutions with privacy-preserving characteristics @xcite . our work is an effort towards leveraging fl-based solutions in nlp. in this paper, we focus on nmt to combine it with fl. a review of the nlp literature shows that almost all recent groundbreaking architectures have been first proposed, or at least evaluated, for translation @xcite @xcite . this is an indication of the intricacy of the nmt task. therefore, it is fair to claim that any fl technique that is capable of training high-quality nmt models could also be considered as a trustworthy alternative for other nlp tasks. thus, nmt could be a strong candidate for fl benchmarking purposes. nmt also has other unique features that can directly benefit fl. one key factor in a fair simulation of fl is to mimic data heterogeneity (also referred to as non-iidness) in experimental environments @xcite . different sampling techniques have been proposed to model such data distribution processes @xcite @xcite , however there is no guarantee that what we simulate is what we encounter in real-world applications. nmt, to some extent, solves this problem since parallel datasets by nature have such heterogeneity. there are only a few nlp tasks that can provide as large and diverse training corpora as nmt. for some language pairs, there exist multiple datasets with hundreds of thousands or even millions of parallel sentences. 1 this means, we should have enough data for each fl node. moreover, each node can naturally pick one dataset/language, so we do not have to artificially distribute data. nmt, together with offering rich training data, also has a bi-lingual setting which is a compelling testbed for fl. coping with the complexity of not only different domains, but also different languages, at the same time in a distributed platform is worth investigating. while nmt offers a natural testbed for fl, we argue that the task itself benefits from the offerings of fl. training a mixed-domain translation model is a challenging task. we show that the aggregation phase of fl can greatly help with this challenge, as it efficiently fuses information from different domains. from the privacy perspective, nmt can also benefit from fl. in fact, the necessity of having a private training pipeline is relatively understudied in nmt. it is mostly assumed that all training datasets are available in a homogeneous and compatible format through a central repository, which is usually not attainable in real world. all these reasons make nmt a compelling case to study in the context of fl. ## federated learning for nlp there are several models in the field that have been proposed to leverage fl for nlp. hartmann et al. (2019) studied whether they could improve the ranking of suggestions in the firefox url bar and train a model on user interactions without violating user privacy. they incorporated feedback received from different clients using fedavg which resulted in significant improvements. ji et al. (2019) suggested that the simple averaging strategy used in fedavg might not be sufficient enough, so they improved the aggregation phase by incorporating the significance of each client by using an attention mask to weigh clients. @xcite focused on language modelling and addressed the problem of out-of-vocabulary entries when working with different clients. @xcite investigated the effect of fl in training better and more personalized user/data representations. their results show that when aggregating information via fl, the model quality increases significantly; at the same time, the training pipeline is distributed and private. we also make a similar observation in our experiments, though not at the representation level, but in terms of the final translation quality. ge et al. (2020) proposed a named-entity recognition (ner) model that is trained with fl to work on medical data. their results demonstrate that not only fl preserves privacy, but also outperforms models trained in a centralized fashion. apart from ner, models with similar concerns have been proposed for mobile keyword prediction @xcite , keyword spotting @xcite , and next word prediction @xcite . @xcite utilized fl to pre-train and/or fine-tune bert models @xcite . from a research standpoint, it is worthwhile to understand if it is even feasible to handle such deep models in an fl framework, and whether a simple averaging-based aggregation is enough. they attempted to address these questions and provided supporting results. in that sense, our work is similar to theirs as we also work with complex models. in addition, we also discuss the bandwidth problem to facilitate exchanging large sets of parameters. ## low-cost domain adaptation in fl in our fl setting, we initialize each client with a high-quality but generic nmt engine. clients use local data to fine-tune their internal model and the combination of a pre-loaded model with local data should lead to better quality. clients also connect with the server regularly to transfer local knowledge and contribute to the aggregation phase. in such a process, domain adaptation happens naturally. inspired by @xcite , we implemented this idea and observed a substantial boost in our translation engines. not only are we able to deliver better results but also we train nmt models in a distributed and private fashion. however, we noticed that communication could be quite costly. in the default configuration, for every pull (in algorithm 1) a large nmt engine has to be exchanged, which might not be necessary. in order to clarify why, we designed an experiment whose information is illustrated in figure 1 . in this experiment, we pick three of our engines and train them for 120k steps within the fl pipeline. for each model, we compare tensors from the 120k checkpoint to their 110k peers and measure how much they changed in-between these two checkpoints. more specifically, for a given tensor w, we compute the pair-wise difference between values from the two checkpoints (w @xmath0, then compute the absolute-value norm of the dif-ference tensor (∥w d ∥). the norm value indicates the shift of each tensor during fl rounds. figure 1 provides a histograms of norm values that belong to different tensors from the encoders and decoders of our translation engines (for more information about the engines and datasets see section 5). results from figure 1 , together with our other observations, show that a small subset of tensors diverge substantially (mostly shown in the right half of the figures), but for the rest, there is a heavy concentration around small norms. more interestingly, we realized that this is a pattern that consistently occurs from one round to the next in our fl experiments, namely each tensor either belongs to a set of highly-fluctuating parameters or it only changes marginally and lies in the less active set. the fluctuation threshold can change but tensors almost always stay in their respective clusters between rounds. we used this finding as a basis of our design to improve bandwidth consumption, such that we decided to focus on either highly fluctuating tensors or those in the other cluster and only pull one type of tensors during communication. the strategy is simple but has led to promising results. more formally, a variation of the aforementioned idea can be simply formulated as shown in equation 2: dp c g in the r-th round consists of all tensors (w t r ) that deviate from their previous values in the previous round by θ. dp stands for dynamic pulling and g indicates that the difference norm of candidate tensors should be greater than or equal to the threshold. dp is exclusively computed for each client which is specified with the c superscript. moreover, the dp sets for decoders and encoders are calculated separately as they vary at different scales. based on equation 2, we do not need to pull all tensors but each client decides what to share with s (only highly-fluctuating tensors in this case). figure 2 visualizes this concept. as shown in the figure, c 1 computes the difference between tensors from rounds r and r -1 and decides to only share two tensors (vertices in diagrams) that have the highest norms. in this scenario, the communication cost between c 1 and s is approximately reduced by 78% for the pulling phase as only 2 out of 9 tensors (22%) are transferred. the exact percentage of the bandwidth saved in this communication protocol directly depends on the client's architecture and θ, but the figure uses an imaginary scenario to explain the reduction mechanism. the intuition behind dp c g is that only highlyfluctuating tensors should be involved in the communication process. it assumes that the main adaptation (or out-domain) knowledge lies in those tensors and what each client needs to learn about its external world is only communicated through such active tensors. therefore, clients only need to exchange them with the server. this is an assumption that might either result in effective communication or conversely hurt the client. because, if by any chance local (or in-domain) knowledge is stored in such tensors, dp c g manipulates the most important parameters and overwrites them with external domains' information. in other words, there is a possibility that the reason for observing high fluctuations in active tensors is not because they carry the community knowledge but because they are responsible to learn local data, so they have to vary frequently to adjust and learn local data. if this second assumption is correct, modifying the active set can easily delay the convergence of local models, and thus deteriorate their quality. due to this concern, we propose another alternative, dp c l , which relies on less active tensors for communications. the pulling condition for dp c l is as in equation 3 : in the next section, we compare dp c l and dp c g and show which strategy is more impactful. ## conclusion and future work in this paper, we reported a set of benchmarking result for nmt in the context of fl. we also proposed an effective technique to reduce the communication bandwidth. our solution tries to determine a subset of parameters that are responsible for learning out-domain knowledge and only exchanges them with the server. in future work, we are interested in i) adding more languages to train multilingual engines, ii) improving communication protocols even further, iii) comparing other fl algorithms in the presence of dp, and iv) investigating nmt in cross-device settings.