联邦学习的研究最近从卷积神经网络(CNNs)转向了视觉 Transformer (ViTs),因为 ViTs 具有更优越的能力。由于 ViTs 缺乏 CNN 固有的 2D 归纳偏差,其训练所需的计算资源更高。然而,如何在资源受限的边缘设备上高效地进行 ViTs 的联邦训练尚未得到社区的探索。
在本文中,作者提出了一种基于遮掩图像的分层联邦框架 EFTViT,以在资源受限的边缘设备上实现高效的全参数训练,从而为异构数据的学习带来显著优势。
总体而言,作者将图像分割成块,并随机遮掩一部分块,观察到排除这些块对性能影响极小,而大幅降低了计算成本并增强了数据内容隐私保护。具体来说,EFTViT 由一系列轻量级本地模块和一个更大的全局模块组成,它们分别独立地在客户端和中央服务器上更新。
本地模块在遮掩图像块上进行训练,而全局模块则在从本地客户端上传的中间块特征上进行训练,并通过提出的中位数采样策略来平衡,从而消除客户端数据分布的隐私性。作者分析了 EFTViT 的计算复杂度和隐私保护。
在流行基准上的广泛实验表明,与现有方法相比,EFTViT 可以实现高达 28.17%的准确性提升,将本地训练的计算成本降低多达 2.8 倍,并将本地训练时间缩短多达 4.4 倍。
在注重数据隐私保护的前提下,客户端设备能够有效地执行联邦学习。早期的联邦学习研究主要集中在卷积神经网络(CNNs)上[1, 20, 22]。
最近,研究重点逐渐转向视觉 Transformer (ViTs),其 Self-Attention 机制在捕捉图像内的长距离对应关系方面表现出色,已在诸如物体识别[8]、检测[6, 13]和语义分割[40]等视觉问题上取得了最先进的性能。
尽管 ViTs 具有令人印象深刻的能力,但由于图像中缺乏空间归纳偏差,训练 ViTs 通常会带来显著更高的计算成本和更长的训练时间[3, 30],这使得资源受限的边缘设备面临巨大的挑战。
联邦学习(FL)旨在 enabling 多个数据分布在不同设备或服务器上时实现协作训练。
1. Introduction
在卷积神经网络(CNN)时代,一些研究者探索了资源受限的联邦学习(FL)问题。这些方法的工作流程总结在图 1a 中。通常情况下,模型异构方法[1, 4, 23, 37]根据客户端可用的资源训练不同规模的模型。然而,这些方法并不适合视觉 Transformer (ViT),因为它们无法从根本上减少客户端训练的计算需求。
在本研究中,作者探索是否可以通过不显著牺牲联邦学习(FL)性能的方式来根本性地降低 Vision Transformer(ViT)的训练计算成本。近期自监督学习领域的研究表明,通过随机遮掩输入图像中的像素并重建这些遮掩的像素, Mask 图像建模能够有效地学习到具有泛化能力的视觉表示[13, 32],从而突显了图像中可能对识别不必要的大量冗余信息。为了验证这一假设,作者在无资源约束条件下进行联邦学习实验,使用遮掩图像来考察其对模型性能和训练计算成本的影响。在实验中,图像被均匀分割成非重叠的块,其中有一部分比例为rm的块被随机遮掩。仅使用未遮掩的块进行模型训练。
受这些观察的启发,作者提出了 EFTViT,这是一种分层的联邦学习框架(如图 1b 所示),该框架利用 Mask 图像在资源受限的客户端上高效地训练 ViT 模型,并通过隐藏客户端数据内容来增强隐私保护。EFTViT 包含轻量级的本地模块部署在边缘客户端和较大的全局模块部署在中央服务器上,旨在适应有限的客户端资源。
本地模块在 Mask 图像上进行训练。与从客户端聚合参数不同,全局模块接收来自本地模块的中间 patches 特征,使其能够学习适用于异质数据的通用表示。为了保持客户端数据分布,作者提出了一种中位数采样策略,在上传前将每个类别的 patch 特征计数调整到所有类别的中位数,从而提升性能并提高训练效率。
本项工作的主要贡献总结如下:
- 据作者所知,作者提出了 EFTViT,这是首个利用掩蔽图像在多个资源受限客户端上高效训练 ViT 模型的联邦学习框架,同时也增强了客户端数据内容保护。EFTViT 能够在客户端和中央服务器之间分层训练所有模型参数,展示了对异构数据的巨大优势。此外,作者引入了一种中位数采样策略,在上传到服务器之前模糊中间特征的分布信息。
- 在流行的基准测试中的实验表明,与 Baseline 相比,EFTViT 能够提高准确性最多 28.17%,降低本地训练计算成本最多 2.8 倍,并将本地训练时间缩短最多 4.4 倍,从而建立了新的最佳性能标准。
2. Related Works
2.1. General Federated Learning
联邦学习是一种去中心化的机器学习方法,通过在客户端设备上直接训练模型,并仅将模型参数传输到中央服务器来增强隐私性。大多数研究集中在解决 FL 中的数据异质性[11, 17, 20, 22]和隐私保护问题[2, 5, 27]。例如,FedProx[22]通过添加 proximal 项来优化局部更新,以解决数据异质性问题。
在隐私保护方面,Asad 等人[2]将同态加密应用于联邦学习,使客户端能够使用私钥对其本地模型进行加密。Shi 等人[27]提出了一个结合差分隐私(DP)的联邦学习方法。然而,这些工作依赖于理想假设,即客户端有足够的资源来处理模型训练过程。
2.2. Federated Learning on Edge Devices
资源受限客户端上的联邦学习方法可以分为联邦蒸馏(FD)和部分训练(PT)。FD 方法侧重于从异构客户端模型聚合知识到服务器模型。例如,FedGKT [12] 在客户端训练小型模型,并通过知识蒸馏定期将其知识转移到大型服务器模型中。PT 方法将全局模型分成可以在资源受限客户端上局部训练的小型子模型。
例如,HeteroFL [7] 随机从全局模型中选择子模型并分发给客户端。然而,这些方法是根据客户端的能力调整模型大小,而不是从根本上解决客户端训练的计算负担。
2.3. Parameter-Efficient Fine-Tuning
在处理基于 Transformer 的复杂模型时,参数高效微调(PEFT) 提供了一种实用的方法,可以在各种下游任务中高效地适应预训练模型,通过固定大部分预训练参数并仅微调一个小子集来减少存储和计算成本[10]。一些研究[29, 38]已经探索了使用不同的 PEFT 技术来评估联邦系统中的性能改进和资源节省。
然而,在 PEFT 中有限的参数微调不可避免地限制了预训练模型对新任务的适应性,可能导致在数据异质性较强的联邦系统中出现次优性能。
3. Efficient Federated Learning with Masked Images
3.1. Problem Definition
3.2.Overview
如图 3 所示,EFTViT 在客户端和中央服务器之间采用分层训练机制,以实现隐私保护和高效的协作学习。每个客户端包含一个具有 M 个 Transformer 层的本地模块、一个具有 N 个 Transformer 层的共享全局模块以及一个分类头。本地模块和分类头在客户端上使用未遮掩的图像片段 Xp 进行训练,实现高效的地方训练并生成代表局部知识的片段特征。
为了保护数据分布的隐私,在上传到服务器之前,客户端采用中位数采样策略创建平衡片段特征(BPF)数据集。随后,服务器使用来自客户端的 BPF 数据集对全局模块进行训练,以有效学习所有任务的全局表示。最后,服务器将更新后的全局模块参数传回客户端,以便进行下一轮训练。
3.3. Training with Masked Images
3.4.Data Distribution Protection with Median Sam pling
作者的中位数采样策略通过计算每个客户端上每类样本的数量中位数来区分少数类和多数类,然后通过对少数类进行过采样增加样本数量,对多数类进行下采样减少样本数量。具体来说,对于少数类样本,所有跨多个本地训练周期生成的 patch 特征都会保留;而对于多数类样本,则仅保留最终周期生成的 patch 特征。随后,通过下采样将每个类别的样本数量减少到中位数。实验结果显示,提高采样阈值会增加计算成本,但不会显著提升最终性能。
3.5. Hierarchical Training Paradigm
为了在不牺牲性能的前提下有效减轻客户端的计算负担,作者提出了一种新的层级训练策略,该策略旨在最小化分配给客户端的可训练参数数量。如前所述,作者的 ViT 模型由一系列轻量级局部模块、一个共享的大规模全局模块以及一个分类头组成。
3.6. Collaborative Algorithms
作者的 EFTViT 的整体工作流程如 Algorithm 1 和 Algorithm 2 所示。在每一轮 t 开始时,服务器会从 K 个客户端中随机选择比例为 P 的客户端参与训练。每个客户端使用从服务器接收到的参数更新其全局模块和分类头,并随后启动本地训练。此外,EFTViT 还会上传对应于图像面积 25%的 Patch 特征,这些特征的上传受到 Mask 率 rm 的控制,即使从理论上讲可以恢复原始图像,但这样做也极大地增加了图像原貌恢复的难度。如果需要进一步增强数据内容隐私性,还可以提高 Mask 率。
数据分布隐私保护。为了保护用户统计信息和偏好,在客户端上作者通过提出的中位数采样策略平衡了 Patch 特征,确保每个类别的样本数量相等。此外,作者的策略与其它隐私保护方法(如差分隐私[9])相互独立,可以无缝集成到 EFTViT 中,从而提供更强的攻击防护能力。
3.7. Privacy & Complexity Analysis
由于您提供的是一张图片链接,而图片中的内容无法直接显示或复制,因此我无法直接为您完成翻译。请您将图片中的文本内容复制或输入到这里,我将根据您提供的英文文本进行翻译。
4. Experiments
4.1. Datasets
为了全面评估 EFTViT 的表现,作者在两个广泛使用的联邦学习数据集中进行了实验,即 CIFAR-10 [18]和 CIFAR-100 [18],以及一个更具挑战性的数据集 UC Merced Land-Use [34],用于遥感应用。CIFAR-10 和 CIFAR-100 数据集各自包含 60,000 张彩色图像。
CIFAR-10 按照 10 个类别组织,每个类别有 6,000 张图像(其中 5,000 张用于训练,1,000 张用于测试),而 CIFAR-100 包含 100 个类别,每个类别有 600 张图像(其中 500 张用于训练,100 张用于测试)。UC Merced Land-Use 数据集包含 21 个土地利用类别,例如农业用地、森林、高速公路、海滩以及其他类别,每个类别有 100 张图像(其中 80 张用于训练,20 张用于测试)。作者根据 Dirichlet 分布 DirN(β)将样本分配给所有客户端,并设置浓度参数 β 为{0.1,1}以模拟高度异质性或低度异质性。
4.2. Implementations
作者使用在 ImageNet-21K [26] 上预训练的 ViT-B [8] 作为框架的主干网络。输入图像被调整为大小为 224 × 224 ,同时使用大小为 16 × 16 的 patch 。
所有实验均在一台单块 NVIDIA GeForce RTX 3090 GPU 上进行。每一轮中,客户端在当地训练 5 个周期,而服务器额外执行 2 个周期。整个框架总共训练 200 轮,大约需要 24 小时。
4.3. Comparison with State-of-the-Art Methods
由于缺乏在资源受限客户端训练 ViTs 的研究,作者采用 FEDBFPT 方法[33]作为强 Baseline ,该方法最初是为自然语言处理任务设计的,并逐步优化较浅层的网络结构,同时选择性地采样较深层以减少资源消耗。为了建立额外的 Baseline ,作者将几种广为人知的 PEFT 方法适应到作者的联邦学习框架中:
(a) Fed-Head:仅训练 Head 层参数;
(b) Fed-Bias:应用偏差调整[36],专注于仅训练偏差项;
(c) Fed-Prompt:结合 Prompt 调整[16],向输入添加可训练的 Prompt 嵌入;以及(d) FedLoRA:通过向 Query 和值层加入 LoRA 模块整合 LoRA 调整[14]。这些方法使用 FedAVG[24]进行参数聚合。除了这一点外,在联邦学习场景中,作者的方法和 Baseline 方法在设置上是相同的。
测试准确性。各种数据集和数据异质性水平下所有方法的测试结果在表 1 中有所呈现。需要注意的是,Fed-Full 是指在没有资源约束的情况下,在客户端训练所有 ViT 参数,这作为比较的参考基准。相比 Baseline 方法,EFTViT 在所有场景中均显示出明显的性能提升。例如,在 β=0.1 的情况下,作者在 UC Merced 土地利用数据集上的表现比第二好的方法高出超过 7.61%。值得注意的是,作者的方法在高和低数据异质性设置中显示了一致的结果,并且在更高的异质性环境下性能甚至更好。
相比之下, Baseline 方法随着数据异质性的增加,其性能显著下降。这些发现强调了作者分层训练策略在有效处理数据异质性方面的重要性。
收敛性。作者在高异质性设置下,在 CIFAR-10、CIFAR-100 和 UC 梅塞德斯土地用途数据集上报告了 EFTViT、FEDBFPT 以及其他 Baseline 在 100 轮训练过程中的测试准确性变化情况,如图 4 所示。在整个训练阶段,作者的方法在三个数据集上始终获得最高的测试准确性,并且收敛速度更快、更稳定。为了定量比较收敛速度,作者将目标准确性设为 85%,并记录达到该阈值所需的训练轮数(## Rounds)。如表 2 所示,在 UC 梅塞德斯土地用途数据集上,EFTViT 显著加速了收敛过程,相比 Fed-Prompt 实现了 27.1 倍的更快收敛。
计算效率。作者从两个方面评估了 EFTViT 在客户端的计算效率:训练过程中前向传播的计算成本以及客户端的最大局部训练时间。计算成本以十亿浮点运算(GFLOPs)为单位进行衡量。在目标准确率为 85%的情况下,作者在三个数据集上报告了 EFTViT 与其他基准方法在本地训练中的最大局部训练时间(Time)。
表 3 的结果显示,作者的方法在这两项指标上都显著提高了计算效率。特别是,在 GFLOPs 方面,EFTViT 至少比其他方法提高了两倍的效率。对于训练时间,与 UC Merced Land-Use 数据集上的 FEDBFPT 相比,EFTViT 将本地训练时间减少了 2.8 倍。这表明,作者的 Mask 图像和分层训练策略有效地减少了客户端的计算量,使得 EFTViT 非常适合在资源受限环境中进行联邦学习。
4.4. Ablation Study
作者进行了广泛的消融实验以探究作者方法的关键组件。
Mask 比例的影响。 Mask 比例(r_{m})决定了被 Mask 的图像片段的数量。较小的(r_{m})减少了输入数据量,从而在模型训练期间降低计算需求。表 4 提供了不同 Mask 率下的 GFLOPs,显示增加 Mask 比例显著减少 GFLOPs。然而,增加 Mask 比例也会影响整体性能。作者评估了 EFTViT 在不同 Mask 率下的效果。图 5 展示了当(\beta=0.1)时,EFTViT 在 CIFAR-100 上的不同 Mask 率结果。结果显示,EFTViT 可以支持广泛的 Mask 比例范围。当 Mask 比例从 0%增加到 75%时,准确率仍超过 90%。然而,当 Mask 比例超过 75%时,性能会大幅下降。因此,作者选择 75%的 Mask 比例,在准确率和计算效率之间取得平衡。局部模块层数(M)的影响。层数(M)决定了客户端与服务器之间的可训练参数分配,影响客户端的计算负载及最终性能
。表 5 列出了不同(M)值下每个客户端的可训练参数数量(## Params)及其对应的模型准确率。结果表明,(M)对测试准确率的影响甚微,展示了 EFTViT 对于客户端资源的优越鲁棒性。鉴于较大的(M)值在客户端带来的更高计算成本以及准确率的下降,作者选择(M=2)作为默认设置。采样阈值的影响。如第 3.4 节所述,采样阈值决定了上传至服务器用于训练的平衡片段特征的数量。因此,较高的阈值会增加服务器训练的成本。作者研究了在 EFTViT 中使用中位数或较高采样阈值的影响,如图 6 所示。结果显示,提高阈值带来的性能改进微乎其微。为了提高服务器端的计算效率,作者在方法中选择了中位数作为阈值。
参考
[0]. EFTViT: Efficient Federated Training of Vision Transformers with Masked Images on Resource-Constrained Edge Devices .
END
作者:小书童
来源:集智书童
推荐阅读
- 使用 Triton 加速 2D 动态块量化 Float8 GEMM 简介
- 武大提出 Point Teacher,两阶段去噪,让小物体点标注检测更可靠 !
- PyTorch 博客 CUTLASS Ping-Pong GEMM Kernel 简介
- 【视觉多模态专栏】理解 Transformer 模型结构
欢迎大家点赞留言,更多 Arm 技术文章动态请关注极术社区嵌入式 AI 专栏欢迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。