FedAvP: Augment Local Data via Shared Policy in Federated Learning

W1. More experiments on FL benchmark on non-IID data

We conduct experiments on comprehensive non-IID settings suggested by the reviewer, specifically the quantity-based label skew settings described in [1]. The table below presents the results across different datasets and partitioning strategies. Here, C is the number of different labels held by each client. In extreme label skew cases, such as C = 1, where data labels are highly partitioned, our algorithm shows slightly lower performance on CIFAR-100 (C = 1). However, in all other cases, our algorithm demonstrates improved performance.

Q1.

This question is related to W1. We conducted experiments in non-IID settings as suggested by the reviewer. We kindly ask you to refer to our response to W1. Additionally, please refer to our responses to Reviewer zQTG W1 and Reviewer WeMd W2, which include comparisons with more non-IID settings.

[1] Li, Qinbin, et al. "Federated learning on non-iid data silos: An experimental study." 2022 IEEE 38th international conference on data engineering (ICDE). IEEE, 2022.

W1. Improvement to the explanation

Thank you for your valuable suggestion to improve our paper. In response, we will provide more detailed explanations to address your points, including the effect of different heterogeneity levels $\alpha$, an additional baseline, the meaning of OOD (Out-of-distribution), and the weighted average metric.

W2. Why does different alpha learn different augmentations?

The alpha controls the degree of data heterogeneity among clients. As the value of alpha decreases, the degree of heterogeneity increases. As the heterogeneity increases, the distribution of labels among the clients becomes more uneven, resulting in the model in each client trained from vastly different distributions even when the same dataset is used. This explains why the optimal augmentations learned by our algorithm vary with the changes in alpha. We will include figures of the dataset distributions in the appendix to further aid readers' understanding.

Q1. The performance of FedAvg + AutoAugment

The AutoAugment paper proposes a method where multiple child networks based on a controller RNN are trained through reinforcement learning. This approach was not developed with FL in consideration, where data privacy is essential. Therefore, applying AutoAugment directly to FL is challenging without further development. The original AutoAugment also requires significant computational resources, as the paper reports 5000 GPU hours for CIFAR-10 and 100 GPU hours for SVHN, making it impractical for FL environments regarding computation time and communication costs.

To facilitate comparison, we applied the “pre-trained” augmentation policies reported in AutoAugment to the FedAvg algorithm. We used the CIFAR-10 and SVHN AutoAugment policies from the paper and compared the results. For CIFAR-100, we used the augmentation policy trained on CIFAR-10, as done in AutoAugment. The “Default” augmentation refers to using Random Crop and Horizontal Flip. As shown in the following table, our algorithm FedAVP shows superior performance across various heterogeneity settings for CIFAR-100, CIFAR-10, and SVHN.

Q2. The meaning of OOD client

The OOD client stands for the out-of-distribution (OOD) client, which refers to the one that does not participate in training but is used to evaluate the global model’s performance during the test. Evaluating OOD clients allows us to determine how well the FL model generalizes to new clients. Even though we briefly mentioned the OOD client in Section 4.1, the actual comparison and experimental results are exhibited in Appendix A; we will clarify this. As shown in Table 5, our approach also outperforms the experiment with the OOD clients, confirming superior generalization capacity.

Q3. The weighted average of each client’s accuracy

We followed the weighted accuracy metric according to the pFL-bench paper [1]. Based on each client's data size, the weighted average metric ensures fair representation and robustness to data imbalances, providing a more accurate overall performance evaluation. It reflects the real-world impact of larger datasets and mitigates the noise from smaller datasets, leading to a more stable assessment in FL.

The label distribution of each client's test dataset is the same as the label distribution of their own training dataset. For example, in the case of the CIFAR-10 dataset, if the number of data samples per label in Client A's training dataset is [100, 30, 50, 70, 10, 80, 20, 90, 40, 60], then the number of data samples per label in Client A's test dataset will be [20, 6, 10, 14, 2, 16, 4, 18, 8, 12].

Additionally, the distribution of the number of test data samples among clients is the same as the distribution of the number of training data samples among clients. Therefore, the weighted average of the global model's accuracy on each client's test dataset by the number of data samples they have is equal to the global model's accuracy on the entire test dataset.

[1] Chen, Daoyuan, et al. "PFL-bench: A comprehensive benchmark for personalized federated learning." NeurIPS, 2022.

[2] Pascanu, Razvan, et al. "On the difficulty of training recurrent neural networks." ICML, 2013.

[3] Mikolov, Tomas, et al. "Empirical evaluation and combination of advanced language modeling techniques." Interspeech, 2011.

[4] Shu, Jun, et al. "Meta-weight-net: Learning an explicit mapping for sample weighting." NeurIPS, 2019.

[5] Zhou, Fengwei, et al. "Metaaugment: Sample-aware data augmentation policy learning." AAAI, 2021.

W1. Comparison with other classic non-IID methods, such as FedNova and Scaffold.

In Table 1 of the manuscript, we compare our model with baselines, including state-of-the-art federated learning and federated data augmentation algorithms. As the review suggested, we conducted an additional experiment with non-IID algorithms, including FedNova [1] and Scaffold [2].

"Default" refers to the default augmentation, which applies Random Crop and Horizontal Flip. Both the FedNova and Scaffold algorithms were tested with default augmentation, as well as with RandAugment and TrivialAugment. FedAvP (Fast Update) was trained using the Fast Update method described in Section 3.3 of the paper. We examined the performance under diverse heterogeneity conditions on highly non-IID environments using CIFAR100, CIFAR10, SVHN, and FEMNIST datasets. Our FedAvP method outperformed across all tested scenarios.

W2. Why the shared policy can alleviate the non-iid?

When the client's data distribution follows non-i.i.d. in federated learning, local models are trained on the different data distributions, which can lead to poor global model aggregation due to distribution shifts. However, our shared data-augmentation policy and meta-policy search strategy can generate additional data samples that help balance the data across clients, thereby reducing the disparity and scarcity in data distributions and improving model convergence and accuracy. As shown in Figure 1(a), under the 'local policy update,' our algorithm updates the augmentation policy so that the aggregated model performs well across the various local data distributions (non-iid, pre-augmentation).

W3. The detailed explanations of gradient clipping and reweighting

We sincerely thank the reviewer for the suggestion regarding the improvement of the readability of our paper. Gradient clipping [3,4] is performed using the method $g \leftarrow \frac{c}{||g||}$ when $||g|| > c$, and it is applied to the gradients of $D_{k}^{val}$ and $D_{k,n-1}^{train}$ in Eq. (6). Reweighting [5,6] is performed using the sigmoid output of the neural network $P_{\theta}$ in Eq. (3). Specifically, when the search space is a 17x17-dimensional neural network output $P_{\theta}$, each dimension $P_{\theta}(1), ..., P_{\theta}(17 \times 17)$ corresponds to a two-operations augmentation, where a higher value indicates a higher probability of selecting that augmentation. During training, when there is a batch size of $B$, $B$ samples are drawn from a multinomial distribution based on the value of $P_{\theta}$, resulting in $P_{\theta}^1,..., P_{\theta}^B$. Here, each $P_{\theta}^i$ corresponds to one of the 17x17 neural network output dimensions. For each data sample, two-operations augmentation is applied based on $P_{\theta}^i$, and the weight of the augmented data sample is determined by the corresponding $P_{\theta}^i$. We will include this explanation in the Appendix.

Q1.

Please refer to our response to W2.

Q2. The advantage of the reweighting in Eq.(3).

We applied the corresponding two-operations to each data sample in the batch using the 17x17-dimensional neural network output, where the weight of the augmented sample was determined by $P_{\theta}^i$. After training the local model, our policy loss tries to improve the performance on the validation data. In this process, two-operations that hinder the performance on the validation data are updated through backpropagation using the weight $P_{\theta}^i$, causing the weight $P_{\theta}^i$ to decrease.

Our Federated Meta-policy (FMPL) approach updates a shared data augmentation policy for each client’s unique environment to mitigate data heterogeneity. The sampling operation is non-differentiable in the typical data-augmentation approach. However, the reweighting approach in Eq.(3) enables differentiation, which is a key advantage.

[1] Wang, Jianyu, et al. “Tackling the objective inconsistency problem in heterogeneous federated optimization.” NeurIPS, 2020.

[2] Karimireddy, Sai Praneeth, et al. “Scaffold: Stochastic controlled averaging for federated learning.” ICML, 2020.

[3] Pascanu, Razvan, et al. “On the difficulty of training recurrent neural networks.” ICML, 2013.

[4] Mikolov, Tomas, et al. “Empirical evaluation and combination of advanced language modeling techniques.” Interspeech, 2011.

[5] Shu, Jun, et al. “Meta-weight-net: Learning an explicit mapping for sample weighting.” NeurIPS, 2019.

[6] Zhou, Fengwei, et al. “Metaaugment: Sample-aware data augmentation policy learning.” AAAI, 2021.

[7] Li, Qinbin, et al. “Federated learning on non-iid data silos: An experimental study.” ICDE. IEEE, 2022.

W1. The scalability of FedAvP

When using two operations, the neural network output $P_{\theta}$ in the paper has a 17x17 dimension, represented as $P_{\theta}(1), ..., P_{\theta}(17 \times 17)$. From this output, we sample $P_{\theta}^{1}, ..., P_{\theta}^{B}$ based on the batch size $B$ and apply the augmentation. Since the augmentation sampling is performed from the output $P_{\theta}$ obtained through a single forward pass, even if the search space expands to 17x17x17 for three operations, the computational load of FedAvP does not increase exponentially due to only requiring one forward pass of the neural network. To assess whether our algorithm can effectively search within this expanded search space, we conducted the following experiment. We extend the layers in FedAvP to experiment with 17x17x17 possible combinations of three operations on the SVHN dataset under different heterogeneity conditions, with $\alpha = 5.0$ and $\alpha = 0.1$. The results, recorded as test accuracy per training round, are presented below.

Table: Test accuracies with different search spaces on the SVHN dataset.

As shown in the table, despite the 17-fold increase in the search space when using three layers on the SVHN/0.1 dataset, FedAvP’s performance significantly improved compared to the two-layer setup reported in the main table of the paper. This demonstrates that FedAvP can effectively search the expanded search space. On the SVHN/5.0 dataset, the performance remained consistent with that of the two-layer setup. We observed that even with a 17-fold increase in the search space due to the use of three layers, FedAvP effectively searches for policies that either improve performance or maintain consistent results, depending on the data distribution.

W2. More experimental results of non-IID scenarios

We conducted additional experiments on CIFAR-10/0.1 and SVHN/5.0, which were not included in the previous paper version. The alpha in FEMNIST is not further explored due to the absence of an adjustable heterogeneity parameter. The results are provided in the official comment.

Q1.

This question is related to W2. We conducted additional experiments in non-IID environments. For a detailed explanation, we kindly ask you to refer to our response to W2. Additionally, please refer to our response to Reviewer 866Z W1, which addresses quantity-based label skew settings[7], where labels are highly partitioned.

Q2. The application of our methods to larger visual model like ViT

We reported the performance of a larger model, specifically the VGG11-s model with approximately 3x more parameters in Appendix A.2. To further investigate, we conduct experiments using the ViT-T model[1,2] on the CIFAR100/5.0 and CIFAR100/0.1, and the results are presented below.

In the table, "Default" refers to the default augmentation, which applies Random Crop and Horizontal Flip. Both the FedAvg and FedExp algorithms were tested with default augmentation, as well as with RandAugment and TrivialAugment. FedAvP (Fast Update) was trained using the Fast Update method described in Section 3.3 of the paper. Our algorithm demonstrates improved performance over the CNN-based models reported in the original paper. For baseline algorithms, we observe overfitting due to the increased size of model parameters.

We also compare the computation time on the client side needed to reach the target accuracy, similar to the methods reported in the paper. The results are as follows.

In the first CIFAR-100/5.0 table, we recorded the number of rounds and computation time required for the global model to reach a target accuracy of 30% when trained with ViT-T model. In the CIFAR-100/0.1 table, we recorded the rounds and computation time needed to reach a target accuracy of 25%. We compared the FedAvP (Fast Update) to the FedAvg with each of the Default Augmentation, RandAugment, and TrivialAugment applied, as explained above. Compared to the FedAvg, FedAvP (Fast Update) takes slightly longer on the CIFAR-100/5.0 dataset. However, on the CIFAR-100/0.1 dataset, FedAvP (Fast Update) actually achieves faster learning with the ViT-T model. Overall, applying our algorithm to the ViT-T model does not result in a significant increase in computation time.

[1] Dosovitskiy, Alexey, et al. "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale." ICLR, 2021.

[2] Qu, Liangqiong, et al. "Rethinking architecture design for tackling data heterogeneity in federated learning." CVPR, 2022.