Assortment of Attention Heads: Accelerating Federated PEFT with Head Pruning and Strategic Client Selection
This addresses efficiency and privacy issues for deploying large language models on resource-constrained devices in federated settings, though it is incremental as it builds on existing PEFT and FL methods.
The paper tackles the challenge of efficiently performing Parameter Efficient Fine-Tuning (PEFT) in Federated Learning (FL) for language models by proposing head pruning, weighted aggregation, and client selection, achieving up to 90% sparsity, 1.8x communication advantage, and 3.9x reduction in training operations with under 2% accuracy drop on benchmarks.
Parameter Efficient Fine-Tuning (PEFT) has become the de-facto approach in adapting Large Language Models (LLMs) for downstream tasks in Natural Language Processing. However, its adoption in privacy-preserving distributed learning frameworks, such as Federated Learning (FL), remains relatively limited. This is mainly due to challenges specific to FL, such as resource-constrained devices and diverse data distributions among clients. In this paper, we propose an efficient method to perform PEFT within the FL framework for Multi-Head Attention (MHA) based language models. We address the challenges through head pruning, a novel head-specific weighted aggregation mechanism, and a client selection strategy. Head pruning minimizes training complexity within the clients, guided by the importance score computed based on the confidence of the attention head. Weighted aggregation of heads ensures the global model captures crucial updates from diverse clients complementing our client selection strategy. We show results on the MultiNLI benchmark along with 20 Newsgroups, XL-Sum, and E2E NLG datasets. We use the MultiNLI dataset and T5-small model with LoRA as our PEFT method, attaining sparsity levels of up to 90%, resulting in a communication advantage of up to 1.8x and a reduction in training OPs of 3.9x while maintaining the accuracy drop under 2%.