FedNLP: Benchmarking Federated Learning Methods for Natural Language Processing Tasks

Bill Yuchen Lin; Chaoyang He; Zihang Zeng; Hulin Wang; Yufen Huang; Christophe Dupuy; Rahul Gupta; Mahdi Soltanolkotabi; Xiang Ren; Salman Avestimehr

Source-linked AI summary

FedNLP: Benchmarking Federated Learning Methods for Natural Language Processing Tasks

Bill Yuchen Lin, Chaoyang He, Zihang Zeng, Hulin Wang, Yufen Huang, Christophe Dupuy, Rahul Gupta, Mahdi Soltanolkotabi, Xiang Ren, Salman Avestimehr

arXiv:2104.08815v3Apr 18, 2021cs.CLcs.AIcs.LG

TL;DR

Privacy constraints and fragmented NLP data motivate federated learning, yet systematic comparisons across NLP tasks are limited. FedNLP provides a benchmark linking Transformer models, FL algorithms, and non-IID partitions across four task formulations, revealing both method-dependent performance and persistent challenges for federated NLP.

Problem
Systematic comparisons of federated-learning methods across NLP tasks are lacking despite privacy-driven data decentralization.
Method
FedNLP benchmarks FL methods with Transformer models across four NLP task formulations and multiple non-IID data partitions.
Results
FedOPT generally outperforms FedAvg and FedProx, except on seq2seq generation, while performance depends on task structure and partition heterogeneity.
Takeaways & Limitations
Federated optimization algorithms may need to be tailored to NLP tasks, with FL-friendly architectures or loss functions as promising directions.
Takeaways & Limitations
The study does not analyze privacy leakage, which remains an open problem beyond its scope.

Abstract

from arXiv · show

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, a benchmarking framework for evaluating federated learning methods on four different task formulations: text classification, sequence tagging, question answering, and seq2seq. We propose a universal interface between Transformer-based language models (e.g., BERT, BART) and FL methods (e.g., FedAvg, FedOPT, etc.) under various non-IID partitioning strategies. Our extensive experiments with FedNLP provide empirical comparisons between FL methods and helps 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.

1 Introduction

FedNLP addresses the lack of systematic evaluation for federated learning in NLP, where privacy constraints distribute data across clients. It provides a benchmarking framework for comparing methods across task formulations, models, and non-IID partitions.

Privacy regulations and fragmented data silos prevent transferring users’ data to centralized servers or other clients.This motivates decentralized NLP training across organizations and personal devices.
Federated learning enables clients to train jointly while keeping raw data decentralized and private through model-parameter exchanges.Clients send model parameters to a server, which updates and redistributes a global model.
Prior NLP federated-learning studies used task-specific experimental setups, making fair comparison across methods and tasks difficult.The paper identifies the absence of a universal benchmarking platform in the literature.
FedNLP compares popular federated-learning methods across text classification, sequence tagging, question answering, and seq2seq generation.Existing datasets are partitioned into non-IID client distributions to simulate realistic shifts.
FedNLP connects Transformer-based language models with multiple federated-learning algorithms and supports varied non-IID partitioning strategies.The framework is modular across task formulations, NLP models, algorithms, and partitions.
The paper contributes systematic evaluation and analysis, an open-source framework, and decentralized NLP datasets for future research.The analysis reports a gap between centralized and decentralized training and examines efficiency across methods and model sizes.

2 Federated Learning for NLP

The paper frames federated learning as decentralized optimization in which clients retain local data while collaboratively updating a global model. FedNLP instantiates this framework with generalized optimization interfaces connecting Transformer models, FL algorithms, NLP tasks, and partitioning strategies.

Federated learning keeps raw client data local and exchanges focused model updates coordinated by a central server.This design is intended to reduce data movement and address privacy-related storage constraints.
The global objective averages client-specific losses while allowing local data distributions to vary across clients.This formulation explicitly captures data heterogeneity.
Federated averaging: FedAvg trains in rounds: the server broadcasts a global model, clients perform local SGD updates, and the server aggregates their changes.The procedure repeats until convergence.
Federated averaging: In the cross-silo setting, all clients participate in every round, enabling a shared global model while preserving client data privacy.The cohort is the entire client population in this setting.
FedOPT framework: FedOPT generalizes FedAvg by using separate client and server gradient-based optimizers, with the server treating aggregated local changes as a pseudo-gradient.FedNLP uses this universal optimization framework to simplify aggregation-based algorithm design.
FedNLP framework: FedNLP connects multiple FL algorithms with Transformer models and supports diverse NLP tasks, models, non-IID partitions, and secure aggregation methods.The framework is designed for benchmarking and future development in federated NLP.

3 Benchmarking Setup with FedNLP

FedNLP benchmarks federated learning across four NLP task formulations using representative datasets, Transformer models, and synthetic or natural non-IID client partitions. Its setup varies label, quantity, feature, and natural distribution shifts to evaluate decentralized learning methods under shared experimental conditions.

Task Formulations, Datasets, and Models: FedNLP covers text classification, sequence tagging, question answering, and seq2seq generation with representative datasets including 20Newsgroup, OntoNotes, MRQA, and Gigaword.The framework supports a wider range of tasks, but these representative tasks are used in the reported experiments.
Task Formulations, Datasets, and Models: The experiments use Transformer-based pretrained language models, including BART for seq2seq tasks.The setup also uses DistilBERT and BART-base, balancing model performance with computational cost for federated applications.
Non-IID Partitioning Strategies: Non-IID label partitions assign clients categorical label distributions sampled from Dirichlet(αp), where α controls similarity across clients.As α increases, client distributions approach the prior; as α approaches zero, clients tend toward examples from a single class.
Non-IID Partitioning Strategies: Quantity shift is modeled by allocating dataset examples according to a client-level Dirichlet distribution z ∼ Dir_N(β), optionally combined with label shift.A uniform allocation can be represented by z ∼ U(N), or by β approaching infinity.
Non-IID Partitioning Strategies: Feature clustering extends synthetic partitioning beyond classification, while MRQA provides a natural cross-silo shift through dataset source and annotating style.The feature-based method is intended for non-classification NLP tasks such as summarization.

4 Experimental Results and Analysis

FedNLP compares federated learning methods across NLP tasks, data partitions, model choices, and communication-saving strategies. The experiments show that method performance depends on task structure and data heterogeneity, while compact models and selective freezing can reduce costs with qualifications.

Experimental setup: FedNLP uses DistilBERT and BART-base for most experiments, citing DistilBERT’s 7x speed improvement over BERT-base on mobile devices and BART-base’s performance–computation trade-off.The setup targets federated-learning scenarios where communication and device computation matter.
FL method comparison: FedOPT generally outperforms FedAvg and FedProx, except in seq2seq generation, while FedAvg exceeds FedProx in sequence tagging.The comparison uses the same data partition, communication rounds, and other settings for each task formulation.
FL method comparison: FedOPT initially outperforms the other methods but gradually becomes worse over time, suggesting task-dependent optimization behavior.The paper conjectures that differing loss functions, especially in structured prediction tasks, affect federated performance.
Non-IID partitions: Smaller α values degrade performance and increase variance, whereas quantity skew is less challenging when label distributions are closer to uniform.Uniform and quantity-skew partitions produce smoother curves, and larger α reduces variance.
Model efficiency: Freezing Transformer layers can reduce communication costs, but FedOPT must unfreeze the last three layers to approach full-model performance.Centralized training achieves its largest gain after unfreezing the last layer, so centralized experience may not generalize to FL.
Model efficiency: DistilBERT is not significantly worse than BERT-base in the reported comparison, while BERT-base is almost 2x larger.The paper therefore presents DistilBERT as a more cost-effective choice for experiments and realistic applications.

5 Related Work

Prior work applies federated learning to selected privacy-oriented NLP applications, but those studies are typically tied to particular tasks and experimental settings. FedNLP provides a shared environment for comparing methods across broader NLP applications and treats prior applications as supported instances.

Federated learning in NLP: Prior federated NLP applications include keyboard systems, intent classification, medical relation extraction, and medical named entity recognition.These applications use federated learning to avoid exchanging sensitive data across platforms.
Federated learning in NLP: Existing applications often use task-specific experimental setups, making fair comparison across FL methods and NLP tasks difficult.This motivates a universal benchmarking platform for federated learning in NLP.
FedNLP’s positioning: FedNLP offers a unified platform for studying varied NLP applications in a shared environment, with prior federated NLP applications represented as supported settings.The platform is intended to help researchers design task-specific or general-purpose FL methods.

6 Conclusion and Future Directions

FedNLP contributes a comprehensive empirical analysis and an open-source framework for studying federated learning with NLP models. The paper identifies performance, efficiency, scalability, privacy, and personalization as future directions.

Contributions: FedNLP compares FL algorithms across four NLP task formulations and multiple non-IID partitions, revealing both promise and challenges for FL in NLP.The framework connects Transformer models with popular FL methods and partition strategies.
Contributions: The open-source FedNLP resources support future research on federated learning for NLP.The paper emphasizes a maintained codebase and a unifying framework for connecting Transformer models, FL methods, and non-IID partitions.
Future directions: Future directions include narrowing the performance gap, improving system efficiency and scalability, and developing trustworthy, privacy-preserving, and personalized NLP methods.These directions are stated as priorities for advancing FL in NLP.

Ethical Considerations and Limitations(*)

FedNLP addresses privacy-sensitive, decentralized NLP settings but leaves privacy leakage analysis outside its scope. Applying FL to NLP also involves heterogeneous task formats, realistic non-IID partitions, and integration challenges between FL methods and Transformer models.

Limitations: FedNLP does not analyze privacy leakage, because privacy-centric measures are treated as orthogonal and fair analysis remains an open problem.The authors state that privacy leakage can be studied when a platform such as FedNLP exists.
Privacy context: Federated learning keeps user data local while clients jointly train a shared model, addressing settings where private data cannot be transferred.The paper describes personal devices and organizational data silos as examples of this setting.
Research gap: A standardized FL-for-NLP platform is needed because existing platforms often lack connections among pretrained language models, NLP datasets, and federated methods.The paper identifies this absence as a barrier to systematic research.
Technical challenges: A comprehensive platform must support varied task input and output formats, realistic non-IID partitions, and flexible integration of FL methods with Transformer models.The conventional Transformer trainer also requires modification for efficient and safe federated communication.

B Basic Formulations of NLP Tasks

The paper organizes mainstream NLP applications into four task formulations—text classification, sequence tagging, question answering, and sequence-to-sequence generation—while also discussing language modeling and extensibility to other tasks.

Four formulations cover many mainstream NLP applications: text classification, sequence tagging, question answering, and sequence-to-sequence generation.
Text Classification (TC): Text classification maps a word sequence to one label from a fixed label set, supporting tasks such as topic and sentiment classification.
Sequence Tagging (ST): Sequence tagging assigns one label to each input token, enabling syntactic analysis and span extraction such as named entity recognition.
Question Answering (QA): Question answering takes a passage and question as input and returns the start and end token indices of the answer span.
Sequence-to-Sequence Generation: Sequence-to-sequence generation uses word sequences as both input and output, supporting applications such as summarization, dialogue response generation, and machine translation.
Language Modeling (LM): Language modeling predicts the most plausible next token from an incomplete sequence and can benefit downstream applications through fine-tuning.
Others: FedNLP also allows users to implement task formulations beyond the four basic categories through its extensible platform.

C Implementation Details

The implementation addresses label-distribution non-IID partitioning by dynamically reallocating examples when clients lack enough data for particular labels.

Label-distribution partitioning can leave some clients without enough examples to sample particular labels.
The proposed dynamic reassignment fills label vacancies with examples from other labels according to the current ratio of remaining unassigned examples.

C.1 The FedNLP Training Pipeline: Security and Efficiency

FedNLP is designed as a modular FL research platform supporting diverse algorithms, secure aggregation, and distributed execution for NLP benchmarking.

Supporting diverse FL algorithms: FedNLP supports classical FL algorithms including FedAvg, FedOPT, and FedProx through modular algorithmic APIs.
Enabling secure benchmarking with lightweight secure aggregation: The platform integrates lightweight secure aggregation because client model weights may still leak private information during federated training.
Enabling secure benchmarking with lightweight secure aggregation: Secure aggregation masks client models so the server sees only the aggregated model rather than each client’s raw model.
Realistic evaluation with efficient distributed system design: FedNLP uses application, algorithm, and infrastructure layers to support distributed training across edge servers and devices.
Federated Learning Methods: The broader FL literature addresses statistical heterogeneity, trustworthiness, system optimization, personalization, and fairness through multiple algorithmic directions.
Existing FL algorithms and systems had not been fully evaluated on the realistic NLP tasks introduced by this paper.

E Future Directions

FedNLP’s future directions emphasize closing the centralized–decentralized performance gap, improving efficiency and scalability, strengthening privacy analysis, and enabling personalization; its implementation provides layered interfaces for these extensions.

Minimizing the performance gap: Federated fine-tuning retains a large accuracy gap relative to centralized fine-tuning on non-IID data, motivating algorithms for Transformer-based NLP.
Improving the system efficiency and scalability: Efficient FL methods must address large Transformer models, resource-constrained edge devices, stragglers, user selection, and convergence speed.
Trustworthy and privacy-preserving NLP: Privacy-preserving capability remains an open research problem, while this paper focuses on implementing and analyzing FL methods for NLP.
Trustworthy and privacy-preserving NLP: FedNLP provides flexible interfaces for future analysis and improvement of privacy preservation in FL methods for NLP and beyond.
Personalized FedNLP: Personalized FedNLP could target accuracy, fairness, and heterogeneous client architectures, including adaptation of small parameter subsets through prefix tuning.
Implementation architecture: The platform separates application, algorithm, and infrastructure layers, with modules for data management, model definition, training, FL algorithms, and distributed learning.
Implementation workflow: Federated training initializes devices, loads client data, creates the model, and invokes the selected FL algorithm through the workflow API.
Application layer: FedNLP supports reusable Transformer and LSTM models, HuggingFace-compatible components, customizable data managers, and single-process task trainers.

F.5 Enhancing Security with Secure Aggregation (SA)

FedNLP incorporates secure aggregation algorithms for privacy-preserving federated NLP. LightSecAgg preserves the model privacy guarantees of SecAgg and SecAgg+ while reducing aggregation complexity and runtime.

LightSecAgg provides the same model privacy guarantees as SecAgg and SecAgg+ while substantially reducing aggregation complexity and runtime.The paper presents this comparison as a performance result illustrated in a figure.
Each user protects its local model with a locally generated random mask that is encoded and shared with other users.
The server can directly reconstruct the aggregate mask from any sufficiently large set of surviving users.
FedNLP integrates secure aggregation algorithms, optimizes system performance, and provides user-friendly APIs compatible with NLP models and federated learning algorithms.