[Paper Review] 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 Data38 references4 citations

TL;DR

This paper introduces FL-NTK, a novel Neural Tangent Kernel-based framework for analyzing convergence and generalization in Federated Learning with overparameterized ReLU neural networks. By leveraging the NTK theory, it proves linear convergence to global optimum under non-convex, non-smooth objectives and establishes a data-dependent generalization bound, overcoming limitations of prior gradient-based analyses in FL.

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.

Motivation & Objective

• To address the lack of theoretical convergence guarantees for federated learning with deep neural networks, especially under non-convex, non-smooth objectives and multi-step local updates.
• To overcome the challenge that standard gradient-based convergence analysis fails in FL due to gradient heterogeneity and non-gradient-based global updates.
• To develop a unified, generalizable framework for analyzing convergence and generalization in FL using overparameterized neural networks.
• To establish explicit convergence rates and generalization bounds without requiring convexity or smoothness assumptions on the loss function.

Proposed method

• Proposes a Federated Learning Neural Tangent Kernel (FL-NTK) framework inspired by the Neural Tangent Kernel (NTK) theory for overparameterized networks.
• Analyzes the evolution of model weights in FL using an asymmetric Gram matrix, which captures the dynamics of local updates and global aggregation.
• Introduces a data-dependent generalization bound via Rademacher complexity and empirical process theory, accounting for client data heterogeneity.
• Employs probabilistic concentration bounds and union bounds over failure events to control deviations in model weights and generalization error.
• Uses random initialization and spectral assumptions on the infinite-width NTK Gram matrix to ensure convergence and stability.
• Applies tools from statistical learning theory, including covering numbers and empirical process bounds, to derive generalization error with high probability.

Experimental results

Research questions

• RQ1 Can we establish a convergence analysis framework for federated learning with ReLU neural networks that does not rely on convexity or smoothness of the loss function?
• RQ2 How does gradient heterogeneity across clients affect the convergence of FL when using overparameterized models?
• RQ3 Can we derive a generalization bound for FL that depends on data distribution and model capacity in the overparameterized regime?
• RQ4 What is the impact of the number of clients and local update steps on the convergence rate in FL with overparameterized networks?
• RQ5 How can we handle the asymmetric structure of the FL-NTK Gram matrix, which differs from the symmetric NTK in centralized training?

Key findings

• The training loss across all clients converges to zero at a linear rate under proper learning rate tuning, when the network is sufficiently wide.
• The population loss is upper bounded by $ \sqrt{\frac{2y^{\top}(H^{\infty})^{-1}y}{n}} + O\left(\sqrt{\frac{\log(n/\lambda\delta)}{2n}}\right) $ with high probability.
• The framework achieves convergence without assuming convexity or smoothness of the objective function.
• Generalization error is bounded using Rademacher complexity, with the bound depending on data distribution and model width.
• The number of communication rounds increases with the number of clients, which aligns with empirical observations.
• The analysis holds under a non-degeneracy condition on the data distribution, ensuring the infinite-width NTK Gram matrix is invertible with positive minimum eigenvalue.