Weighted Correlation Attention Network-Based Dual-Phase Federated Learning Framework for Image Recognition

The recognition task of privacy-sensitive images requires effective federated learning (FL) models. However, there are challenges such as inconsistent data distributions, as well as the impact of different class correlations and class imbalance for classification across regions. Existing research on explicitly addressing class correlation and class imbalance within FL frameworks remains limited. Additionally, due to the heterogeneity and class imbalance across different clients, the global model’s parameter aggregation process faces issues of inconsistency, resulting in a significant discrepancy between the global model parameters and parts of local models’ parameters, which leads to poor classification performance on these clients’ local models. To address these challenges, we propose a novel dual-phase federated learning framework, Fd-WCAT, which incorporates a weighted correlation attention network for multi-label electricity image recognition. In this framework, each client trains a local model, constructing masked label correlation graphs. Then, the extracted label correlation embedding features are extracted from the graphs for a class imbalance-weighted classifier. To tackle the parameter inconsistency between local and global models during training, Fd-WCAT designs a dual-phase weighted loss function based on global-local parameter regularization. Each client computes its class imbalance coefficient and sends the local model parameters to the server. On the server side, local models are clustered into <italic xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">T</i> groups to ensure similarity in model parameters within each group. Intra-group aggregation is performed to generate <italic xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">T</i> temporary models with minimal inconsistencies. Then, imbalance weights for the prototype models are calculated based on the average imbalance coefficient of each group, and a global model is generated through imbalance weighted aggregation. Experimental results on multi-label benchmarks show that Fd-WCAT outperforms existing baseline models.