Edge Computing Model based on Federated Learning for COVID-19 Clinical Outcome Prediction in the 5G Era

Abstract

As 5G and AI continue to develop, there has been a significant surge in the healthcare industry. The COVID-19 pandemic has posed immense challenges to the global health system. This study proposes an FL-supported edge computing model based on federated learning (FL) for predicting clinical outcomes of COVID-19 patients during hospitalization. The model aims to address the challenges posed by the pandemic, such as the need for sophisticated predictive models, privacy concerns, and the non-IID nature of COVID-19 data. The model utilizes the FATE framework, known for its privacy-preserving technologies, to enhance predictive precision while ensuring data privacy and effectively managing data heterogeneity. The model's ability to generalize across diverse datasets and its adaptability in real-world clinical settings are highlighted by the use of SHAP values, which streamline the training process by identifying influential features, thus reducing computational overhead without compromising predictive precision. The study demonstrates that the proposed model achieves comparable precision to specific machine learning models when dataset sizes are identical and surpasses traditional models when larger training data volumes are employed. The model's performance is further improved when trained on datasets from diverse nodes, leading to superior generalization and overall performance, especially in scenarios with insufficient node features. The integration of FL with edge computing contributes significantly to the reliable prediction of COVID-19 patient outcomes with greater privacy. The research contributes to healthcare technology by providing a practical solution for early intervention and personalized treatment plans, leading to improved patient outcomes and efficient resource allocation during public health crises.

1. Introduction

The COVID-19 pandemic has caused profound devastation across diverse regions globally. While the majority of patients have recovered, those with compromised constitutions or underlying conditions have tragically succumbed to the disease [1]. The novel coronavirus continues to undergo mutations. The situation calls for the urgent development of sophisticated predictive models to enhance outbreak tracking, early detection, and rapid deployment of medical resources, underscoring the critical need for advanced data analysis and machine learning in combating these challenges.

The pandemic's response has been hampered by several hurdles, notably the safeguarding of user privacy and the collection of comprehensive, real-time data sets for predictive modeling. Traditional machine learning approaches, while foundational for predicting clinical outcomes, grapple with significant issues [2]-[9]. Centralizing sensitive COVID-19 patient data raises serious privacy and security concerns, risking patient confidentiality and exposing data to potential breaches.

Additionally, the non-IID nature of COVID-19 data, shaped by diverse factors like regional healthcare practices and genetic variations, challenges traditional models that rely on IID assumptions. This can lead to biased or inaccurate predictions. Moreover, the sheer volume of data required for robust predictions may overwhelm these models, particularly in regions lacking the digital infrastructure to support such data analysis [10].

In this context, the study introduces a novel model supported by Federated Learning (FL) and edge computing, leveraging decentralized machine learning to address these challenges. This model, utilizing the FATE framework known for privacy-preserving technologies like homomorphic encryption and secure multi-party computation, is designed to enhance predictive precision while ensuring data privacy and managing data heterogeneity effectively.

This research, by integrating FL with edge computing, contributes significantly to predicting COVID-19 patient outcomes more reliably and with greater privacy. It presents a robust framework that not only facilitates broader participation in model training but also remains precise across varying data conditions.

The remainder of this paper is structured as follows: Section 2 provides an overview of the applications of machine learning and edge computing in the prevention, diagnosis, and treatment of COVID-19. Section 3 presents a comprehensive overview of the FL-supported edge computing model, including its structure and specific details. Section 4 presents the results obtained from applying the proposed model to a real-world dataset. Finally, Section 5 concludes the paper with a summary of the key findings and contributions.

2. Related Work

The advent of AI technology has had a profound impact on the healthcare industry, particularly in the areas of prevention, diagnosis, and treatment of COVID-19. Traditional machine learning models, such as Decision Trees (DT), Support Vector Machines (SVM), and Logistic Regression (LR), have been widely adopted due to their low algorithmic complexity and cost-effective hardware implementation [11]-[13]. These models have been crucial in preventing further infections by predicting the number of confirmed and death cases of COVID-19, predicting mortality and severity by combining cardiac markers with demographic and clinical features, and providing early risk assessment and accurate prediction of clinical endpoints [2][3][11].

On the other hand, deep learning models have shown exceptional capability in various medical scenarios, including the detection and diagnosis of COVID-19 through imaging [14]-[16]. The integration of these models with 5G and edge computing has led to the development of frameworks that leverage low latency, scalability, data protection, and reliable local edge servers for effective management of epidemics. However, it is crucial to acknowledge that deep learning models typically require massive amounts of data and high hardware requirements for training, which can be a limiting factor, especially when dealing with multimodal features in cloud-based scenarios.

The revolutionary integration of 5G technology into healthcare has marked a new epoch, particularly with the advent of edge computing. This technology stands out for its ability to provide low-latency and highly reliable data processing capabilities, which are imperative in critical healthcare applications [17]-[20]. One of the salient features of 5G edge computing in healthcare is its transformative impact on medical imaging. It has significantly expedited the process of image processing and data transmission, ensuring that high-quality medical images are readily available to healthcare professionals in real-time. This has not only augmented diagnostic precision but also dramatically increased the efficiency of the entire diagnostic process, resulting in faster and more reliable patient care [21].

Moreover, edge computing has paved the way for the implementation of innovative healthcare solutions such as remote patient monitoring and telesurgery [22][23]. These applications require Ultra-Reliable Low-Latency Communication (URLLC), which edge computing is uniquely positioned to provide. By processing data closer to the source, it ensures minimal delay, thereby facilitating real-time decision-making and interventions that could be life-saving.

However, it is imperative to acknowledge the challenges that come with these innovations. The consistency of network stability is crucial, especially when dealing with critical healthcare applications where any delay or disruption can have dire consequences. Additionally, as healthcare data is extremely sensitive, the decentralized nature of edge computing networks opens up potential security vulnerabilities. Protecting patient data from unauthorized access and ensuring the integrity of the data becomes paramount, and this is an area that requires continuous attention and improvement [24].

FL has emerged as a pivotal innovation within distributed machine learning, facilitating the cooperative training of models across an extensive network of devices while preserving the privacy of sensitive data. This revolutionary approach has found applications across a diverse range of sectors, from finance [25][26] to industrial Internet of Things (IoT) systems [27]-[29], demonstrating its versatility and effectiveness in harnessing decentralized data sources for machine learning. Among these applications, the healthcare sector stands out as one of the most critical areas where FL is making a significant impact.

In healthcare, where the protection of patient confidentiality is of utmost importance, FL offers a transformative solution. By enabling the application of FL to electronic health records, significant advancements have been made in predictive analytics, leading to the development of more personalized and precise treatment plans. This method not only enhances the accuracy and effectiveness of medical care but also ensures the protection of patient data, preventing any breaches of privacy. The widespread adoption of FL in various domains underscores its potential to revolutionize data analysis and machine learning, with the healthcare industry benefiting immensely from its ability to safeguard patient information while improving treatment outcomes [30][31].

In addition to preserving data confidentiality, FL also addresses the critical issue of data silos in healthcare. It enables the utilization of diverse datasets from various institutions for model training without needing to centralize the data, thereby enhancing the robustness and generalizability of the models. This is crucial in a field like healthcare where patient populations can vary significantly across different regions and institutions.

The unprecedented challenges posed by the COVID-19 pandemic necessitate innovative and technologically advanced solutions. The integration of FL with edge computing stands out as a strategic approach, synergistically harnessing the capabilities of both technologies to establish a system that is both responsive and efficient in crisis management.

By capitalizing on the low-latency and reliable data processing features of edge computing, this integration markedly enhances the performance of FL models, especially in scenarios that demand real-time responsiveness. Ensuring the promptness and precision of insights garnered from these models is imperative, as it underpins informed decision-making, a critical component in effectively combating the COVID-19 pandemic.

3. Materials and methods

3.1 Datasets

In this study, a real-world COVID-19 dataset was obtained from a hospital in Jiangsu province, China. Participants were diagnosed with COVID-19 using RT-PCR detection of SARS-CoV-2 nucleic acid or serum SARS-CoV-2-specific IgM and IgG antibodies. Severe cases were defined based on the criteria outlined in the "COVID-19 Diagnosis and Treatment Scheme (Trial Version 6)", which included the following: (1) blood oxygen saturation≤93%, (2) respiratory rate≥30 times / minute, and (3) arterial blood oxygen partial pressure (PaO2) / oxygen concentration (FiO2)≤300 mmHg

The study conducted a thorough investigation, encompassing a diverse array of variables such as demographic details, symptoms, clinical test indicators (inclusive of CT indices), disease and treatment specifics, clinical attributes, and pathological factors. The utilization of SHapley Additive exPlanations (SHAP) values facilitated an evaluation of the individual contributions of features towards the model's efficacy, akin to assessing each player's impact on the final outcome in a game scenario. The SHAP values of the top 20 features, identified through the highest mean absolute SHAP values, are depicted as an extensive distribution in Fig. 1. In this representation, each dataset entry is denoted as a data point for its corresponding feature. The x-axis illustrates the feature's directional influence—positive or negative—on the model's predictive capability. The color scheme denotes feature values, scaled according to the range observed within the dataset.

Fig. 1. SHAP values for the top 20 features.

3.2 FL-supported edge computing model

The preceding introduction underscores the significant potential of 5G in preventing and treating COVID-19 across various scenarios, with edge computing playing a pivotal role in realizing its unique capabilities. Cloud computing is distant from terminal devices, resulting in considerable bandwidth consumption and notable transmission delays. In addition to the considerable expenses associated with transmission and delays, edge devices must transmit sensitive data to cloud servers for processing, posing potential risks to data confidentiality. The concept of edge computing has emerged as a response to the growing need for enhanced privacy and cost-effective communication. Through the utilization of distributed computing, edge computing facilitates the transfer of computing power to the data collection point, leading to decreased data transmission delays, improved user experience, and heightened levels of security and privacy. In comparison to cloud computing, edge computing offers the following advantages: 1) It reduces bandwidth consumption during transmission by performing small transaction fragmentation data processing tasks at the forefront. 2) It ensures data security through local processing of privacy-sensitive data. 3) Edge computing enhances system robustness, enabling uninterrupted local data computing and transmission even in the presence of abnormal external network conditions.

The concept of FL has made a significant contribution to addressing the challenge of preserving user data privacy during machine learning training. By adopting a decentralized approach with a centralized server and distributed clients, FL enables collaborative model training through the transmission of model parameters. Each data island maintains its original data locally, ensuring data privacy throughout the training process, while the centralized server enhances model performance by aggregating and sharing the model parameters from the distributed clients. The emergence of FL has significantly lowered the barriers to data sharing, enabling a broader participation of data owners in the model training process. Additionally, it effectively mitigates the risks associated with data leakage and substantially reduces the costs associated with data concentration. The specific process of integrating FL and edge computing is outlined below.

In the first step, task initialization is performed. Initially, the cloud defines the training objectives and data format for the task, such as predicting the outcomes of hospitalized COVID-19 patients. Based on the characteristics of the edge nodes, the FL training process and model parameters, including the model parameter weights 𝑃^𝑖 , where 𝑖 = {0, … , 𝑀}, are determined. Here, 𝑖 denotes the current iteration round, and when 𝑖 = 0, 𝑃⁰ represents the initialized parameters. The number of iterations 𝑀, the number of all participating nodes 𝑁 and local minibatch size 𝐵 in the training process are also determined. Once these specifications are finalized, the cloud distributes the task requirements and model parameters to all edge nodes.

The second step involves edge node training and model updating. Let 𝑛 = {1,2, … , 𝑁} represent the edge nodes. The cloud randomly selects 𝑛′ (where 𝑛′ ≤ 𝑁) edge nodes to distribute the model parameters. Upon receiving the global model parameters 𝑃^𝑖 , each edge node updates its local data to generate model parameters with local characteristics, denoted as 𝑃^𝑖_𝑛 . The training objective is to minimize the loss function 𝐿(𝑃^𝑖_𝑛).

\(\begin{align}P_{n}^{i^{*}}=\underset{P_{n}^{i}}{\arg \min } L\left(P_{n}^{i}\right)\end{align}\)       (1)

The local dataset 𝐷_𝑛 is initially partitioned into 𝑗 subsets 𝐵_𝑗 = {1,2, . . . ,𝑗}, based on local minibatch size 𝐵. For each subset 𝐵_𝑗, iterate 𝑇 rounds, updating the parameter 𝑃^𝑖_𝑛 for each 𝑏 ∈ 𝐵_𝑗.

𝑃^𝑖_𝑛 ← 𝑃^𝑖_𝑛 − 𝜂∇𝐿(𝑃; 𝑏)       (2)

Where 𝜂 represents the learning rate, and ∇L represents the gradient. After completing the training, each edge node will upload the trained parameters 𝑃^𝑖_𝑛 to the cloud.

In the third step, the cloud integrates and updates the parameters. The cloud consolidates the parameters uploaded by 𝑛′ edge nodes and updates the global parameters. At the cloud, the optimization objective is to minimize the global loss function 𝐿(𝑃^𝑖 ).

\(\begin{align}L\left(P^{i}\right)=\frac{1}{n^{\prime}} \sum_{i=1}^{n^{\prime}} L\left(P_{n}^{i}\right)\end{align}\)       (3)

The cloud aggregates the parameters 𝑃^𝑖_𝑛 from all edge nodes to compute the new global parameter 𝑃^𝑖+1.

\(\begin{align}P^{i+1}=\frac{1}{\sum_{n \in N} D_{n}} \sum_{n=1}^{n^{\prime}} D_{n} P_{n}^{i}\end{align}\)       (4)

After multiple iterations of steps 2 and 3, the loss function converges or reaches the desired objective, signaling the completion of model training.

As illustrated in Fig. 2, the FL-supported edge computing model distinguishes itself from traditional machine learning models by effectively leveraging the advantages of both edge computing and FL. By utilizing the computational capabilities of edge nodes for local data model training, the system optimizes resource utilization due to each node's limited data volume. Additionally, FL's parameter synchronization aligns the computational capabilities of local and cloud models, enhancing the delivery of AI services directly at the edge. This approach not only significantly reduces system latency and bandwidth consumption but also ensures rapid and efficient services for COVID-19 patients. Importantly, with patient data remaining local during model training, privacy concerns are mitigated, encouraging greater participation from medical practitioners and thereby enriching the dataset for developing more robust models.

Fig. 2. Framework of FL-supported edge computing model for COVID-19.

Building upon this foundation, the study meticulously recreated the technological environment of various departments within a hospital in Jiangsu Province, China. The deployment of the FATE framework on cloud-based virtual machines within a designated sandbox environment enabled the analysis of partial COVID-19 case data. This setup demonstrated the FL-supported edge computing model's capability for sophisticated data analysis and machine learning model training within a secure, decentralized healthcare setting, highlighting its potential to revolutionize data-driven decision-making in critical sectors.

The integration of FATE, a leading collaborative learning platform designed for industrial-scale applications, further empowers organizations and institutions to develop machine learning models while prioritizing data privacy. Developed by Webank's AI department and made open-source in 2019, FATE now contributes to the Linux Foundation under the Apache 2.0 license. Its key features include a distributed architecture that supports various FL configurations, advanced security protocols for user privacy protection, and performance optimization tools for efficient large-scale model training. These attributes make FATE an ideal framework for implementing our FL-supported edge computing model, reinforcing the model's capabilities in managing complex, privacy-sensitive datasets across healthcare and other sectors.

4. Experiment results

In this study, the proposed model was compared with seven independent machine learning models: SVM, LR, DT, Catboost, Random Forest, Extra Trees, and Ada Boost. This rigorous approach was underpinned by clinical data meticulously gathered from a specified hospital in Jiangsu Province, China, ensuring that the process adhered to the highest clinical standards. The data collection process received rigorous evaluation and approval by the hospital's ethics committee, under the ethical approval designation number 2020-SR-106, with all pertinent documents related to this clearance submitted as part of the supplementary materials.

The dataset, comprising data from 303 patients confirmed to have COVID-19 through diagnostic methods such as the RT-PCR test and the assay for serum SARS-CoV-2-specific IgM and IgG antibodies, was categorized into two groups based on their admission detection indicators and initial severity classification: severe outcome (n=132) and mild outcome (n=171). These diagnostic approaches ensure a comprehensive identification of infected individuals, encompassing both active infections and previous exposure to the virus. Furthermore, the classification of severe COVID-19 cases adheres rigorously to the criteria outlined in the "COVID-19 Diagnosis and Treatment Scheme (Trial Version 6)", which include detailed clinical indicators such as respiratory distress, oxygen saturation levels, and other critical parameters. This methodology allows for a precise segmentation of the study population based on the severity of their condition, facilitating a nuanced analysis of the disease's impact and the effectiveness of various treatment approaches.

Fig. 3(a) to (d) comparatively assessed the predictive precision of the machine learning models in edge computing environments, including a detailed simulation that evaluated the efficacy of the FL-supported model. Additionally, datasets from all nodes were aggregated to enhance the proposed model, and its performance was compared against the averaged results of the seven models, as depicted in Fig. 4. The assessment used performance metrics such as the F1 score, Area Under Curve (AUC), and precision, employing a 5-fold cross-validation method to ensure robust evaluation.

Fig. 3. F1 score, precision, and AUC comparison of DT, SVM, LR, Catboost, Random Forest, Extra Trees, Ada Boost and Proposed model. (a) shows the validation results after training all models on edge node 1. (b) shows the validation results after training all models on edge node 2. (c) shows the validation results after training all models on edge node 3. (d) shows the validation results after training all models on edge node 4.

Fig. 4. Average validation results of all models trained on edge nodes 1 to 4 and the proposed model's validation on aggregated data from all four nodes.

The proposed model exhibited remarkable performance, notably outperforming SVM, LR, and DT in precision metrics. Its performance was competitive against sophisticated ensemble methods such as Catboost. This demonstrated its efficacy in balancing sensitivity and specificity, particularly in the nuanced context of COVID-19 severity classification.

A critical observation was the proposed model’s exceptional efficacy when leveraging the amalgamated dataset from all nodes. It achieved an F1 score of 0.87, precision of 0.78, and an AUC of 0.91, highlighting its scalability and adaptability, as well as its enhanced predictive performance over both traditional and ensemble classifiers. Further analysis revealed the varied performance of models like SVM, LR, and DT across different nodes could be attributed to their distinct reactions to the feature space and distribution. For instance, SVM and LR may struggle with non-linear separations in high-dimensional spaces, while DT is prone to overfitting amidst complex structures. Conversely, ensemble methods, by aggregating decisions from multiple weak learners, exhibit superior performance, indicating robust handling of data’s complexity and diversity. The standout performance of the proposed model in an aggregated setting is likely attributed to the FL framework’s capacity to harness diverse datasets without compromising privacy. Its design, tailored for the decentralized nature of edge computing, facilitates more generalized representations.

Model efficacy variations across nodes (E1 to E4) reflect the heterogeneity in data structure and distribution, impacting model suitability. The complexity of data structures, including feature count, correlation, and non-linear relationships, significantly influences model performance. Ensemble models like Catboost and Random Forest excel in managing complex, high-dimensional data structures through their aggregate learning approach.

The primary advantage of the proposed model lies in its ability to enhance data utilization efficiency while safeguarding privacy. By processing data locally and sharing only model updates, this approach reduces privacy concerns, promoting data sharing among individuals and medical entities. This is particularly crucial in addressing global health emergencies like COVID-19, where diverse data is essential in developing effective strategies.

Building on the extensive performance evaluation of the proposed model alongside seven independent machine learning models, the next phase of the analysis delves into the computational complexity. The comparison, based on training times, is summarized in the Table 1:

Table 1. Comparative analysis of training times across models

The results indicate that the proposed model, which incorporates privacy-preserving technologies such as homomorphic encryption and secure multi-party computation within the FATE framework, incurs a significant computational overhead in terms of encrypted data communication and computation during the training process. However, it is important to note that the overhead associated with these data processing techniques does not scale linearly. The distributed nature of the FATE platform can effectively mitigate this overhead, especially when handling large datasets. This mitigation strategy significantly narrows the gap in training times between the proposed model and other models, underscoring its efficiency and feasibility in processing large-scale datasets.

Based on an analysis from some of the authors' experiences treating COVID-19 patients on the front lines, it was noted that during the initial outbreak, inconsistencies in medical conditions led to variations in the feature data collected by different hospitals (edge nodes). For instance, community hospitals, due to limited resources, struggled to provide CT scans for patients. Consequently, building predictive models with limited feature data posed a significant challenge. This study attempted to integrate edge computing with FL, taking into account the use of different features by edge nodes for training and prediction. The dataset was divided into two edge nodes, E1 and E2, with only age and gender as overlapping features, as shown in Table 2. Such a division not only caters to edge nodes with varying medical capabilities but is also applicable to model training scenarios across different countries. Given that different countries often adopt varied medical standards, it was challenging to rapidly standardize all features during the early stages of the outbreak, resulting in significant feature disparities. After dividing into E1 and E2 nodes, 80% of the data was used for training the models, while the remaining 20% was allocated for testing.

Table 2. Feature in E1 and E2

Fig. 5 demonstrates the model's performance on the test set, where the "proposed model" represents the performance of the model after federated training on the global dataset, while "proposed model on E1" indicates the performance after training and prediction on the E1 dataset alone. Similarly, "proposed model on E2" signifies the model's performance after training and prediction exclusively on the E2 dataset. Fig. 5 illustrates that, despite the presence of diverse features, the proposed model maintains performance, resulting in a substantial enhancement in the performance of the cloud-based model. Compared to training the E1 and E2 nodes separately, the F1 scores for the E1 and E2 nodes increased by 0.05 and 0.16, respectively, with a corresponding 0.03 improvement in precision. The performance of the proposed model tends to improve with the incorporation of more training data, ensuring privacy preservation. This contrasts with traditional machine learning models, which often suffer from poor generalization due to insufficient training data. In such scenarios, the performance of the proposed model surpasses that of traditional machine learning models. However, the proposed model's drawback is the bandwidth cost incurred from transmitting model parameters.

Fig. 5. F1 score, AUC, and precision comparison of proposed model on E1 and E2

In the proposed method, the FL framework utilizes XGBoost as the local algorithm, which constructs a model comprising 80 trees, with each tree having a depth and leaf node count of 3. The resultant single model approximates a size of 300KB. Considering the architecture comprises four sub-nodes, the communication cost for a complete round of parameter updates per node is calculated as 300KB * 4 * 2, equating to 2.4MB. Factoring in the 78 iterations required for model training, the cumulative communication cost is 187.2MB. When evaluated in the context of 5G technology, which is characterized by high data transfer rates and minimal latency, this incurred cost is relatively negligible, affirming the transmission's feasibility and efficiency.

5. Conclusion

In conclusion, the study has successfully demonstrated the potential of integrating 5G technology with FL to create a robust predictive model for COVID-19 clinical outcomes. This integration not only enhances the precision of patient prognosis during hospitalization but also addresses critical privacy concerns by ensuring data remains localized, a feature that is particularly valuable in the sensitive healthcare domain.

The model's ability to generalize across diverse datasets, even in the presence of varying node features, underscores its adaptability and resilience in real-world clinical settings. This adaptability is further amplified by the use of SHAP values, which streamline the training process by identifying influential features, thus reducing computational overhead without compromising predictive precision.

The research contributes to the field of healthcare technology by providing a practical solution that can aid healthcare professionals in making informed decisions for early intervention and personalized treatment plans. This, in turn, can lead to improved patient outcomes and more efficient resource allocation during public health crises.

Furthermore, the study offers policymakers a powerful tool for developing targeted strategies to combat COVID-19 and future pandemics, by providing insights into the effectiveness of various treatment approaches based on predictive analytics.

However, the study's limitations, particularly the challenges posed by data disparities among nodes, highlight the importance of future research in data balance and training weight optimization. Looking forward, the work paves the way for more sophisticated, efficient, and patient-focused healthcare solutions, marking a significant step in the integration of 5G and FL in healthcare.