[论文解读] Compare Where It Matters: Using Layer-Wise Regularization To Improve Federated Learning on Heterogeneous Data

Ha Min Son, Moon Hyun Kim+1 more|arXiv (Cornell University)|Dec 1, 2021

Privacy-Preserving Technologies in Data参考文献 20被引用 3

一句话总结

FedCKA 提出了一种用于联邦学习的逐层正则化框架，通过使用中心核对齐（CKA）度量表示相似性，仅选择性地对最相关且天然相似的层（特别是前两层）进行正则化，从而在非独立同分布（non-IID）数据上提升性能。该方法在深度模型（如 ResNet-50）上实现了与 FedAvg 相当的训练效率，同时在 CIFAR-10 上达到最高 97.58% 的准确率，性能达到当前最先进水平。

ABSTRACT

Federated Learning is a widely adopted method to train neural networks over distributed data. One main limitation is the performance degradation that occurs when data is heterogeneously distributed. While many works have attempted to address this problem, these methods under-perform because they are founded on a limited understanding of neural networks. In this work, we verify that only certain important layers in a neural network require regularization for effective training. We additionally verify that Centered Kernel Alignment (CKA) most accurately calculates similarity between layers of neural networks trained on different data. By applying CKA-based regularization to important layers during training, we significantly improve performance in heterogeneous settings. We present FedCKA: a simple framework that out-performs previous state-of-the-art methods on various deep learning tasks while also improving efficiency and scalability.

研究动机与目标

• 解决在非独立同分布数据分布下联邦学习性能下降的问题。
• 识别在异构设置中神经网络中哪些层对正则化最为关键。
• 通过避免对所有层进行正则化，提升训练效率与可扩展性。
• 评估 CKA 作为比较客户端之间表示相似性的度量指标的有效性。
• 开发一种可扩展、高性能的正则化框架，超越现有最先进方法。

提出的方法

• 仅对被识别为天然相似且对性能影响最大的前两层应用逐层正则化。
• 使用中心核对齐（CKA）计算跨客户端本地模型激活表示之间的相似性。
• 引入基于 CKA 的正则化项 ℓcka，在本地训练过程中对前两层的表示进行对齐。
• 通过三模型前向传播（客户端、全局模型和目标模型）高效计算 CKA 相似性，无需增加层数。
• 通过仅在关键层上进行正则化来优化计算成本，避免在深层架构中进行逐层操作。
• 将 FedCKA 设计为与 FedAvg 兼容的即插即用框架，对现有联邦学习流程的修改极少。

实验结果

研究问题

• RQ1 在神经网络中，哪些层对重新初始化最敏感，因此在联邦学习中正则化最为关键？
• RQ2 在非独立同分布数据上训练的模型之间，CKA 与其它相似性度量（如 l2、余弦）相比，度量表示相似性的表现如何？
• RQ3 仅对天然相似的层进行选择性正则化，是否能在不降低训练效率的前提下提升性能？
• RQ4 FedCKA 在模型深度增加时的可扩展性如何，尤其与对所有层进行正则化的其他方法相比？
• RQ5 使用基于 CKA 的正则化是否能在 CIFAR-10 和 Tiny ImageNet 等标准基准上带来更好的收敛性和准确率？

主要发现

• 在 α = 5.0 时，FedCKA 在 CIFAR-10 上达到 97.58% 的准确率，显著优于 FedAvg（54.82%）及其他 SOTA 方法。
• 使用 CKA 进行正则化可获得最高性能，其中核 CKA 略优于线性 CKA，但计算成本更高。
• 在 ResNet-50 和 Tiny ImageNet 上，FedCKA 的训练时间与 FedAvg 相当（750.97s），而 FedProx、SCAFFOLD 和 MOON 的训练时间呈指数级增长。
• 仅对前两层进行正则化可降低训练开销并提升可扩展性，尤其在具有大量层的深层模型中优势明显。
• 使用 CKA 作为相似性度量可实现更精确的表示对齐，仅惩罚真正不相似的更新。
• FedCKA 表明，将正则化聚焦于天然相似的层，比对所有层进行均匀正则化更为有效。