[论文解读] Gradient Masked Averaging for Federated Learning

Irene Tenison, Sai Aravind Sreeramadas+4 more|arXiv (Cornell University)|Jan 28, 2022

Privacy-Preserving Technologies in Data被引用 6

一句话总结

该论文提出了一种新型聚合方法——梯度掩码平均（Gradient Masked Averaging, GMA），通过强调在客户端间方向一致的梯度分量，同时降低不一致分量的权重，从而在非独立同分布（non-i.i.d.）设置下提升联邦学习的泛化能力。GMA在数据异构性和客户端参与不平衡的情况下优于标准平均方法，且在多个数据集和算法上均表现出一致的性能提升。

ABSTRACT

Federated learning (FL) is an emerging paradigm that permits a large number of clients with heterogeneous data to coordinate learning of a unified global model without the need to share data amongst each other. A major challenge in federated learning is the heterogeneity of data across client, which can degrade the performance of standard FL algorithms. Standard FL algorithms involve averaging of model parameters or gradient updates to approximate the global model at the server. However, we argue that in heterogeneous settings, averaging can result in information loss and lead to poor generalization due to the bias induced by dominant client gradients. We hypothesize that to generalize better across non-i.i.d datasets, the algorithms should focus on learning the invariant mechanism that is constant while ignoring spurious mechanisms that differ across clients. Inspired from recent works in Out-of-Distribution generalization, we propose a gradient masked averaging approach for FL as an alternative to the standard averaging of client updates. This aggregation technique for client updates can be adapted as a drop-in replacement in most existing federated algorithms. We perform extensive experiments on multiple FL algorithms with in-distribution, real-world, feature-skewed out-of-distribution, and quantity imbalanced datasets and show that it provides consistent improvements, particularly in the case of heterogeneous clients.

研究动机与目标

• 解决因客户端间数据异构性导致的联邦学习模型性能下降问题。
• 克服标准平均方法在非独立同分布设置下引入的偏差和信息损失。
• 通过聚焦于客户端间一致的机制，抑制特定客户端的虚假模式，从而提升泛化能力。
• 开发一种可插拔的聚合方法，与现有联邦学习算法兼容，无需修改模型架构即可增强鲁棒性。
• 在包括分布内、分布外以及数量不平衡设置在内的多种数据集上，实现一致的性能提升。

提出的方法

• 计算所有客户端中每个梯度分量的符号，以确定更新方向上的分量级一致性。
• 基于一致性得分构建软掩码矩阵，一致性越高则在最终更新中权重越大。
• 在平均前对客户端梯度应用掩码，突出客户端一致的方向，同时抑制冲突的更新。
• 将掩码后的梯度平均结果作为全局模型更新，替代标准的FedAvg式平均方法。
• 将GMA作为即插即用的替代方案集成到现有联邦学习算法（如FedAvg）中，无需修改客户端训练流程。
• 掩码机制具有可微性和稳定性，其理论收敛性与稳定性分析已在补充材料中提供。

实验结果

研究问题

• RQ1 与标准平均相比，梯度掩码是否能在非独立同分布数据分布下提升联邦学习的泛化能力？
• RQ2 在现实联邦学习场景中常见的特征偏移和数量不平衡数据分布下，GMA的性能表现如何？
• RQ3 在客户端部分参与的轮次中，GMA是否能为未参与客户端提供更强的性能增益？
• RQ4 随着客户端数量和本地训练轮次的增加，GMA的性能表现如何变化？
• RQ5 在逻辑回归等凸优化设置下的MNIST数据集上，GMA能否提升收敛性和鲁棒性？

主要发现

• 在10个客户端的FMNIST数据集上，非独立同分布设置下，GMA使参与客户端的测试准确率提升2.4%，未参与客户端提升4.6%，优于标准平均方法。
• 在100个客户端的CIFAR-10数据集上，采用Dirichlet分布的数量偏移（β=0.5）设置下，GMA优于标准平均方法，表现出在极端数据不平衡下的鲁棒性。
• 在MNIST数据集的凸优化设置下，使用逻辑回归时，GMA在非独立同分布数据下的测试准确率达到88.5%，而标准平均方法为87.0%。
• GMA在多种算法和数据集上均表现出一致的性能提升，尤其在异构性和不平衡设置下增益更显著。
• 消融实验表明，随着客户端数量增加和本地训练轮次增多，GMA的优势进一步增强，表明其对客户端漂移和数据异构性的更强鲁棒性。
• 该方法展现出更高的稳定性和收敛性，理论分析支持其鲁棒性与收敛特性。