FedMAGE: Enhancing Federated Self-Supervised Learning under Heterogeneous Data with Generative Learning

Federated self-supervised learning enables representation learning from decentralized unlabeled data, but non-IID client distributions can induce inconsistent local representations and weak global coordination. Existing methods mitigate this issue by communicating compact representation summaries, such as prototypes or centroids, yet such compression retains only partial information about client data. Generative learning offers a promising alternative through richer sample-level surrogates, but its representation-level effect at both clients and the server remains unclear. In this paper, we theoretically characterize generative enhancement in federated masked image modeling and show that local generative modeling strengthens the recovery of globally relevant semantic directions under heterogeneous data. We further show that semantically consistent generated data improve global representability beyond conventional parameter aggregation. Building on these insights, we propose FedMAGE, a theory-guided framework that combines masked reconstruction and variational generation through a shared client encoder, while using quality-aware client generators to refine the aggregated global model. Extensive experiments show consistent gains over representative baselines across varying levels of data heterogeneity, while ablation studies confirm the contribution of the main components of FedMAGE.