An Agnostic Approach to Federated Learning with Class Imbalance

1 INTRODUCTION . Class imbalance is ubiquitous in real world supervised learning problems . Examples span from medical applications ( Lee & Shin , 2020 ; Roy et al. , 2019 ; Choudhury et al. , 2019 ) , fraud detection ( Yang et al. , 2019 ; Chan et al. , 1999 ) , to consumer based applications ( Wang et al. , 2021b ; Wu et al. , 2020 ; Long et al. , 2020 ) . In these scenarios , data belonging to a subset of classes constitute a great proportion of the population while data from the minority classes , generated by uncommon events , are scarce ( He & Garcia , 2009 ) . Having a non-uniform number of samples per class deteriorates the performance of the classifier in the minority class ( Huang et al. , 2016 ) , resulting in low training and testing accuracy . More importantly , unintended consequences of mistreating the minority class can be catastrophic ( Van Hulse et al. , 2007 ) if the problem is not handled appropriately . From the perspectives of data heterogeneity and privacy , the issue of class imbalance is even more significant in the setting of Federated Learning ( FL ) . Due to the heterogeneity of the local data distributions , there can be a significant mismatch between the local and global imbalance , i.e . the class that is a minority locally can actually be a majority class globally . Moreover , for the purpose of privacy protection in FL , one should tackle the class-imbalance problem in an agnostic manner , i.e . the proposed algorithms should not require the minority class to be identified . In the centralized training setting , techniques like balanced sampling , loss re-weighting , and gradient tuning have achieved many successes ( Johnson & Khoshgoftaar , 2019 ) . However , due to the extra difficulty of handling the class imbalance problem in FL , previous arts that rely on the explicit identification of the minority class do not directly apply . While research along this line is quite limited , a commonly used heuristic is to estimate the portion of a class using the norm of the gradient per class . When used as proxies , these quantities are then utilized to reweight the losses corresponding to different classes ( Wang et al. , 2021a ; Yang et al. , 2020 ) . A key drawback of ( Wang et al. , 2021a ) is that a subset of the client ’ s local dataset needs to be shared , which is not ideal for privacy preserving purposes . Moreover , the effectiveness of such approaches degrades when there is a notable mismatch between the local imbalance and the global imbalance ( Wang et al. , 2021a ) . In this work , we target the most difficult yet possibly most interesting setting in FL , where there is a significant mismatch between the local and global imbalance . Concretely , on certain clients , a minority class from a global perspective is in the majority locally . Practitioners often encounters such a mismatch due to the highly heterogeneous data configuration in FL and the principle of locality ( Wang et al. , 2021a ) . For example , when the local data are produced by a minority user , the corresponding minority data could occupy a large portion of the local distribution . In this scenario , due to the scarceness of the minority data from a global perspective , the corresponding underrepresented client will usually experience poor service quality from the trained model , giving rise to potentially severe fairness issue . To overcome the aforementioned challenges , we propose a constrained learning formulation to handle the class imbalance issue in FL while accounting for both heterogeneity and privacy . In brief , we impose constraints on the standard FL formulation so that the empirical loss on every client should not overly exceed the average empirical loss . Such constraints are shown to force the classifier to account for all classes equally and hence mitigate the detrimental effect of class imbalance , under a type of heterogeneous data configuration that captures the mismatch between the local and global imbalance . The advantages of our formulation are threefold : First , in contrast to previous arts which usually rely on some heuristics to reweight the loss functions for different classses , our approach is principled , yielding a simple optimization interpretation . Second , unlike existing methods , our formulation requires no additional information compared to the original FL formulation and it treats data from different classes agnostically , and hence is less likely to leak the private information of the clients . Third , from our extensive empirical study , our formulation can significantly improve the testing accuracy on the minority class , without compromising the overall performance . Contribution . We summarize our contributions as follows . • Problem Formulation : To address the challenge of class imbalance in FL with data heterogeneity , we propose a novel constrained FL formulation with explicit enforcement on the similarity between the local empirical losses . Using only information from the standard FL objective , our approach is completely agnostic to class distribution of the client data and its proxies , as opposed to the existing literature ( which mostly violate privacy requirements of FL ) . • Meta-Algorithm : We solve our constrained optimization problem via a primal-dual approach . For a fixed dual variable , the corresponding Lagrangian function enjoys a similar structure as the standard FL loss , but with non-uniform weights on the local objectives . Accordingly , we propose an efficient ( meta- ) algorithm called CLIMB that can use any FL optimization method as a subroutine , with nearly negligible communication overhead . Furthermore , we analyze the convergence properties of CLIMB under certain oracle assumptions . • Accuracy Improvement in the minority class : On benchmark datasets , CLIMB with Fed-Avg as base solver achieves significant enhancement of the accuracy in the minority class without compromising the overall accuracy under various data heterogeneity and class imbalance scenarios . 2 RELATED WORKS . Class Imbalance in Centralized Setting . In the centralized learning setting , the number of samples per class is known . We categorize existing solutions that exploit such information as follows . Balanced sampling . In order to balance the data used for gradient calculation , the most common approaches are either to under sample from the majority classes ( Liu et al. , 2008 ) , or to augment the minority class ( Chawla et al. , 2002 ; Guo & Viktor , 2004 ) . Both methods seek to generate an artificial uniform distribution of data ( Buda et al. , 2018 ; Pouyanfar et al. , 2018 ) . Loss reweighting . Methods in this category focus on re-weighting the loss functions for different classes , with extra emphases on the mistakes made in the minority class ( Cui et al. , 2019 ; Ling & Sheng , 2008 ; Sun et al. , 2007 ; Wang et al. , 2016 ) . There also exist works like ( Lin et al. , 2017 ) that adjust the scale of the loss sample-wise , without exploiting the global data distribution . Gradient tuning . In the context of neural networks , some other works focus on tuning the gradient per class ( Anand et al. , 1993 ) showing faster rates of convergence in the class imbalance setting . To do this , the gradient is reshaped in order to have an equal magnitude when projected to all the gradients per class , while keeping the same norm as the original gradient . Class Imbalance in Federated Learning . Combating with the issue of class imbalance is more challenging in FL since the data composition in such a setting is generally unknown and the minority classes are often difficult to identify Li et al . ( 2020 ) ; Wang et al . ( 2020b ) . While research in this direction are quite limited , to the best of our knowledge , we classify them into the following two classes based on whether a proxy of the data composition is explicitly built . ( i ) A common strategy in existing methods is to explicitly build proxies of the data composition based on some heuristics . These proxies are dynamically adjusted during the learning procedure and allow the aforementioned three techniques , balanced sampling , loss re-weighting and gradient tuning , to be utilized in the FL setting : Wang et al . ( 2021a ) suggests that there exists a proportional relation between the magnitude of the gradients the corresponds to the last layer of the neural network and the sample quantity . Based on such observation , the proposed the Ratio-Loss , a class-wise re-weighted version of the standard cross entropy loss . We note that to define the Ratio-Loss , one needs to collect a subset of client data as “ auxiliary data ” , which may not be ideal in the setting of FL . On the same vein , Yang et al . ( 2020 ) argues that the gradient per class can be used as a proxy to infer the imbalances in the distributions of the clients , and thus clients should be selected according to how uniform the magnitude of gradient per class is . We also notice the work Astraea which introduces extra virtual components called mediators between the FL server and clients ( Duan et al. , 2019 ) . The mediator is assumed to have access to the local data distributions of the clients so that client rescheduling and data augmentation can be carried out accordingly . Through the prepossessing on the mediator , the gradients communicated to the server are made balanced class-wise . However , when privacy is taken into consideration , methods like Astraea may not be appropriate as in essence it is directly built on the global data distribution . ( ii ) Without explicitly constructing an estimation of the data composition , there are works that resort to techniques like active learning and reinforcement learning to implicitly learn the data composition as the optimization progresses . Works along this research line usually rely on client selection to mitigate the effects of class imbalances . In ( Goetz et al. , 2019 ) , to address the class imbalance problem , there will a higher probability to sampling clients that posses the global minority class . Other works , leverage client selection as a multi-armed bandit problem , and design a client selection policy to balance the gradients ( Xia et al. , 2020 ) . Other arts leverage Q-learning techniques to select clients in order to minimize the overall loss Wang et al . ( 2020a ) . However , we believe in the most practical setting of FL , clients that are available per communication round is incoercible and hence these strategies have limited applicability . Comparison with Previous Works . Our method significantly differs from the above strategies in the following ways : ( 1 ) Our method requires no knowledge of the data composition , nor any proxy of such information . Consequently , our approach is agnostic to the minority class and better preserves clients ’ privacy . Such a property draws clear distinction between our work and the works listed in class ( i ) above . ( 2 ) As will be more clear in the following section , we only perform clientwise re-weighting as opposed to the class-wise re-weighting schemes used in previous works , which further emphasize the agnostic nature of our approach . ( 3 ) In contrast to the methods in class ( ii ) above , our approach does not require active client selection , enjoying a broader applicability in the most practical and interesting settings .

Summary: This paper focuses on an interesting and challenging setting in federated learning, where number of classes and samples are imbalanced across different local workers. Compared to other methods which might require subset of local data or prior knowledge of the minor classes, the algorithm proposed in this paper does not need neither local data nor prior knowledge, which maintains the privacy principle in federated learning and aligned to real-world settings. The main contributions are: 1. Proposes a new algorithm solving the non-iid issue in federated learning which only adds constraints on local losses, without sacrificing local data privacy; 2. Utilizes a primal-dual optimization mechanism to solve the problem, provides theoretical guarantees on convergence; 3. Based on the experiments, there is a significant performance increase compared to FedAvg.