[论文解读] TCT: Convexifying Federated Learning using Bootstrapped Neural Tangent Kernels

Yaodong Yu, Alexander Wei+3 more|arXiv (Cornell University)|Jul 13, 2022

Privacy-Preserving Technologies in Data被引用 5

一句话总结

本文提出 TCT 框架，该框架首先通过联邦训练使用 FedAvg 提取特征，然后在通过经验神经正切核（eNTK）凸化后的模型上应用梯度校正优化。该方法在 CIFAR10 上实现了高达 37% 的准确率提升，并在极端数据异构性条件下完全弥合了集中式训练与联邦学习之间的性能差距。

ABSTRACT

State-of-the-art federated learning methods can perform far worse than their centralized counterparts when clients have dissimilar data distributions. For neural networks, even when centralized SGD easily finds a solution that is simultaneously performant for all clients, current federated optimization methods fail to converge to a comparable solution. We show that this performance disparity can largely be attributed to optimization challenges presented by nonconvexity. Specifically, we find that the early layers of the network do learn useful features, but the final layers fail to make use of them. That is, federated optimization applied to this non-convex problem distorts the learning of the final layers. Leveraging this observation, we propose a Train-Convexify-Train (TCT) procedure to sidestep this issue: first, learn features using off-the-shelf methods (e.g., FedAvg); then, optimize a convexified problem obtained from the network's empirical neural tangent kernel approximation. Our technique yields accuracy improvements of up to +36% on FMNIST and +37% on CIFAR10 when clients have dissimilar data.

研究动机与目标

• 为解决在数据异构性条件下联邦学习性能下降的问题，特别是针对非凸深度神经网络。
• 探究为何现有联邦优化方法尽管早期网络层已学习到有用特征，却仍无法有效利用这些特征。
• 开发一种方法，在保持已学习特征的同时，实现在非凸性条件下最终层的稳定、高精度训练。
• 证明通过 eNTK 的凸化可使异构联邦设置下的性能恢复至接近集中式训练的水平。
• 表明所提出的 TCT 框架在超参数和通信效率方面具有鲁棒性，即使在极端数据异构性条件下亦然。

提出的方法

• 首先，使用 FedAvg 训练深度神经网络，以联邦方式从客户端数据中提取有用特征。
• 其次，计算已训练模型的经验神经正切核（eNTK）近似，以线性化最后几层。
• 第三，对 eNTK 近似后的特征，应用梯度校正优化（如 SCAFFOLD）求解所得凸问题。
• 对 eNTK 特征进行归一化，以提升第二阶段的收敛速度和通信效率。
• 使用第一阶段的模型检查点初始化基于 eNTK 的凸优化，确保特征一致性。
• 允许最终层在固定 eNTK 特征上进行端到端训练，从而有效学习一个基于预训练表征的线性头。

实验结果

研究问题

• RQ1 为何当前联邦优化方法在数据非独立同分布时仍无法实现与集中式训练相当的性能？
• RQ2 在联邦设置下，神经网络的早期层在多大程度上学习到了有用特征？为何最终层无法有效利用这些特征？
• RQ3 通过经验神经正切核（eNTK）进行凸化，是否能恢复在数据异构性条件下的非凸联邦学习性能？
• RQ4 第一阶段预训练持续时间的选择如何影响 TCT 框架的最终模型准确率？
• RQ5 对 eNTK 特征进行归一化是否能提升第二阶段凸优化的收敛速度和通信效率？

主要发现

• 当客户端数据分布差异显著时，TCT 在 FMNIST 上实现了高达 36% 的准确率增益，在 CIFAR10 上实现了 37% 的增益。
• 当第一阶段使用预训练特征提取器时，TCT 完全弥合了集中式与联邦学习之间的性能差距。
• 使用 FedAvg 训练模型提取的 eNTK 特征，TCT 在测试准确率上比标准 FedAvg 提高约 20%。
• 对 eNTK 特征进行归一化可显著提升收敛速度，使 TCT 在约 40 次通信轮次内达到接近 100% 的训练准确率。
• TCT 在第一阶段通信轮数（T₁ ≥ 60）变化时性能保持稳定，表明其对超参数选择具有鲁棒性。
• FedAvg 的性能随本地训练步数增加而下降（M=1000），而 TCT 中使用的 SCAFFOLD 则保持或提升收敛性能，凸显了梯度校正的重要性。