FedAvP: Augment Local Data via Shared Policy in Federated Learning

Authors: minui  hong
junhyeog  yun
insu  jeon
gunhee  kim
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 07:24:08

Abstract: Federated Learning (FL) allows multiple clients to collaboratively train models without directly sharing their private data. While various data augmentation techniques have been actively studied in the FL environment, most of these methods share input-level or feature-level data information over communication, posing potential privacy leakage. In response to this challenge, we introduce a federated data augmentation algorithm named FedAvP that shares only the augmentation policies, not the data-related information. 
For data security and efficient policy search, we interpret the policy loss as a meta update loss in standard FL algorithms and utilize the first-order gradient information to further enhance privacy and reduce communication costs. Moreover, we propose a meta-learning method to search for adaptive personalized policies tailored to heterogeneous clients. Our approach outperforms existing best performing augmentation policy search methods and federated data augmentation methods, in the benchmarks for heterogeneous FL.

Main paper: 
 introduction: 
 Federated Learning (FL) is a collaborative learning approach that allows multiple clients to learn without sharing their private information [1] [2] [3] [4] [5] [6] . A central server coordinates the training process across multiple devices and aggregates the locally trained models into a global one; thus reducing the communication cost of exchanging raw data and mitigating the risk of privacy leakage associated with data sharing [7] .However, the limited accessibility of data in FL still poses many challenges, such as insufficient training data and local data bias. To address these challenges, there has been a growing interest in federated data augmentation techniques [8] [9] [10] [11] [12] . They aim to increase the diversity and volume of data available at each client, thereby improving the overall robustness and performance of the federated models. For example, FedMix [8] improves performance and privacy by averaging multiple images to facilitate data mixup among clients. Similarly, FedFA [11] utilizes feature statistics to mitigate local data biases, to improve model generalization. Despite their benefits, these methods often apply the sharing of input-level [8, 10, 12] or feature-level [13, 11] information. Such information sharing poses additional privacy concerns since malicious attackers could potentially reconstruct original data by applying gradient matching loss [14] on the additional input and feature information.We propose a novel federated data augmentation algorithm named FedAvP (Augmet Local Data via Shared Policy in Federated Learning), which shares only the augmentation policies during training. Thus, each client does not need to share its own data or data-related information directly but obtains collective knowledge on how to augment the dataset at local learning. Previously, AutoAugment [15] utilizes reinforcement learning (RL) to automatically find the optimal data augmentation policy for a target dataset. While AutoAugment requires extensive GPU resources, more efficient and faster policy search has been studied [16] [17] [18] [19] [20] . However, these methods are not designed for FL environments but they perform policy search for public batch datasets. The data scarcity and heterogeneity in FL could fail these standard policy search frameworks.To improve the robustness and generalization of model training in the heterogeneous FL setting, we first introduce a Federated Meta-Policy Loss (FMPL) specifically designed to compute a gradient of augmentation policy that updates a shared data augmentation policy for each client's unique environment. Our approach guides the policy gradient to account for the effects of data augmentation on the unseen local data. The policy gradient utilizes higher-order information; the impact of the data augmentation is estimated by its effect on the validation loss observed after a few gradient descent steps with the initial augmented data. However, computing the direct meta-policy gradient in FL requires an additional communication step between the server and clients. To bypass this, we also develop an alternative meta-policy search method that utilizes a first-order approximation. We further demonstrate that the adaptive policy search technique can adapt to heterogeneous data distributions among clients in the FL environment.In the experiments, our FedAvP demonstrates superior performance on CIFAR-10/100 [21] , SVHN [22] , and FEMNIST [23] datasets within an FL context, compared to existing federated learning algorithms, including FedAvg [2] , FedProx [4] , FedDyn [5] , FedExP [6] , and federated data augmentation algorithms, including FedGen [9] , FedMix [8] , and FedFA [11] . Moreover, to further leverage the potential performance of these algorithms, we also conducted experiments applying data augmentation techniques such as RandAugment [17] and TrivialAugment [24] to these algorithms for comparison. We also evaluate our algorithm in environments where data is non-i.i.d. with heterogeneous clients [25] . We further compare the effectiveness of the utility of sharing this policy across clients for searching, in contrast to conducting a local policy search.Our primary contributions are as follows.1. We propose FedAvP(Augment Local Data via Shared Policy in Federated Learning) as the first algorithm in federated learning that facilitates shared augmentation policies among clients for federated policy search, to the best of our knowledge.2. We introduce the federated meta-policy loss for effective policy search, and further propose a first-order approximation to this loss to enhance privacy and reduce communication costs.3. Enabling meta-learning, our algorithm allows for rapid adaptation of a personalized policy by each client, addressing the challenge of highly heterogeneous data distributions among clients in the FL environment. 
related work: 
 Automated Data Augmentation. AutoAugment [15] utilizes RL to find an optimal data augmentation policy for a target dataset automatically. FastAA [16] proposes a more efficient search strategy by training small NNs in parallel without iterative training, using the density matching method. RandAugment [17] suggests a simplified search method composed of two hyper-parameters, which find the augmentation policy without a separate search process. TrivialAugment [24] further simplifies the algorithm and applies a single augmentation to each image as a parameter-free method. MetaAugment [19] proposes a sample-aware augmentation policy network to capture the variability of training samples more accurately than previous dataset-based search methods. Deep AutoAugment [18] proposes a fully automated search method that builds a multi-layer data augmentation pipeline from scratch by stacking augmentation layers. All of these recent data augmentation methods were developed under the assumption that all training data is accessible on a server. This assumption is invalid in FL since data privacy is a significant concern and it is challenging to tune these algorithms for each of the numerous heterogeneous local datasets. Therefore, our study aims to develop a new data augmentation policy search algorithm that takes into account the distributed FL process while preserving data security. RandAugment [17] and TrivialAugment [24] can be applied simultaneously to many clients in a federated learning environment, so we compared these methods in our experiments.Federated Data Augmentation. The standard Federated Learning (FL) framework, such as Fe-dAvg [2] , typically performs iterative local model updates at each client and a global update at a server. Since the clients and server only communicate through the model parameters instead of raw data, it enables secured and decentralized learning [1, 3] . Despite its benefits, FL still has challenges such as convergence degradation and model overfitting due to heterogeneity and sparsity among the data caused by the differences in client's actions and preferences [25] . To address these challenges, federated data augmentation (FDA) employs a data augmentation approach instead of a model-centric approach. FedMix [8] applies Mixup [26] to FL, augmenting data by linear interpolation of two random training examples and their labels. For privacy reasons, FedMix transfers mixup data to the server by averaging multiple images from the local device. In FedGen [9] , the server learns a lightweight generator to ensemble user information, which is then broadcasted to users to regulate local training using the learned knowledge. FedFA [11] assumes that the data distribution of the clients can be summarized by the statistics of the latent features (i.e., mean and standard deviation). This allows learning local models by regularizing the gradients of the latent representations, weighted by the variances of the feature statistics estimated from the entire client federation. StatMix [13] sends the mean and standard deviation information of local client images to the server and makes it available for learning for each client. These methods use input-level data (image) averaging or feature-level statistics to prevent direct data transfer. Unlike previous methods, our methodology focuses on transmitting only the policy information optimized for local datasets from each client.ATSPriavacy [27] has demonstrated that searching for transformation policies can also protect against reconstruction attacks in Federated Learning (FL), while preserving performance. We compare our algorithm with this approach in terms of both vulnerability to reconstruction attacks and performance in the experimental section 3 Approach: FedAvPWe introduce FedAvP (Augment Local Data via Shared Policy in Federated Learning), which performs data augmentation search by sharing policies among clients in a federated learning (FL) environment.Starting from the problem formulation ( §3.1), we address the challenges of heterogeneous clients with a proposal for adaptive policy search ( §3.2). Finally, we extend its applicability through integration with the FedAvg algorithm [1] and joint learning ( §3.3). 
problem formulation: 
 Our approach is based on the idea of centralized FL [7] ; it is not possible to share personal data neither between the server and clients nor among clients. This is different from traditional automated policy search algorithms, which find the optimal policy in a single batch dataset [15, 16, 20, 19, 18] .Objective. We begin with the standard FL algorithm, FedAvg [1] , aiming to find augmentation policies that minimize a given objective function as follows:Model:EQUATIONK is the number of clients used to train the global model w, and D k is the local data of client k. The transformation policy p θ is used to augment the data denoted by t p θ (D train k ) with coefficients α k satisfying α k = 1 and α k ≥ 0. Assuming client k has n k data samples, α k is defined asα k = n k n where n = k n k .Under a global policy assumption, all clients share the same p θ . Alternatively, if we assume that each client has its own transformation policy p local θ k , we representp local θ = {p local θ 1 , p local θ 2 , . . . , p local θ K }. Search Space.We use the augmentation space having a sequence of two operations, following the search space from previous studies on automated policy search [17, 19] . We examine a set of 17 operations in total, including {Identity, ShearX/Y, TranslateX/Y, Rotate, AutoContrast, Equalize, Solarize, Posterize, Contrast, Color, Brightness, Sharpness, RandFlip, RandCutout, RandCrop}. Each operation includes a random magnitude, rescaled and uniformly sampled from the normalized interval [0, 1]. Assuming each of the K clients has individual policy, the entire search space consists of K joint distributions, each having a size of 17 × 17. For a global policy, we learn a single 17 × 17 dimensional joint distribution corresponding to the two operations. Specifically, p θ has a value in [0, 1] using the sigmoid function on θ. For augmentation sampling, we normalize the policy parameter vector p θ whose sum to be one, and sample from a joint distribution with a regularization term ϵ. Operation pairs (op1, op2) are drawn from a mixed distribution: The server sends global model parameters w gr and policy parameters θ gr to clients. Clients train local models with augmented data, and the server aggregates them to compute w gr+1 . Clients update policies on w gr+1 using validation data, and the server aggregates these policies. (b) Clients update the model and policy parameters via first-order approximation. The server aggregates client updates to form the updated global model w gr+1 and policy parameters θ gr+1 . This sampling strategy utilizes a uniform probability across all operation pairs for a balanced exploration and exploitation, as adopted in previous work [19] .(op1, op2) ∼ (1 -ϵ) • p θ (p θ ) + ϵ • 1 17 2 (2)To search the policy in a differentiable manner, we adopt the concept from meta-learning [28] , where a model is trained on training data n times by SGD algorithm (n inner steps), followed by validation of one outer step. In our context, training is performed on augmented data, followed by evaluation on validation data [29, 19] . However, directly applying this concept of inner and outer steps to FL is quite challenging due to the added complexity of FL, which involves training local models and aggregating them to update the global model.Since our goal is to optimize the global augmentation policy, we redefine the inner and outer steps from the FL perspective. In a standard FL setting, the server sends a global model to the clients for each round, which is trained with their local training data. Afterward, the aggregation process occurs on the server to update the global model. We can regard one round of local training and aggregation as one inner step. After r rounds, where the global model is updated r times, validating the final updated model on each client can be considered one outer step. We set r = 1, meaning validation occurs after each round.Federated Meta-Policy Loss. In a single round of client updates, we first perform local training on each client and aggregate the local models to update the global weights, which are then redistributed to the clients for computing the validation loss. Figure 1 EQUATIONwhere P i θ k n is the unnormalized probability p θ k n (op1, op2) for the i-th transformation, when the sampled transformation t i p θ k n is the operations (op1, op2). This reweighting strategy is inspired by recent sample reweighting [29, 19] . Following the local updates with the augmentation policy p θ k , we aggregate the results to obtain the new global weights w gr+1 , which are redistributed to the clients for validation loss assessment. At the start of round r, after distributing the global weights w gr to each participating client, the procedure for the federated meta-policy loss can be summarized as follows:1. Each client k performs local updates by optimizing their local weights w k n using the augmented data with the local loss in Eq. (3). The updated local weights are then aggregated according to w gr+1 = k α k w k N . 2. The aggregated global weights w gr+1 are then sent back to the same clients.3. The Federated Meta-Policy Loss (FMPL) is computed on each client's validation set.EQUATIONFor this loss, the gradient with respect to the augmentation policy p θ is computed using backpropagation. Nonetheless, this approach presents two significant challenges. Firstly, it necessitates access to validation gradient information of other clients, which poses privacy concerns. Secondly, it requires revisiting the same clients for additional updates, thereby doubling the communication overhead.First-order Approximation. To ensure security by preventing access to other clients' gradients and to reduce communication costs, we derive an approximation for the policy gradient with respect to θ k n-1 at the local step n of the client k by a Taylor expansion, as in the following proposition. See Appendix B for more details. Proposition 1. Consider the federated meta-policy loss derived from the updated weight w k n for client k at step n using a first-order Taylor expansion:EQUATIONWhen computing the policy gradient of the loss with respect to θ k n-1 , the first-order gradient approximation isEQUATIONwherew k n = w gr -lr • g aug w k 0 -. . . -lr • g aug w k n-1and α k is a coefficient proportional to the client's data size.In Proposition 1, g augw k n = ∇ℓ tp θ k n (D train k,n ) (w k n )is the gradient obtained from the local loss in Eq. (3) at step n for client k. We calculate the gradient within the same client to prevent gradient leakage across clients. We optimize a policy that maximizes the inner product of the gradient obtained from sampling the validation data and the gradient obtained through augmentation using p θ k n-1 . The validation data are not used separately; instead, they are replaced by sampling the next batch, inspired by Reptile [30] . Proposition 1 also indicates that policy gradients can be computed concurrently with local updates. We utilize this approach to joint training ( §3.3), which enables simultaneous model and policy training.Gradient Clipping. From the policy gradient of Eq. ( 6), it becomes evident that policy search aims to maximize the dot product of the gradients derived from both augmentation and validation data. This approach, however, can introduce a bias toward augmentations with larger gradient magnitudes due to the nature of the dot product. Inspired by [31, 18] , we mitigate the influence of gradient magnitude by gradient clipping [32, 33] . We apply gradient clipping to the gradients from both validation and augmentation data in Eq. ( 6) using a regularizer hyperparameter c.Adaptive Policy Search. Our first-order approximation computes policy gradients at each local update step n without aggregation. These gradients are then averaged across all steps for each client to update the policy after a single communication round. However, this method can slow the policy search, since it only permits one gradient descent update per communication round. Drawing inspiration from meta-learning [28, 30, 34] , which quickly adapts to various tasks using neural network training, we propose an adaptive policy search. The augmentation policy is represented as a vector parameter p θ = sigmoid([θ 1 , θ 2 , . . . , θ 289 ]), where [θ 1 , θ 2 , . . . , θ 289 ] denotes an 17 × 17 joint distribution. We use a neural network comprising the following dense layers: one dummy embedding layer, two 100-dim hidden layers, and an output layer shaped to the 17 × 17 distribution = θg r w k * , θ k * ← LOCALUPDATE(w k 0 , θ k 0 ) wg r+1 ← Σ k α k w k * ∆θ ← Σ k α k (θ k * -θg r ) θg r+1 ← θg r + η∆θ procedure LOCALUPDATE(w, θ) Initialize w k 0 = w, θ k 0 = θ for Each local step n from 0 to E do Sample transformations tp θn from p θn Sample batch data tp θn (D train k,n ) from D k and tp θn Compute ∇ℓ tp θn (D train k,n ) (w k n ) using Eq.(3) w k n+1 ← w k n -γ∇ℓ tp θn (D train k,n ) (w k n ) Sample batch data D val k,n from D k Compute ∇ℓ D val k,n (w k n+1 ) at w k n+1 θ k n+1 ← θ k n -λ∇ℓ FMPL D val k,n (θ k n ) using Proposition 1 Set w k * = w k n+1 and θ k * = θ k n+1Send w k * and θ k * to the server size. We update our policy as done in Reptile [30, 35] . That is, the local policy updates on each client correspond to the inner steps in the Reptile, while the global policy updates on the server are analogous to the outer steps, as detailed in Algorithm 1. We train the policy neural network by increasing the dot-product between policy gradients on each client as follows:EQUATIONwhereL k,j = ℓ FMPL D val k,j (θ k 0 )is the federated meta-policy loss in Eq. ( 4) computed on the client k's j-th validation data batch using the global policy parameters θ k 0 . ⟨∇L k,j • ∇L k,s ⟩ is the dot-product between policy gradients on the client k. See Appendix C for more details. This process enables the policy neural network to learn in a direction that enhances the dot-product between the policy gradients of each client, thereby facilitating efficient policy search and enabling personalized policy search. We incorporate this strategy in the joint training in §3.3. An experiment result regarding this will be conducted in §4.As done in AutoAugment, one could employ a pretrained network to perform policy search [16, 18] . However, such methods require a separate training phase and present complications in FL environments, since training a pretrained network beforehand is cumbersome, and the separate phase is disadvantageous for parallel training. Note that we previously adopt an adaptive policy search that can simultaneously train the model and policy search. At each local update step n, the weights w k n+1 and gradients concurrently update the model and policy by comparing the gradient on validation data at w k n+1 . Algorithm 1 illustrates this joint training of our FedAvP. It begins by initializing the global model parameter w gr and the policy parameter θ gr , which are then sent to the clients from the server. Each client generates augmented data t p θn (D train k,n ) using the policy parameter p θn at every step n and updates the model accordingly. Subsequently, the gradient ∇ℓ D val k,n (w k n+1 ) is computed using the newly updated w k n+1 . Following Proposition 1, the policy parameter θ k n is updated to maximize the gradient on both validation and augmentation data ∇ℓ tp θn (D train k,n ) (w k n ). This process of model and policy updates is repeated in every local update. Afterwards, the model and policy parameters from each client are aggregated at the server using α k .Fast Update One limitation of joint learning is that it requires backpropagation for the policy gradient at every step. To reduce the computation load on local clients, the policy can be updated periodically instead of at every local step n, specifically when n mod τ == 0. In all our experiments, we set τ = 5. See the variant of the algorithm in Appendix A.3. We also reduced the hidden size of the policy neural network to two 25-dimensional hidden layers. In our experiments in §4, we will compare the performance of the Fast Update. Previous FL studies have demonstrated that the standard algorithms are effective and converged when the data is i.i.d. [36] [37] [38] . To evaluate robust performance in non-i.i.d. data, we set up our experimental environment by distributing the CIFAR-10/100 [21] and SVHN [22] datasets with different levels of data heterogeneity among clients. We assign the data to 130 clients based on a Dirichlet distribution with different hyperparameters of α = [5.0, 0.1], as done in pFL-Bench [25] . The smaller α is, the higher the degree of heterogeneity is. Among these clients, only 100 randomly selected clients participate in the training, while the remaining 30 are nominated as out-of-distribution (OOD) clients. In each communication round, only 10 clients are sampled. We employed a standard CNN model, consistent with those in previous studies [39] [40] [41] , for the global model. Experiments involving the larger model and OOD clients are provided in Appendix A.2. For the FEMNIST dataset [23] , we introduced variability in data size by distributing data based on the writers [8, 11] . Further environments details are provided in Appendix A. Our code is available at https://github.com/alsdml/FedAvP.Baselines. For a comprehensive evaluation, we compared our method with state-of-the-art federated learning algorithms such as FedAvg [2] , FedProx [4] , FedDyn [5] , FedExP [6] , and federated data augmentation algorithms including FedGen [9] , FedMix [8] , and FedFA [11] . To further explore the potential performance of these algorithms, we conducted experiments applying data augmentation techniques such as RandAugment [17] and TrivialAugment [24] . We also compared the results with those using default augmentations (random crops and horizontal flipping). Additionally, for comparison with our proposed model, we included FedAvP (W/ Local Policy), which trains each Settings. Table 1 reports the test accuracy, which measures the accuracy on the test dataset of the 100 participating clients. The test accuracy is calculated as the weighted average of each client's accuracy by the number of data points they have. We compared each baseline with three data augmentations: +Default (random crops and horizontal flipping), +RandAugment, and +TrivialAugment. For FedAvP (W/ Local Policy), each client has its own transformation policy and policy aggregation is removed from our algorithm (Algorithm 1). For FedAvP (Fast Update), we used periodic local updates (Algorithm 2). For FedAvP, we used full local updates (Algorithm 1).Results. The application of automated data augmentation algorithms such as RandAugment and TrivialAugment within a federated learning framework does not consistently enhance performance across all cases. In contrast, our algorithm learned distinct augmentations for each dataset, as depicted in Figure 2 . When compared to FedAvP (W/ Local Policy) and FedAvP, notable performance improvements were observed in highly non-i.i.d. scenarios such as α = 0.1 in CIFAR-100 and SVHN. For instance, the standard deviation of local dataset sizes in CIFAR-100 with α = 5.0 was relatively small at 19.43, while it was significantly higher at 118.00 for CIFAR-100 and 533.60 for SVHN with α = 0.1, indicating that imbalanced datasets in non-i.i.d. settings pose challenges for training local policies with smaller local datasets. Although FedAvP (Fast Update) experienced some performance declines compared to FedAvP, it generally achieved higher performance than baseline methods across most datasets. In our FedAvP, at the beginning of each round, the participating clients receive a global policy from the server, which is then optimized into a local policy using each client's local data. That is, this optimization adapts the global policy into a personalized policy on the local data of each client. The clients use the personalized policy to train their local model to achieve high performance and to aggregate well with other local models in the server. Figure 3 shows the statistics on the Euclidean distances between the personalized policies of clients participating in each round and the global policy for that round. With α = 5.0 in CIFAR-100, the data is sufficiently i.i.d., and thus the personalized policies of clients tend not to deviate from the global policy. On the other hand, with α = 0.1, the deviation from the global policy is initially high but decreases as training progresses, particularly after about 100 rounds. The variance of the Euclidean distances also follows this pattern. 
reconstruction attack: 
 Settings. In collaborative learning systems, it has been reported that gradient leakage attacks can occur [14, 42, 43] , leveraging gradients to reconstruct the original training data. We conducted these reconstruction attack [14] experiments to evaluate whether our algorithm, which shares policies in a federated learning setting, provides enhanced privacy compared to FedGen, FedMix, and FedFA. Specifically, FedGen shares information at the generator and label distribution levels, FedMix shares at the input level, and FedFA shares at the feature level. For computational comparison, we measured the time taken to reach a target accuracy of 35% on the CIFAR-100 dataset with α = 0.1. Regarding communication comparison in FedMix, the method involves sending an average image to the server prior to training, which is detailed in the "Before" section of the table. For FedAvP (Fast Update) using small neural networks, although it is higher than these baselines, the increase of 0.38MB from the cost of FedAvg is about 2.48% compared to the gradient transmission cost of the model. See Appendix A.4 for additional results. 
conclusion: 
 We proposed a novel federated data augmentation method named FedAvP (Augment Local Data via Shared Policy in Federated Learning). It shares the augmentation policies during training rather than preprocessed or encoded data such as the average of data or statistics of features. Direct exposure of personal information was constrained, yet clients still benefited from the policies learned and shared across the clients to augment their local data. We also proposed Federated Meta-Policy Loss (FMPL) and used the first-order gradient information to enhance privacy with reduced communication costs. A potential limitation of our algorithm is the introduction of Joint Training. This approach requires consideration when applying our federated data augmentation method to existing federated learning algorithms. Also, our algorithm assumed the use of FedAvg, which does not account for model personalization [44, 45] in the computation of FMPL. Investigating a policy loss that aligns with model personalization algorithms would be interesting. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: Our research proposes a federated data augmentation method that does not have a high risk for misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: We have cited all the assets used in the paper.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: Model Consistent with prior studies, we employed a standard CNN model for all experiments as referenced in [39] [40] [41] . The global model consists of three convolutional layers with 64 filters and 3x3 kernels, followed by three fully-connected layers of 256, 128, and the final classification layer. All experiments are run on a cluster of 32 NVIDIA GTX 1080 GPUs.Datasets To evaluate robust performance in non-i.i.d. data, we set up our experimental environment by distributing the CIFAR-10/100 [21] and SVHN [22] datasets with varying levels of data heterogeneity among clients. We allocated the data to 130 clients based on a Dirichlet distribution with different hyperparameters of α = [5.0, 0.1], following the procedure in pFL-Bench [25] . The lower the α, the higher the degree of heterogeneity. Among these clients, only 100 were randomly selected to participate in the training, while the remaining 30 were designated as out-of-distribution (OOD) clients. For the CIFAR10/100, the number of training rounds R was set to 1000. For the FEMNIST, the number of training rounds R was set to 500. For the SVHN, the number of training rounds R was set to 100, 300, and 500, as reported in Table 5 . The results reported in the main paper, Table 1 , are for the 500 round. In each communication round, only 10 clients are sampled, and the remaining 30 clients serve as out-of-distribution (OOD) clients.Hyperparameters For the FedAvP algorithm, the hyperparameters include the server policy learning rate η, client policy learning rate λ, gradient clipping threshold c, and a regularization term ϵ. In our experimentation, we tuned η within [0.4, 0.9], λ within [0.1 ∼ 0.9], c within [0. Baselines In all baseline experiments, Random Crop and HorizontalFlip were applied as default augmentations. For the baseline algorithms' hyperparameter settings, specifically for RandAugment (RA) [17] , we leveraged a PyTorch implementation [47] . Hyperparameter values were determined in alignment with the procedures described by the authors. For the CIFAR-100 dataset, the hyperparameter M was investigated within the range of [2, 6, 10, 14] and N within [1, 2] . For CIFAR-10, the hyperparameter M was investigated within [4, 5, 7, 9, 11] and N within [2, 3] . For FEMNIST, the hyperparameter M was investigated within [4, 5, 7, 9, 11] and N within [1, 2, 3] . For SVHN, the hyperparameter M was investigated within [5, 7, 9, 11] and N within [3] . For TrivialAugment (TA) [24] , we utilized the PyTorch library's built-in algorithm [47] . For RandAugment and TrivialAugment, Random Crop and HorizontalFlip were applied first, followed by RandCutout [48] . As for FedMix [8] , the hyperparameter λ was investigated within We performed experiments using reconstruction attacks [14] to assess if our policy-sharing algorithm offers better privacy in a federated learning context, in comparison to FedMix [8] and FedFA [11] , which share data at the input and feature levels respectively. Additionally, we conducted experiments to compare our algorithm with FedGen [9] , which shares the label distribution of the client's training data and a generative model that produces a latent from learned latent feature space over the clients' local training data. Despite our algorithm not being specifically designed as a defense against reconstruction attacks, we included ATSPrivacy [27] in our performance evaluation for comparison. ATSPrivacy aims to identify the best policy to counter reconstruction attacks by conducting a policy search. Although ATSPrivacy's primary goal isn't to enhance performance through policy sharing in collaborative learning environments, we chose it as a baseline because it also involves policy search in the context of collaborative learning. We have described below the attack algorithms that utilize input-level, feature-level, and policy gradient information.It has been reported that in Collaborative learning systems, Gradient leakage attacks are feasible [14, 42, 43] . These attacks use gradients to reconstruct the training data. We consider a scenario with a central server and clients where learning occurs through the exchange of gradients. Assume we have a given gradient ∇W (x, y), then we can optimize for a dummy data and label pair (x ′ , y ′ ) by minimizing the following objective:EQUATIONwhere || • || denotes a norm distance measure, ∇W (x, y) is the given gradient, (x, y) is the client's sample data and label, and (x ′ , y ′ ) are the targets of optimization. Following [42] , we utilize cosine similarity as a cost function instead of norm distance, like in the case of x * , y * = arg minEQUATIONwhere ℓ(x, y) = ⟨x, y⟩ /(∥x∥ • ∥y∥).Here, we assume that input-level information is available, specifically the mean image of the client's data xmean in the FedMix [8] . The method for reconstruction attack utilizing input-level information is as follows:x * , y * = arg minEQUATION)where αinput is a hyperparameter for the additional term, it is fixed at 0.1 in experiments where input-level information is available and set to 0 in scenarios without input-level information. xmean represents the mean image of the client's data. Cosine similarity was utilized to measure gradient distance, as described in [42] .Similarly, in cases where the client provides feature-level information to the server, as in the FedFA [11] , this information can also be used to help with reconstruction. The method for reconstruction attack utilizing feature-level information is as follows:x * , y * = arg minEQUATIONwhere μk and σk are the momentum updated feature statistics of layer k in the client model, µ k ′ and σ k ′ are the feature statistics of layer k with regard to x ′ , α f eat and β f eat are hyperparameters for the additional terms fixed at 0.1 and 0.0, respectively.In the FedGen algorithm, clients transmit the label distribution of their local training data and the generative model. This information allows us to improve the cost function for reconstruction, which is formulated as follows:EQUATION)where G is a generative model to generate the latent distribution, W p refers to the final layer in the client model that takes a latent as an input and outputs prediction logits, c is the label distribution of the client's local training data, αgen and βgen are hyperparameters for the additional terms fixed at 0.0001 and 0.0001, respectively. Also, we can initialize y ′ with c.For the FedAvP algorithm, clients send the policy gradient information. Such policy gradient information can also be utilized as a term of cost function for reconstruction in the following manner:EQUATIONwhere α policy is a hyperparameter for the additional term fixed at 0.1.When additional information, such as input-level or feature-level, is available in addition to gradient information, it can be utilized, potentially posing a security risk. On the other hand, we have confirmed that even when utilizing the policy gradient information additionally, the reconstruction risk does not increase. The results of the experiments regarding this are summarized in Table 2 .Our reconstruction attack experiments were conducted on CIFAR-100 with α = 0.1, featuring heterogeneous data distribution. Specifically, after 1000 rounds of training, we selected clients for the reconstruction attack.The selection criteria were the client with the most training data (895 data points), Client(L), and the one with the least (156 data points), Client(S). For both clients, we randomly sampled 150 training data points to perform the reconstruction attack. In all experiments, the batch size and reconstruction step for the reconstruction attack were set to 1 and 2400, respectively. We used the Adam optimizer [49] . Gradients obtained from 1 round after 1000 rounds of training were used for the reconstruction attack.For ATSPrivacy, we trained the 6 policies with FedAvg used in the CIFAR-100 dataset of the author's paper [27] , namely 3-1-7, 43-18-18, (3-1-7, 43-18-18), 21-13-3, 7-4-15, and (21-13-3, 7-4-15) , in the CIFAR-100 with (α = 0.1) environment. We performed the reconstruction attack on the policies (7-4-15) and (21-13-3, 7-4-15) , which showed the best accuracy performance. The (21-13-3, 7-4-15) policy is a hybrid policy described in the author's paper [27] , and we applied it by randomly sampling one of the two policies.We conducted experiments on a model with more parameters than those used in the main experiments. The model used in these experiments is a simplified version of VGG11 [50, 51] , where all dropout and batch normalization layers are removed, and the filters and the size of all fully-connected layers are reduced by a factor of 2. This model contains about three times more network parameters than the networks used in the main paper. We verify the performance of our model on non-IID datasets, specifically SVHN-10 (α = 0. In the main paper, we followed the weighted accuracy metric as described in the pFL-Bench [25] .Here, we provide additional results for non-i.i.d. experiments using an equally-weighted metric on CIFAR-100.Method CIFAR-100 / 5.0 CIFAR-100 / 0. Table 16 : Test accuracy (%) on CIFAR-100 with α = 5.0 and α = 0.1 using an equally-weighted metric.The table below presents the results across different datasets and partitioning strategies, specifically the quantity-based label skew settings described in [56] . 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.Method CIFAR-100 SVHN C = 3 C = 2 C = 1 C = 3 C = 2 C =Consider the federated meta-policy loss derived from the updated weight w k n for client k at step n using a first-order Taylor expansion:EQUATIONWhen calculating the policy gradient of this loss with respect to θ k n-1 for client k, the first-order gradient approximation is as follows:EQUATIONwhere w k n = w gr -lr • g augEQUATIONEQUATIONEQUATIONProof.EQUATIONFirst, let's calculate starting from the left term in Eq. ( 19)∂ℓ D val k (w k n ) ∂θ k n-1 (20) = ∇ℓ D val k (w k n ) T • ∂w k n ∂θ k n-1(using the chain rule) (21)= ∇ℓ D val k (w k n ) T • ∂(w k n-1 -lr • g aug w k n-1)∂θ k n-1(using the definition of w k n in ( 16))(22) = -lr • ∇ℓ D val k (w k n ) T ∂(g aug w k n-1)∂θ k n-1 (using ∂w k n-1 ∂θ k n-1 = 0 ) (23) = -lr • ∂(∇ℓ D val k (w k n ) T • ∇ℓ tp θ k n-1 (D train k,n-1 ) (w k n-1 )) ∂θ k n-1(using the definition of g augEQUATIONNext, we calculate the right term in Eq. ( 19)∂(∇ℓ D val k (w k n ) T (w gr+1 -w k n )) ∂θ k n-1 = ∂w k n ∂θ k n-1 • ∂∇ℓ D val k (w k n ) T ∂w k n • (w gr+1 -w k n ) + ∇ℓ D val k (w k n ) T • ∂(w gr+1 -w k n ) ∂θ k n-1 (using the chain rule and (f g) ′ = f ′ g + f g ′ ) (25) = ∇ℓ D val k (w k n ) T • ∂(w gr+1 -w k n ) ∂θ k n-1 (treating ∇ℓ D val k (w k n ) as a constant for first-order) (26) = ∇ℓ D val k (w k n ) T • ∂(α k • w k N -w k n ) ∂θ k n-1(using the definition of w gr+1 in ( 18))(27) = ∇ℓ D val k (w k n ) T • ∂((α k -1) • w k n ) ∂θ k n-1(using the definition of w k N in ( 16)) and treating g aug w k n as a constant for first-order)EQUATION)∂θ k n-1(using the definition of w k n in ( 16))EQUATION(using the definition of g augEQUATION) By combining Eq. ( 24) and Eq. (33) ,EQUATIONQuestion: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We provide the code for our paper and its instructions in a zip file. Guidelines:• The answer NA means that paper does not include experiments requiring code. • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: Our research proposes a federated data augmentation method that does not have negative impacts on society. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).38th Conference on Neural Information Processing Systems (NeurIPS 2024).https://optuna.org/ 


Decision: This paper proposes FedAvP, a data augmentation methods by sharing policies among clients in a federated learning environment. It focuses on an important and challenging aspect in federated learning. The effectiveness of the proposed solution is verified by extensive experiments. We recommend accept it.

Reviews: review: paper_summary: paper_summary: This paper proposes FedAvP, which performs data augmentation search by sharing policies among clients in a federated learning (FL) environment. They introduce federated meta-policy loss and utilize the first-order gradient information to further enhance privacy and reduce communication costs. The proposed algorithm allows for rapid adaptation of a personalized policy by each client, relieving the challenge of non-iid.
strength_weakness: strength_weakness: 1. If the focus of this paper is on addressing non-iid using data augmentation, then the baselines used for comparisons are actually lacking. Classic methods for addressing non-iid, such as FedNova and Scaffold, have not been compared.
2. Additionally, to address the non-iid problem, it would be helpful to have some more vivid and intuitive explanations of why the shared policy can alleviate the non-iid issue.
3. One weakness affecting readability is that when the article references some inspired works to address certain issues, it opts for direct citations without more detailed explanations, such as gradient clipping and the reweighting in Eq. (3).
questions: questions: 1. Regarding W2, could you please add more intuitive explanations of the shared policy strategy?

2. Regarding W3, could you please explain what is the advantage of the reweighting in Eq.(3)?
limitations: limitations: The authors provide discussions on the limitations of the work in Sections 5.
paper_summary: paper_summary: The paper introduces a novel federated data augmentation algorithm, FedAvP, designed for data augmentation of the client-side without the need to share client data information. Specifically, the authors propose a meta-learning method that allows multiple clients to collaboratively learn data augmentation policies and design a Federated Meta-Policy Loss (FMPL) for the optimization of data augmentation policies. Moreover, to prevent the gradient updates of the data augmentation policies from leaking local data privacy, they propose using first-order approximation optimization to protect privacy and reduce communication costs. Finally, the authors conducted experiments in different data heterogeneity scenarios across multiple datasets to validate the effectiveness of the method.
strength_weakness: strength_weakness: 1) The augmentation policy search method in the paper only supports image data augmentation and includes only two types of augmentation methods. Moreover, as the number of augmentation methods increases, the search space will grow exponentially. The authors should analyze the scalability of FedAvP in this regard.
2) The paper lacks comparisons with the latest federated data augmentation methods, and there are few non-iid scenarios involved in the performance experiments, with only one scenario on some datasets.
questions: questions: 1. Why were the performance experiments on the Cifar-10, SVHN, and FEMNIST datasets only conducted in one scenario?
2. The gradient calculation of the policy loss introduces significant additional overhead. Could this be a hindrance to applying this method to larger visual pre-training models like ViT?
limitations: limitations: The method proposed in the paper seems to be applicable only for federated training of small-scale visual models. Its performance on tasks involving text, speech, or multi-modalities and larger-scale models still needs to be further validated.
paper_summary: paper_summary: This paper points out that the shared input-level and feature-level information poses potential privacy leakage. FedAvP only shares the augmentation policy which is not related to the data. This method leverages the first-order information as a replacement to reduce privacy leakage and communication costs. Comprehensive experiments demonstrate the efficacy and efficiency of FedAvp.
strength_weakness: strength_weakness: - The writing of this paper should be improved to express ideas clearly.
- A deeper analysis is needed to reveal why different $\alpha$ learn different augmentation methods under the same data as depicted in Figure 1. There is also no in-depth insight about how the different augmentation influences the performance in different FL setting
questions: questions: - As FedAvp is also a learning policy for augmentation strategy. So whats the performance of FedAvg+AutoAugment which is also an automation policy?
- What is the meaning of OOD client is not mentioned.
- This paper claims that the reported accuracy is calculated as the weighted average of each client’s accuracy by the number of data points they have. However, most of the methods report the accuracy of global model. How the test dataset is partitioned of each client?
limitations: limitations: See weakness.
paper_summary: paper_summary: This paper proposes FedAvP, a novel algorithm to perform data augmentation via shared policy in federated learning (FL). Extensive experiments verify the effectiveness of the proposed algorithm.
strength_weakness: strength_weakness: W1: The authors can conduct more experiments on more comprehensive non-IID settings. For example, in a previous FL benchmark on non-IID data [1], the quantity-based label skew settings are challenging. The authors are suggested to experiment on such settings to further support the effectiveness of the proposed algorithm. 

[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.
questions: questions: Please see weaknesses. The authors are suggested to conduct more comprehensive experiments. I will adjust the score based on the author response.
limitations: limitations: Please see weaknesses.

Comments: title: Thank you
Response to authors
comment: Thanks to the efforts of the authors, I will raise my score.
Dear Reviewer. We sincerely appreciate your thoughtful advice.

However, FedAvP does not learn just one strategy. As shown in Figure 2, it clearly learns different data augmentation strategies depending on the heterogeneity levels and distributions.
As described in Section 4.3, the global policies are then adapted into a local policy using each client's local data. Figure 3 also demonstrates the statistics of personalized policies among different clients.

As you suggest, we also evaluated the accuracy of the global model on CIFAR-100 ($\alpha = 5.0$ and $\alpha = 0.1$) using equally-weighted metric.
We will include the results in the appendix. In this experiment, FedAvP still outperforms other baselines.

|Dataset/heterogeneity degree α|CIFAR100/5.0|CIFAR100/0.1|
|------------------------------|------------|------------|
|FedAvg+Default|40.04|36.98|
|FedAvg+RandAugment|47.3|43.17|
|FedAvg+TrivialAugment|46.61| 42.04|
|FedProx+Default|40.56|37.61|
|FedProx+RandAugment|45.95|41.25|
|FedProx+TrivialAugment|46.59|41.67|
|FedDyn+Default|42.11|38.23|
|FedDyn+RandAugment|45.68|42.08|
|FedDyn+TrivialAugment|46.84|40.92|
|FedExp+Default|42.78|38.22|
|FedExp+RandAugment|46.14|41.97|
|FedExp+TrivialAugment|48.54|42.01|
|FedGen+Default|42.12|38.05|
|FedGen+RandAugment|47.11|42.96|
|FedGen+TrivialAugment|47.73|40.62|
|FedMix+Default|39.59|38.46|
|FedMix+RandAugment|46.67|42.7|
|FedMix+TrivialAugment|46.62|42.49|
|FedFA+Default|43.68|41.18|
|FedFA+RandAugment|48.87|43.26|
|FedFA+TrivialAugment|47.86|43.36|
|FedAvP(local)|49.05|43.64|
|FedAvP(Fast Update)|49.94|45.09|
|FedAvP|**50.59**|**45.93**|


We respectfully hope that the reviewer re-examines the evaluation of our work. Our approach to facilitating shared data augmentation policy is a novel direction in federated learning research.
We have addressed data scarcity and heterogeneity, preserving security, and also provided a theoretical analysis that elucidates the role of meta-policy updates in distributed learning.
Thanks to the efforts of the authors. However, I am still concerned about the insight of this paper. e.g., why FedAvP strategies tend to learn one strategy and not the other when heterogeneity is severe? This may provide new insight into addressing heterogeneity in FL. At the same time, there is a lot of work to test the global model accuracy, and the metric for this task is also needed. I  hope this paper can inspire a new view in FL to solve the heterogeneity. In summary, I will continue to maintain my grade
Thank you for the detailed responses. Most of my concerns have been well addressed.
Thank you for the responses. Most of my concerns have been addressed.
I appreciate the authors' efforts on additional experiments. Given the rebuttal and other reviews, I decide to raise the rating.