SlowGuess/NAACL_2022_OCR huggingface dataset:
Bill Yuchen Lin $^{1*}$ , Chaoyang He $^{1*}$ , Zihang Zeng $^{1}$ , Hulin Wang $^{1}$ , Yufen Huang $^{1}$ , Christophe Dupuy $^{2}$ , Rahul Gupta $^{2}$ , Mahdi Soltanolkotabi $^{1}$ , Xiang Ren $^{1*}$ , Salman Avestimehr $^{1*}$
University of Southern California Amazon Alexa AI2
{yuchen.lin,chaoyang.he,saltanol,xiangren,avestime}@usc.edu gupra@amazon.com
Increasing concerns and regulations about data privacy and sparsity necessitate the study of privacy-preserving, decentralized learning methods for natural language processing (NLP) tasks. Federated learning (FL) provides promising approaches for a large number of clients (e.g., personal devices or organizations) to collaboratively learn a shared global model to benefit all clients while allowing users to keep their data locally. Despite interest in studying FL methods for NLP tasks, a systematic comparison and analysis is lacking in the literature. Herein, we present the FedNLP<sup>1</sup>, a benchmarking framework for evaluating federated learning methods on four common formulations of NLP tasks: text classification, sequence tagging, question answering, and seq2seq generation. We propose a universal interface between Transformer-based language models (e.g., BERT, BART) and FL methods under various non-IID partitioning strategies. Our extensive experiments with FedNLP provide empirical comparisons between FL methods and help us better understand the inherent challenges of this direction. The comprehensive analysis points to intriguing and exciting future research aimed at developing FL methods for NLP tasks.
Fine-tuning large pre-trained language models (LMs) such as BERT (Devlin et al., 2019) often leads to state-of-the-art performance in many realistic NLP applications (e.g., text classification, named entity recognition, question answering, summarization, etc.), when large-scale, centralized training datasets are available. However, due to the increasing concerns and regu
Figure 1: The FedNLP benchmarking framework.
lations about data privacy (e.g., GPDR (Regulation, 2016)) emerging data from realistic users have been much more fragmented and distributed, forming decentralized private datasets of multiple "data silos" (a data silo can be viewed as an individual dataset) — across different clients (e.g., organizations or personal devices).
To respect the privacy of the users and abide by these regulations, we must assume that users' data in a silo are not allowed to transfer to a centralized server or other clients. For example, a client cannot share its private user data (e.g., documents, conversations, questions asked on the website/app) with other clients. This is a common concern for organizations such as hospitals, financial institutions, or legal firms, as well as personal computing devices such as smartphones, virtual assistants (e.g., Amazon Alexa, Google Assistant, etc.), or a personal computer. However, from a machine learning perspective, models trained on
*Bill and Chaoyang contributed equally; Xiang and Salman are equal advisors for this work.
https://github.com/FedML-AI/FedNLP
Findings of the Association for Computational Linguistics: NAACL 2022, pages 157 - 175
a centralized dataset that combine the data from all organizations or devices usually result in better performance in the NLP domain. Therefore, it is of vital importance to study NLP problems in such a realistic yet more challenging scenario — i.e., training data are distributed across different clients and cannot be shared for privacy concerns.
The nascent field of federated learning (et al, 2019; Li et al., 2020a) (FL) aims to enable many individual clients to train their models jointly while keeping their local data decentralized and completely private from other users or a centralized server. A common training schema of FL methods is that each client sends its model parameters to the server, which updates and sends back the global model to all clients in each round. Since the raw data of one client has never been exposed to others, FL is promising as an effective way to address the above challenges, particularly in the NLP domain, where many user-generated text data contain sensitive and/or personal information.
Despite the growing progress in the FL domain, research into and application for NLP has been rather limited. There are indeed several recent works on using FL methods for processing medical information extraction tasks (Sui et al., 2020). However, such prior work usually has its experimental setup and specific task, making it difficult to fairly compare these FL methods and analyze their performance in other NLP tasks. We argue that future research in this promising direction (FL for NLP) would highly benefit from a universal benchmarking platform for systematically comparing different FL methods for NLP. To the best of our knowledge, such a benchmarking platform is still absent from the literature.
Therefore, our goal in this paper is to provide comprehensive comparisons between popular FL methods (e.g., FedAvg (McMahan et al., 2017a), FedOPT (Reddi et al., 2021), FedProx (Li et al., 2020b)) for four mainstream formulations of NLP tasks: text classification, sequence tagging, question answering, and seq2seq generation. Although there are few available realistic FL datasets for NLP due to privacy concerns, we manage to use existing NLP datasets to create various non-IID data partitions over clients. These non-IID partitions simulate various kinds of distribution shifts (e.g., label, features, quantities, etc.) over the
clients, which often happen in real-world NLP applications. As for the base NLP models, we use the Transformer architecture (Vaswani et al., 2017) as the backbone and support a wide range of pre-trained LMs such as DistilBERT (Sanh et al., 2019), BERT (Devlin et al., 2019), BART (Lewis et al., 2020), etc. To conduct extensive experiments, we need to support the experiments with multiple options on dimensions such as (1) task formulations, (2) NLP models, (3) FL algorithms, and (4) non-IID partitions. Therefore, we propose FedNLP, a modular framework with universal interfaces among the above four components, which is thus more extensible for supporting future research in FL for NLP.
We aim to unblock the research of FL for NLP with the following two-fold contributions:
- Evaluation and analysis. We systematically compare popular federated learning algorithms for mainstream NLP task formulations under multiple non-IID data partitions, which thus provides the first comprehensive understanding. Our analysis reveals that there is a considerably large gap between centralized and decentralized training in various settings. We also analyze the efficiency of different FL methods and model sizes. With our analysis, we highlight several directions to advance FL for NLP.
- Resource. The implementation of our experiments also forms a general open-source framework named FedNLP, which is capable of evaluating, analyzing, and developing FL methods for NLP. We also provide decentralized NLP datasets of various task formulations created by various non-IID partitioning strategies for future research.
The remainder of this paper is structured as follows. We introduce the background knowledge of federated learning and several typical FL algorithms in §2. Then, we present the proposed non-IID partitioning strategies to create synthetic datasets for different task formulations in §3. Our results, analysis, and findings are in §4. Finally, we discuss related work (§5) and conclusions (§6).
2 Federated Learning for NLP
In this section, we first introduce the background knowledge of federated learning (FL) in the con
text of NLP tasks. Then, we illustrate a unified FL framework that we used to study typical FL algorithms. Based on this, we build our research framework, a general pipeline for benchmarking and developing FL methods for NLP.
2.1 Federated Learning Concepts
Federated learning (FL) is a machine learning paradigm where multiple entities (clients) collaborate in solving a machine learning problem under the coordination of a central server or service provider. Each client's raw data is stored locally and not exchanged or transferred; instead, focused updates intended for immediate aggregation are used to achieve the learning objectives (Kairouz et al., 2019). Therefore, federated learning has been seen as a promising direction to decrease the risk of attack and leakage, reduce the difficulty and cost of data movement, and meet the privacy-related data storage regulations.
In the basic conception of federated learning, we would like to minimize the objective function,
$$
F (\boldsymbol {x}) = \mathbb {E} _ {i \sim \mathcal {P}} [ F _ {i} (\boldsymbol {x}) ],
$$
where  $F_{i}(\pmb {x}) = \mathbb{E}_{\xi \sim \mathcal{D}_{i}}[f_{i}(\pmb {x},\xi)]$  (1)
$\pmb{x} \in \mathbb{R}^d$  represents the parameter for the global model,  $F_i: \mathbb{R}^d \to \mathbb{R}$  denotes the local objective function at client  $i$ , and  $\mathcal{P}$  denotes a distribution on the collection of clients  $\mathcal{I}$ . The local loss functions  $f_i(\pmb{x}, \xi)$  are often the same across all clients, but the local data distribution  $\mathcal{D}_i$  will often vary, capturing data heterogeneity.
Federated averaging (FedAvg) (McMahan et al., 2017a) is a common algorithm to solve (1) by dividing the training process into rounds. At the beginning of the  $t$ -th round ( $t \geq 0$ ), the server broadcasts the current global model  $\boldsymbol{x}^{(t)}$  to a cohort of participants: a random subset of clients from  $S^{(t)}$  which includes  $M$  clients in total. Then, each sampled client in the round's cohort performs  $\tau_{i}$  local SGD updates on its own local dataset and sends the local model changes  $\Delta_{i}^{(t)} = \boldsymbol{x}_{i}^{(t,\tau_{i})} - \boldsymbol{x}^{(t)}$  to the server. Finally, the server uses the aggregated  $\Delta_{i}^{(t)}$  to update the global model:  $\boldsymbol{x}^{(t+1)} = \boldsymbol{x}^{(t)} + \frac{\sum_{i \in S(t)} p_{i} \Delta_{i}^{(t)}}{\sum_{i \in S(t)} p_{i}}$ . where  $p_{i}$  is the relative weight of client  $i$ . The above procedure will repeat until the algorithm converges. In the cross-silo setting where all clients participate in training on every round (each cohort is the entire population),
we have  $\mathcal{S}^{(t)} = \{1,2,\dots ,M\}$ . Consequently, we can learn a global model to benefit all clients while preserving their data privacy.
2.2 Our Unified Framework for FL
In this work, we propose to use FedOPT (Reddi et al., 2021), a generalized version of FedAvg, to build the FedNLP platform. As the pseudo-code presented in Algorithm 1, the algorithm is parameterized by two gradient-based optimizers: CLIENTOPT and SERVEROPT with client learning rate  $\eta$  and server learning rate  $\eta_{s}$ , respectively. While CLIENTOPT is used to update the local models, SERVEROPT treats the negative of aggregated local changes  $-\Delta^{(t)}$  as a pseudo-gradient and applies it to the global model. This optimization framework generalizes to many aggregation-based FL algorithms and simplifies the system design.
To make our research general, we explore different combinations of SEVEROPT and CLIENTOPT. The original FedAvg algorithm implicitly sets SEVEROPT and CLIENTOPT to be SGD, with a fixed server learning rate  $\eta_{s}$  of 1.0. FedProx (Li et al., 2020b), tackling statistical heterogeneity by restricting the local model updates to be closer to the initial (global) model, can be easily incorporated into this framework by adding L2 regularization for better stability in training. Moreover, given that AdamW (Loshchilov and Hutter, 2019) is widely used in NLP, we set it for ClientOpt and let the ServerOpt be SGD with momentum
to reduce the burden of tuning.
2.3 The Proposed FedNLP Framework
To support our research in this paper and other future work in the area of federated learning for NLP, we build a general research framework named FedNLP, based on the above universal optimization framework. We here briefly highlight its unique features and leave the details in the following content and a detailed design is shown in App. F. First, FedNLP is the very first framework that connects multiple FL algorithms with Transformer-based models, to our best knowledge. Also, we implement a flexible suite of interfaces to support different types of NLP tasks and models, as well as different non-IID partitioning strategies (Sec. 3.2). To study security and privacy guarantees, we incorporate state-of-the-art secure aggregation algorithms such as LightSecAgg (see F.5).
3 Benchmarking Setup with FedNLP
In this section, we introduce the creation of our benchmark datasets from a set of chosen NLP tasks with different non-IID partition methods. We evaluate various FL methods on these datasets.
3.1 Task Formulations, Datasets, and Models
There are numerous NLP applications, but most of them can be categorized based on four mainstream formulations: text classification (TC), sequence tagging (ST), question answering (QA), and seq2seq generation (SS). The formal definition of each formulation is detailed in Appendix §B. To cover all formulations while keeping our experiments in a reasonable scope, we select one representative task for each formulation:
- Text Classification: 20Newsgroup (Lang, 1995) is a news classification dataset with annotations for 20 labels. We showcase our FedNLP with this dataset as it has a larger output space (20 labels) than sentiment-analysis datasets, which is an important factor for the label-distribution shift scenarios.
- Sequence Tagging: OntoNotes (Pradhan et al., 2013) (5.0) is a corpus where sentences have annotations for the entity spans and types. We use it for the named entity recognition task, which is fundamental to information extraction and other applications.
Task	Txt.Cls.	Seq.Tag.	QA	Seq2Seq
Dataset	20News	Onto.	MRQA	Giga.
# Training	11.3k	50k	53.9k	10k
# Test	7.5k	5k	3k	2k
# Labels	20	37*	N/A	N/A
Metrics	Acc.	F-1	F-1	ROUGE
Table 1: Statistics of the selected datasets for our experiments. *37 is the size of the tag vocabulary.
QA: MRQA (Fisch et al., 2019) is a benchmark consisting of 6 popular datasets<sup>2</sup>: SQuAD (Rajpurkar et al., 2016) (8529/431), NewsQA (Trischler et al., 2017) (11877/613), TriviaQA (Joshi et al., 2017) (4120/176), SearchQA (Dunn et al., 2017) (9972/499), HotpotQA (Yang et al., 2018b), and NQ (Kwiatkowski et al., 2019) (9617/795).
- Seq2Seq Generation: Gigaword (DBL, 2012) is a news corpus with headlines that are often used for testing seq2seq models as a summarization task. Other tasks such as dialogue response generation and machine translation can also be adapted to this format.
We show the basic statistics of the above datasets in Table 1. Note that our FedNLP as a research platform supports a much wider range of specific tasks of each formulation, while we only introduce the ones used in our experiments here with typical settings. Moreover, our contribution is more of a general FL+NLP benchmarking platform instead of particular datasets and partitions.
Base NLP Models. Fine-tuning pre-trained LMs has been the de facto method for NLP research, so we focus on testing Transformer-based architectures in FedNLP. Specifically, we choose to use BART (Lewis et al., 2020), a text-to-text Transformer model similar to the T5 model (Raffel et al., 2020), for seq2seq tasks.
3.2 Non-IID Partitioning Strategies
The existing datasets have been used for centralized training in NLP. As our focus here is to test decentralized learning methods, we need to distribute the existing datasets to a set of clients. It is the non-IIDness of the client distribution that
2We only use part of the data to demonstrate and verify our hypothesis; we show the train/test split in brackets.
makes federated learning a challenging problem. Thus, we extend the common practice widely used in prior works to the NLP domain for generating synthetic FL benchmarks (Li et al., 2021a). We first introduce how we control the label distribution shift for TC and ST, then the quantity distribution shift, and finally how we model the distribution shift in terms of input features for non-classification NLP tasks (e.g., summarization).
Non-III Label Distributions. Here we present how we synthesize the data partitions such that clients share the same (or very similar) number of examples, but have different label distributions from each other. We assume that on every client training, examples are drawn independently with labels following a categorical distribution over  $L$  classes parameterized by a vector  $\mathbf{q}$  ( $q_i \geq 0, i \in [1, L]$  and  $\| \mathbf{q} \|_1 = 1$ ). To synthesize a population of non-identical clients, we draw  $\mathbf{q} \sim \mathrm{Dir}_L(\alpha \mathbf{p})$  from a Dirichlet distribution, where  $\mathbf{p}$  characterizes a prior class distribution over  $L$  classes, and  $\alpha > 0$  is a concentration parameter controlling the identicalness among clients. For each client  $C_j$ , we draw a  $\mathbf{q}_j$  as its label distribution and then sample examples without replacement from the global dataset according to  $\mathbf{q}_j$ . With  $\alpha \to \infty$ , all clients have identical distributions to the prior (i.e., uniform distribution); with  $\alpha \to 0$ , on the other extreme, each client holds examples from only one class chosen at random. In Fig. 2, we show heatmaps for visualizing the distribution differences between each client. Figure 3 shows an example of the concrete label distributions for all clients with different  $\alpha$ . We can see that when  $\alpha$  is smaller, the overall label distribution shift becomes larger.
Controlling non-III Quantity. It is also common that different clients have very different data quantities while sharing similar label distribution. We thus also provide a quantity-level Dirichlet allocation  $z \sim \mathrm{Dir}_N(\beta)$  where  $N$  is the number of clients. Then, we can allocate examples in a global dataset to all clients according to the distribution  $z$  — i.e.,  $|\mathcal{D}_i| = z_i|\mathcal{D}_G|$ . If we would like to model both quantity and label distribution shift, it is also easy to combine both factors. Note that one could assume it is a uniform distribution  $z \sim U(N)$ , (or  $\beta \to \infty$ ) if we expect all clients to share a sim
Figure 2: The  $J-S$  divergence matrix between 100 clients on the 20News dataset when  $\alpha \in \{1,5,10,100\}$ . Each sub-figure is a  $100 \times 100$  symmetric matrix. The intensity of a cell  $(i,j)$ 's color here represents the distance between the label distribution of Client  $i$  and  $j$ . It is expected that when  $\alpha$  is smaller, the partition over clients is more non-IID in terms of their label distributions.
ilar number of examples. A concrete example is shown in Figure 8 (Appendix).
Controlling non-HID Features. Although straightforward and effective, the above label-based Dirichlet allocation method has a major limitation — it is only suitable for text classification tasks where the outputs can be modeled as category-based random variables. To create synthetic partitions for other non-classification NLP tasks and model distribution shifts, we thus propose a partition method based on feature clustering. Specifically, we use SentenceBERT (Reimers and Gurevych, 2019) to encode each example to a dense vector by their text then we apply K-Means clustering to get the cluster label of each example; finally, we use these cluster labels (as if they were classification tasks) to follow the steps in modeling label distribution shift. There are two obvious benefits of this clustering-based Dirichlet partition method: 1) It enables us to easily synthesize the FL datasets for non-classification tasks (i.e., ST, QA, SS) as they do not have discrete labels as output space; 2) The BERT-based clustering results naturally imply different sub-topics of a dataset, and thus feature shift can be seen as a shift of latent labels — we can reuse the same method for the label-based Dirichlet partition method.
Natural Factors For datasets like MRQA, we consider a cross-silo setting where each client is associated with a particular sub-dataset (out of the six datasets of the same format), forming a natural distribution shift based on the inherent factors such as data source and annotating style.
Task	Dataset	Partition	Clients	FedAvg	FedProx	FedOPT	# Rounds
Text Classification	20news	α=1 (label shift)	100	0.5142	0.5143	0.5349	22
Sequence Tagging	OntoNotes	α=0.1 (label shift)	30	0.7382	0.6731	0.7918	17
Question Answering	MRQA	natural factor	6	0.2707	0.2706	0.3280	13
Seq2Seq Generation	Gigaword	α=0.1 (feature shift)	100	0.3192	0.3169	0.3037	13
Table 2: The comparisons between different FL methods under the same setting on different NLP tasks. The number of workers per round are 10, expect for the MRQA task, which uses 6.
Figure 3: Visualizing the non-HID label distributions on 20News with  $\alpha$  being  $\{1,5,10,100\}$ . Each subfigure is a  $100 \times 20$  matrix, where 100 is the number of clients, and 20 is the number of labels. The intensity of a cell here represents the ratio of a particular label in the local data of a client. When  $\alpha$  is smaller (1, 5, 10), each client has a relatively unique label distribution, thus the differences between clients are larger; when  $\alpha = 100$ , every client has a nearly uniform label distribution.
In this section, we aim to analyze typical federated learning methods (introduced in our benchmark datasets with multiple dimensions with the base NLP models listed previously. We put more implementation details and additional results in Appendix. We organize our extensive experimental results and findings from the analysis as a collection of research questions with answers.
Experimental Setup and Hyper-parameters.
We use DistilBERT and BART-base for most of our experiments, as the former is a distilled version of the BERT model and has a 7x speed improvement over BERT-base on mobile devices — a common scenario for FL applications; the BART-base model is the most suitable option considering the trade-off between performance and computation cost. We leave our implementation details and the selected hyper-parameters in the
submitted supplementary materials.
Our experiments cover both cross-device and cross-silo settings. As shown in Table 2, in the cross-device setting, we use uniform sampling to select 10 clients for each round when the client number in a dataset is very large (e.g., 100). For the cross-silo setting, each round will select the same number of clients (we use 6 for the QA task). The local epoch number is set to 1 for all experiments. To make our results reproducible, we use wandb.ai to store all experiment logs and hyperparameters as well as running scripts.
Q1: How do popular FL methods perform differently under the same setting?
We compare the three typical FL methods under the same setting (i.e., data partition, communication rounds, etc.) for each task formulation. As shown in Table 2, we report the results of FedAvg, FedProx, and FedOPT. We can see that overall FedOPT performs better than the other two methods, with the only exception being in the seq2seq generation task. FedAvg and FedProx perform similarly with marginal differences, but FedAvg outperforms FedProx in sequence tagging. These two exceptions are surprising findings, as many prior works in the FL community show that FedOPT is generally better than FedProx and FedAvg on vision tasks and datasets.
We conjecture that such inconsistent performance across tasks suggests the difference in terms of the loss functions has a great impact on FL performance. Seq2seq and sequence tagging tasks usually have more complex loss landscapes than text classification, as they are both typical structured prediction tasks, while the text classification has a much smaller output space. From Fig. 4, we see that the FedOPT outperforms the other two methods at the beginning while gradually becoming worse over time.
Figure 4: The learning curves of the three FL Methods on four different task formulations. The metrics used for these tasks are accuracy, span-F1, token-F1, and ROUGE respectively; The x-axis is the number of rounds.