分享自:

基于联邦学习的BERT模型高效训练框架

期刊:软件学报

本研究的主要作者为王鑫澳、陈珂、寿黎但、骆歆远和陈刚,所有作者均来自区块链与数据安全全国重点实验室(浙江大学)。这项名为《基于联邦学习的bert模型高效训练框架》的研究成果,以中文发表于《软件学报》2025年第36卷第9期(第4110–4133页),其英文标题为“Efficient framework for BERT model training based on federated learning”。

研究的学术背景主要涉及两个快速发展的领域:联邦学习(Federated Learning, FL)和预训练语言模型(Pre-trained Language Models, PLMs),特别是像BERT(Bidirectional Encoder Representations from Transformers)这样的基于Transformer架构的大型模型。近年来,人工智能模型,尤其是PLMs的性能取得了显著提升,但其成功严重依赖于大量高质量的训练数据。然而,在许多专业领域(如医疗、金融),数据因包含敏感隐私信息而分散存储于不同机构,无法集中收集用于模型训练。联邦学习作为一种分布式机器学习范式,允许在数据不离开本地设备(客户端)的前提下协作训练模型,从而为隐私保护下的数据利用提供了解决方案。

然而,一个核心矛盾由此产生:联邦学习中的客户端设备(如个人手机、边缘设备)通常计算资源有限、网络带宽受限;而BERT等大型PLMs的参数规模庞大,训练它们需要巨大的计算和通信开销,远超一般客户端的承载能力。现有的一些解决方案,如信息压缩(量化、稀疏化)主要针对通信开销,模型剪枝和知识蒸馏则更多关注于模型推理阶段的轻量化,对于在资源受限客户端上高效训练完整大型模型这一根本问题效果有限。因此,如何在联邦学习场景下,使资源受限的客户端能够高效参与BERT模型的训练(包括在专业语料上的进一步预训练和针对具体任务的下游微调),同时不牺牲模型精度,成为一个亟待解决的关键挑战。本研究旨在攻克这一难题,其核心目标是:设计一个名为FedBT的高效训练框架,在保护数据隐私的前提下,显著降低客户端在联邦学习过程中训练BERT模型的计算成本和通信成本,并达到与传统联邦学习(训练完整模型)接近的模型性能。

研究的工作流程详细且系统,主要分为以下几个环节:首先是问题的形式化定义与理论分析,然后是FedBT框架针对两种不同场景(进一步预训练和下游任务微调)的具体算法设计,最后通过大量实验进行验证。具体而言,研究流程包括: 1. 问题定义与开销建模:研究者首先将联邦学习场景下BERT模型高效训练的问题进行了严格的形式化定义。他们明确了两种核心应用场景:(1)进一步预训练场景:服务端拥有一个通用BERT模型,客户端拥有同领域的无标注隐私语料,目标是利用客户端语料增强模型在该领域的表现。(2)下游任务微调场景:服务端拥有一个预训练好的BERT模型和少量测试数据,客户端拥有同领域的带标注隐私任务数据,目标是利用客户端数据使模型适配具体下游任务。研究进一步构建了计算开销和通信开销的数学模型(如公式2、3),将开销与客户端模型规模、更新层数、训练策略等变量关联起来,并确立了优化目标(公式5、6):在保证模型精度的前提下,将客户端的计算和通信开销至少降低至传统联邦学习的一半以下。此外,研究还对BERT模型各组成部分(嵌入层、Transformer层等)的计算和空间复杂度进行了理论分析,为框架设计提供了理论基础。 2. FedBT框架核心算法设计:这是研究的核心创新部分,针对上述两种场景分别提出了不同的训练策略。两种策略共享一个核心理念:基于神经网络(包括BERT)训练过程中“浅层先稳定,深层后习得”的特性,以及BERT模型中浅层Transformer(捕捉语法特征)和深层Transformer(捕捉语义特征)在不同任务中重要性不同的先验知识,选择性地在客户端训练和更新全局模型的关键部分参数,而非全部参数。 * 高效进一步预训练框架:针对进一步预训练场景,FedBT设计了“浅层渐进式训练算法”和“深层采样映射算法”。 * 浅层渐进式训练算法:由于预训练阶段浅层Transformer更为关键,该算法决定客户端在每一轮联邦学习中训练哪一层Transformer。其核心是从第0层开始,将总训练轮次的前一半用于训练第0层,剩余轮次的一半用于训练第1层,以此类推,渐进式地、重点地训练浅层参数。例如,总轮次为10时,训练层序为[0,0,0,0,0,1,1,1,2,2]。 * 深层采样映射算法:客户端构建一个规模小于全局模型的本地模型(例如6层Transformer,而全局模型为12层)。对于本地模型中需要训练的关键浅层(如上例中的第0、1、2层),其参数直接使用或更新;对于本地模型中更深层的、不用于训练的参数,则通过从全局模型深层Transformer参数中随机采样映射的方式来构建。这样做的好处是,让浅层参数在训练时也能“感知”到深层模型的知识表示,从而保证训练出的浅层参数与完整的全局模型架构兼容,提升训练效果。 * 工作流程:在每一轮联邦学习中,服务端指定一个训练层索引l并下发相关参数。客户端根据l和本地参数池(存储着历史聚合参数)构建本地小模型,然后仅用本地隐私数据训练该模型的第l层Transformer和输出层,其他层冻结。训练完成后,客户端仅上传第l层Transformer和输出层的更新参数到服务端。服务端使用联邦平均(FedAvg)等算法聚合所有客户端的更新,更新全局模型的对应层,并决定下一轮要训练的新层索引l。 * 高效下游任务微调框架:针对下游微调场景,FedBT设计了“循环递减训练算法”。 * 核心思想:下游任务微调更依赖于深层Transformer的特征表示。因此,FedBT在该场景下让客户端使用与全局模型规模相同的模型,但采用从深层向浅层循环递减的方式选择训练层。即优先训练和更新最深层(如第11层)的Transformer参数,经过一定轮次后,再开始同时训练次深层(第10层),以此类推,逐步扩大训练范围。 * 工作流程:客户端从服务端下载完整的全局模型。在本地训练时,并非更新所有参数,而是根据算法确定的当前训练层集合(例如初期只更新第11层和输出层),对这些选中的深层参数进行有监督的微调训练。同样,客户端只上传这些被更新层的参数到服务端进行聚合。这样既大幅减少了通信量,又由于深层参数离输出层近、反向传播路径短,也降低了计算开销。 3. 实验设计与验证:为了验证FedBT框架的有效性,研究者在多个专业领域数据集上进行了充分的实验对比。实验对象包括不同领域的文本语料。实验主要比较了FedBT框架与传统联邦学习(在客户端训练完整BERT模型)在两种场景下的性能。实验指标主要包括:(1)模型精度:在目标任务(如下游分类、NER)上的准确率、F1值等;(2)计算开销:客户端训练时间的占比;(3)通信开销:客户端上传/下载参数量的占比。实验设置了相同的总训练轮次、学习率等超参数,以确保公平比较。数据处理的流程是标准的自然语言处理流程,包括分词、构建掩码语言模型(MLM)任务用于预训练、构建分类任务用于微调等。分析数据的工作流程主要是收集各轮次训练后的模型在测试集上的性能指标,并记录每个客户端在训练过程中的实际计算时间和通信数据量,最后进行统计分析。

研究的主要结果令人印象深刻,充分证明了FedBT框架的优越性。在进一步预训练场景下,实验结果显示:与传统联邦学习训练完整模型相比,FedBT框架能够将客户端的训练开销(计算成本)降低至原来的34.31%,将通信开销降低至原来的7.04%。这意味着计算成本节约了近三分之二,而通信成本节约了超过90%。更重要的是,在如此大幅的成本削减下,使用FedBT框架训练出的模型,在后续下游任务微调后达到的精度,与传统方法训练出的模型精度非常接近。这一结果逻辑上验证了“重点训练浅层Transformer”策略在预训练场景下的有效性,以及“深层采样映射”算法在保证模型兼容性方面的作用。

在下游任务微调场景下,实验结果同样显著:FedBT框架将客户端的训练开销降低至原来的48.26%,将通信开销降低至原来的20.19%,即计算和通信成本均降低了一半以上。同时,模型在目标下游任务上达到的精度也与训练完整模型的方法基本持平。这证实了“循环递减训练深层参数”策略在下游微调场景下的高效性。这些结果(成本大幅降低,精度基本保持)直接支撑了研究的核心结论:FedBT框架成功地解决了资源受限客户端无法高效训练大型BERT模型的问题。

研究的结论是明确且有力的。本文成功设计并实现了一种适用于联邦学习环境的高效BERT模型训练框架——FedBT。该框架通过选择性训练关键参数(预训练重浅层、微调重深层)、结合渐进式/循环式训练算法以及参数采样映射技术,在保护数据隐私的前提下,显著降低了联邦学习客户端在参与大型预训练语言模型训练时的计算与通信开销,并且保持了与传统方法相当的模型精度。这项工作不仅为在隐私敏感的专业领域部署和利用大型AI模型提供了切实可行的技术方案,具有重要的应用价值(可应用于医疗诊断、金融风控、法律咨询等需要数据隐私又依赖专业模型的领域),也丰富了联邦学习与大型模型高效训练相结合的理论与方法,具有较高的科学价值。它展示了一种不同于单纯压缩或剪枝的新思路:通过剖析模型训练的内在规律,动态地、智能地分配有限的训练资源,从而实现全局效率的最优。

本研究的亮点突出体现在以下几个方面:首先,在研究问题上,精准抓住了联邦学习与大型预训练模型结合的核心痛点——客户端资源瓶颈,具有明确的现实意义。其次,在方法论上具有显著的新颖性:①提出了“分场景差异化训练”的核心思想,根据预训练和微调的不同特性设计截然不同的参数选择策略,而非一刀切;②创新性地提出了“浅层渐进式训练”、“深层采样映射”和“循环递减训练”等一套组合算法,将渐进式学习思想创造性地应用于联邦学习的参数更新调度中;③与传统的压缩或蒸馏方法不同,FedBT专注于训练过程的优化,在推理阶段仍使用完整的原始大模型,从而最大限度地保留了模型能力。最后,在实证效果上,取得了非常突出的量化成果,特别是在进一步预训练中将通信开销降至7.04%,这一数据极具说服力,证明了框架的高效性。此外,研究还包含了严谨的理论开销分析,并公开了前期探索工作(IJCAI 2023),体现了研究的系统性和延续性。

上述解读依据用户上传的学术文献,如有不准确或可能侵权之处请联系本站站长:admin@fmread.com