Understanding the Robustness of Distributed Self-Supervised Learning Frameworks Against Non-IID Data

2026-07-02Machine Learning

Machine Learning
AI summary

The authors studied how different ways of training AI on decentralized data handle situations where the data is not evenly mixed (non-IID). They found that a method called Masked Image Modeling (MIM) is better at dealing with this uneven data than another method called Contrastive Learning (CL). They also showed that networks with better connectivity improve training robustness, meaning federated learning works just as well as other decentralized methods. To apply their findings, the authors created a new training loss called MAR loss, which helps align local and global information. They tested their ideas thoroughly and confirmed both their theory and the usefulness of MAR loss.

Distributed Self-Supervised LearningNon-IID DataMasked Image ModelingContrastive LearningFederated LearningDecentralized LearningNetwork ConnectivityPre-trainingLocal-to-Global AlignmentMAR Loss
Authors
Xuanyu Chen, Nan Yang, Shuai Wang, Dong Yuan
Abstract
Recent research has introduced distributed self-supervised learning (D-SSL) approaches to leverage vast amounts of unlabeled decentralized data. However, D-SSL faces the critical challenge of data heterogeneity, and there is limited theoretical understanding of how different D-SSL frameworks respond to this challenge. To fill this gap, we present a rigorous theoretical analysis of the robustness of D-SSL frameworks under non-IID (non-independent and identically distributed) settings. Our results show that pre-training with Masked Image Modeling (MIM) is inherently more robust to heterogeneous data than Contrastive Learning (CL), and that the robustness of decentralized SSL increases with average network connectivity, implying that federated learning (FL) is no less robust than decentralized learning (DecL). These findings provide a solid theoretical foundation for guiding the design of future D-SSL algorithms. To further illustrate the practical implications of our theory, we introduce MAR loss, a refinement of the MIM objective with local-to-global alignment regularization. Extensive experiments across model architectures and distributed settings validate our theoretical insights, and additionally confirm the effectiveness of MAR loss as an application of our analysis.