Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper proposes an algorithm for federated learning that leverages recent advances in the NTK framework. Participants share sample-wise jacobian matrices instead of model weights or gradients, based on the intuition that the NTK is able to capture useful statistical information when learning under statistical heterogeneity. The paper presents numerical results showing that the proposed approach maintains efficiency across different heterogeneity levels. Further, the authors introduce a practical implementation that reduces communication, and compare this approach with FedAvg and FedNova. 
main_review: This paper combines two popular relevant machine learning frameworks, NTK and federated learning, with the aim of addressing statistical heterogeneity in federated settings. The paper proposes a practical implementation that addresses potential communication burdens. The experimental section provides interesting experimentation showing this framework outperforms FedNova and FedAvg on three datasets. 

The paper is well organized and clear. Previous work is discussed and necessary background is introduced. The experimental section also presents interesting results. Plots and experiment details are discussed and explained. 

Weaknesses / clarifications.  
- My main criticism is that throughout the paper it is mentioned that the NTK “inherently solves the non-IID data problem”, when learning global model under statistical heterogeneity of workers. Yet, from the formulation of the NTK this is not evident and it is not clearly developed in the paper, for example by  explaining why the Jacobian enhances generalization while gradients don’t. Why is this framework more robust? I acknowledge that at the end of section 4.1 there is a paragraph on this, however it does not respond to my question as it affirms that the global kernel $H^{(k)}$ is more expressive than gradients without an explanation of why. 

- It is not very clear for me how this algorithm is different from Algorithm 1 in [HLSY21]. Further, it would also help to differentiate the paper from this previous to include a discussion on how the presented convergence results are different from  the ones in [HLSY21]. 

- Given that several $t^{k}$ are tested at each round, how does this affect the convergence result?


- Throughout the paper a”privacy” is used in different contexts. I would suggest a clarification of what privacy means in each context and what attacks are prevented by taking a specific action. For example, the random projections do not protect from an "honest but curious" server, since all workers share the same projection matrix. This step does not provide Differential privacy. Some differential privacy guarantees can be provided from the shuffling model, but again this is not very clearly explained, and what it is protecting from. 


Minor:
- I understand NTK intrinsic difficulties to try on larger models, but given that the contribution of this paper is a practical algorithm it would be interesting to see results in more realistic FL settings, perhaps a small language task. 
- Given that FL is an applied field, It would be good to have a small discussion on practical aspects of FL like what to expect under user intermittency, privacy (see above), computation and memory costs, and formal communication costs, e.g. $O(poly(M, N_m, \beta, ...))$ and how it compares with SotA algorithms. 

summary_of_the_review: This paper proposes a practical algorithm for FL based on recent advances on the NTK, and FL-NTK. It provides interesting preliminary and promising results. However, there are several aspects related to clarifying intuitions and differentiating this paper from previous work that still need to be addressed.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper presents a new federated learning (FL) framework that employs the neural tangent kernel (NTK) method to replace the widely used gradient descent algorithm for optimization. To improve communication efficiency and privacy-preserving features, data sampling and random projection techniques are used in the proposed FL-NTK framework. Experiments are conducted to demonstrate the advantages of the proposed FL-NTK in the robustness to data heterogeneity and communication efficiency, as compared to the baseline FedAvg. 
main_review: Strengths:
+ The idea of using NTK for model optimization without gradient descent in the FL paradigm is interesting and somewhat novel. 
+ Inherent issues such as communication cost and data privacy of the proposed method are properly discussed. 

Weaknesses:

A. Method
- In a way, the proposed method shifts much of the computation to the server-side (e.g., weight update and selecting the best one). The various integer update steps for the weight evolution are not clearly specified in the experiments (reproducibility). And the time cost for the grid search of $t^{(k)}$ should be discussed. 
- To address the risk of data leakage and the high communication cost of FL-NTK, random data projection, and Jacobian shuffling are used. Although it seems that these tricks will do the job, additional resources (i.e., the trusted key server and the shuffling server) are required, which may not be available for some applications. 

B. Experiment evaluation
- The network used in the experiments, i.e., a multilayer perception with 100 hidden nodes, is too simple. More standard network architectures like ResNets (e.g., ResNet-18) should be employed for evaluation and comparison. 
- The datasets used in the experiments are also small, and not diverse (basically MNIST and its variants). It would be more convincing to show results and comparison on larger datasets (more samples and classes), such as CIFAR-100 and Tiny-ImageNet.
- In the data heterogeneity experiments, a few popular baselines such as FedProx and SCAFFOLD are not compared. 
- In Fig. 5, the performance of FL-NTK does not increase as the level of data non-IIDness goes down. An accuracy drop is observed when $\alpha=0.3$. Can you provide an explanation for this? 
- For CP-NTK-FL, communication with the trusted key server and the shuffling server is also required, in addition to the communication between clients and server. Therefore, a breakdown of the comm. cost is helpful information. 
summary_of_the_review: Overall, the proposed NTK approach for FL is interesting and has novelty. However, there are a few key limitations/weaknesses: (1) FL-NTK has a higher risk of data leakage and communication cost compared to FedAvg. The proposed CP-FL-NTK may address these issues but at the cost of requiring additional resources (i.e., the trusted key server and the shuffling server), which could limit the practical usage of the method. (2) As a general approach, a more comprehensive and rigorous experimental evaluation is expected to fully demonstrate the effectiveness of the proposed method.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper proposes a federated learning (FL) paradigm empowered by the neural tangent kernel (NTK) framework. The NTK under the infinite-width regime allows us to analyze a dynamic of the corresponding neural network without a gradient descent algorithm. The authors utilize it for predicting the best parameters aggregated from multiple workers. However, this framework requires transmitting Jacobian matrices between workers and the aggregation server, which results in increasing computational overhead and exposing more private information. To this end, they adopt dimensionality reduction via randomness-sharing projection, zeroing out compression as well as shuffling. They also address that their NTK-based FL scheme has a faster convergence rate compared to that of FedAvg, for a two-layer network under specific assumptions. Finally, empirical results support that the proposed FL method performs better than other FL algorithms and achieves similar test accuracy to the ideal centralized case.
main_review: - This paper uses the NTK method as a tool for predicting the output of neural networks. As in eq (6), the dynamics of the gradient flow can be expressed by the neural tangent kernel (NTK). However, the NTK does not change over time under the **infinite-width regime** (please see [1]). The paper does not mention anything about the width limit of networks. Moreover, the neural network used in the proposed FL framework cannot have infinite width, hence, the evolution dynamics (Eq. 8, 10(b)) may be not valid to predict the network outputs.

    [1] Chizat et al., On Lazy Training in Differentiable Programming. NeurIPS. 2019


- The paradigm requires each worker to send a Jacobian tensor with respect to its assigned dataset, which increases communication overhead. Indeed, the transmitted parameters are used for updating the global parameter governed by Eq. (10b). Can the communication bottleneck improve by transmitting only essential factors for the updates? For example, instead of sending J_m^{(k)}, f^{(k)}(X_m), Y_m, one can transmit the corresponding sub-matrix of w^{(k,t)}, which leads to reduce the communication overhead. Is there any reason that workers send the Jacobian matrix?


- CP-FL-NTK variant uses dimension reduction via a random projection with a shared seed. A larger compression ratio (beta) reduces the communication overhead better but the performance also does so. How does degradation relate to the compression ratio? Is there any provable guarantee or practical guidance to choose a proper value of beta? And what is the exact random projection? Is it Gaussian random projection? or random sketching? Moreover, can the corresponding privacy guarantee be analyzed?


- In Remark 1, the authors argue that Eq. (19) implies that convergence in Eq. (17) is faster than that in Eq. (18). But, precisely, it is not true because of $1- t \cdot (\eta \lambda_0 / (2N)) \leq 1- (\eta \lambda_0 / (2N))^t$ for integers $t \geq 1$.

- Minor Issues:

  - The big-O notations usually do not include constants, hence writing $\Omega(N^2 / \lambda^2 * \log(N/d))$ in Theorem 1 seems more natural.

  - It is good to briefly describe FedAvg algorithms for better understanding.

  - A d-dimensional Dirichlet distribution requires a d-dimensional parameter vector (alpha), but it is chosen by scalar values in section 6. Does it mean that the vector with same values? Do the authors have a chance to try other parameter (alpha) settings?

  - It would be great if graph colors and markers in Figures 4 and 5 match each other.

summary_of_the_review: The paper is well-written and the methodology is very clear. However, it seems that the paper needs to address the usage of NTK more concretely, e.g., clarifying the width limit. In addition, the second algorithm with communication efficiency and privacy protection also needs more rigorous analyses. Overall, this paper is under the bar of acceptance.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper introduces new algorithms for the federated learning (FL) farmwork using neural tangent kernel (NTK) paradigm. Two algorithms are referred to as NTK-FL and CP-NTK-FL, where the latter is a variant of the former for improved communication efficiency and privacy preserving. The proposed algorithms are aimed to address statistical heterogeneity across the workers. An important point is that unlike typical training algorithms for the FL setting, here the workers upload the labels and sample Jacobian matrices (representing NTK) to the server. The server then uses tools from NTK to obtain the trained neural network (instead of using gradient descent). Theoretical and empirical validations (on MNIST, Fashion-MNIST and EMNIST), and comparisons with FedAvg are provided.  
main_review: Strengths:

- The paper provides a novel approach to FL using NTK paradaigm
- The text is well written and clear to follow

Weaknesses:

- The notations can improve to make both the algorithm and the analysis more clear. It seems that $X_m$ and $Y_m$ are used for worker m's samples, and $X^{(k)}$ and $Y^{(k)}$ are used for the sample in round $k$. In the analysis and equation (11) $X$ is used without subscript/superscript. Could you clarify if $X^{(k)}$ and $Y^{(k)}$ are concatenations of $X_m$ and $Y_m$ at round $k$? What do X and y in Theorem 1 represent?  

- I may be missing something here, but it seems that the closed form expressions for the trained network using NTK, for example, the ones given in equations (8), (10) and (11) of (Lee et al., 2019) also depend on the inverse kernel matrix. That seems to be standard in kernel methods in general. In fact, kernel methods are known to be computationally prohibitive due to the matrix inverse step. It is not clear to me why and how the closed form solution given in equation (8) does not depend on the kernel inverse. I augest to use clear references for all the results which are not proven in this paper, to make the paper self-contained. 
It seems X in equation (8) is only the sample. Has this become possible because (8) does not actually represent a closed form solution, but rather only the output values at the sample?

- While the novel framework is nicely contextualized through a literature review, the information that is uploaded by the workers in this framework violates the privacy preserving feature of FL. Noting that privacy preserving is the main motivation for the FL framework, this can be a problematic aspect of the novel framework. Although CP-NTK-FL is introduced to improve the privacy by compression and shuffling, without a clear discussion and characterization of privacy guarantee, the framework may not be effective. I think authors should include a clear discussion on this issue.  
summary_of_the_review: In summary, I appreciate the novel framework introduced in this paper for FL problem. Some of the results are however not completely clear to me. I also suggest a clear discussion on privacy issue that is central to the FL problem.
review_title: null

Meta-review: This paper proposes a novel Federated Learning (FL) framework that leverages the Neural Tangent Kernel (NTK), to replace the gradient-descent algorithm for optimization. Specifically, the workers upload the labels and the Jacobian matrices to the server, and the server uses the tools from the NTK to obtain a trained neural network. However since this could lead to increased communication cost and compromise of data privacy, the authors propose data sampling and random projection techniques to alleviate the problem. The authors provide a theoretical analysis that the proposed scheme has a faster convergence than FedAvg under specific assumptions, and experimentally validate that it significantly outperforms previous FL algorithms, achieving similar test accuracy to ideal centralized cases.

Pros
- The idea of using NTK for model optimization without gradient descent and use of it in the FL setting is both interesting and novel.
- The paper properly discusses and tackles the new challenges posed by the introduction of the new method.
- The paper is well-organized and clearly written, with sufficient discussion of related works and backgrounds. 

Cons

- The proposed method puts heavy computational burdens on the server-side.
- The method violates the privacy preserving feature of FL by its nature, and while the proposed compression shuffling alleviates the concern, more discussion is necessary.
- Missing comparison against popular baselines such as FedProx and SCAFFOLD. 
- The faster convergence of the proposed method in comparison to FedAvg depends on the learning rate and is not always true.
- There is a gap between the theory and practice, which makes the practicality of the algorithm still questionable. 

Although the reviewers found the idea as novel, the proposed techniques for alleviating communication cost and privacy concerns convincing, and considered both the theoretical analysis and experimental validation thorough, all reviewers leaned toward rejection due to critical concerns unanswered. During the discussion period, the authors alleviate many of the minor concerns from the reviewers, but there were still remaining concerns on the gap between the theory and practice on its convergence behavior, and insufficient discussion of the privacy-preserving feature of the proposed method, as well as shifting of computation burdens to the server. Thus, the reviewers reached a consensus that the paper is not yet ready for publication. 

Despite the low average score, the novelty of the idea and the quality of the paper is much higher than those of the accepted papers in my batch, and I strongly believe that this will become a high impact paper, if remaining concerns from the reviewers are properly resolved.
Decision: Reject