NTK-DFL: Enhancing Decentralized Federated Learning in Heterogeneous Settings via Neural Tangent Kernel

Gabriel Thompson, Kai Yue, Chau-Wai Wong, Huaiyu Dai

Abstract: Decentralized federated learning (DFL) is a collaborative machine learning framework for training a model across participants without a central server or raw data exchange. DFL faces challenges due to statistical heterogeneity, as participants often possess data of different distributions reflecting local environments and user behaviors. Recent work has shown that the neural tangent kernel (NTK) approach, when applied to federated learning in a centralized framework, can lead to improved performance. We propose an approach leveraging the NTK to train client models in the decentralized setting, while introducing a synergy between NTK-based evolution and model averaging. This synergy exploits inter-client model deviation and improves both accuracy and convergence in heterogeneous settings. Empirical results demonstrate that our approach consistently achieves higher accuracy than baselines in highly heterogeneous settings, where other approaches often underperform. Additionally, it reaches target performance in 4.6 times fewer communication rounds. We validate our approach across multiple datasets, network topologies, and heterogeneity settings to ensure robustness and generalization. Source code for NTK-DFL is available at https://github.com/Gabe-Thomp/ntk-dfl}{https://github.com/Gabe-Thomp/ntk-dfl

Lay Summary:

Collaborative training of machine learning models is gaining traction as large, ChatGPT-like models take an unprecedented amount of computing resources to train. Particularly, individuals may want to collaboratively train models without explicitly sharing their personal data, and without the need for a larger, central server. In other words, models may be trained in a decentralized fashion.

This introduces the problem of data heterogeneity: different users have different kinds of data! Your mom may love pictures of cats, while your brother only takes pictures of dogs and parrots. If we were training an animal classifier on these images, the variation of data across devices can hinder training.

In our paper, we investigate a training algorithm that replaces the typical training approach with a new one. We use a mathematical tool called the neural tangent kernel. In plain English, this tool allows the user to share more expressive data in the training process. By sharing better data, model training is enhanced despite the data variation discussed prior. Also, users must communicate with over fewer rounds than in previously proposed algorithms. Lastly, we provide open-source code for the research community to test and build upon our algorithm.

Link To Code: https://github.com/Gabe-Thomp/ntk-dfl

Primary Area: General Machine Learning

Keywords: Federated Learning, Decentralized Federated Learning, Neural Tangent Kernel

Submission Number: 7782