Mitigating Dimensional Collapse and Model Drift in Non-IID Data of Federated Learning

One of the key challenges in federated learning is addressing the non-independent and identically distributed (Non-IID) data among parties, which can lead to divergence of local model parameters and decreased convergence accuracy of the global model, along with serious dimension collapse issues. In this paper, we introduce VICON (Variance-Invariance-Covariance model Contrastive Learning), a method to prevent dimension collapse issues. In local training, particularly, a regularization technique is employed to foster orthogonal feature representations across dimensions and to sustain the variance of individual embedded dimensions above a predefined level. Complementary to this, contrastive learning methodologies are utilized to cluster-like instances while dissociating divergent ones, further enhancing discriminative capabilities. It helps control the model's parameter norm and adapt it to high-dimensional data, reducing information loss and aligning local models with the global optimization objective in federated learning to minimize bias and collapse. Extensive experiments show that VICON performs better than other algorithms such as federated averaging (FedAvg), federated proximal optimization (FedProx), and model-agnostic federated learning (Moon). Compared to the Moon algorithm, it not only improves accuracy by 2.2% to 3.7%, but also enables efficient communication, achieves higher accuracy, and remains robust when dealing with imbalanced data and uncertain local updates.