A Multi-Level Approach for Class Imbalance Problem in Federated Learning for Remote Industry 4.0 Applications

TL;DR

The paper tackles class imbalance in federated learning for remote Industry 4.0 applications by proposing FedBal, a two-level approach that (i) uses a Tversky loss at the local level to mitigate false negatives and (ii) employs a dynamic, performance-based worker selection mechanism to robustly aggregate models at the global level. Through extensive experiments on non-IID and unbalanced data, FedBal demonstrates improved and more stable mean intersection-over-union ($mIoU$) compared to baseline FL methods, with gains up to 3–5% in final rounds and stronger resilience under varying class imbalance intensities. The work highlights practical strategies for robust FL in privacy-sensitive, resource-constrained, remote environments and sets the stage for deployment as a service plugin atop existing FL frameworks. Overall, FedBal offers a principled mechanism to harmonize local loss design and global client selection to enhance segmentation performance in oil spill detection and similar industrial remote sensing tasks.

Abstract

Deep neural network (DNN) models are effective solutions for industry 4.0 applications (\eg oil spill detection, fire detection, anomaly detection). However, training a DNN network model needs a considerable amount of data collected from various sources and transferred to the central cloud server that can be expensive and sensitive to privacy. For instance, in the remote offshore oil field where network connectivity is vulnerable, a federated fog environment can be a potential computing platform. Hence it is feasible to perform computation within the federation. On the contrary, performing a DNN model training using fog systems poses a security issue that the federated learning (FL) technique can resolve. In this case, the new challenge is the class imbalance problem that can be inherited in local data sets and can degrade the performance of the global model. Therefore, FL training needs to be performed considering the class imbalance problem locally. In addition, an efficient technique to select the relevant worker model needs to be adopted at the global level to increase the robustness of the global model. Accordingly, we utilize one of the suitable loss functions addressing the class imbalance in workers at the local level. In addition, we employ a dynamic threshold mechanism with user-defined worker's weight to efficiently select workers for aggregation that improve the global model's robustness. Finally, we perform an extensive empirical evaluation to explore the benefits of our solution and find up to 3-5% performance improvement than baseline federated learning methods.

Figures (8)

• Figure 1: A federated learning setup in fog federation. Multiple company share their fog systems to train oil spill detection DNN model where data security is preserved by federated learning.
• Figure 2: Comparison of minimum training loss for single machine (fog) learning and federated learning. The first four bars from the left represents single machine learning with four different data sources (S1, S2, S3, and S4), fifth bar is single machine learning with combined data source as centralized learning, and right most bar represent the federated learning. Minimum training loss for all cases are considered, and federated learning has the minimum loss comparing other single machine learning.
• Figure 3: The training loss of Unet model with oil spill data set in both central machine learning, and federated learning setup. The convergence trends of training losses are similar in both setup.
• Figure 4: A typical federated learning (FL) setup with different worker clients participating in training phase. Global model is broadcast or downloaded to the worker clients. Then workers train the global model with their local data and send the updated model to the central server for updating the global model.
• Figure 5: Federated learning training considering class imbalance and global convergence. Tversky loss is used in the training considering class imbalance. After training of each epoch, mean intersection over union (mIoU) is checked with a dynamic threshold for global convergence.