FedMAO: Federated Learning With Model Alignment and Two-Stage Optimization

Federated Learning (FL) aims to collaboratively train a global model using decentralized data residing on edge device clients, thereby effectively preserving user data privacy. The inherent challenge of non-independent and identically distributed (Non-IID) client data directly impairs global model performance. While existing methods have made progress in mitigating Non-IID issues, However, existing methods often overlook scenarios with significant global class imbalance, which constrains the global model’s performance. To address this gap, we propose FedMAO: a novel FL approach based on model alignment and two-stage optimization. FedMAO introduces two key innovations:(1) A local model alignment mechanism: This dynamically adjusts the classification decision boundaries of local models according to client-specific class distribution disparities, effectively mitigating output inconsistency across local models.(2) A two-stage optimization strategy with guide models: In the first stage, local models are updated under the constraint of a guide model to alleviate update divergence caused by client class imbalance. In the second stage, the guide model constraint is removed, enabling models to fully adapt to local data characteristics and leverage client-specific knowledge.Extensive experiments conducted on three public datasets demonstrate that FedMAO significantly outperforms state-of-the-art baselines, achieving accuracy improvements of up to 7.4%. These results validate the effectiveness of FedMAO addressing the dual challenges of global class imbalance and Non-IID client data.