Context-aggregator: An approach of loss- and class imbalance-aware aggregation in federated learning

Abstract

Federated Learning (FL) is an emerging distributed learning paradigm which offers data privacy to contributing nodes in the collaborating environment. By exploiting the Individual datasets of different hospitals in FL setting could be used to develop reliable screening, diagnosis, and treatment predictive models to tackle major challenges such as pandemics. FL can enable the development of very diverse medical imaging datasets and thus provide more reliable models for all participating nodes, including those with low quality data. However, the issue with the traditional Federated Learning paradigm is the degradation of generalization power due to poorly trained local models at the client nodes. The generalization power of the FL paradigm can be improved by considering the relative learning contribution of client nodes. Simple aggregation of learning parameters in the standard FL model faces a diversity issue and results in more validation loss during the learning process. This issue can be resolved by considering the relative contribution of each client node participating in the learning process. The class imbalance at each site is another significant challenge that greatly impacts the performance of the aggregated learning model. This work considers Context Aggregator FL based on the context of loss-factor and class-imbalance issues by incorporating the relative contribution of the collaborating nodes in FL by proposing Validation-Loss based Context Aggregator (CAVL) and Class Imbalance based Context Aggregator (CACI). The proposed Context Aggregator is evaluated on several different Covid-19 imaging classification datasets present on participating nodes. The evaluation results show that Context Aggregator performs better than standard Federating average Learning algorithms and FedProx Algorithm for Covid-19 image classification problems.