1 Introduction↩︎

Deep learning advancements have been driven by large-scale datasets, as seen in the training of LLMs, which require billions of data points [1], [2]. However, real-world data is often decentralized, such as surveillance footage from distributed security cameras. This abundance of unlabeled distributed data has spurred interest in distributed self-supervised learning (D-SSL) [3], [4], which extends self-supervised learning (SSL) to decentralized settings. Existing D-SSL frameworks can generally be distinguished in two aspects: differing by the adopted self-supervised learning (SSL) method or by the applied distributed framework. Self-supervised learning (SSL) is a widely used technique to learn representations without human-labeled annotations by solving pretext tasks that generate supervisory signals from raw data [5]. Depending on the approach used to generate supervisory signals, SSL methods are broadly categorized into Contrastive Learning (CL) and Masked Image Modeling (MIM) [6], [7], with representative methods like SimSiam [8] and MAE [9]. On the other hand, federated learning (FL) and decentralized learning (DecL) are two main frameworks in training models with distributed data [10], [11]. FL aggregates local models via a central server [3], [12], while DecL enables direct inter-client communications for aggregating models, enhancing privacy and avoiding the dependence on the central server [13], [14].

One unique challenge of D-SSL research is handling highly heterogeneous data on clients. Distributed data among multiple clients are normally non-independent and identically distributed (non-IID), leading to performance degradation [15]. To tackle this challenge, previous works proposed advanced D-SSL algorithms with robustness to heterogeneous data. Notable examples include FedU [3], Orchestra [16], and L-DAWA [17]. However, despite continuous algorithmic innovation, there is still a lack of theoretical understanding of this heterogeneity problem. For example, FedU was designed within the FL framework, but how would its robustness to non-IID data change if deployed in a DecL framework without coordination from the server? Similarly, state-of-the-art D-SSL algorithms are primarily based on CL, while the adaptation of MIM methods to distributed settings remains under-explored. Could D-SSL based on MIM offer greater robustness to non-IID data than CL-based methods? These confusions converge into a fundamental research question affecting the advancement of D-SSL:

To address this question, this paper aims to provide a theoretical understanding of how different D-SSL frameworks behave under heterogeneous data. We construct mathematical models in a simplified non-IID setting and rigorously analyze the representability of local and global representations learned by these algorithms. Our analysis reveals two key insights: (i) D-SSL algorithms based on Masked Image Modeling (MIM) are inherently more robust than those based on Contrastive Learning (CL), although their robustness still degrades under severe divergence between local and global distributions; and (ii) the robustness of decentralized SSL improves with the average connectivity of the network, which suggests that decentralized SSL is only as robust as federated D-SSL in the limited case of full connectivity (i.e., a fully connected network). Building on these insights, we also explore how theoretical results can inform algorithmic design. As an illustration, we refine the MIM objective with the additional alignment regularization, which we call MAR loss, to encourage local-to-global representation consistency. Finally, we conduct extensive experiments on ResNet [18] and Vision Transformer (ViT) [19] across a variety of distributed settings and benchmark datasets to validate our theoretical findings and to demonstrate the usefulness of MAR loss as a practical example.

In summary, our main contributions are listed below:

1. We develop a rigorous theoretical analysis of distributed self-supervised learning (D-SSL) under non-IID data, showing that MIM-based D-SSL is inherently more robust than CL-based D-SSL.

2. We establish the relationship between network connectivity and robustness, proving that decentralized SSL benefits from higher connectivity and that federated SSL is no less robust than decentralized SSL.

3. We introduce MAR loss as an illustrative case study demonstrating how our theoretical results can guide algorithmic design, by refining the MIM objective with alignment regularization.

4. We conduct extensive experiments across model architectures and distributed settings, which validate our theoretical insights and further confirm the effectiveness of MAR loss.

2 Related Work↩︎

Self-Supervised Learning. Self-supervised learning (SSL) leverages unlabeled data by generating pseudo labels from raw inputs to learn meaningful representations [5]. Vision-based SSL methods are typically categorized into contrastive learning (CL) and masked image modeling (MIM) [6], [7]. CL learns representations by maximizing the similarity between positive pairs (i.e., similar data points created by data augmentation) and minimizing it between negative pairs (i.e., data pairs created by other data points) [20], [21]. Recent methods like SimSiam [8] and BYOL [22] advance the original contrastive loss by removing terms related to negative pairs, which improves stability and reduces batch size dependence. MIM, in contrast, randomly masks out patches of input images and predicts the missing parts, learning representations through a reconstruction loss [9], [23]–[25]. Although different in formulation, recent studies have shown that many MIM methods have close connections to CL (i.e., their objectives can be directly re-formulated as contrastive loss [7], [26]). In this work, we aim to figure out which SSL paradigm is inherently more robust against data heterogeneity.

Distributed Learning. Distributed learning enables collaborative model training across multiple clients without sharing data. Two dominant frameworks in this area are: federated learning (FL), which uses a central server to coordinate and aggregate models [12], and decentralized learning (DecL), where clients exchange models locally with neighbors [13], [14]. While FL is more widely adopted [27] for better convergence and training effectiveness, DecL offers benefits in scalability and privacy. Recent studies have started comparing these two frameworks [28], [29]. For example, Sun et al. explored which leads to better generalization and the impact of network architecture on generalization [11]. However, the relationship between network architecture and the non-IID robustness in distributed settings is still unclear. Our work addresses this gap by providing both theoretical analysis and empirical findings to clarify this relationship.

Distributed SSL. Distributed SSL (D-SSL) integrates SSL with distributed frameworks to leverage unlabeled, decentralized data while preserving privacy [3], [30]. A core challenge is learning robust representations under data heterogeneity [15]. Prior work has primarily focused on algorithmic solutions such as FedU [3] and L-DAWA [17]. Although some studies also provide theoretical analyses, their purpose is to demonstrate the validity of the proposed algorithms rather than to advance the understanding of the robustness variance between different D-SSL frameworks [16], [31]. The most relevant theoretical work is by Wang et al., who showed that SSL is more robust than supervised learning in distributed settings [4]. Unfortunately, their study only analyzed a specific case of D-SSL where CL is combined with FL and did not extend it to other types of D-SSL frameworks. In contrast, our work delves deeper into these differences, shedding light on understanding the insensitivity of various D-SSL approaches under heterogeneous conditions.

3 Problem Setup↩︎

To provide theoretical insights on understanding this central question, we first introduce our problem setup about distributed training and D-SSL with heterogeneous data.

3.1 Distributed Training↩︎

Distributed Setting. Consider a distributed scenario consisting of a connected network of N clients, represented as a graph \(\mathcal{G} = (\mathcal{V}, \mathcal{E})\), where \(\mathcal{V}\) is the set of clients and \(\mathcal{E}\) is the set of edges denoting direct communication links between clients. The connectivity of the graph is captured by a matrix \(A \in \mathbb{R}^{N \times N}\), referred to as the adjacency matrix, where \(A_i\) denotes the set including client \(i \in [N]\) itself and its neighbors shown by \(\mathcal{E}\), \(|A_i|\) represents the size of this neighborhood set or the connectivity of client \(i\), and \(|\bar{A}| = \frac{1}{n} \sum_{i=1}^{n}(|A_i|)\) is the average connectivity. Hence, distributed training conducted through DecL satisfies \(\forall i \in [N], 2 \leq |A_i| \leq N\). In contrast, FL relies on a central server that aggregates local models from all clients and broadcasts the global model back to them in each round, as in FedAvg [12]. This architecture effectively enables every client to communicate with all others through the server, which corresponds to a fully connected decentralized topology where \(\forall i \in [N], |A_i| = N\). A more formal specification of the graph structure and the mixing-weight conditions for this distributed setting is provided in Appendix 8.6.1.

Objective of Distributed Optimization. To utilize different clients to learn useful representations, distributed training generally optimizes the below global objectives: \[\label{eq95distributed95global} W^*_{Dec} = \min_W \frac{1}{N} \sum_{i=1}^N \frac{1}{|A_i|} \sum_{j \in A_i} \mathcal{L}_j(W_j); \quad W^*_{Fed} =\min_W \frac{1}{N} \sum_{i=1}^N \mathcal{L}_i(W_i)\tag{1}\] where \(\mathcal{L}_j\) is the objective of local SSL on client \(j\), \(W^*_{Dec}\) and \(W^*_{Fed}\) denote the global objective of DecL and FL, respectively. In particular, at each iteration of DecL, each client conducts local updates using the local dataset and aggregates the updated local model with those from neighbors [13]. For generating the global model for downstream tasks, there will be an additional aggregation on all local models after all iterations. Differently, the optimization of FL involves each round of model aggregation only on the central server [12]. Then, the server broadcasts the global model to all clients for the next round of training. Note that the FL framework does not need another aggregation between all local models since the updated global model on the server can be used directly for fine-tuning.

3.2 Rigorous Analysis of D-SSL on a Simplified Non-IID setting↩︎

Non-IID Client Data. D-SSL involves all clients collaboratively training a global model by leveraging their local unlabeled datasets \(\{D_i\}^{N}_{i=1}\) and communicating over the graph \(\mathcal{G}\). Since sharing data is prohibited to protect privacy, the heterogeneity across these distributed data sources generally leads to a performance drop in many distributed applications [3], [12]. Two common types of data heterogeneity are: feature heterogeneity and label heterogeneity [15]. In this paper, we follow previous works [4], [32] to model a simplified but formal label non-IIDness between local datasets as follows.

The global data distribution \(D = \bigcup_{i=1}^N D_i\) across clients is assumed to contain unlabeled data from \(2N\) classes. For the dataset on client \(i\), the local data distribution \(D_i\) is constrained and imbalanced on three classes, with most samples belonging to classes \(2i-1\) and \(2i\), while the remaining very few samples come from the class \(h_i \in [2N] \setminus \{2i-1, 2i\}\). Specifically, for a sufficiently large positive integer \(d>0\), let \(x \in \mathbb{R}^d \sim D_i\) be the data points in the local dataset and \(e_1, \ldots, e_d\) be the standard unit-norm vectors of the d-dimensional Euclidean space. For class \(2i-1\), we set \(x^{(2i-1)} = e_i - \Sigma^{N}_{k \neq i, k=1}q^{(2i-1,k)}\tau e_k + \mu \xi^{(2i-1)}\), where \(\tau\) and \(\mu\) are two positive hyperparameters, \(q\) is sampled uniformly from \(\{0,1\}\) and \(\xi \sim \mathcal{N}(0, I)\) from Gaussian distribution. Likewise, for class \(2i\), we define \(x^{(2i)} = -e_i - \Sigma^{N}_{k \neq i, k=1}q^{(2i,k)}\tau e_k + \mu \xi^{(2i)}\). The size of the data from classes \(2i-1\) and \(2i\) are equal and both grow in polynomials of \(d\). For infrequent class \(h_i\), the samples are generated as: \(x^{(h_i)} = e_{h_i} + \mu \xi^{(h_i)}\) and the amount of data is sublinear in \(d\), denoted as \(O(d^{\alpha})\) with \(\alpha \in (0,1)\)). Furthermore, we assume all \(N\) local datasets to have an equal total number of samples, i.e., \(|D_1| = |D_2| = \ldots = |D_N|\). To facilitate understanding, we provide an overview of this non-IID data distribution in Figure 1. Next, we consider CL and MIM as two main paradigms of SSL and formulate CL and MIM, respectively.

CL Formulation. For CL, we adopt the more advanced Simsiam [8] which trains with only the positive pairs \((g_a(x), g_b(x))\), where \(g_a(\cdot)\) and \(g_b(\cdot)\) are random augmentations drawn from SimSiam’s augmentation policy (e.g., Gaussian noise, flipping). Consider a linear embedding function \(f_{W}(x) = Wx\), where the weight matrix \(W\) satisfies \(W \in \mathbb{R}^{c \times d}\) and \(c \geq 2N\) according to the distributed settings, the local objective on client \(i\) is defined as: \[\label{eq95simsiam95loss} \begin{align} \mathcal{L}_{CL} &= -\mathbb{E}_{x \sim D_i} [|(W(g_a(x)))^{\intercal}(W(g_b(x)))] + \frac{1}{2}||W^{\intercal}W||^2_F. \end{align}\tag{2}\] Eq.(2 ) captures the SimSiam loss by utilizing the negative inner product \(\langle a,b \rangle\) to measure the distance between the positive pairs. This objective also excludes a feature predictor for simplicity and includes a regularization term \(||W^{\intercal}W||^2_F\) for more mathematically tractable, similar to previous works [4], [32]. Note that Eq.(2 ) stands for a general form of Simsiam loss due to the wide class of augmentation functions [5]. For a detailed and tractable theoretical exploration, we consider the linear formulation of data augmentation and further differ CL by the similarity between \(g_a(\cdot)\) and \(g_b(\cdot)\). In particular, for the case where the positive pairs are generated by similar augmentations, the objective becomes \(\mathcal{L}_{CL} = -\mathbb{E}_{x \sim D_i} [(W(x+\xi))^{\intercal}(W(x+\xi'))] + \frac{1}{2}||W^{\intercal}W||^2_F\), where \(\xi, \xi' \sim \mathcal{N}(0,I)\) are random noise sampled IID from the Gaussian distribution. On the other hand, when \(g_a(\cdot)\) and \(g_b(\cdot)\) are different, we define the loss as \(\mathcal{L}_{CL}' = -\mathbb{E}_{x \sim D_i} [(W(x+\xi))^{\intercal}(W(Hx))] + \frac{1}{2}||W^{\intercal}W||^2_F\), where \(H \in \mathbb{R}^{d \times d}\) denotes a linear image transformation (e.g., rotation, translation, etc.). The formal conditions on \(H\) are given in Appendix 8.6.1.

MIM Formulation. For MIM, a random binary mask \(m \in \{0,1\}^d\) (created by uniformly sampling 0 with probability \(p\), i.e., mask ratio) is applied to partition the input \(x\) into two complementary views: the unmasked part \(x_1 = x \odot m\) and the masked part \(x_2 = x \odot (1-m)\) satisfying \(x_1 + x_2 = x\). Then, we train an encoder-decoder model \(f = f_d \circ f_e\), where the encoder \(f_e\) encodes the input \(x_1\) to a latent representation \(z=f_e(x_1)\), and the decoder \(f_d\) decodes \(z\) back to pixel space to reconstruct the masked part \(x_2\). Hence, considering a linear encoder and decoder with embedding matrix \(W_e \in \mathbb{R}^{c\times d}\) and \(W_d \in \mathbb{R}^{d\times c}\), the local objective of MIM is given by \[\label{eq95MIM95loss} \begin{align} \mathcal{L}_{MIM} &= \mathbb{E}_{x \sim D_i} \mathbb{E}_{x_1, x_2 | x} ||f_d(f_e(x_1)) - x_2||^2 = \mathbb{E}_{x \sim D_i} ||W_d W_e (x \odot m) - (x \odot (1-m))||^2, \end{align}\tag{3}\] where the mean square error (MSE) loss is utilized to enforce the reconstructed image to be similar to the original image, and \(\odot\) denotes the Hadamard product. Recent studies have focused on the connection between MIM and contrastive losses and found that the MIM reconstruction objective admits an alignment between the masked and unmasked parts [7], [26]. Based on these results, we adopt an alignment-style formulation of Eq.(3 ) with \(W := W_e \in \mathbb{R}^{c\times d}\): \[\label{eq95new95MIM95loss} \begin{align} \mathcal{L}_{MIM} &= - \mathbb{E}_{x \sim D_i} [(W (x \odot m))^{\intercal}(W(x \odot (1-m)))] + \frac{1}{2}||W^{\intercal}W||^2_F, \end{align}\tag{4}\] which implicitly aligns the masked and unmasked views in the embedding space. The regularization term \(||W^{\intercal}W||^2_F\) is also introduced to ensure a well-posed quadratic form and improve the traceability.

4 Theoretical Insights↩︎

In this section, we use the above problem setup to model different D-SSL frameworks and compare their robustness to heterogeneous data. Differences in applied SSL and network architecture lead to distinctions in learned representations, which can be further explored to determine variance in robustness. Due to page limitations, the complete proof for our analysis is provided in Appendix 8.6.

4.1 Analysis of Representations Learned by D-SSL↩︎

We begin our analysis with the definition of the representability of the learned representation.

The intuition behind this definition is that for any input vectors \(x \in \mathbb{R}^d\), the learned feature space should have a good representation of the standard basis vectors, \(e_1, \ldots, e_d\), to perform well. In particular, these basis vectors should have large projections onto the feature space. The introduction of the representability vector allows us to quantitatively assess the feature space learned by different D-SSL frameworks. Similar definitions and notations have also been used in previous works studying the feature space of SSL [4], [32]. Based on this definition and the above problem setup, we establish the following theorem for D-SSL based on MIM pre-training.

Theorem 1 shows the status of the feature space learned by distributed MIM with different objectives (i.e., local vs decentralized global vs federated global). Note that for each provided representability vector, we find a unique lower bound and a shared upper bound (considering \(\sum^d_{j=1} (e_k^{\intercal}e_j)^2 = 1\)). The distance between the lower and upper bound states how much the learned representation fluctuates in the c unit directions, \(e_1, \ldots, e_c\), associated with data generation. Therefore, the smaller the distance, the less sensitive the representation space is to the non-IID distribution of local datasets on clients. In other words, the corresponding D-SSL is more robust to heterogeneity.

By a similar proof, we derive the representability vectors for D-SSL with CL pre-training as follows.

Theorem 2 demonstrates that the local and global feature spaces learned by distributed CL are distinct from those learned by distributed MIM. However, it is not obvious which feature spaces hold a smaller gap between the lower and upper bounds. To determine which type of pre-training is less sensitive to data heterogeneity, we further compare their global feature spaces learned in DecL and FL framework, respectively, and summarize the results in the following theorem.

4.2 MIM is inherently more Robust than CL with Heterogeneous Data↩︎

The main intuition for the greater robustness (or smaller sensitivity) of distributed MIM is that CL learns representations from aligning features of the positive pair generated from the original data through data augmentation, whereas MIM aligns features of the reconstructed and the raw data to learn representations. Although the applied augmentation generally does not lead to a change in data labels [8], [20], the output is still a different image. In contrast, the masking operation splits the original image into the masked and unmasked parts, but a portion of the original data is retained in both parts. As a result, CL learns a local representation with greater randomness, and that additional randomness is also biased by local labels. Considering that data heterogeneity already exists among clients, the global representation learned by distributed CL is less uniform than that learned by distributed MIM.

4.3 Impact of the Average Client Connectivity on Non-IID Robustness↩︎

Next, we shift our focus to another dimension that distinguishes D-SSL algorithms and address the question: how does the network architecture affect the robustness of the feature space learned by D-SSL? The tool for solving this question is again the bounds of the representability vector. For the DecL setup where clients directly communicate with their direct neighbors, Theorem 1 and 2 have implicitly shown the answer.

Corollary 1 also implies that the robustness of D-SSL conducted in a federated setup should be no worse than in a fully decentralized network. Consider the best case of the network topology, where each client can communicate with all other clients in the network. In this case, each client receives a model aggregated by the local models from all clients, which is exactly the global model distributed by the server in the federated setup. We can continue exploring to verify that this intuition is correct. Theoretically, combining Theorem 1, Theorem 2, and Corollary 1, we arrive at another main theorem addressing the question introduced at the beginning of this section.

This theorem further demonstrates the robustness trade-off between applying SSL in federated and decentralized frameworks. For less concern about the impact of data heterogeneity, we should conduct distributed SSL in a federated setup (often also referred to as federated self-supervised learning [3], [16], [17], [33]). However, the decentralized case is more common in reality, as it is challenging to provide a central server that can be trusted by all clients and has stable communication with them. Then, we can consider increasing the average connectivity between clients to minimize the negative impact of heterogeneous data on training (e.g., identifying under-connected clients and creating additional direct communication links).

5 MAR Loss: An Illustrative Case Study in Enhancing Robustness↩︎

The preceding analysis has addressed the main focus of this paper by establishing theoretical insights into the robustness of different D-SSL frameworks under heterogeneous data. As a further step, we illustrate how these insights can guide a more robust algorithmic design. In particular, our results show that although distributed MIM is fundamentally more robust than CL, its training dynamics are dominated by the client-specific covariance, causing local encoders to drift toward different directions before aggregation gradually mitigates this effect. This observation motivates us to refine the MIM objective with an additional term that explicitly and dynamically promotes consistency between local and global masked representations, which we term MAR loss. The integration of MAR into both federated and decentralized frameworks is summarized in Algorithm 3 and Algorithm 4

Formally, MAR loss augments the MIM objective with an alignment regularization term: \[\label{eq95mar95loss} \begin{align} \mathcal{L}_{MAR} &= \mathbb{E}_{x \sim D_i}\, \mathbb{E}_{x_1, x_2 | x} \Big[ \|f_d(f_e(x_1)) - x_2\|^2 + \gamma_t^{(i)} \cdot \text{A-MMD}(z_i, \bar{z}) \Big], \end{align}\tag{5}\] where \(z_i = f_e(x_1)\) and \(\bar{z}\) denote the local masked and global representations, and \(\gamma_t^{(i)} > 0\) is a dynamic weight for alignment. The alignment regularizer is based on Maximum Mean Discrepancy (MMD), a widely used measure of distributional discrepancy in machine learning [34]–[36]. MMD compares whether two distributions \(P\) and \(Q\) differ by mapping samples into a reproducing kernel Hilbert space (RKHS) and evaluating differences in their feature means. Typically, MMD adopts a Gaussian kernel \(k(x, x') = \exp(-\|x-x'\|^2 / 2\sigma^2)\).

In MAR, we employ an adaptive version (A-MMD) to compare the feature spaces of local and global representations more robustly. Unlike prior FL works that use vanilla MMD [37]–[39], A-MMD selects the kernel bandwidth automatically rather than fixing it. Given batches of local and global embeddings of equal size \(B\), A-MMD is computed as: \[\label{eq95A-MMD} \text{A-MMD}(z_i, \bar{z}) = \frac{1}{B(B-1)} \Bigg(\sum_{a \neq b} k(z_{i,a}, z_{i,b}) + \sum_{a \neq b} k(\bar{z}_a, \bar{z}_b)\Bigg) - \frac{2}{B^2}\sum_{a=1}^{B}\sum_{b=1}^{B} k(z_{i,a}, \bar{z}_b),\tag{6}\] with the adaptive kernel defined as \(k(z,z') = exp(-\frac{||z-z'||}{2 (\text{mean}_{a \neq b}||z_a- z_b||)^2})\) . This data-driven choice ensures stability across non-IID clients by scaling the kernel to the observed embedding distribution.

Finally, to balance early-stage consensus and late-stage efficiency, we design the regularization weight \(\gamma_t^{(i)}\) to decay smoothly from \(\gamma_{\max}\) to \(\gamma_{\min}\). We adopt a cosine schedule based on client participation: \[\label{eq95dynamic95gamma} \gamma_t^{(i)} = \gamma_{\min} + (\gamma_{\max} - \gamma_{\min}) \cdot \tfrac{1}{2}\Big(1 + \cos \tfrac{\pi \cdot \omega_t^{(i)}}{\Omega}\Big),\tag{7}\] where \(\omega_t^{(i)}\) counts the number of times client \(i\) has been selected up to round \(t\), and \(\Omega\) controls the decay horizon. In DecL, where all clients participate every round, one can simply set \(\Omega = T\). In FL with partial participation, a practical choice is the expected number of selections per client, or \(T\) as a default. This schedule applies stronger alignment when client divergence is most pronounced, and gradually relaxes toward \(\gamma_{\min}\) as training progresses, ensuring dynamic robustness gains.

6 Experiments↩︎

In this section, we conduct extensive experiments to validate the correctness of our derived theoretical insights and evaluate the effectiveness of the MAR loss in improving the robustness of distributed MIM against data heterogeneity. We first introduce the experimental setup. Then we assess our results in different datasets, model backbones, and distributed settings.

6.1 Experimental Setup↩︎

Datasets and Distributed Simulation. We pre-train our models on the Mini-ImageNet dataset [40], which contains 60,000 images extracted from ImageNet [41]. To simulate a distributed scenario with label non-IIDness, the dataset is partitioned by sampling the class priors of the Dirichlet distribution [42]. More heterogeneous division can be made with a smaller Dirichlet parameter \(\alpha\) during sampling, while the IID case is simulated with a very large \(\alpha\). Besides, we follow prior works to simulate feature heterogeneity by uniformly dividing datasets and applying unique data augmentation for each client [4], [15]. Hence, the local labels are kept the same but features are skewed into different domains before training. Furthermore, to simulate DecL, we use the Erdős-Rényi model [43] to initialize a connected network with the number of clients and the average connectivity as inputs and return the adjacency matrix \(A\). For FL, we additionally assume a central server connecting with all clients. After pre-training, the models’ backbones are fine-tuned on benchmark datasets, including CIFAR-10, CIFAR-100 [44], and ImageNet. The fine-tuning accuracies are used for analysis.

Implementation Details. For our experiments, we use ResNet [18] and Vision Transformer (ViT) [19] as the model architecture. Following the problem setup in theoretical analysis, we select Simsiam [8] and MAE [9] as the representatives of CL and MIM pre-training, respectively. In original works, Simsiam is used to pre-train ResNet models, while MAE is used to pre-train ViTs. We implement two new SSL baselines to show that our theoretical insights apply to any model architecture. One uses Simsiam to pre-train ViTs, and the other one pre-trains ResNet through MAE. Furthermore, we follow the classical distributed algorithms, D-PSGD [45] and FedAvg [12], to implement the DecL and FL frameworks, and then implement our FedMAR and DecMAR algorithms based on these frameworks. All our codes are implemented in Python using the Pytorch framework and executed on a server with 4 \(\text{NVIDIA\textregistered}\) RTX 3090 GPUs. The detailed training setup and server configuration can be found in Appendix 8.2. Our codes are also available at https://github.com/xuanyuLawrence/FedMAR-DecMAR.

6.2 Empirical Study↩︎

Insensitivity Superiority of Distributed MIM. Table 1 compares the impact of data heterogeneity on the pre-training effectiveness between distributed MIM and CL. With highly heterogeneous data, the learned local feature space will be significantly different across clients, resulting in a greater divergence between local and global feature space and a larger drop in the performance compared to the IID setup [3], [16], [33]. Across various datasets and backbone architectures, we observe that distributed MIM consistently exhibits a smaller gap between IID and non-IID settings compared to distributed CL. The experimental results align with Theorem 3, verifying that MIM is less sensitive than CL when handling heterogeneous data in distributed scenarios. To further substantiate this theoretical insight, we also visualize the local and global feature spaces learned by distributed MIM and CL and compute the \(l_2\)-norm weight distance between their local and global models. Please see Appendix 8.3.1 for these external experimental results.

3.5pt

Impact of Average Connectivity on Non-IID Robustness. We verify our second insight by setting up decentralized networks with different average connectivity \(|\bar{A}|\). For the same \(|\bar{A}|\), we consider two cases: (1) a general case where the number of neighbors \(|A_i|\) varies across clients, and (2) a uniform case where all clients have the same connectivity, i.e., \(\forall i \in [N], |A_i| =|\bar{A}|\). Additionally, we set up an FL scenario with 20 clients training in parallel per round. Figure 2 (a) shows that Corollary 1 is correct. We can observe that the fine-tuning accuracy of decentralized SSL increases with \(|\bar{A}|\). Moreover, Figure 2 (a) provides empirical evidence for Theorem 4. We find that pre-training in FL is no less robust than in DecL against heterogeneous data. Additional results using alternative consensus matrices for DecL are given in Appendix 8.3.2, and confirm the same robustness ordering.

Effectiveness of MAR loss. To further illustrate the practical relevance of our analysis, we evaluate MAR loss against the standard MIM objective in both FL and DecL frameworks under varying degrees of heterogeneity. Figure 2 (b) shows that, as the level of non-IIDness increases (i.e., the Dirichlet parameter \(\alpha\) decreases from 1 to 0.001), the fine-tuning accuracy of all methods declines. Nevertheless, models pre-trained with MAR loss consistently outperform those trained with the vanilla MIM loss across all non-IID levels. This trend holds in both FL and DecL settings, suggesting that MAR loss can effectively reduce the sensitivity of distributed MIM to heterogeneity.

Besides, to provide a more comprehensive evaluation, we extend the comparison to recent federated SSL baselines. We evaluate the effectiveness of our FedMAR by comparing it against several state-of-the-art (SOTA) federated self-supervised learning (F-SSL) baselines in a non-IID distributed setting. The SOTA baselines involve: 1) FedU [3]: Using the divergence-aware predictor module for dynamic updates within the self-supervised BYOL network [22]; 4) FedEMA [33]: Employing EMA of the global model to adaptively update online networks; 5) Orchestra [16]: Combining clustering algorithms with Federated Learning for better model aggregation. 6) FeatARC [4]: Combing clustering techniques with feature alignment; 7) LDAWA [17]: Smartly aggregating models according to the angular divergence between local models; and 8) Fed\(U^2\) [46]: Optimizing training with the flexible uniform regularizer and efficient unified aggregator. Following prior works [3], [17], we simulate a highly heterogeneous scenario with 100 clients sampled from a Dirichlet distribution with \(\alpha=0.1\). In each round, 5 clients are randomly selected and each conducts 10 epochs of local training for 200 rounds in total.

Since most baselines employ ResNet-18 [18] as the backbone, we first implement FedMAR with ResNet-18 for a direct comparison. As shown in Table 2, FedMAR employed on ResNet-18 achieves higher accuracy on CIFAR-10 and CIFAR-100 while obtaining comparable results on ImageNet. This indicates that MAR loss can provide tangible improvements even when using the same CNN backbone as prior methods. To further examine the generality of MAR, we also evaluate FedMAR with a lightweight Vision Transformer backbone (Tiny-ViT). Importantly, this model has a comparable number of parameters and GFLOPs to ResNet-18, ensuring fairness in comparison. In this setting, FedMAR employed on Tiny-ViT achieves superior performance on all three benchmarks, surpassing CNN-based baselines while maintaining lower computational cost. These results suggest that MAR loss is not limited to convolutional architectures and can be particularly effective when applied to transformer-based models in federated self-supervised learning.

Extended Study on MAR. In addition to the comparative experiments, we also conduct ablation studies on the main components of MAR loss, including the alignment term and its dynamic weighting. The detailed results of these analyses are provided in 8.4.1, and 8.4.2, respectively. The reported results have confirmed that each component contributes to the performance improvement led by MAR loss. Finally, we assess the practical feasibility of MAR by analyzing its privacy concerns and communication overhead. Since MAR communicates only masked embeddings, the additional communication overhead remains modest, and raw data are never shared. Privacy is therefore also preserved through masking and can be further strengthened with differential privacy. A detailed discussion is reported in Appendix 8.5.

7 Conclusion↩︎

In this paper, we investigated the robustness of distributed self-supervised learning (D-SSL) under heterogeneous data. Our theoretical analysis shows that MIM-based frameworks achieve greater robustness than CL-based ones, and that the degree of robustness in decentralized learning is closely tied to the average network connectivity, with federated learning being no less robust than decentralized learning. These findings provide a principled foundation for understanding how algorithmic choices and network structures affect distributed learning with unlabeled and heterogeneous data. Beyond the theory, we also illustrated how such insights can inform practical design. As a case study, we introduced MAR loss, a refinement of the MIM objective with alignment regularization, which serves to demonstrate the applicability of our analysis. Extensive experiments across model architectures and distributed settings validate our theoretical predictions, and further confirm the utility of MAR loss in practice. We hope that our results can serve as a theoretical grounding and guiding framework for future developments in distributed self-supervised learning.

Acknowledgments↩︎

This work is partly supported by the Australian Research Council Linkage Project (Grant No.LP220200893).

8 Appendix↩︎

8.1 Full Pseudocode of D-SSL with MAR loss↩︎

8.2 Details about Experiment Setup↩︎

In this section, we have provided two tables to present our experiment setup. Table ¿tbl:tab1? shows the experiment details, which include the specific settings for the model architecture, dataset, scenario, and training. Table ¿tbl:tab2? demonstrates the setup of the running environment, including the configuration of our test server.

\begin{table}[h!]

	Details
Model Architecture	ResNet, Vision Transformer (ViT)
Number of layers in ResNet	18
Number of blocks in ViT	5
Pre-train Method	MAE, Simsiam
Pre-train Dataset	Mini-ImageNet
Fine-tune Dataset	CIFAR-10/100, ImageNet
Non-IID Options (i.e. the value of \(\alpha\))	\(\{1e5 \text{ (IID)}, 1, 0.1, 0.01, 0.001\}\)
Options for the \(\gamma\) used in MAR loss	\(\{1, 0.1, 0.01, 0.001\}\)
For Federated Learning (FL):
Number of clients	\(100\)
Number of sampled clients per round	\(5\)
Number of local training epochs	\(2\)
Number of total training rounds	\(100\)
For Decentralized Learning (DecL):
Number of clients	\(20\)
Options for average connectivity	\(3, 5, 10, 20 \text{ (equals to FL)}\)
Number of local training epochs	\(1\)
Number of total training rounds	\(25\)
Fine-tuning Epochs	50/100 (CIFAR-10/100), 20/100 (ImageNet)
Pre-train Batch Size	128
Fine-tune Batch Size	256 (CIFAR-10/100), 1024 (ImageNet)
Base Learning Rate	1.5e-4

\vskip

-0.1in \end{table}

\begin{table}[h!]

Config	Details
Server GPU Count	4
Server GPU Type	RTX 3090 (24GB)
Server CPU Type	AMD EPYC 7282 16-Core
CUDA	12.4
Framework	PyTorch

\vskip

-0.1in \end{table}

8.3 External Experiments↩︎

8.3.1 Feature Space Visualization and Model Difference↩︎

Besides Table 1 demonstrating the non-IID robustness of distributed CL and MIM by the gap in fine-tuning accuracy, we further explore the differences in their learned features empirically. Specifically, we simulate a heterogeneous setting with 100 clients using a Dirichlet sampling with \(\alpha=0.1\). For each D-SSL framework, we obtain three pre-trained ViT backbones: (1) a global model trained using FL across all clients; and (2) two local models trained solely on data from client 1 and client 100, respectively. To compare their learned feature spaces, we extract the encoder features of each model. These high-dimensional features are first projected to 20 dimensions using principal component analysis (PCA) and then embedded into 2D space using Umap [47] for visualization.

Figure 5 presents the feature of local and global models learned by each D-SSL method. Each column corresponds to one method, while each row shows features from a specific model (client 1, client 100, and the global model). We observe that for distributed MIM methods, the local features are more aligned with each other and also closer to the global features, suggesting more consistent representations across heterogeneous clients. In contrast, distributed CL exhibits greater divergence between local and global features, indicating that it is inherently more sensitive to data heterogeneity.

To provide a more quantitative comparison, we also show the weight differences between local and global models in Table 3. In particular, we compute the layer-wise \(\ell_2\)-norm difference between local and global model weights and report the sum across all layers. The results show that distributed MIM methods (MAE and MAR) yield significantly lower weight distances compared to distributed CL, reinforcing the observation that MIM leads to more stable and consistent model updates in the presence of non-IID data.

8.3.2 Effect of Different Consensus Matrices on Robustness Theory↩︎

This experiment examines whether the robustness findings in Section 4.3 remain valid under different choices of consensus matrices in decentralized learning. The theoretical bounds link robustness to the average connectivity of the network graph, while the average connectivity is closely related to how efficiently information mixes across clients. If various consensus rules produced qualitatively different mixing behavior, they could in principle affect robustness. To test this, we conduct two groups of experiments on decentralized networks with 20 clients under strong non-IID conditions using \(\alpha=0.1\). In the first group, the network has an average connectivity of around 5, while in the second group, connectivity is increased to about 10. Within each group, we pre-train distributed MIM using four commonly adopted consensus schemes, including data size weighting, degree normalized averaging, doubly stochastic matrices, and push sum, and compare all results against a federated learning baseline with the same number of clients. The results in Table 4 show a consistent pattern. When connectivity is low, all decentralized variants suffer a noticeable accuracy loss relative to federated learning, and the specific consensus rule makes only minor differences. When connectivity increases, all decentralized variants recover to a level that is close to the federated baseline yet never surpass it. These findings confirm that the qualitative ordering predicted by the theory persists. The choice of consensus matrix influences only constant factors in mixing but does not overturn the robustness relation that federated learning is at least as robust as decentralized learning.

8.4 Ablation Studies on MAR↩︎

8.4.1 Ablation on Alignment Metric↩︎

Our MAR loss (Eq. 5 ) involves two key components: the dynamic regularization weight \(\gamma_t\) and the A-MMD distributional penalty used to align local and global representations. To understand their impact, we perform ablation studies on each component. We first evaluate the contribution of the alignment metric.

For baselines, we consider two commonly used choices in prior work: cosine similarity, which has been widely adopted in federated SSL studies for enforcing alignment between local and global feature spaces [4], and vanilla MMD with a fixed kernel bandwidth, which has also been explored in recent federated learning works [37]–[39]. On top of these, we evaluate our adaptive variant A-MMD, where the kernel bandwidth is chosen automatically based on either the median or mean of pairwise distances. As shown in Table 5, A-MMD consistently outperforms cosine similarity and vanilla MMD across different \(\gamma\) values. Between the two adaptive variants, using the mean of pairwise distances provides slightly better performance, and we adopt this as our default design.

8.4.2 Ablation on Regularization Weight↩︎

Next, we analyze the impact of the regularization weight \(\gamma\) by fixing the alignment metric to A-MMD. Results in Table 6 show that using a large weight (\(\gamma=1\)) degrades performance, as the alignment term overwhelms the reconstruction objective. Conversely, very small weights such as \(\gamma=1\mathrm{e}{-3}\) reduce MAR to a near-vanilla MIM objective and fail to deliver sufficient robustness gains. Moderate fixed values such as \(\gamma=1\mathrm{e}{-2}\) and \(\gamma=1\mathrm{e}{-1}\) yield stronger results, but still remain below our proposed dynamic schedule.

Notably, the cosine decay schedule that smoothly decreases \(\gamma\) from \(1\mathrm{e}{-1}\) to \(1\mathrm{e}{-3}\) achieves the best performance (54.91%). This validates our intuition behind dynamic weighting: stronger alignment is most beneficial in the early stage when client divergence is high, while gradually relaxing the weight avoids excessive penalty in later stages. These findings highlight the importance of the dynamic design in MAR loss, which not only achieves higher accuracy but also improves training stability.

8.5 Discussion on Privacy and Communication Overhead of MAR↩︎

When deploying MAR loss in practice, natural concerns arise regarding the potential privacy risks and the additional communication associated with sharing local representations. We provide both quantitative and qualitative analyses below to show that these costs remain modest and manageable.

Privacy considerations. The information communicated by MAR is limited to local representations \(z_i=f_e(x_1)\) derived from the unmasked portion of the input. Because MIM typically adopts a high masking ratio (e.g., 75% in MAE [9]), most raw content remains hidden and the embedding dimensionality is substantially reduced, which mitigates potential leakage. For stronger guarantees, MAR can be further combined with standard Differential Privacy (DP) mechanisms [48], [49] by perturbing embeddings before transmission, e.g., \(z_i \leftarrow f_e(x_1) + \mathcal{N}(0,\sigma^2 I)\) with \(\sigma\) calibrated to satisfy \((\epsilon,\delta)\)-DP.

Communication overhead. In addition to the standard model updates (e.g., gradients or weights), MAR transmits compact masked embeddings computed from the unmasked portion of each input. This is the sole extra payload introduced by MAR. For instance, in the MAE (ViT-B/16) setting on ImageNet with a 75% masking ratio, each image has 196 patches, of which 49 remain visible. With hidden size 768 and batch size 256, this yields about \(49 \times 768 \times 256\) float values (\(\approx 36.8\) MB in float32). By contrast, a full model with 86M parameters is \(\approx 328\) MB, so the additional cost from MAR is only \(\sim\!11\%\) under this configuration. Crucially, in cross-device settings where small batches are common, this extra cost decreases proportionally with the batch size: at \(B{=}128\) it is \(\approx 18\) MB (\(\sim\!5\%\)), at \(B{=}64\) it is \(\approx 9\) MB (\(\sim\!3\%\)), and at \(B{=}32\) it drops to around \(\sim\!1\%\). These calculations indicate that the MAR-induced overhead remains acceptable in realistic deployments. Moreover, MAR is optional: when minimal communication is the overriding priority, one can simply use the standard MIM objective, whose effectiveness is explained by our theory, at zero additional cost. When a small extra cost is acceptable, MAR offers corresponding robustness gains while keeping the overhead low.

8.6 Full proof for Theoretical Analysis↩︎

8.6.1 Formal Assumptions↩︎

To make the theoretical analysis fully transparent and self-contained, we first summarize here all assumptions used in deriving the main results. These assumptions complement the problem setup in Section 3 and reflect the standard modeling choices commonly adopted in theoretical studies of distributed and self-supervised learning.

8.6.2 Learned Representability for Decentralized MIM↩︎

This section provides the full proof of Theorem 1.

8.6.3 Learned Representability for decentralized contrastive learning↩︎

This section provides the full proof of Theorem 2.

Then, we start to prove Theorem 2 as follows.

8.6.4 Proof of First Theoretical Insight↩︎

This section provides the full proof of Theorem 3.

8.6.5 Proof of Second Theoretical Insight↩︎

This section provides the full proof of Corollary 1 and Theorem 4.

References↩︎

---

1. Corresponding Authors: Nan Yang and Dong Yuan.↩︎