Paper: FedWon: Triumphing Multi-domain Federated Learning Without Normalization (ICLR-2024)

Reviewer comment (Theory):
Insufficient Theoretical Analysis: The paper lacks a theoretical analysis of the effectiveness of the proposed Scaled Weight convolution (WSConv) layer. The authors should provide a more rigorous theoretical foundation to explain why the WSConv layer is suitable for addressing the challenges in federated multi-domain learning.

Author response:
We are also keen to provide more theoretical guarantees of our proposed FedWon. However, as the first work that considers FL without normalization, we have demonstrated multiple important properties of this direction through extensive experiments. We will provide more theoretical guarantees in the extension of the paper.

Reviewer comment (Experiments):
Limited Applicability to Convolutional Neural Network (CNN) Models: The proposed method is limited to CNN-based deep learning models, which restricts its applicability to a specific class of models, such as recurrent neural networks (RNNs), graph neural networks (GNNs) or transformer.

Author response:
Thank you for raising this interesting perspective. As the first work considering FL without normalization, we primarily focus on CNN models and discover many interesting properties of via extensive experiments, such as versatility to both cross-silo and cross-device FL and applicability to batch sizes as small as 1. We will further explore FL without normalization for the other networks mentioned by the reviewer, such as RNN, GNN, or transformers, in the extension of the paper.

Reviewer comment (Novelty):
The paper mainly discusses the challenges of normalization in federated multi-domain learning and proposes to eliminate normalization. However, other directions exist to solve the federated multi-domain problem, such as PartialFed and FMTDA, which do not follow the normalization idea. What are the advantages and characteristics of FedWon compared with these methods?

Author response:
Thank you for pointing out these two works. We have included them in the discussion in the paper. The table below compares our performance with PartialFed on the Office-Caltech-10 dataset. FedWon generally achieves better performance than PartialFed. Besides, PartialFed may be also not scalable to cross-device FL like FedBN as the clients are stateful and load the loading strategy from the previous round. In contrast, our proposed FedWon is versatile for both cross-silo FL and cross-device FL. The setting of FMTDA is different, FMTDA assumes source domain data in the server and aims to adapt to target domains in clients. Our proposed method, as well as FedBN and PartialFed, assumes there is no data in the server. Thus, it is not fair to compare the methods directly. 

| Methods | Amazon | Caltech | DSLR | Webcam | Average |
| --- | --- | --- | --- | --- | --- |
| PartialFed-Fix | 58.3 | 44.9 | 88.1 | **91.2** | 70.6 |
| PartialFed-Adaptive | 63.4 | 45.4 | 85.6 | 90.5 | 71.3 |
| FedWon | **67.0** | **50.4** | **95.3** | 90.7 | **75.6** |

Reviewer comment (Novelty):
The proposed method lacks innovation; it essentially directly applies the weight standardization and gradient clipping from the NF-Net series [1, 2] to the federated learning setting. It does not offer targeted improvements to address the unique challenges of the federated learning setting.

Author response:
We would like to clarify that the proposed method of FL without normalization is mainly motivated by the unique challenges of the multi-domain FL setting. As shown in Figure 1b, the BN statistics between clients and between a client and the server differ significantly, which hinders the performance of multi-domain FL. We are the first work that considers FL without normalization. We believe that this new perspective offers important insights for the community via extensive empirical studies and provides useful lessons for the community on the intriguing properties of FL without normalization.

Reviewer comment (Experiments):
The experiments for cross-device FL in the paper are not sufficient for the proposed method's effectiveness. The cross-device FL experiments only include a single dataset and 100 clients.

Author response:
Thank you for the comment. We would like to clarify that we actually provide experiments of cross-device FL on both Digits-Five and CIFAR-10 datasets. To further validate its effectiveness in cross-device FL, we further conduct the experiments of a total of 1000 clients, with a selection of only 0.1 clients per round, in the table below. FedWon also generally outperforms FedAvg under this setting. 

| Methods | MNIST | SVHN | USPS | SynthDigits | MNIST-M | Average |
| --- | --- | --- | --- | --- | --- | --- |
| FedAvg | 96.0 | 71.2 | 94.7 | 82.9 | **82.8** | 85.5 |
| FedWon | **96.4** | **73.6** | **95.5** | **83.7** | 81.9 | **86.2** |

Reviewer comment (Presentation):
In the upper image in Figure 1b, it seems that the green points representing the server are entirely invisible, while in the lower image, it appears that there is no information at all.

Author response:
Thank you for raising the concern. We have revised Figure 1b to show the variation of BN statistics between two clients to make it clearer. The full comparison of BN statistics between clients and between the client and the server is provided in Figure 12 in the Supplementary.

Reviewer comment (Evaluation):
The presentation of feature maps in Figure 5 doesn't seem very informative. What information can we extract from the feature maps, and how are they related to the model's performance? Additionally, it appears that there is not much difference between the feature maps of FedAvg and FedBN.

Author response:
Feature maps are commonly used to offer interpretability. Clear and discernible features normally indicate that the model is learning meaningful representations from the data. FedWon has much better feature maps focusing on the object of interest compared to the other methods. As for FedBN and FedAvg, the feature maps between clients differ significantly, indicating large discrepancies in models learned in different domains. The feature maps of FedBN are generally slightly clearer than FedAvg, especially on the Client 1 local models.

Reviewer comment (Experiments):
Why the proposed method is only applied to the convolution layer? From Eq. 1, it seems that it can be applied to any kinds of layers represented by a weight matrix $W$.

Author response:
The equation outlines a general methodology applicable to various layers represented by weight matrices. We initially focus on convolution layers as they are commonly followed by a BN layer (e.g., ResNet and model architectures in Table 5 and 6 in the Supplementary). It has the potential to extend to other layers such as fully connected layers in scenarios where it comes with BNs.

Reviewer comment (Experiments):
In Eq. 1, how is $\gamma$ chosen? In particular, how robust is the training result to the choice of $\gamma$?

Author response:
Thank you for raising this question. Eq. 1 is the standard BN formulation, so we suppose the reviewer is interested in our method's γ in Eq. 3. The value of γ is associated with the activation function used. For ReLUs, $γ = \sqrt{\frac{2}{1-(1-\pi)}}$.

Reviewer comment (Evaluation):
Since the proposed method does not depend of batch statistics, I'm curious why FedWon B=1 better than B=2 (e.g., Figure 4)?

Author response:
Thank you for this interesting question. The results shown in the manuscript use the same learning rate of 0.01 for both batch sizes B = 1 and B = 2 of FedWon. To further analyze this result, we further tune the learning rate of B = 2 to 0.02 and they achieve similar performance, as shown in the table below. We have updated the paper to reflect these results.

| B | MNIST | SVHN | USPS | SynthDigits | MNIST-M | Average |
| --- | --- | --- | --- | --- | --- | --- |
| 1 | 98.7 | 85.8 | 98.3 | 93.7 | 90.6 | 93.4 |
| 2 | 98.6 | 85.6 | 98.2 | 93.7 | 90.8 | 93.4 |

Reviewer comment (Experiments):
In figure 4, I'm curious to see how would FedAvg (B=1) perform.

Author response:
We also intend to report results for FedAvg with B = 1. However, methods such as FedAvg contains BN layers and cannot run with B = 1. This is because the statistics computed in BN (such as mean and variance) rely on batch-wise computations, and with a batch size of 1, there's no variance among samples, making normalization difficult.

Reviewer comment (Theory):
It is claimed in this paper that FedBN can't do cross-device FL. However, there is not enough evidence, as far I can tell, from the paper that supporting this argument. Can the authors elaborate more on why FedBN can't do cross-device FL?

Author response:
FedBN stores BN statistics locally in clients and reloads these BN statistics after receiving the global model aggregated in the server. It assumes that the clients are stateful -- each client may participate in each round of the computation, carrying these BN statistics from round to round. 

However, cross-device FL described in [1]  involves a vast number of clients, with only a small subset participating in FL training per round. In cross-device FL, clients are mostly stateless -- each client will likely participate only once in the whole training. 

When applying FedBN to cross-device FL, the randomly selected clients have a high possibility of not carrying the BN statistics from the previous round. Thus, FedBN is largely unsuitable for cross-device FL.

[1] Kairouz, Peter, et al. "Advances and open problems in federated learning." Foundations and Trends® in Machine Learning 14.1–2 (2021): 1-210.

Reviewer comment (Theory):
The paper brought the Scaled Weight Standardization (SWS) technique to handle the multi-domain FL problem. However, there is less analysis about the SWS’s impacts to the FL process, e.g., will it lead to a better convergence bound?

Author response:
Thank you for raising the question. We provided empirical analysis on the convergence of the FL process in Figure 4 (right), Figure 6 (right), and Figure 10 (left). As the first work that considers FL without normalization, we will provide more theoretical analysis on the convergence bound in the extension of the paper.

Reviewer comment (Presentation):
Many methods are compared in the paper, some of them have BN, some of them are suitable for cross-silo FL only, and some of them are suitable for cross-device FL, it would be clearer to have a structured summarization to help understand the scenarios where these methods are suitable for.

Author response:
Thank you for this constructive comment. Below is the table to summarize and compare these methods. We have also added this table in Table 7 in the Supplementary.

| Method | Has no BN | Multi-domain FL | Skewed Labeld Distribution | Cross-silo FL | Cross-device FL |
| --- | --- | --- | --- | --- | --- |
| FedAvg | ✖ | ✓ | ✓ | ✓ | ✓ |
| FedAvg+GN | ✓ | ◯ | ✓ | ✓ | ✓ |
| FedAVG+LN | ✓ | ◯ | ✓ | ✓ | ✓ |
| FedBN | ✖ | ✓ | ◯ | ✓ | ✖ |
| SiloBN | ✖ | ◯ | ✓ | ✓ | ✖ |
| FedWon | ✓ | ✓ | ✓ | ✓ | ✓ |

Reviewer comment (Experiments):
Only one dataset is evaluated for the skewed label distribution problem.

Author response:
Thank you for raising the concern. As the paper primarily focuses on multi-domain FL, we conduct experiments on one dataset to demonstrate the adaptability to the skewed label distribution problem. We will conduct further analysis on more datasets in the extension of the paper.

Reviewer comment (Experiments):
Sec.4 claims that the FedWon will achieve competitive results even at a batch size of 1 while [https://arxiv.org/abs/1602.05629] shows the FedAvg will degenerate to the FedSGD and be less effective. It is an interesting topic, and do you mind to report more details about the learning curves (communication round v.s accuracy) at different batch sizes or on different datasets?

Author response:
Thank you for raising this interesting point. We revised the manuscript and provided the learning curves in Figure 14 in the Supplementary. Different batch sizes tend to have a similar trend of convergence.

Reviewer comment (Experiments):
What is the impact of local epochs on the proposed method?

Author response:
We provided analysis on the impact of local epochs in Table 14 in the Supplementary. Our proposed FedWon is robust on different local epochs and consistently outperforms FedBN under different settings.