[论文解读] SPATL: Salient Parameter Aggregation and Transfer Learning for Heterogeneous Clients in Federated Learning

Sixing Yu, Phuong Nguyen+4 more|arXiv (Cornell University)|Nov 29, 2021

Privacy-Preserving Technologies in Data参考文献 25被引用 5

一句话总结

SPATL 提出了一种联邦学习框架，通过选择性聚合关键模型参数并利用客户端特定的预测器进行知识迁移，在异构、非独立同分布（non-IID）环境下有效降低了通信开销并提升了训练稳定性。该方法将通信成本最高降低86.45%，推理所需的浮点运算量（FLOPs）减少39.7%，同时在收敛速度和性能方面优于当前最先进方法。

ABSTRACT

Federated learning~(FL) facilitates the training and deploying AI models on edge devices. Preserving user data privacy in FL introduces several challenges, including expensive communication costs, limited resources, and data heterogeneity. In this paper, we propose SPATL, an FL method that addresses these issues by: (a) introducing a salient parameter selection agent and communicating selected parameters only; (b) splitting a model into a shared encoder and a local predictor, and transferring its knowledge to heterogeneous clients via the locally customized predictor. Additionally, we leverage a gradient control mechanism to further speed up model convergence and increase robustness of training processes. Experiments demonstrate that SPATL reduces communication overhead, accelerates model inference, and enables stable training processes with better results compared to state-of-the-art methods. Our approach reduces communication cost by up to $86.45\%$, accelerates local inference by reducing up to $39.7\%$ FLOPs on VGG-11, and requires $7.4 imes$ less communication overhead when training ResNet-20.

研究动机与目标

• 解决因频繁传输模型权重而导致联邦学习中通信开销过高的挑战。
• 提升在非独立同分布（non-IID）、异构客户端环境下的训练稳定性和收敛性。
• 通过减少资源受限边缘设备上处理的参数数量，降低本地推理成本。
• 通过共享编码器向客户端特定预测器进行有效知识迁移，以应对数据异构性问题。
• 开发一种可扩展、高效的联邦学习框架，适用于大规模去中心化人工智能部署。

提出的方法

• 基于图神经网络（GNNs）和多层感知机（MLPs）引入一个关键参数选择代理，用于识别并仅传输模型编码器中最关键的参数。
• 将模型分解为共享编码器和客户端特定的预测头，通过微调的本地预测器实现对非IID数据的知识迁移。
• 使用在线强化学习在每个客户端上微调关键参数选择代理，使其以极低开销适应本地数据分布。
• 在本地和服务器环境中均采用控制变量（control variates）机制实现梯度控制，以纠正梯度漂移，并使本地更新与全局优化方向对齐。
• 在强化学习代理的预训练阶段应用模型压缩与剪枝技术，确保边缘设备部署时推理成本极低（V100上为26 KB内存、0.36 ms）。
• 通过重用预训练的关键参数选择代理，在不同模型架构（如ResNet-18、ResNet-56）之间实现迁移学习，仅需极少微调。

实验结果

研究问题

• RQ1 选择性聚合关键参数是否能显著降低联邦学习中的通信成本，同时不损害模型准确性？
• RQ2 通过客户端特定预测器进行知识迁移，在非IID、异构数据分布上是否能有效提升模型性能？
• RQ3 轻量级强化学习代理是否能高效地在边缘设备上微调，以适应本地数据的关键参数选择？
• RQ4 梯度控制在去中心化、非IID联邦学习中能在多大程度上提升训练稳定性和收敛速度？
• RQ5 SPATL在不同模型架构上（尤其是VGG和ResNet等过参数化模型）的表现如何？

主要发现

• 与最先进方法相比，SPATL将通信成本最高降低86.45%，在ResNet-20上最大降低达7.4倍。
• 该方法通过减少VGG-11上高达39.7%的FLOPs，显著提升了边缘设备上的本地推理效率。
• 通过客户端特定预测器的迁移学习，模型准确率相比非迁移学习基线最高提升19.86%。
• 梯度控制机制显著提升了训练稳定性，并在非IID设置下高梯度方差场景中实现更快收敛。
• 预训练的关键参数选择代理在微调后，可在约40轮强化学习更新内快速收敛，并在ResNet-18和ResNet-56等不同架构上达到相当的性能。
• 强化学习代理的推理成本极低（V100上为26 KB内存、0.36 ms），使其在边缘设备上具备实际部署可行性。