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

## Abstract

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 different data 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. The NTK-based update mechanism is more expressive than typical gradient descent methods, enabling more efficient convergence and better handling of data heterogeneity. 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-model variance and improves both accuracy and convergence in heterogeneous settings. Our model averaging technique significantly enhances performance, boosting accuracy by at least 10% compared to the mean local model accuracy. 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 generalizability.

## 1 Introduction

Federated learning (FL) is a machine learning paradigm in which multiple clients train a global model without the explicit communication of training data. In most FL scenarios, clients communicate with a central server that performs model aggregation. In the popular federated averaging (FedAvg) algorithm (McMahan et al., ), clients perform multiple rounds of stochastic gradient descent (SGD) on their own local data, then send this new weight vector to a central server for aggregation. As FL gains popularity in both theoretical studies and real-world applications, numerous improvements have been made to address challenges, including communication efficiency, heterogeneous data distributions, and security concerns (Sattler et al., ; Li et al., ; Zhu et al., ).
To handle the performance degradation caused by data heterogeneity, many works have proposed mitigation for FedAvg (Karimireddy et al., ; Li et al., ).
Notably, some researchers have introduced the neural tangent kernel (NTK), replacing the commonly-used SGD in order to improve the model convergence (Yu et al., ; Yue et al., ).

Despite these advancements, the centralized nature of traditional FL schemes introduces the possibility for client data leakage, computational bottlenecks at the server, and high communication bandwidth demand (Kairouz et al., ). Decentralized federated learning (DFL) has been proposed as a solution to these issues (Martínez Beltrán et al., ). In DFL, clients may communicate with each other along an undirected graph, where each node represents a client and each edge represents a communication channel between clients.
While DFL addresses some of the issues inherent to centralized FL, both frameworks grapple with the challenge of statistical heterogeneity across clients.
Although mixing data on a central server could readily resolve this issue, privacy concerns and the burden of extensive communication make FL and DFL approaches necessary to address this challenge.
This paper focuses on the following research question: How can we design a DFL approach that effectively addresses statistical heterogeneity?

We propose a DFL method that exploits the NTK to evolve weights. We denote this paradigm NTK-DFL. Our approach combines the advantages of NTK-based optimization with the decentralized structure of DFL. The NTK-DFL weight evolution scheme makes use of the communication of client Jacobians, allowing for more expressive updates than traditional weight vector transmissions and improving performance under heterogeneity. Complementing this NTK-based evolution, we utilize a model averaging step that exploits inter-model variance among clients, creating a global model with much better generalization than any local model.
The contributions of this paper are threefold.

- 1.
The proposed NTK-DFL method achieves convergence with 4.6 times fewer communication rounds than existing approaches in heterogeneous settings. To the best of our knowledge, this is the first work leveraging NTK-based weight evolution for decentralized federated training.
- 2.
The effective synergy between NTK-based evolution and DFL demonstrates superior resilience to data heterogeneity with model averaging.
- 3.
The NTK-DFL aggregated model achieves at least 10% higher accuracy than the average accuracy of individual client models. This aggregated model exhibits robust performance across various network topologies, datasets, and data distributions.

## 2 Related Work

Federated Learning (FL) FL was introduced by McMahan et al. () as a machine learning approach that enables training a model on distributed datasets without sharing raw data. It attempts to address key issues such as data privacy, training on decentralized data, and data compliance for more heavily regulated data (e.g., medical imaging) (Zhang et al., ). Despite its advantages, the centralized topology of FL introduces several challenges. These include potential privacy risks at the central server, scalability issues due to computational bottlenecks, and high communication overhead from frequent model updates between clients and the server (Mothukuri et al., ).

Decentralized Federated Learning (DFL) DFL aims to eliminate the need for a central server by connecting clients in a fully decentralized topology.
Sun et al. () adapted the FedAvg approach of multiple local SGD iterations to the decentralized setting, leveraging momentum to improve model convergence and weight quantization to reduce total communication cost. Our NTK-DFL method may be viewed as building on this foundation, using the neural tangent kernel for more effective weight updates.
Dai et al. () proposed a method of DFL where each client possesses their own sparse mask personalized to their specific data distribution.
Shi et al. () employed the sharpness-aware minimization optimizer to reduce the inconsistency of local models, whereas we tackle this issue through per-round averaging and final model aggregation. DFL approaches can aim to train one global model, such as the case of many hospitals training a model for tumor classification with local, confidential images (Shiri et al., ). They may also aim to train a personalized model for each client in order to perform better on the local data distribution. For example, different groups of mobile phone users may use different words or emojis and would benefit from a personalized model (Tan et al., ). Our method focuses on training a high-performing global model that generalizes well across all clients, offering improved convergence and resilience to data heterogeneity compared to existing DFL approaches.

Neural Tangent Kernel (NTK) NTK has primarily been used for the analysis of neural networks (Golikov et al., ), though it has recently seen use in the training of neural networks for FL (Yue et al., ). Introduced by Jacot et al. (), it shows that the evolution of an infinitely wide neural network converges to a kernelized model. This approach has enabled the analytical study of models that are well approximated by this infinite width limit (Liu et al., ). The NTK has also been extended to other model types, such as the recurrent neural network (Alemohammad et al., ) and convolutional neural network (Arora et al., ). We instead use the linearized model of the NTK approximation as a tool for weight evolution.
Some studies have explored the integration of NTKs with FL.
For instance, Huang et al. () applied the NTK analysis framework to study the convergence properties of FedAvg, while Yu et al. () extended NTK applications beyond theoretical analysis by training a convex neural network. Moreover, Yue et al. () replaced traditional SGD-based optimization with NTK-based evolution in a federated setting, where clients transmit Jacobian matrices to a central server that performs weight updates using the NTK.

## 3 Proposed Method: NTK-Based Decentralized Federated Learning

### 3.1 Problem Statement

We begin with a brief overview of centralized FL. The goal of centralized FL is to train a global model ${\bm{w}}$ across $M$ clients with their private, local data
$\mathcal{D}_{i}=\{({\bm{x}}_{i,j},{\bm{y}}_{i,j})\}_{j=1}^{N_{i}}$, where $N_{i}$ is the number of training examples of the $i$th client. FL algorithms aim to
numerically solve the sample-wise optimization problem of $\operatorname*{min}_{\bm{w}}F({\bm{w}})$, where
$F({\bm{w}})=\frac{1}{M}\sum_{i=1}^{M}N_{i}F_{i}({\bm{w}})$ and
$F_{i}({\bm{w}})=\frac{1}{N_{i}}\sum_{j=1}^{N_{i}}\mathcal{\ell}({\bm{w}},{\bm{
x}}_{j},{\bm{y}}_{j})$,
where $M$ denotes the number of clients.
In the decentralized setting, an omnipresent global weight ${\bm{w}}$ is not available to clients in each communication round. Rather, each client possesses their own model ${\bm{w}}_{i}$ that is trained in the update process. Following related DFL work (Shi et al., ; Sun et al., ), we seek a global model ${\bm{w}}$ that benefits from the heterogeneous data stored locally across clients and generalizes better than any individual client model.
A global or aggregated model may take the form ${\bm{w}}=\frac{1}{M}\sum_{i=1}^{M}N_{i}{\bm{w}}_{i}$.

Notation Formally, we have a set of clients $\mathcal{C}=\{C_{1},C_{2},\ldots,{C_{M}}\}$. Each client is initialized with its weight ${\bm{w}}_{i}^{(0)}\in\mathbb{R}^{d}$,
where $d$ is the size of the parameter vector and the superscript in ${\bm{w}}_{i}^{(0)}$ denotes the initial communication round. Model training is done in a series of communication rounds denoted $k\in\{1,2,...,K\}$. Let the graph at round $k$ be $\mathcal{G}^{(k)}=(\mathcal{C},E^{(k)})$, where $E^{(k)}$ is the set of edges representing connections between clients.
Furthermore, the neighborhood of client $i$ at round $k$ is denoted $\mathcal{N}_{i}^{(k)}=\{C_{j}\mid(C_{i},C_{j})\in E^{(k)}\}$. This graph is specified before each communication round and can take an arbitrary form.

Figure: Figure 1: NTK-DFL process: ① Clients exchange weights, ② Average weights with neighbors, ③ Compute and exchange Jacobians, labels, and function evaluations, ④ Construct local NTK and evolve weights [Eq. (5)]. This decentralized approach enables direct client collaboration and NTK-driven model evolution without a central server.
Refer to caption: x1.png

### 3.2 Communication Protocol

In the following sections, we present the NTK-DFL paradigm (Figure 1). We describe the key components of the algorithm, including the communication protocol and weight evolution process.

Per-round Parameter Averaging At the beginning of each communication round $k$, each client $C_{i}$ both sends and receives weights. Every client sends their model ${\bm{w}}_{i}^{(k)}$ to all neighbors $C_{j}\in\mathcal{N}_{i}^{(k)}$. Simultaneously, $C_{i}$ receives the weight vectors ${\bm{w}}_{j}^{(k)}$ from all neighbors $C_{j}\in\mathcal{N}_{i}^{(k)}$. Each client then aggregates its own weights along with the neighboring weights to form a new weight as follows:

$$ $\bar{{\bm{w}}}_{i}^{(k)}=\frac{1}{|\mathcal{N}_{i}^{(k)}|+1}\big{(}{\bm{w}}_{i }^{(k)}+\sum_{j\in\mathcal{N}_{i}^{(k)}}{\bm{w}}_{j}^{(k)}\big{)}.$ (1) $$

The client must then send this aggregated weight $\bar{{\bm{w}}}_{i}^{(k)}$ back to all neighbors $C_{j}\in\mathcal{N}_{i}^{(k)}$. This step enables each client to
construct a local NTK, comprised of inner products of Jacobians from both neighboring clients and their own Jacobians. See Algorithm 2 in Appendix A for a detailed description of this process.

Local Jacobian Computation At this point, each client possesses its own aggregated weight $\bar{{\bm{w}}}_{i}^{(k)}$ as well as an aggregated weight $\bar{{\bm{w}}}_{j}^{(k)}$ for each of their neighbors $C_{j}\in\mathcal{N}_{i}^{(k)}$. The clients use these weights $\bar{{\bm{w}}}_{j}^{(k)}$ and local data $\mathbf{X}_{i}$ to compute the Jacobian of $f(\mathbf{X}_{i};\bar{{\bm{w}}}_{j}^{(k)})$ with respect to the neighboring model parameters $\bar{{\bm{w}}_{j}}^{(k)}$ and their local data $\mathbf{X}_{i}$. We can denote this neighbor-specific Jacobian as

$$ ${\bm{J}}^{(k)}_{i,j}\equiv[\nabla_{\bar{{\bm{w}}_{j}}}{\bm{f}}(\mathbf{X}_{i}; \bar{{\bm{w}}}_{j}^{(k)})]^{\top}.$ (2) $$

For a given client, the gradient is taken with respect to its neighbor’s aggregated weight $\bar{{\bm{w}}}_{j}^{(k)}$, but the function is evaluated on the client’s local data $\mathbf{X}_{i}$. Each client sends every neighbor their respective Jacobian tensor ${\bm{J}}^{(k)}_{i,j}$, true label $\mathbf{Y}_{i}$, and function evaluation ${\bm{f}}(\mathbf{X}_{i};\bar{{\bm{w}}}_{j}^{(k)})$. Note the order of the indices in the Jacobian: the client sends ${\bm{J}}_{i,j}^{(k)}$, an evaluation on the client’s data and the neighbor’s weights. In contrast, the client receives ${\bm{J}}_{j,i}^{(k)}$ from each of its neighbors, an evaluation on the neighbor’s data and the client’s weights. Algorithm 3 in Appendix A describes this process.

### 3.3 Weight Evolution

After all inter-client communication is completed, the clients begin the weight evolution phase of the round (see Algorithm 4 in Appendix A). Here, all clients act in parallel as computational nodes. Each client possesses their own Jacobian tensor ${\bm{J}}_{i,i}^{(k)}$ as well as their neighboring Jacobian tensors ${\bm{J}}_{j,i}^{(k)}$ for each $C_{j}\in\mathcal{N}_{i}^{(k)}$.

We denote the tensor of all Jacobian matrices possessed by a client at round $k$ as $\bm{\mathcal{J}}_{i}^{(k)}$, which is composed of matrices from the set $\{{\bm{J}}_{i,i}^{(k)}\}\cup\{{\bm{J}}_{j,i}^{(k)}\ |\ C_{j}\in\mathcal{N}_{i}
^{(k)}\}$ stacked along a third dimension. We denote the matrix of true labels and function evaluations stacked in the same manner as $\bm{\mathcal{Y}}_{i}$ and ${\bm{f}}(\bm{\mathcal{X}}_{i})$, respectively. Here, $i$ denotes a client index $C_{i}\in\mathcal{C}$. Each tensor $\bm{\mathcal{J}}_{i}$, $\bm{\mathcal{Y}}_{i}$, and ${\bm{f}}(\bm{\mathcal{X}}_{i})$ is a stacked representation of the data from each client and its neighbors. Explicitly, we have $\bm{\mathcal{J}}_{i}^{(k)}\in\mathbb{R}^{\tilde{N_{i}}\times d_{2}\times d}$, $\bm{\mathcal{Y}}_{i}^{(k)}\in\mathbb{R}^{\tilde{N_{i}}\times d_{2}}$, and ${\bm{f}}(\bm{\mathcal{X}}_{i})\in\mathbb{R}^{\tilde{N_{i}}\times d_{2}}$. $\tilde{N_{i}}$ denotes the total number of data points between client $C_{i}$ and its neighbors, and $d_{2}$ is the output dimension.

From here, each client performs the following operations to evolve its weights. First, compute the local NTK $\mathbf{H}_{i}^{(k)}$ from the Jacobian tensor $\bm{\mathcal{J}}_{i}^{(k)}$ using the definition of the NTK

$$ $\mathbf{H}^{(k)}_{i,mn}=\frac{1}{d_{2}}\langle\bm{\mathcal{J}}_{i}^{(k)}({\bm{ x}}_{m}),\bm{\mathcal{J}}_{i}^{(k)}({\bm{x}}_{n})\rangle_{F}.$ (3) $$

Each element of the NTK is a pairwise Frobenius inner product between Jacobian matrices, where the indices $m$ and $n$ correspond to two different data points. Second, using $\mathbf{H}_{i}^{(k)}$, the client evolves their weights as follows:

$$ ${\bm{f}}^{(k,t)}(\bm{\mathcal{X}}_{i})=(\mathbf{I}-e^{\frac{\eta t}{\tilde{N_{ i}}}\mathbf{H}_{i}^{(k)}})\bm{\mathcal{Y}}_{i}^{(k)}+e^{\frac{\eta t}{\tilde{N _{i}}}\mathbf{H}_{i}^{(k)}}{\bm{f}}^{(k)}(\bm{\mathcal{X}}_{i}),$ (4) $$

$$ ${\bm{w}}_{i}^{(k,t)}=\sum_{j=1}^{d_{2}}(\bm{\mathcal{J}}^{(k)}_{i,:j:})^{\top} \mathbf{R}_{i,:j}^{(k,t)}+\bar{{\bm{w}}}_{i}^{(k)},\qquad\mathbf{R}_{i,:j}^{(k ,t)}\equiv\frac{\eta}{\tilde{N_{i}}d_{2}}\sum_{u=0}^{t-1}[\bm{\mathcal{Y}}_{i} ^{(k)}-{\bm{f}}^{(k,u)}(\bm{\mathcal{X}}_{i})].$ (5) $$

Third, the client selects the weight ${\bm{w}}_{i}^{(k,t)}$ for a timestep $t$ with the lowest loss according to the evolved residual ${\bm{f}}^{(k,t)}(\bm{\mathcal{X}}_{i})-\bm{\mathcal{Y}}_{i}$. This is used as the new weight ${\bm{w}}_{i}^{(k+1,0)}$ for the next communication round.

Final Model Averaging  Throughout the paper, we study the convergence of the aggregated model ${\bm{w}}=\frac{1}{{M}}\sum_{i=1}^{M}N_{i}{\bm{w}}_{i}$. In the decentralized setting, clients would average all models to create ${\bm{w}}$ after all training is completed. This may be done through a fully-connected topology, sequential averaging on a ring topology, or in a secure, centralized manner. Clients may also connect in a denser topology than that of training, and average with a desired number of neighbors. In practice, we observe that the aggregated model is more accurate than any individual client model. We study the impact of client averaging order on model performance with a client selection algorithm and show the results in Figure 6. Each client that opts in to model averaging contributes a portion of its data to a global validation set before training begins. Our client selection algorithm selects clients in the order of their accuracy on the validation set. We will demonstrate that in the practical setting, with a proper selection of clients, not all nodes must opt into final model averaging in order for the aggregate model to benefit from improved convergence. We note a difference between model consensus, often discussed in the DFL literature (Savazzi et al., ; Liu et al., ), and the proposed final model averaging approach. Model consensus refers to the gradual convergence of all client models to a single, unified model over numerous communication rounds. In contrast, our approach implements final model averaging as a distinct step performed after the completion of the training process.

Lastly, while memory efficiency is not the primary focus of this paper, we briefly note a technique to address potential memory constraints in NTK-DFL implementations. For scenarios involving dense networks or large datasets, we introduce Jacobian batching. This approach allows clients to process their local datasets in smaller batches, reducing memory complexity from $O(N_{i}d_{2}d)$ to $O(N_{i}d_{2}d/m)$, where $m$ is the number of batches. Clients compute and transmit Jacobians for each batch separately, evolving their weights multiple times per communication round. This complexity reduction allows clients to connect in a denser network for the same memory cost.

## 4 Experiments

### 4.1 Experimental Setup

Datasets and Model Specifications Following Yue et al. (), we experiment on three datasets: Fashion-MNIST (Xiao et al., ), FEMNIST (Caldas et al., ), and MNIST (Lecun et al., ). Each dataset contains $C=10$ output classes. For Fashion-MNIST and MNIST, data heterogeneity has been introduced in the form of non-IID partitions created by the symmetric Dirichlet distribution (Good, ). For each client, a vector ${\bm{q}}_{i}\sim\text{Dir}(\alpha)$ is sampled, where ${\bm{q}}_{i}\in\mathbb{R}^{C}$ is confined to the $(C-1)$-standard simplex such that $\sum_{j=1}^{C}q_{ij}=1$. This assigns a distribution over labels to each client, creating heterogeneity in the form of label-skewness. For smaller values of $\alpha$, clients possess a distribution concentrated over fewer classes. We test over a range of $\alpha$ values in order to simulate different degrees of heterogeneity. In FEMNIST, data is split into shards based on the writer of each digit, introducing heterogeneity in the form of feature-skewness. For the model, we use a two-layer multilayer perceptron with a hidden width of 100 neurons for all trials.

Network Topologies A sparse, time-variant $\kappa$-regular graph with $\kappa=5$ was used as the standard topology for experimentation, where for each communication round $k$, a new random graph $\mathcal{G}^{(k)}$ with the same parameter $\kappa$ is created. Various values of $\kappa$ were tested to observe the effect of network density on model convergence. We also experimented with various topologies to ensure robustness to different connection settings. We used a network of 300 clients throughout our experiments.

Baseline Methods We compare our approach to various state-of-the-art baselines in the DFL setting. These include D-PSGD (Lian et al., ), DFedAvg, DFedAvgM (Sun et al., ), and DisPFL (Dai et al., ). We also compare with the centralized baseline NTK-FL (Yue et al., ).
The upper bound NTK-FL would consist of a client fraction of $1.0$ where the servers constructs an NTK from all client data each round, which is infeasible due to memory constraints. Instead, we conducted a comparison following Dai et al. (), equalizing the degree of the busiest node in DFL with that of the server in centralized FL. Additional details regarding baselines and NTK-FL results can be found in Appendix B.

Performance Metrics We evaluate the performance of the various DFL approaches by studying the aggregate model accuracy on a global, holdout test set. This ensures that we are measuring the generalization of the aggregate model from individual, heterogeneous local data to a more representative data sample. Our approach is in line with the goal of training a global model capable of improved generalization over any single local model (Section 3.1), unlike personalized federated learning where the goal is to fine-tune a global model to each local dataset (Tan et al., ). When evaluating the selection algorithm in Figure 6, we split the global test set in a 50:50 ratio of validation to test data. We use the validation data to sort the models based on their accuracy, and report the test accuracy in the figure.

### 4.2 Experimental Results

Test Accuracy & Convergence Our experiments demonstrate the superior convergence properties of NTK-DFL compared to baselines.
Figure 2 illustrates the convergence trajectories of NTK-DFL and other baselines on Fashion-MNIST. We see that NTK-DFL convergence benefits are enhanced under increased heterogeneity.
Under high heterogeneity with $\alpha=0.1$, NTK-DFL establishes a 3–4% accuracy lead over the best-performing baseline within just five communication rounds and maintains this advantage throughout the training process. Additionally presented are the number of communication rounds necessary for convergence to 85% test accuracy, where NTK-DFL consistently outperforms all baselines.
For the $\alpha=0.1$ setting, NTK-DFL achieves convergence in 4.6 times fewer communication rounds than DFedAvg, the next best performing baseline.
Figure 8 demonstrates a similar convergence advantage for NTK-DFL on both FEMNIST and non-IID MNIST datasets.

Figure: Figure 2: Convergence of different methods on Fashion-MNIST for (left) highly non-IID with $\alpha=0.1$ and (middle) IID settings. (Right) The table displays the communication rounds required to reach 85% test accuracy on Fashion-MNIST. We observe increased improvement in NTK-DFL convergence over baselines for more heterogeneous settings.
Refer to caption: x2.png

Factor Analyses of NTK-DFL We evaluate NTK-DFL’s performance over various factors, including the sparsity and heterogeneity levels, and the choices of the topology and weight initialization scheme.
Figure 3 illustrates the test accuracy of NTK-DFL and other baselines as functions of the sparsity and heterogeneity levels, respectively.
We observe a mild increase in convergence accuracy with decreasing sparsity. NTK-DFL experiences stable convergence across heterogeneity values $\alpha$ ranging from $0.1$ to $0.5$.

Figure: Figure 3: Performance of NTK-DFL vs. (left) sparsity level and (right) heterogeneity level (smaller $\alpha$ $\to$ more heterogeneous). NTK-DFL outperforms the baselines and the gains are stable as the factors vary.
Refer to caption: x4.png

The left plot reveals that NTK-DFL consistently outperforms baselines by 2–3% across all sparsity levels.
The right plot demonstrates NTK-DFL’s resilience to data heterogeneity—while baseline methods’ performance deteriorates with decreasing $\alpha$, NTK-DFL maintains stable performance.
Figure 10 illustrates the impact of network topology on NTK-DFL convergence. The dynamic topology accelerates convergence compared to the static topology, likely due to improved information flow among clients. Figure 10 demonstrates the effect of weight initialization on NTK-DFL performance. While random per-client initialization slightly slows convergence compared to uniform initialization, NTK-DFL exhibits robustness to these initialization differences.

Figure: Figure 4: Performance gains of model averaging on convergence, trained on Fashion-MNIST. Solid lines correspond to accuracy of the aggregated global model, whereas dotted lines correspond to the mean accuracy across client models. NTK-DFL’s aggregated model maintains high performance, whereas mean client accuracy declines significantly with increased heterogeneity.
Refer to caption: x6.png

Gains Due to Final Model Aggregation
Figure 4
demonstrates the dramatic effect of final model aggregation on final test accuracy. Though the individual client models decrease in accuracy as the level of heterogeneity increases, the final aggregated model remains consistent across all levels of heterogeneity (as seen in Figures 2 and 3).
In the most heterogeneous setting $\alpha=0.1$ that we tested, the difference between the mean accuracy of each client and the aggregated model accuracy is nearly 10%. A similar phenomenon is observed as the client topology becomes more sparse. For the same heterogeneity setting with a sparser topology of $\kappa=2$, the difference between these accuracies is nearly 15%. Though the individual performance of local client models may suffer under extreme conditions, the inter-client variance created by such unfavorable settings is exploited by model averaging to recuperate much of that lost performance.

Figure 7 suggests
that inter-model variance enhances the performance of model averaging in DFL.
While extreme dissimilarity in model weights would likely result in poor performance of the averaged model, we observe that a moderate degree of variance can be beneficial. We posit that the NTK-based update steps generate a more advantageous level of variance compared to baseline approaches, contributing to improved overall performance.

Figure: Figure 5: Final model test accuracy on Fashion-MNIST vs. the number of clients averaged for a highly heterogeneous setting with $\alpha=0.1$. The histogram shows the distribution of individual client model accuracies. Three client selection criteria are tested: (red) high-to-low (proposed), (green) random, and (blue) low-to-high.
Refer to caption: x7.png

Figure: Figure 7: Relationship between model variance and final test accuracy on Fashion-MNIST. Each point represents a trial with distinct hyperparameters. The plot reveals a positive correlation between model variance and accuracy, suggesting that higher variance may benefit model averaging in DFL to a certain extent. Notably, the NTK-DFL approach demonstrates both higher accuracy and greater model variance compared to other methods.
Refer to caption: x9.png

Selection Algorithm
Figure 6 demonstrates the results of the selection algorithm for the final model aggregation. The selection algorithm is highly effective in the heterogeneous setting. The effect is most notable in the $\alpha=0.1$ setting, where the performance with the selection algorithm significantly outperforms a random averaging order and the lower-bound averaging order. The proposed selection criterion requires the fewest clients to be averaged to achieve the same level of accuracy in this highly heterogeneous setting. In practical deployments, this has implications for final-round averaging in a fully decentralized setting. For example, clients may connect in a denser final topology and prioritize averaging with neighbors possessing a higher validation accuracy. This approach could optimize the efficacy of the final aggregation step while maintaining the decentralized nature of the system.

Per-round Averaging Ablation Study In Figure 6, we perform an ablation study in which we remove the per-round parameter averaging that is a part of the NTK-DFL process. Here, clients forego the step of averaging their weight vectors with their neighbors during each communication round. Instead, clients compute Jacobians with respect to their original weight vector and send these to each of their neighbors (see Algorithm 3 in Appendix A). A massive distribution shift can be seen in the figure, where the distribution in the ablated setting is clearly skewed into lower accuracies. In contrast, NTK-DFL with per-round averaging demonstrates a much tighter distribution around a higher mean accuracy, effectively eliminating the long tail of low-performing models. Per-round averaging in NTK-DFL serves as a stabilizing mechanism against local model drift, safeguarding clients against convergence to suboptimal solutions early in the training process.
In other words, client collaboration in the form of per-round averaging with neighbors ensures that no client lags behind in convergence. This is a particularly valuable feature in decentralized federated learning scenarios where maintaining uniformity across a diverse set of clients with heterogeneous data is a major challenge (Martínez Beltrán et al., ).

## 5 Conclusion and Future Work

In this paper, we have introduced NTK-DFL, a novel approach to decentralized federated learning that leverages the neural tangent kernel to address the challenges of statistical heterogeneity in decentralized learning settings. Our method combines the expressiveness of NTK-based weight evolution with a decentralized architecture, allowing for efficient and privacy-preserving collaborative learning without a central server. We reduce the number of communication rounds needed for convergence, which may prove advantageous for high-latency settings or those with heavy encoding/decoding costs.

There are promising unexplored directions for NTK-DFL. For instance, extending the algorithm to training models such as CNNs, ResNets (He et al., ), and transformers (Vaswani et al., ). Additionally, future research could explore the application of NTK-DFL to cross-silo federated learning scenarios, particularly in domains such as healthcare, where data privacy concerns and regulatory requirements often necessitate decentralized approaches. Lastly, NTK-DFL may serve as a useful paradigm for transfer learning applications in scenarios where a single, centralized source of both compute and data is not available.

Figure: (a) FEMNIST
Refer to caption: x10.png

Figure: Figure 9: The effect of static vs. dynamic topology on NTK-DFL. Solid lines correspond to a dynamic topology, whereas dotted lines correspond to a static topology. Both methods benefit from the dynamic topology and NTK-DFL outperforms DFedAvg under both topologies. Other baselines are not drawn but perform similarly to DFedAvg.
Refer to caption: x13.png

## Appendix A NTK-DFL Algorithms

Figure: Algorithm 1 Consolidated Federated Learning Process

Figure: Algorithm 2 Per-Round Parameter Averaging

Figure: Algorithm 3 Local Jacobian Computation and Sending Jacobians

Figure: Algorithm 4 Weight Evolution

## Appendix B Additional Experimental Details

### B.1 Baselines

NTK-FL is the only centralized baseline that we compare with. We choose a part rate which ensures that the busiest node in the centralized setting is no busier than the busiest decentralized setting. By busier, we mean the degree of the node or number of clients communicating with it. We note that NTK-FL is not an upper bound in this case due to the comparison being founded on node busyness. Evaluating NTK-FL in the same setting as the table in Figure 2, NTK-FL converges to threshold accuracy in 73, 85, and 180 communication rounds for heterogeneity settings IID, $\alpha=0.5$, and $\alpha=0.1$ respectively. D-PSGD is one of the first decentralized, parallel algorithms for distributed machine learning that allows nodes to only communicate with neighbors. DFedAvg adapts FedAvg to the decentralized setting, and DFedAvgM makes the use of SGD-based momentum and extends DFedAvg. Both use multiple local epochs between communication rounds, like vanilla FedAvg. DisPFL is a personalized federated learning approach that aims to train a global model and personalize it to each client with a local mask. In order to make the comparision fair, we report the accuracy of the global model on our test set.

### B.2 Hyperparameters

We perform a hyperparameter search over each baseline and select the hyperparameters corresponding to the best test accuracy. We use the $\alpha=0.1$ Fashion-MNIST test accuracy at communication round 30 as the metric for selection. This is done because the majority of comparisons take place on Fashion-MNIST in the non-IID setting. For D-PSGD, we use a learning rate of $0.1$, and a batch size of $10$ (local epochs are defined to be one in this approach). For DFedAvg, we use a learning rate of $0.1$, a batch size of $25$, and $20$ local epochs. For DFedAvgM, we use a learning rate of $0.01$, a batch size of $50$, $20$ local epochs, and a momentum of $0.9$. For DisPFL, we use a learning rate of $0.1$, a batch size of $10$, and $10$ local epochs. Following (Dai et al., ), we use the sparsity rate of $0.5$ for DisPFL. As for the NTK-DFL, we use a learning rate of 0.01 and search over values $t\in\{100,200,\ldots,800\}$ during the weight evolution process.

### B.3 Inter-Model Variance

We describe the inter-model variance between NTK-DFL clients in Figure 7. Here, the variance is calculated as follows

$$ $V=\frac{1}{d_{2}}\sum_{j=1}^{d_{2}}\sqrt{\sum_{i=1}^{M}[\bar{{\bm{w}}}-{\bm{w} }_{i}]_{j}^{2}}$ (6) $$

where we investigate the average per-parameter variance to normalize for the scale of each parameter among clients.

### B.4 Further Experimental Results

Figure: Figure 11: Convergence of NTK-DFL across different dynamic topologies, trained on Fashion-MNIST. NTK-DFL is evaluated with (blue) a $\kappa=5$ regular graph, (yellow) an Erdos-Renyi random graph with five mean neighbors, and (green) a ring topology, where each client is connected to two neighbors. We observe that NTK-DFL demonstrates steady convergence across different topology classes.
Refer to caption: x15.png

Figure: Figure 12: NTK-DFL model accuracy as a function of neighbor count ($\kappa$), trained on Fashion-MNIST. Notably, (blue) the aggregated model accuracy across NTK-DFL clients remains consistent, even as network sparsity varies. This stability persists despite a significant decline in (yellow) mean individual client test accuracy as the number of neighbors decreases.
Refer to caption: x16.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 3
presentation: 2
questions: 1. The derivative notation should be: $\boldsymbol{J}_{i, j}^{(k)} \equiv\big[\nabla _{\bar{w}_j} \boldsymbol{f} (\mathbf{X}_i; \overline{\boldsymbol{w}} _j^{(k)}) \big]^{\top} \Rightarrow$ 

$
\boldsymbol{J}_{i, j}^{(k)} \equiv\big[\nabla _{w} \boldsymbol{f}(\mathbf{X}_i ; \overline{\boldsymbol{w}}_j^{(k)})\big]^{\top}
$

2. The expression “A global or aggregated model may take the form $\boldsymbol{w}=\frac{1}{M} \sum _{i=1}^M N_i \boldsymbol{w}_i$” should be revised to: $\boldsymbol{w}=\frac{1}{N} \sum _{i=1}^M N_i \boldsymbol{w}_i$, where $N=\sum _{i=1}^M N_i.$

3. Consider simplifying the set of clients using the notation $\mathcal{C}=${$1,..i,..M$} instead of $\mathcal{C}=${$C_1,...C_M$}. This way, subsequent neighborhoods can be expressed as $\mathcal{N}_i^{(k)}= ${$ j \mid(i, j) \in E^{(k)}$ }, making later expressions more concise.



4. The true label $Y_i$ in Figure 1 does not match the definition $\mathbf{Y}_i$ in the main text.

5. Given the definition of the aggregated model as $\boldsymbol{w}=\frac{1}{N} \sum_{i=1}^M N_i \boldsymbol{w}_i$, Eq. (1) should be modified to:
   $
   \overline{\boldsymbol{w}}_i^{(k)}=\frac{1}{N_i+\sum _{j \in \mathcal{N}_i^{(k)}}N_j}(N_i\boldsymbol{w}_i^{(k)}+\sum _{j \in \mathcal{N}_i^{(k)}} N_j\boldsymbol{w}_j^{(k)}).
   $

6. It is currently unclear how NTK-DFL differs from the previous work by Yue et al. (2022); it seems to extend their work to decentralized FL. It would be helpful for the paper to add more discussion on the unique aspects of NTK-DFL in the context of DFL.

7. NTK-DFL requires each client to share local true labels, function evaluations, Jacobians, and other information with their neighbors. This appears to contradict the fundamental privacy-preserving features of FL. More discussion on privacy protection is recommended.
soundness: 2
strengths: - The idea of using the NTK paradigm for DFL without gradient descent is both interesting and novel.
- This paper contributes a practical algorithm for DFL.
summary: This paper proposes a novel framework for decentralized federated learning (DFL) that leverages a Neural Tangent Kernel (NTK)-based update mechanism instead of typical gradient descent methods. Specifically, clients send labels and Jacobians to their neighbors, who then use tools from NTK to obtain the trained neural network instead of relying on gradient descent. The paper provides empirical results validating that this approach significantly outperforms previous baselines in highly heterogeneous settings, achieving 4.6 times fewer communication rounds.
weaknesses: - The proposed method requires clients to share their respective Jacobians, true labels, and function evaluations with their neighbors. This seems to violate the privacy-preserving feature of FL. More discussion is needed.
- The experiments were only validated on simple datasets (MNIST, Fashion-MNIST, and EMNIST); it is necessary to test on more complex datasets, such as CIFAR-100.
- The notations could be improved to make the method and algorithm clearer.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - How to prevent privacy data leakage in NTK-DFL's 3rd step? Is there any technique to enhance this? For example more discussion on differential privacy.
- What's the comparison to other methods about network overhead? Is there a comparison of total bytes transferred per round for NTK-DFL v.s. centralized approaches, or an analysis of how the overhead scales with number of clients and model size.
-  How about other network topologies and protocols in DFL?
soundness: 3
strengths: - Presentation is easy to read.
- Straightforward motivation and clear validation: NTK works well in Centralized FL; it's not used in Decentralized FL yet; a vanilla method NTK-DFL is proposed; it empirically works well.
- Experimental results about comparison and ablation seem good.
summary: This paper introduce Neural Tangent Kernal into Decentralized Federated Learning, method of which is called NTK-DFL. The corresponding framework, NTK-DFL weight evolving method, empirically works on 3 common vision datasets in federated learning literatures.
weaknesses: - Privacy concerned. In the 3rd step of proposed framework NTK-DFL, ground truth label are exchanged among clients, which would lead to more serious bad impacts. This potention privacy leakage breaks the claim made in Line 45-48. More discussion about privacy preserving should be considered.
- Additional overhead. centralized federated learning system do have a heavy bandwidth issue on the path from clients to the server. However, the additional ones in NTK-DFL exchanged on the entire network system is far larger than that of centralized federated learning system. Further analyses on network overhead should be involved and make the paper more solid.
- The content and the main claim are supported by the experiment. However, more experiments will be considered for higher rating. More detailed ablation study, analyses on overhead through the whole DFL system and NTK on other network topologies and protocols (e.g., multi-server) are recommended.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: 1.See the weakness above.
soundness: 3
strengths: 1.The experiments demonstrate that the proposed method is useful.
2.The study with NTK approach on decentralized FL is meaningful.
summary: This work proposes an approach leveraging the NTK to train client models in the decentralized setting while introducing a synergy between NTK-based evolution and model averaging. It exploits intermodel variance and improves both accuracy and convergence in heterogeneous settings. The empirical results show the proposed approach consistently achieves higher accuracy than baselines in highly heterogeneous settings and reaches target performance in fewer communication rounds with multiple datasets, network topologies, and heterogeneity settings.
weaknesses: 1. A significant concern revolves around the novelty of the proposed method. It seems that the proposed method may appear to be an extension of NTK from FL to DFL with some effective trick methods. It is essential for the authors to underscore their distinctive contributions in a more prominent manner.
2. In terms of experimental baselines, it is recommended that the authors include the most recent decentralized federated learning method (DFedSAM (Shi et al. (2023)) for a comprehensive comparison. This will enhance the paper's completeness and relevance in the context of the current state of the field.
3.The analysis in line 310-325 “Gains Due to Final Model Aggregation” is not clear enough. Where is the 10% in line 318 and 15% in line 321?

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: - How are implement (4) and (5) implemented in practice? To exactly calculate the exponential map, eigen-decomposition on $\tilde{N}_i\times\tilde{N}_i$ matrix $H$ is needed. It would be better to explicitly present the implementation in practice.

- Can authors also present the communication cost in terms of bits rather than communication rounds?

- Can the method scale to larger scale problems like TinyImagenet?
soundness: 2
strengths: The paper is motivated by an important problem in decentralized federated learning: heterogeneity. The idea to reconstruct empirical NTK from neighbors is novel and interesting. Experiments also show great potential of the NTK-DFL in tacking heterogeneous datasets.
summary: The paper proposes to leverage the Neural Tangent Kernel (NTK) to guide decentralized federated learning. More specifically, in each communication round, clients send their weights to neighbors and ask neighbors to calculate the Jacobian on neighbors' datasets. Then, clients receive the Jacobian from neighbors, and construct the empirical NTK matrix with the data from all neighbors. Next, each client updates the weights with reconstructed NTK. The proposed method NTK-DFL is shown to outperform standard DFL algorithms on heterogeneous datasets.
weaknesses: - My major concern is the communication overhead of NTK-DFL. Conventional decentralized gradient descent only communicates gradients with neighbors, so the communication cost is $O(d)$. However, NTK-DFL needs to communicate Jacobian with the cost of $O(N_i\times d)$, where $N_i$ is the number of data in one client. This approach does not seem scalable to big datasets. Authors mentioned Jacobian batching in line 216. It would be interesting to see the trade-off between communication cost and performance for different batch sizes.

- The presentation of weight evolution can be made clearer. It seems negative signs are missing in the exponents in (4).

Meta-review: The authors introduce a novel framework for decentralized federated learning (DFL) which replaces the typical gradient descent method with a Neural Tangent Kernel (NTK)-based update mechanism. Particularly, clients send labels and Jacobians to their neighbors, who will leverage NTK to obtain the trained neural net. This is an interesting contribution and the idea of using NTK-based update instead of gradient descent for DFL is interesting and novel. 

However, the reviewers raised critical concerns regarding (1) potential violation of privacy concern when clients have to share labels and Jacobians; (2) lack of experiments on more complicated datasets and comparison with important baselines.

The authors have made substantial efforts to address the concerns raised by the reviewers but have not provided any extra experiments following multiple reviewers' suggestions. Some of the suggested baseline comparisons, as pointed out by Reviewer wkMx and Ump1, are important, as acknowledged by the authors. 

The reason stated by the authors (in response to wkMx) for not providing comparison during rebuttal is that it would require intricate implementation. While I understand that the authors might not have enough time to do that, it cannot be used as an excuse because without those comparison, the empirical studies are not sufficiently conclusive. This basically means this paper needs more time and polishing to reach an acceptable state.

Thus, unfortunately, lacking strong empirical studies, I cannot recommend acceptance for this paper. The authors are encouraged to strengthen the paper's empirical studies following the suggestions from the reviewers for future submission.

Decision: Reject