FedMix: Approximation of Mixup under Mean Augmented Federated Learning

Federated learning ( FL ) allows edge devices to collectively learn a model without directly sharing data within each device , thus preserving privacy and eliminating the need to store data globally . While there are promising results under the assumption of independent and identically distributed ( iid ) local data , current state-of-the-art algorithms suffer from performance degradation as the heterogeneity of local data across clients increases . To resolve this issue , we propose a simple framework , Mean Augmented Federated Learning ( MAFL ) , where clients send and receive averaged local data , subject to the privacy requirements of target applications . Under our framework , we propose a new augmentation algorithm , named FedMix , which is inspired by a phenomenal yet simple data augmentation method , Mixup , but does not require local raw data to be directly shared among devices . Our method shows greatly improved performance in the standard benchmark datasets of FL , under highly non-iid federated settings , compared to conventional algorithms . 1 INTRODUCTION . As we enter the era of edge computing , more data is being collected directly from edge devices such as mobile phones , vehicles , facilities , and so on . By decoupling the ability to learn from the delicate process of merging sensitive personal data , Federated learning ( FL ) proposes a paradigm that allows a global neural network to learn to be trained collaboratively from individual clients without directly accessing the local data of other clients , thus preserving the privacy of each client ( Konečný et al. , 2016 ; McMahan et al. , 2017 ) . Federated learning lets clients do most of the computation using its local data , with the global server only aggregating and updating the model parameters based on those sent by clients . One of the standard and most widely used algorithm for federated learning is FedAvg ( McMahan et al. , 2017 ) , which simply averages model parameters trained by each client in an element-wise manner , weighted proportionately by the size of data used by clients . FedProx ( Li et al. , 2020b ) is a variant of FedAvg that adds a proximal term to the objective function of clients , improving statistical stability of the training process . While several other methods have been proposed until recently ( Mohri et al. , 2019 ; Yurochkin et al. , 2019 ; Wang et al. , 2020 ) , they all build on the idea that updated model parameters from clients are averaged in certain manners . Although conceptually it provides an ideal learning environment for edge devices , the federated learning still has some practical challenges that prevent the widespread application of it ( Li et al. , 2020a ; Kairouz et al. , 2019 ) . Among such challenges , the one that we are interested in this paper is the heterogeneity of the data , as data is distributed non-iid across clients in many real-world settings ; in other words , each local client data is not fairly drawn from identical underlying distribution . Since each client will learn from different data distributions , it becomes harder for the model to be trained efficiently , as reported in ( McMahan et al. , 2017 ) . While theoretical evidence on the convergence of FedAvg with non-iid case has recently been shown in ( Li et al. , 2020c ) , efficient algorithms suitable for this setting have not yet been developed or systematically examined despite some efforts ( Zhao et al. , 2018 ; Hsieh et al. , 2020 ) . In addition to non-iid problem , another important issue is that updating model parameters individually trained by each client is very costly and becomes even heavier as the model complexity increases . Some existing works Smith et al . ( 2016 ) ; Sattler et al . ( 2019 ) target this issue to decrease the amount of communications while maintaining the performance of FedAvg . A more practical approach to reduce communication cost is to selectively update individual models at each round , rather than having all clients participate in parameter updates . This partial participation of clients per round hardly affects test performance in ideal iid settings but it can exacerbate the heterogeneity of weight updates across clients and as a result , the issue of non-iid ( McMahan et al. , 2017 ) . In order to mitigate the heterogeneity across clients while protecting privacy , we provide a novel yet simple framework , mean augmented federated learning ( MAFL ) , in which each client exchanges the updated model parameters as well as its mashed ( or averaged ) data . MAFL framework allows the trade-off between the amount of meaningful information exchanged and the privacy across clients , depending on several factors such as the number of data instances used in computing the average . We first introduce a naive approach in our framework that simply applies Mixup ( Zhang et al. , 2018 ) between local data and averaged external data from other clients to reduce a myopic bias . Here , we go further in our framework and ask the following seemingly impossible question : can only averaged data in our framework that has lost most of the discriminative information , bring the similar effect as a global Mixup in which clients directly access others ’ private data without considering privacy issues ? Toward this , we introduce our second and more important approach in our framework , termed Federated Mixup ( FedMix ) , that simply approximates the loss function of global Mixup via Taylor expansion ( it turns out that such approximation only involves the averaged data from other clients ! ) . Figure 1 briefly describes the concept of our methods . We validate our method on standard benchmark datasets for federated learning , and show its effectiveness against the standard federated learning methods especially for non-iid settings . In particular , we claim that FedMix shows better performance and smaller drop in accuracy with more heterogeneity or fewer clients update per communication round , further increasing difficulty of federated learning . Our contribution is threefold : • We propose a simple framework for federated learning that averages and exchanges each local data . Even naive approach in this framework performing Mixup with other clients ’ mashed data shows performance improvement over existing baselines on several settings . • We further develop a novel approximation for insecure global Mixup accessing other clients ’ local data , and find out that Taylor expansion of global Mixup only involves the averaged data from other clients . Based on this observation , we propose FedMix in our framework approximating global Mixup without accessing others ’ raw data . • We validate FedMix on several FL benchmark datasets especially focusing on non-iid data settings where our method significantly outperforms existing baselines while still preserving privacy with minimal increases in communication cost . 2 RELATED WORK . Federated learning Federated learning was first proposed in Konečný et al . ( 2016 ) where the prevalent asynchronous SGD ( Dean et al. , 2012 ) is used to update a global model in a distributed fashion . A pioneering work in this field proposed the currently most widely used algorithm , FedAvg ( McMahan et al. , 2017 ) , which is also the first synchronous algorithm dedicated to federated setting . Shortly after , Li et al . ( 2020b ) proposed a variant of FedAvg , named FedProx , where the authors claimed to overcome statistical heterogeneity and increase stability in federated learning . Recent studies attempt to expand federated learning with the aim of providing learning in more diverse and practical environments such as multi-task learning ( Smith et al. , 2017 ) , generative models ( Augenstein et al. , 2020 ) , continual learning ( Yoon et al. , 2020 ) , semi-supervised learning ( Jeong et al. , 2020 ) , and data with noisy labels ( Tuor et al. , 2020 ) . Our paper focuses on general federated settings , but it could be considered in such various situations . However , these algorithms may obtain suboptimal performance when clients participating in FL have non-iid ( Zhao et al. , 2018 ; Hsieh et al. , 2020 ) distributions . While the convergence of FedAvg on such settings was initially shown by experiments in McMahan et al . ( 2017 ) and later proved in Li et al . ( 2020c ) , it does not guarantee performance as good as it would have been for iid setting . Existing algorithms that pointed out this issue have major limitations , such as privacy violation by partial global sharing of local data ( Zhao et al. , 2018 ) or no indication of improvement over baseline algorithms such as FedAvg ( Hsieh et al. , 2020 ) . Our method aims to improve performance particularly on these non-iid situations , without compromising privacy . Mixup Mixup ( Zhang et al. , 2018 ) is a popular data augmentation technique that generates additional data by linear interpolation between actual data instances . Mixup has been usually applied to image classification tasks and shown to improve test accuracy on various datasets such as CIFAR10 and ImageNet-2012 ( Russakovsky et al. , 2015 ) , and , on popular architectures such as ResNet ( He et al. , 2016 ) and ResNeXt ( Xie et al. , 2017 ) , for various model complexity . It is also reported in Zhang et al . ( 2018 ) that Mixup helps with stability , adversarial robustness ( Zhang et al. , 2018 ) , calibration , and predictive certainty ( Thulasidasan et al. , 2019 ) . Mixup is expanding from various angles due to its simplicity and popularity . First , beyond image classification tasks , its effectiveness has been proven in various domains such as image segmentation ( Eaton-Rosen et al. , 2020 ) , speech recognition ( Warden , 2018 ) , and natural language processing ( Guo et al. , 2019 ) . Also , several extensions such as Manifold Mixup ( Verma et al. , 2018 ) , which performs Mixup in latent space , or CutMix ( Yun et al. , 2019 ) , which replaces specific regions with others patches , have been proposed . In most of the previous studies on federated learning , Mixup was partially ( or locally ) used as a general data augmentation technique . Some recent studies ( Oh et al. , 2020 ; Shin et al. , 2020 ) proposed to send blended data to server using Mixup , but they require sending locally- and linearly-mixed ( mostly from two instances ) data to server at every round , therefore being susceptible to privacy issues with huge communication costs . Our work properly modifies Mixup under the restrictions of federated learning and mitigates the major challenges of federated learning such as non-iid clients . 3 MEAN AUGMENTED FEDERATED LEARNING ( MAFL ) AND FEDMIX . We now provide our framework exchanging averaged data for federated learning and main method approximating insecure global Mixup under our framework , after briefly introducing the setup . 3.1 SETUP AND BACKGROUND . Federated learning and FedAvg Federated Averaging ( FedAvg ) ( McMahan et al. , 2017 ) has been the most popular algorithmic framework for federated learning . For every communication round t = 0 , . . . , T − 1 , a client k ∈ 1 , . . . , N selected for local training sends back its model wkt ( or only difference to reduce communication cost ) to a global server . For every round , K number of clients are selected to locally update and send model parameters . The server simply averages parameters received , so that the global model wt after t rounds of communications becomes wt = ∑n k=1 pkw k t where pk is the importance of client k based on the relative number of data in k among all selected clients at t. The updated global model is sent back to clients for the next round , which undergoes the following E local updates via stochastic gradient descent ( SGD ) : wkt+1 , i+1 ← wkt+1 , i − ηt+1∇ ` ( f ( xki ; wt+1 , i ) , yki ) for i = 0 , 1 , . . . , E − 1 , batch size B , and local learning rate η . Here , ` is the loss function for learning and f ( x ; wt ) is the model output for input x given model weight wt . Mixup Mixup ( Zhang et al. , 2018 ) is a simple data augmentation technique using a linear interpolation between two input-label pairs ( xi , yi ) and ( xj , yj ) to augment x̃ = λxi + ( 1 − λ ) xj and ỹ = λyi + ( 1 − λ ) yj . The variable λ ∈ [ 0 , 1 ] is a hyperparameter that is chosen from the beta distribution for each training step .

Summary: The paper proposed MAFL, a novel approach to conduct Mixup under the federated learning setting whiling preserving data privacy. The proposed FedMix scheme is inspired by Taylor’s expansion of the global Mixup formulation. The effectiveness of MAFL is justified via empirical studies over a simulated federated learning environment, which indicates that Mixup achieves better test accuracies on various machine learning tasks.