在联邦学习中通过剪枝注意力头加速低资源语言模型微调。
Assortment of Attention Heads: Accelerating Federated PEFT with Head Pruning and Strategic Client Selection
- 剪枝重要性高的注意力头,降低客户端计算负担。
- 90%剪枝率下通信量提升1.8倍,训练算力减少3.9倍,精度损失<2%。
- 适合资源受限设备上的隐私保护模型部署场景。
参数高效微调(PEFT)已成为自然语言处理中适配大语言模型(LLM)下游任务的主流方法。然而,在如联邦学习(FL)这类保护隐私的分布式学习框架中,其应用仍相对有限,主要受制于设备资源受限和客户端数据分布差异大等问题。本文提出一种在联邦学习框架内对基于多头注意力(MHA)的语言模型进行高效PEFT的方法。通过注意力头剪枝、一种新的头级加权聚合机制以及客户端选择策略来应对上述挑战。头剪枝依据注意力头置信度计算的重要性分数,降低客户端训练复杂度;加权聚合确保全局模型捕捉来自多样化客户端的关键更新,与客户端选择策略相辅相成。在MultiNLI基准及20 Newsgroups、XL-Sum、E2E NLG数据集上验证效果。以T5-small模型和LoRA作为PEFT方法,最大实现90%稀疏度,通信开销提升1.8倍,训练操作数减少3.9倍,精度下降保持在2%以内。
原文摘要 · Abstract (English)
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%.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。