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

Authors: Gabriel Thompson, Kai Yue, Chau-Wai Wong, Huaiyu Dai
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: Federated Learning, Decentralized Federated Learning

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. The source code will be available as a link on the discussion forum once it is open.

Decision: reject
Reviewer scores: 5, 6, 5, 5

Reviews and discussion:
<thread>
creation date:
1730492528937

thread_type:
official_review

<thread_object>
creator: reviewer_ump1

<content>
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.

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.

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).

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730512862731

thread_type:
official_review

<thread_object>
creator: reviewer_wkmx

<content>
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.

strengths:
1.The experiments demonstrate that the proposed method is useful.
2.The study with NTK approach on decentralized FL is meaningful.

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?

questions:
1.See the weakness above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730657376079

thread_type:
official_review

<thread_object>
creator: reviewer_45hy

<content>
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.

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730733153880

thread_type:
official_review

<thread_object>
creator: reviewer_ob3c

<content>
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.

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731675507380

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
NTK-DFL Code

comment:
Hello! Here is the link for the NTK-DFL code. The README has a description of prerequisites for the code and example usage.

https://anonymous.4open.science/r/ntk-dfl-07DF/
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.