[论文解读] FL-NTK: A Neural Tangent Kernel-based Framework for Federated Learning Convergence Analysis

Baihe Huang, Xiaoxiao Li+2 more|arXiv (Cornell University)|May 11, 2021

Privacy-Preserving Technologies in Data参考文献 38被引用 4

一句话总结

该论文提出了一种基于神经正切核（Neural Tangent Kernel, NTK）的新框架——FL-NTK，用于分析过参数化ReLU神经网络在联邦学习中的收敛性与泛化性能。通过利用NTK理论，该工作在非凸、非光滑目标函数下证明了线性收敛至全局最优解，并建立了依赖于数据的泛化边界，克服了联邦学习中以往基于梯度分析方法的局限性。

ABSTRACT

Federated Learning (FL) is an emerging learning scheme that allows different distributed clients to train deep neural networks together without data sharing. Neural networks have become popular due to their unprecedented success. To the best of our knowledge, the theoretical guarantees of FL concerning neural networks with explicit forms and multi-step updates are unexplored. Nevertheless, training analysis of neural networks in FL is non-trivial for two reasons: first, the objective loss function we are optimizing is non-smooth and non-convex, and second, we are even not updating in the gradient direction. Existing convergence results for gradient descent-based methods heavily rely on the fact that the gradient direction is used for updating. This paper presents a new class of convergence analysis for FL, Federated Learning Neural Tangent Kernel (FL-NTK), which corresponds to overparamterized ReLU neural networks trained by gradient descent in FL and is inspired by the analysis in Neural Tangent Kernel (NTK). Theoretically, FL-NTK converges to a global-optimal solution at a linear rate with properly tuned learning parameters. Furthermore, with proper distributional assumptions, FL-NTK can also achieve good generalization.

研究动机与目标

• 为解决深度神经网络在联邦学习中缺乏理论收敛保证的问题，特别是针对非凸、非光滑目标函数以及多步本地更新的情形。
• 克服标准基于梯度的收敛性分析在联邦学习中失效的挑战，原因在于梯度异质性以及非基于梯度的全局更新机制。
• 开发一种统一且可泛化的框架，用于分析过参数化神经网络在联邦学习中的收敛性与泛化性能。
• 在不假设损失函数具有凸性或光滑性的前提下，推导出显式的收敛速率与泛化边界。

提出的方法

• 提出一种受神经正切核（NTK）理论启发的联邦学习神经正切核（FL-NTK）框架，适用于过参数化网络。
• 通过使用非对称Gram矩阵分析联邦学习中模型权重的演化过程，该矩阵能够捕捉本地更新与全局聚合的动力学特性。
• 利用Rademacher复杂度与经验过程理论，引入一种依赖于数据的泛化边界，以考虑客户端间数据异质性的影响。
• 采用概率浓度不等式与失败事件的联合界，控制模型权重与泛化误差的偏差。
• 通过随机初始化以及对无限宽度下NTK Gram矩阵的谱假设，确保收敛性与稳定性。
• 应用统计学习理论中的工具，包括覆盖数与经验过程边界，以高概率推导出泛化误差。

实验结果

研究问题

• RQ1 我们能否建立一个适用于ReLU神经网络联邦学习的收敛性分析框架，且该框架不依赖于损失函数的凸性或光滑性？
• RQ2 在使用过参数化模型时，客户端之间的梯度异质性如何影响联邦学习的收敛性？
• RQ3 我们能否在过参数化设置下，推导出依赖于数据分布与模型容量的联邦学习泛化边界？
• RQ4 客户端数量与本地更新步数对过参数化网络在联邦学习中收敛速率的影响是什么？
• RQ5 我们如何处理FL-NTK Gram矩阵的非对称结构，该结构与集中式训练中的对称NTK存在差异？

主要发现

• 当网络足够宽且学习率适当调节时，所有客户端的训练损失以线性速率收敛至零。
• 以高概率，总体损失被上界控制为 $ \sqrt{\frac{2y^{\top}(H^{\infty})^{-1}y}{n}} + O\left(\sqrt{\frac{\log(n/\lambda\delta)}{2n}}\right) $。
• 该框架在不假设目标函数具有凸性或光滑性的情况下实现了收敛。
• 通过Rademacher复杂度对泛化误差进行有界控制，其边界依赖于数据分布与模型宽度。
• 通信轮数随客户端数量的增加而上升，这与经验观察结果一致。
• 在数据分布非退化（non-degeneracy）的条件下分析成立，确保无限宽NTK Gram矩阵具有正的最小特征值且可逆。