Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: This paper focuses on difficulty introduced by non-convexity and data heterogeneity in federated learning.
Authors first show that, with data heterogeneity, linear models and convex optimization problem can be trained efficiently with gradient correction techniques such as SCAFFOLD, while non convex problem cannot. 
Then, in order to sidestep the non-convexity for neural network, authors substitute original model with a linear approximation and original loss with a quadratic loss, such that non-convex optimization problem turns into a convex regression problem based on NTK. 
Since feature learning is necessary for NTK align with data to provide good result, authors introduce two-stage training, where feature learning only happens in first stage, and convex optimization with gradient correction happens in second stage. 
Authors conduct experiments to show that such two-stage optimization can fit the data faster and produce better test accuracy.
strengths_and_weaknesses: Strength:

This paper introduced a novel two-stage scheme to combine the feature learning capacity of neural network and for efficient optimization linear model.

Weakness:

The convexified problem is introduced mainly due to an optimization consideration: SCAFFOLD perform well on convex method. However, the convex formulation clearly sacrifice the model capacity. The feature learning can only happen in first stage (or non-convex part), thus BookNTK learns less feature than centralized training, which doesn't involve two-stage training.


questions: If I understand correctly, the final training result for BookNTK is a neural network plus a linear model. Any inference requires a backward propagation to derive the feature for empirical NTK to apply a linear model. This is different from common workflow. Is there a way to compile the final network and the linear model into a single model that support inference with a single forward pass?

Section 3 mentioned "the train and test accuracies in Figure 1(b) match quite closely, suggesting that the failure lies in optimization". However Figure 2 shows a different picture, where train acc and test acc have a large gap. What's the reason behind this discrepancy?

Why quadratic loss instead of cross-entropy is used in (1)?

How much worse is BookNTK than centralized training? Authors only mentioned one centralized baseline in section 3, where they show that "21% (out of the 35% gap in accuracy) may be attributed to a failure to optimize the linear output layer". However, for BookNTK, none of results in section 5 and appendix contains centralized training. I think this question is as important as "How much better is BookNTK than FedAvg". I strongly suggest authors to add centralized training result for a comparison.
limitations: Yes.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: This work proposes BooNTK; a two stage approach for FL that allows for better performance in cross-silo settings. The main idea is to 1) perform standard FedAvg training on a non-convex model, such as a neural network, then 2) approximate it with a first order Taylor approximation around the parameters found and finally 3) optimize the parameters of this linear model with federated training using optimisers with gradient correction, such as SCAFFOLD. The authors motivate this through the optimization difficulties of non-convex models in the non-iid setting; empirically, the loss of performance when moving to non-idd data is larger for non-convex model compared to convex.

The authors describe several further approximations that they do to make the method more efficient (i.e., removing the bias term of the Taylor approximation and only considering a subset of the coordinates of the gradient vector) and then demonstrate BooNTK’s performance on three label skew non-iid settings on FMNIST, CIFAR10 and CIFAR100.

strengths_and_weaknesses: Strengths
- Good results on the cross-silo setting
    - There is significant improvement by the BooNTK pipeline on all tasks involving a varying amount of label skew.
- Intuitive ablations and method exposition
    - The method is explained / motivated well and I liked the layer importance investigation. The ablation studies are also useful and highlight the sensitivity of the method on the choice of (some) hyper parameters.
- Simple method
    - The method is quite simple and straightforward. As a bonus, the second step is also computationally more efficient than training the original neural network.

Weaknesses:
- No discussion about how different non-iid ness settings affect the method
    - Given the claims of the work about the negative effects of data heterogeneity in the non-convex setting, I would have expected that the authors experimented with more diverse non-iid settings (i.e., not just label skew). As the current label-skew experiments concern mostly different marginals over the labels, $p(y)$, at each client, a better adjusted classification layer is important. I believe this is one of the reasons that BooNTK is effective in these scenarios. However, not all non-iidness is just label skew; for example, you could consider an image recognition system running in different mobile phones. As each phone comes with its own camera sensor, covariate shift (i.e., different $p(x)$ across clients) can also manifest as a source of non-iidness. In this case, I would intuitively expect that the distribution of the features also differs among clients, therefore, adjusting just the classifier might not be enough.
- No discussion about how the linear model is less flexible 
    - The linear model is less flexible than the non-convex neural network. As a result, its improved performance (on a given feature set) could be just because it has less degrees of freedom to adapt to the non-i.i.d. peculiarities. I believe a discussion around this is missing.
- There are several approximations involved in the computation of the eNTK features, the impact of which is unclear
    - In section 4.1 the authors describe a series of approximations to eNTK in order to reduce the computational burden, however there is no (empirical) evaluation on how each one affects the final performance. 

questions: Most of my questions revolve around the weaknesses described above.

1. The authors do not take into account that the linear model is less flexible by design, therefore it is harder to fit the non-iid peculiarities (and could explain its improved performance). I believe a control experiment would highlight whether this can be an issue in practice. For example, one could do a heavily regularised non-convex model to see if the gap is shrinking in Figure 1. How would, e.g., the performance of a pipeline similar to BooNTK fare if instead of the linear model fitting step (i.e., stage 2), one just switched from FedAvg to FedProx with a strong regularisation strength for stage 2?

2. I believe that the effect of various sources of non-iidness on BooNTK should be investigated, as the experiments consider only label skew (where generally adjusting the classifier only is sufficient). What happens when there is, e.g., covariate shift? Intuitively, this should affect the earlier layers more.

3. I believe some more investigation on the eNTK part of BooNTK is required. For example, how does finetuning just the classification layer (while keeping the rest of the network frozen) with SCAFFOLD work, relative to the eNTK approach (which requires further approximations)? Furthermore, how does BooNTK perform without the approximations to the eNTK (i.e., do the approximations have beneficial regularisation effects or are they detrimental for performance)?

4. The authors consider classification tasks, however for the second stage of BooNTK, they consider an MSE loss on an one-hot representation of the targets. What is the motivation of the MSE loss? Intuitively, you should be able to apply a softmax on the linear model of the first eq. at sec. 4.1 to get a logistic regression model, which is more appropriate for classification. 

limitations: The authors could have spent a bit more time discussing potentially negative aspects of their method (e.g., the linearity of the model in the second stage).
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: This paper makes two main contributions: 1) empirically observing that the early layers of neural networks trained by FedAvg learn useful features even in heterogeneous data settings, while the performance degradation due to data heterogeneity is due to ineffective later layers, and 2) based on this observation, proposing a novel two-stage cross-silo federated learning method, BooNTK, that first uses FedAvg to learn early-layer features, applies a particular transformation based on these features to all data points, then learns a linear classifier on top of the transformed data using SCAFFOLD. 

To verify that FedAvg learns useful early-layer features, FedAvg is run on a heterogeneous image dataset, then the last $\ell$-many layers are retrained in a centralized manner while the earlier layers are held fixed. It is observed that using the FedAvg-pretrained early layers leads to huge performance improvement over using random weights for the early layers when the last layer is retrained centrally. Also, there is some improvement due to using the FedAvg-pretrained vs random weights when the last $\ell$-many layers are retrained centrally for all $\ell$, suggesting that the early layers learned by FedAvg are extracting useful information. This observation inspires the first stage of the proposed method: FedAvg pretraining to learn the early layer weights. The paper also observes that SCAFFOLD is robust to data heterogeneity when the client losses are convex (but is not robust when the losses are nonconvex). This inspires the second stage of the proposed method: applying SCAFFOLD to convex losses to learn the last linear layer of the network. In particular, the second stage approach consists of computing the neural tangent kernel (NTK) representation of each data point using the pretrained weights in the NTK computation. Then, SCAFFOLD is employed to solve a multi-output linear regression with the transformed data points as inputs and the one-hot encoded labels as the target vectors. Empirical results are provided showing large improvement in training and testing accuracy of the proposed method over FedAvg, FedProx and SCAFFOLD on CIFAR100, FEMNIST, and MNIST with very heterogeneous partitions.

strengths_and_weaknesses: Strengths

- The paper makes interesting empirical observations that are relevant to NeurIPS and may inspire future work. In particular the observation that FedAvg learns useful features even in data heterogeneous settings is novel and helps to explain the empirical success of FedAvg plus fine-tuning.

- These observations motivate a new federated learning method with promising empirical performance.

- The experimental evaluation is mostly thorough, with multiple datasets tested and helpful ablations.

- The writing is clear.

Weaknesses

1. Not enough intuition or empirical evidence is provided for why the second stage of the algorithm should consist of linear regression on the NTK-transformed data points rather than simply fixing the first L-1 layers and running lear regression to learn the paramters of the last layer with MSE loss, or learn it with multi-class logistic regression and cross entropy loss as is conventional. Both of the latter approaches maintain convexity of the loss functions, are much simpler to implement, and involve an optimization over far fewer parameters (presumably the output of the last layer has dimension far less than p=100,000 as in the NTK approach). These approaches should be compared against as baselines and intuition should be provided as to why they are not used. 

2. On a similar note, results on the computational cost of the proposed method should be included. It seems to be much larger than the baselines due to mapping every data point to a high dimension via parameter derivative computation.

3. Personalized FL has been shown to be an effective alternative to learning a single global model in data heterogeneous settings. As such, some personalized FL methods should also be compared against. Also, the experimental results would be strengthened by comparison against more recent FL methods that learn a single global model, e.g. [23,63]. Comparing with only 3 baselines is low for an empirical paper.

Minor notes
- Missing related works: Huang et al., 2021 and Yue  et al., 2021 employ the NTK for FL.

- Intuition on why it makes sense to run linear regression for classification problems rather than logistic regression would be helpful.

- Since BootNTK gets to pretrain using FedAvg for T_1 rounds, fair comparison with the other methods should allow them to train for an extra T_1 rounds.

-Cross-device FL may still have data heterogeneity

- SCAFFOLD has higher communication and computational cost than fedavg by constant factor due to computing and communicating the gradient correction terms (footnote b).

- “By taking an eNTK approximation, BooNTK optimizes a convex approximation while using information from all layers of the model.” - besides the last layer

Huang et al., FL-NTK: A Neural Tangent Kernel-based Framework for Federated Learning Convergence Analysis, https://arxiv.org/pdf/2105.05001.pdf, 2021.
Yue et al., Neural Tangent Kernel Empowered Federated Learning, https://arxiv.org/pdf/2110.03681.pdf, 2021.


questions: Why is the proposed approach limited to the cross-silo setting, with a small number of clients? It seems to me that it should also work for settings with many clients. Or, if there are many clients and fewer samples per client, does locally optimizing the high-dimensional linear regression diverge?

limitations: The authors should discuss the computational complexity of the proposed method.
review_title: null

Meta-review: metareview: This paper introduced a novel two-stage scheme to combine the feature learning capacity of neural network and for efficient optimization linear model. It makes interesting empirical observations that are relevant to NeurIPS and may inspire future work. In particular, the observation that FedAvg learns useful features even in data heterogeneous settings is novel and helps to explain the empirical success of FedAvg plus fine-tuning. The experimental evaluation is mostly thorough, with multiple datasets tested and helpful ablations.

confidence: Certain
Decision: Accept