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. 本文提出通过头部剪枝和加权聚合机制，结合客户端选择策略，来优化联邦学习中的PEFT过程。
  3. 在MultiNLI等多个数据集上，使用T5-small模型和LoRA方法，取得了显著的性能提升和资源节约。

📝 摘要（中文）

参数高效微调（PEFT）已成为适应大型语言模型（LLMs）于自然语言处理下游任务的主要方法。然而，其在隐私保护的分布式学习框架（如联邦学习）中的应用仍然有限，主要由于资源受限设备和客户端数据分布多样性等挑战。本文提出了一种在联邦学习框架内高效执行PEFT的方法，针对基于多头注意力（MHA）的语言模型，通过头部剪枝、头部特定加权聚合机制和客户端选择策略来解决这些挑战。实验结果表明，该方法在多个基准数据集上表现出色，达到了高达90%的稀疏性，通信效率提升1.8倍，训练操作减少3.9倍，同时保持准确率下降低于2%。

🔬 方法详解

问题定义：本文旨在解决在联邦学习框架中应用PEFT时面临的挑战，包括设备资源限制和客户端数据分布多样性。现有方法在这些方面表现不佳，导致效率低下和准确性下降。

核心思路：论文提出通过头部剪枝来减少训练复杂度，并引入头部特定的加权聚合机制和客户端选择策略，以确保全球模型能够有效捕捉来自不同客户端的重要更新。

技术框架：整体框架包括三个主要模块：头部剪枝模块、加权聚合模块和客户端选择模块。头部剪枝模块根据注意力头的重要性分数进行剪枝，减轻客户端的训练负担；加权聚合模块确保重要更新被有效整合；客户端选择模块优化参与训练的客户端。

关键创新：最重要的创新在于头部剪枝和头部特定加权聚合机制的结合，这与传统的PEFT方法不同，后者通常不考虑头部的重要性和客户端的多样性。

关键设计：在参数设置上，采用LoRA作为PEFT方法，确保模型的稀疏性达到90%。损失函数设计上，关注于保持准确性，同时优化通信效率和训练操作数。

📊 实验亮点

实验结果显示，采用该方法在MultiNLI基准上实现了高达90%的稀疏性，通信效率提升1.8倍，训练操作数减少3.9倍，同时准确率下降保持在2%以内，展现出显著的性能优势。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理中的联邦学习场景，尤其是在隐私保护和资源受限的环境下。通过优化PEFT过程，该方法能够在多样化的数据分布中有效提升模型性能，具有广泛的实际价值和未来影响。

📄 摘要（原文）

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