Assortment of Attention Heads: Accelerating Federated PEFT with Head Pruning and Strategic Client Selection

📄 arXiv: 2506.00743v1 📥 PDF

作者: Yeshwanth Venkatesha, Souvik Kundu, Priyadarshini Panda

分类: cs.CL, cs.AI, cs.DC

发布日期: 2025-05-31

💡 一句话要点

提出基于注意力头剪枝与策略性客户端选择的联邦PEFT加速方法

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 联邦学习 参数高效微调 注意力头剪枝 客户端选择 大型语言模型

📋 核心要点

  1. 联邦学习中参数高效微调（PEFT）面临资源约束设备和数据异构性的挑战。
  2. 通过注意力头剪枝降低客户端计算负担，并设计加权聚合机制融合不同客户端的更新。
  3. 实验表明，该方法在保持精度的情况下，显著降低了通信成本和计算复杂度。

📝 摘要（中文）

本文提出了一种高效的方法，用于在联邦学习（FL）框架内对基于多头注意力（MHA）的语言模型进行参数高效微调（PEFT）。该方法通过注意力头剪枝、一种新颖的头特定加权聚合机制以及客户端选择策略来应对FL中的挑战，例如资源受限的设备和客户端之间不同的数据分布。注意力头剪枝通过基于注意力头置信度计算的重要性得分来最小化客户端内的训练复杂度。头的加权聚合确保全局模型捕获来自不同客户端的关键更新，从而补充了我们的客户端选择策略。在MultiNLI基准以及20 Newsgroups、XL-Sum和E2E NLG数据集上的结果表明，使用MultiNLI数据集和T5-small模型以及LoRA作为PEFT方法，可以达到高达90%的稀疏度，从而实现高达1.8倍的通信优势和3.9倍的训练OPs减少，同时保持精度下降在2%以下。

🔬 方法详解

问题定义：论文旨在解决联邦学习（FL）环境下，对大型语言模型（LLM）进行参数高效微调（PEFT）时面临的挑战。现有方法在FL中应用PEFT时，由于客户端资源有限和数据分布差异大，导致训练效率低下和模型性能下降。尤其是在多头注意力（MHA）机制中，如何选择性地更新和聚合不同客户端的注意力头，是一个关键问题。

核心思路：论文的核心思路是通过注意力头剪枝来降低客户端的计算负担，并设计一种头特定的加权聚合机制，以更好地融合来自不同客户端的更新。此外，还引入了客户端选择策略，以进一步提高训练效率和模型性能。这种方法旨在在保证模型性能的前提下，最大限度地减少通信成本和计算复杂度。

技术框架：整体框架包含以下几个主要阶段：1) 客户端训练：每个客户端使用本地数据对模型进行PEFT，并计算每个注意力头的重要性得分。2) 注意力头剪枝：根据重要性得分，每个客户端对注意力头进行剪枝，减少计算量。3) 加权聚合：服务器接收来自客户端的更新，并使用头特定的权重对注意力头进行聚合，生成全局模型。4) 客户端选择：服务器根据一定的策略选择参与下一轮训练的客户端。

关键创新：论文的关键创新在于以下几个方面：1) 注意力头剪枝：通过重要性得分来选择性地剪枝注意力头，降低了客户端的计算负担。2) 头特定加权聚合：根据注意力头的重要性，对来自不同客户端的更新进行加权聚合，提高了模型性能。3) 客户端选择策略：通过选择合适的客户端参与训练，提高了训练效率。

关键设计：论文的关键设计包括：1) 重要性得分计算：基于注意力头的置信度计算重要性得分，用于指导注意力头剪枝。2) 加权聚合权重：根据注意力头的重要性，设计加权聚合的权重，以更好地融合来自不同客户端的更新。3) 客户端选择策略：具体选择策略未知，但目标是选择对全局模型提升最大的客户端。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，该方法在MultiNLI数据集上，使用T5-small模型和LoRA作为PEFT方法，可以达到高达90%的稀疏度，从而实现高达1.8倍的通信优势和3.9倍的训练OPs减少，同时保持精度下降在2%以下。这表明该方法在降低通信成本和计算复杂度方面具有显著优势。

🎯 应用场景

该研究成果可应用于各种隐私保护的分布式学习场景，例如医疗健康、金融风控等。通过在本地设备上进行模型微调，可以保护用户数据的隐私，同时利用联邦学习的优势，提升模型的泛化能力。该方法有望加速LLM在资源受限设备上的部署和应用。

📄 摘要（原文）

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%.