D OES LEARNING FROM DECENTRALIZED NON -IID UNLABELED DATA BENEFIT FROM SELF SUPERVISION ?
A BSTRACT
The success of machine learning relies heavily on massive amounts of data, which are usually generated and stored across a range of diverse and distributed data sources.
Decentralized learning has thus been advocated and widely deployed to make efficient use of the distributed datasets, with an extensive focus on supervised learning (SL) problems.
Unfortunately, the majority of real-world data are unlabeled and can be highly heterogeneous across sources.
In this work, we carefully study decentralized learning with unlabeled data through the lens of self-supervised learning (SSL), specifically contrastive visual representation learning.
We study the effectiveness of a range of contrastive learning algorithms under decentralized learning setting, on relatively large-scale datasets including
ImageNet-100, MS-COCO, and a new real-world robotic warehouse dataset.
Our experiments show that the decentralized SSL (Dec-SSL) approach is robust to the heterogeneity of decentralized datasets, and learns useful representation for object classification, detection, and segmentation tasks, even when combined with the simple and standard decentralized learning algorithm of Federated Averaging ( FedAvg ).
This robustness makes it possible to significantly reduce communication and to reduce the participation ratio of data sources with only minimal drops in performance.
Interestingly, using the same amount of data, the representation learned by Dec-SSL can not only perform on par with that learned by centralized SSL which requires communication and excessive data storage costs, but also sometimes outperform representations extracted from decentralized SL which requires extra knowledge about the data labels.
Finally, we provide theoretical insights into understanding why data heterogeneity is less of a concern for
Dec-SSL objectives, and introduce feature alignment and clustering techniques to develop a new Dec-SSL algorithm that further improves the performance, in the face of highly non-IID data.
Our study presents positive evidence to embrace unlabeled data in decentralized learning, and we hope to provide new insights into whether and why decentralized SSL is effective and/or even advantageous.
I NTRODUCTION
The success of machine learning hinges heavily on the access to large-scale and diverse datasets.
In practice, most data are generated from different locations, devices, and embodied agents, and stored in a distributed fashion.
Examples include a fleet of self-driving cars collecting a massive amount of streaming images under various road and weather conditions during everyday driving, or individuals using mobile devices to take photos of objects and scenery all over the world.
Besides being largescale, these datasets have two salient features: they are heterogeneous across data sources, and mostly unlabeled .
For instance, images of road conditions, which are expensive to label, vary across cars driving on highways vs. rural areas, and under sunny vs. snowy weather conditions (Figure 19).
Methods that can make the best use of these large-scale distributed datasets can significantly advance the performance of current machine learning algorithms and systems.
This has thus motivated a surge of research in decentralized learning/learning from decentralized data 2 (Koneˇcn`y et al., 2016;
Hsieh et al., 2017; McMahan et al., 2017; Kairouz et al., 2021; Nedic, 2020), where usually a global model is trained on the distributed datasets using communication between the local data sources and a centralized server, or sometimes even only among the local data sources.
The goal is typically to reduce or eliminate the exchanges of local raw data to save communication costs and protect data privacy.
How to mitigate the effect of data heterogeneity remains one of the most important research questions in this area (Zhao et al., 2018;
Hsieh et al., 2020; Karimireddy et al., 2020;
Ghosh et al., 2020; Li et al., 2021a), as it can heavily downgrade the performance of decentralized learning.
Moreover, most existing decentralized learning studies focused on supervised learning (SL) problems that require data labels (McMahan et al., 2017; Jeong et al., 2020; Hsieh et al., 2020).
Hence, it remains unclear whether and how decentralized learning can benefit from large-scale, heterogeneous, and especially unlabeled datasets typically encountered in the real world.
On the other hand, people have developed effective methods of learning purely from unlabeled data and demonstrated impressive results.
Self-supervised learning (SSL), a technique that learns representations by generating supervision signals from the data itself, has unleashed the power of unlabeled data and achieved tremendous successes for a wide range of downstream tasks in computer vision (He et al., 2020; Chen et al., 2020;
He et al., 2021b), natural language processing (Devlin et al., 2018; Sarzynska-Wawer et al., 2021), and embodied intelligence (Sermanet et al., 2018; Florence et al., 2018).
These SSL algorithms, however, are usually trained in a centralized fashion by pooling all the unlabeled data together, without accounting for the heterogeneous nature of the decentralized data sources.
Very recently, there have been a few contemporaneous/concurrent attempts (He et al., 2021a; Zhuang et al., 2021; 2022;
Lu et al., 2022; Makhija et al., 2022) that bridged unsupervised/self-supervised learning and decentralized learning, with focuses on designing better algorithms that mitigate the data heterogeneity issue.
In contrast, we revisit this new paradigm and ask the question:
Does learning from decentralized non
IID unlabeled data really benefit from SSL?
We focus on understanding the use of SSL in decentralized learning when handling unlabeled data.
We aim to answer whether and when decentralized SSL (Dec-SSL) is effective (even combined with simple and off-the-shelf decentralized learning algorithms, e.g., FedAvg (McMahan et al., 2017)); what are the unique inherent properties of Dec-SSL compared to its SL counterpart; how do the properties play a role in decentralized learning, especially with highly heterogeneous data?
We also aim to validate our observations on large-scale and practical datasets.
We defer a more detailed comparison with these most related works to §A.
In this paper, we show that unlike in decentralized (supervised) learning, data heterogeneity can be less concerning in decentralized SSL, with both empirical and theoretical evidence.
This leads to more communication-efficient and robust decentralized learning schemes, which can sometimes even outperform their supervised counterpart that assumes the availability of label information.
Among the first studies to bridge decentralized learning and SSL, our study provides positive evidence to embrace unlabeled data in decentralized learning, and provides new insights into this setting.
We detail our contributions as follows.
Contributions.
(i)
We show that decentralized SSL, specifically contrastive visual representation learning, is a viable learning paradigm to handle relatively large-scale unlabeled datasets, even when combined with the simple FedAvg algorithm.
Moreover, we also provide both experimental evidence and theoretical insights that decentralized SSL can be inherently robust to the data heterogeneity across different data sources.
This allows more local updates, and can significantly improve the communication efficiency in decentralized learning.
(ii) We provide further empirical and theoretical evidences that even when labels are available and decentralized supervised learning (and associated representation learning) is allowed, Dec-SSL still stands out in face of highly non-IID data.
(iii)
To further improve the performance of Dec-SSL, we design a new Dec-SSL algorithm, FeatARC , by using an iterative feature alignment and clustering procedure.
Finally, we validate our hypothesis and algorithm in practical and large-scale data and task domains, including a new real-world robotic warehouse dataset.
P RELIMINARIES AND O VERVIEW
Consider a decentralized learning setting with K different data sources, which might correspond to different devices, machines, embodied agents, or datasets/users that can generate and store data locally.
The goal is to collaboratively solve a learning problem, by exploiting the decentralized data from all data sources.
More specifically, consider each data source k ∈
[ K ] has local dataset
D k
= { x k,i } | D k | i =1 , and x k,i ∈
X ⊆
R d are identically and independently distributed (IID) samples from probability distribution D k , i.e., x k,i ∼ D k .
Note that the distributions D k is in general different across data sources k , yielding an overall heterogeneous (i.e., non-IID) data distribution for the data from all the sources.
Let D =
(cid:83)
k ∈
[ K ] D k denote the set of all data samples.
Moreover, we are interested in situations where no label is provided alongside the data x .
To effectively utilize the large-scale unlabeled data, we resort to self-supervised learning approaches.
Specifically, SSL approaches extract representations from these unlabeled data, by finding an embedding function f w : X → R m , where w is the parameter of the embedding function.
z = f w ( x ) is the representation vector that can be useful for downstream tasks, e.g., classification or segmentation.
We summarize several popular SSL approaches here that will be used later in the paper.
Self-supervised representation learning.
Now consider a given data source k ∈ [ K ] .
There are two popular methods in the SSL community.
In contrastive learning (Chen et al., 2020;
He et al., 2020) specifically, a sample x is used to provide supervision signals along with two generated positive samples x + and x (overloaded for notational simplicity) and (possibly multiple) negative samples x − sampled from the training batch.
The goal of SSL is to find an embedding f w that makes x and x + close, while keeping x and x − s apart, if negative samples are used.
One commonly used loss for SSL is the InfoNCE loss (Oord et al., 2018), which has been used in popular SSL approaches as SimCLR (Chen et al., 2020) and MoCo (He et al., 2020):
[Equation] [Equation]
where τ > 0 is a temperature hyperparameter, j is the index for negative samples, D ( · , · ) is a distance function such as the cosine distance, i.e., D ( z 1 , z 2 ) = − z 1 · z 2 ||
z 1 |||| z
|| .
Some other effective SSL approaches, such as BYOL (Grill et al., 2020) and SimSiam (Chen & He, 2021), remove the terms related to negative samples in (2.1).
These methods also add an additional function g , the feature predictor , which only applies to x to create an asymmetry and to avoid the collapsed solutions.
This usually leads to the following objective: L k ( w ) :=
[Equation]
In our experiments, we make use of both losses and the SSL approaches associated with them.
Decentralized SSL.
To exploit the heterogeneous data distributed at different locations/devices, decentralized SSL optimizes the following global objective:
[Equation]
which can be solved using many existing decentralized learning algorithms.
For instance, FedAvg (McMahan et al., 2017) is one of the most representative, easy-to-implement, and communicationefficient decentralized learning algorithms which optimizes this objective without data-sharing among data sources.
At each iteration t , the server first samples a set of data sources M t with size |M t | = ρK and run δ local update steps on each of the local dataset.
Then, each data source k ∈ M t sends back the updated local model weight w t,δk to the central server, and the server averages them to be the global model w t +1 = 1 |M t | (cid:80) k ∈M t w t,δk for the next round t + 1 .
The server then broadcasts the global model to each data source to reset w t +1 , 0 k as w t +1 .
The number of local updates ( δ ) determines the communication efficiency (larger δ means less communication); in the experiments, we use E to denote the number of epochs of local updates (as a surrogate for δ ).
Both
E and the participation rate ρ are important factors that determine the efficiency of decentralized learning.
The learned representation f w ( x ) can then be used in downstream supervised learning tasks.
There are many real-world applications of decentralized SSL, including self-driving cars, warehouse robots, and mobile devices.
A further discussion can be found in Appendix §D.
2. O VERVIEW OF O UR S TUDY
Terminology & setup.
We separate our experiment pipeline into representation learning (pretraining phase) and downstream evaluation (evaluation phase).
Our main focus is on the aforementioned Dec-SSL approach.
We use FedAvg (McMahan et al., 2017) with SimCLR (Chen et al., 2020) as the default method.
Moreover, we will also compare with settings where the label information is available, i.e., the classical decentralized (supervised) learning, which should be more favorable for learning.
See Figure 1 for a summary of different settings.
The first setting is Dec-
SL : we simply run FedAvg on the decentralized labeled data, for end-to-end classification.
Dec-SL does not learn representations explicitly, and serves as a natural baseline when labels are available.
The second setting is representation learning from Dec-SL, where we train supervised learning with
FedAvg , and then use the feature extractor network as the backbone for downstream tasks.
This way, we can also learn the representation from decentralized labeled data, and make the comparison with Dec-SSL more fair, since both are learning features for various downstream tasks.
We term this setting as Dec-SLRep .
[Table]
+qhYPcviyIMtsA1KwATHoAouQA3UAQYP4Am8gFftUXvW3rT3SWlOy3o2wS9oH9+upZkE</latexit>
The evaluation phase tests the representations from Dec-SSL or Dec-SLRep.
We consider two protocols in the evaluation phase: linear probing for image classification (Zhang et al., 2016) and finetuning for object detection/segmentation (Doersch et al., 2015).
For classification, we train a linear classifier on top of the frozen pretrained network and evaluate the top-1 classification accuracy.
For object detection/segmentation, we finetune the network by using the pretrained weights as initialization and training in an end-to-end fashion, and then we evaluate the mean Average
Precision (mAP) metric.
Downstream tasks are performed on centralized train and test dataset.
Please refer to Appendix §C.1 for implementation details and Table 3 for experiment setups.
Questions of interest.
Through extensive experiments on large-scale datasets, and theoretical analysis in simplified settings, we seek to answer the following questions: (i) How well can decentralized
SSL, even instantiated with the simple FedAvg algorithm, rival the performance of its centralized counterpart, and handle the non-IIDness of decentralized unlabeled data?
(ii) Is there any unique and inherent property of Dec-SSL, compared to its supervised learning counterpart; how and why may the property benefit decentralized learning, even when the label information is available?
(iii)
Is there a way to further improve the performance of Dec-SSL in face of highly non-IID data?
Our hypothesis is that SSL, whose objective is not particularly dependent on the x to y mappings, learns a relatively uniform representation across decentralized and heterogeneous unlabeled datasets, thus leading to more efficient and robust decentralized learning schemes.
We aim to validate this hypothesis and answer these questions in the following sections.
D EC -SSL IS E FFICIENT AND R OBUST TO D ATA H ETEROGENEITY
We first seek to address question (i) in §2.1 – how well decentralized SSL performs, in face of nonIID and decentralized unlabeled data.
To this end, we first introduce the notion of data heterogeneity in decentralized learning, which is usually categorized as input heterogeneity, label distribution heterogeneity, and the heterogeneity in the relationships between the features and labels , respectively (Hsieh et al., 2020).
We create label heterogeneity by distributing each data source with different proportion of classes; we construct the heterogeneity via either sampling from a Dirichlet process with hyperparameter α or via skewness partitioning (Hsieh et al., 2020) with hyperparameter β .
We also create input heterogeneity by leveraging the feature space of a pretrained network on the data.
See §C.2 for more details on how we create data heterogeneity across data sources.
3. E XPERIMENTAL OBSERVATIONS
CIFAR classification under different types of non-IIDness.
In this experiment, we construct input and label non-IIDness using 5 data sources in the CIFAR-10 (Krizhevsky et al., 2009) dataset based on the Dirichlet Process.
The sources of non-IIDness are the feature clusters and labels, respectively.
We control parameter α to create datasets from very IID (each data source has roughly a uniform distribution over 10 classes / 5 feature clusters) to very non-IID (each data source has data from 2 classes / 1 feature clusters).
Recall that E denotes the number of epochs for local updates and ρ denotes the participation ratio of data sources at each round.
We use E = 50 epochs of local updates in this experiment, which is equivalent to around δ = 1000 iterations, i.e., each local data source updates 50 epochs independently before averaging.
The results are shown in Figure 2.
Surprisingly,
[Figure]
the performance of downstream classification, with representations trained using decentralized SSL, is very insensitive to the non-IIDness across the datasets and only bears a slight performance drop.
This robustness over data non-IIDness is encouraging, and stands in sharp contrast with most existing decentralized supervised learning algorithms, which are known to suffer from the data heterogeneity in general (Hsieh et al., 2020).
As a baseline, we consider the classical decentralized SL approach of FedAvg , trained over the same non-IID data, but with label information.
Indeed, the performance of decentralized SL can drop significantly as the non-IIDness increases.
Finally, we note that the simple use of FedAvg in SSL can achieve performance comparable to the centralized
SSL, showing that Dec-SSL is an effective decentralized learning scheme to handle unlabeled data.
Finetuning ImageNet representation for COCO detection.
In this experiment, we finetune the representations learned from ImageNet to COCO detection benchmark (Lin et al., 2014) with the
Detectron pipeline (Girshick et al., 2018).
Specifically, we use ImageNet-100 with ResNet-18 and 1 × training schedule for Mask R-CNN (He et al., 2017) with a ResNet18 FPN being the backbone.
Compared to the contemporary works (Zhuang et al., 2022; Lu et al., 2022) on federated self-supervised learning, our setup is more relevant to real-world applications, as it works on largerscale and more practical datasets and tasks.
We run Dec-SSL on ImageNet-100 dataset with 5 data sources, and with E = 1 epoch of local updates, which corresponds to around δ = 500 local updates, to learn the global representation using FedAvg .
On Table 1 left, we observe that the representation from Dec-SSL almost reaches the performance of the representation from centralized SSL and improves upon baselines that train the model from scratch, i.e., the no pretrain row.
This conveys that SSL can learn useful representations in decentralized settings, avoiding the heavy communication cost of centralized learning.
Decentralized SSL for real-world package segmentation.
The issue of data heterogeneity and communication efficiency is significant for real-world applications such as those in Amazon warehouses, whose fleets of working robots can generate millions of images per day (see Figure 21 for an illustration).
We provide details about the Amazon dataset in §D.1.
We use data from one sample warehouse site at Amazon, and split the data based on the session ID (which is usually a sequence of days).
Each decentralized learner is only allowed to access the local data at one session, which is equivalent to the non-IID case where skewness β = 0 .
We then deploy decentralized self-supervised learning on a subset of the enormous warehouse data, which has around 80000 images with contour labels output by the Amazon work-cells.
We use SimCLR with FedAvg and communication efficiency E = 1 number of local update epochs, as the pretraining method.
On the right subtable of Table 1, we compare different ways to initialize weights for finetuning, and show that the representations learned from decentralized SSL outperforms training from scratch and even matches centralized SSL on the Amazon dataset.
We also experiment with finetuning segmentation task using Mask R-CNN on different fractions of the data, and show that Dec-SSL can further improve the performance of training from scratch, when there is no as much labeled data.
3. T HEORETICAL INSIGHTS
We now provide some theoretical insights into why the objective of Dec-SSL leads to more robust performance in face of data heterogeneity.
In particular, we analyze the property of the solutions to the local and global objectives of Dec-SSL in a simplified setting, and show that the global objective is not affected significantly by the heterogeneity of local datasets.
Our setup is inspired by the very recent work (Liu et al., 2021), where the effect of imbalanced data in centralized SSL was studied in a simplified setting.
In particular, we generalize the centralized and 3 -way classification setting to a decentralized and 2 K -way one, carefully design the generation of data distribution across data sources, and establish analyses for both local and global objectives in decentralized SSL.
We also 𝑒 " 𝑒 # 𝑒 " 𝑒 # 𝑒 $…& 𝑫𝑒 "…)*" 𝑒 )," ,…,𝑒 & 𝑒 ) 𝑫 𝑲 𝑒 $…& 𝑫
SL
SSL
SL
SSL
SL
SSL
[Figure]
[Table]
improve some analysis therein, and design new metrics to characterize the performance adapted to the decentralized setting.
Due to space limitation, we include an abridged introduction here, and defer more details to Appendix §E.
Setup.
Consider a Dec-SSL problem with K data sources.
Similar to the SimSiam approach, we first augment x , an anchor sample from the dataset, by sampling ξ, ξ ′ ∼ N (0 , I ) IID from the
Gaussian distribution.
Consider the linear embedding function f w ( x ) =
wx , where w ∈ R m × d and m ≥ 2 K .
The SSL objective for data source k is given by
[Equation] [Equation]
) where (cid:98) E is taken expectation over the empirical dataset x k,i ∼ D k , and the randomness of ξ k,i and ξ ′ k,i .
Moreover, recall the global objective is given in (2.2).
Note that (3.1) instantiates SimSiam loss with the negative inner-product ⟨ a, b ⟩ as the distance function D ( a, b ) and no feature predictor, and with a regularization term for mathematical tractability, as in Liu et al.
Data heterogeneity.
The K data sources collaboratively solve (2.2) to learn a representation for a 2 K -way classification task.
The K local datasets are generated in a way that for each fixed k ∈ [ K ] , the labels are skewed in that data from classes 2 k − 1 and 2 k constitute the majority of the data, while other classes are rare, or even unseen.
More details on the specifications of data heterogeneity can be found in §E.1.
We visualize the heterogeneity of the data distributions in Figure 3.
To compare the representations learned across data sources and that learned from jointly solving (2.2), we introduce the following definition on the representability of the representation space.
Definition 3.1 (Representability vector) .
Let S ⊆
R
d be the subspace spanned by the rows of the learned feature matrix w ∈ R m × d , where the embedding function f w ( x ) =
wx .
The representability of S is defined as a vector r =
[ r 1 , · · · , r d ] ⊤ ∈ R d , such that r i =
∥
Π S ( e i ) ∥ 22 for i ∈
[ d ] , where Π S ( e i ) ∈
R
d is the projection of standard basis e i onto S , and thus r i = (cid:80) sj =1 ⟨ e i , v j ⟩ 2 where s = dim ( S ) and { v 1 , · · · , v s } is a set of orthonormal bases for S .
The intuition of this definition is that a good feature space should have the property that many standard unit bases among e 1 , · · · , e d , which can be used to represent any vectors in R d , can be represented well by the feature space, i.e., have large projections onto it.
Note that as a vector, r provides a quantitative way to compare the representability of two feature spaces across different directions (i.e., different unit basis).
In the following theorem, we compare the representability learned by local objectives and the global one, for Dec-SSL.
[Figure]
Theorem 3.2 (Representability of local v.s. global objectives for Dec-SSL) .
For decentralized SSL in the setting described above, with high probability, the representability vector learned from any local objective of source k , denoted by r k =
[ r k 1 , · · · , r kd ] ⊤ , satisfies that 1 − O ( d − 4 / 5 ) ≤ r ki ≤for all
i ∈
[ K ] \{ k } .
Moreover, the representability vector learned from the global objective, denoted by ¯ r =
[¯ r 1 , · · · , ¯ r d ] ⊤ , satisfies that 1 − O ( d − 4 / 5 ) ≤ ¯ r i ≤ 1 for all i ∈
[ K ] .
Theorem 3.2 states that the feature spaces learned from local SSL objectives are relatively uniform , in the sense that for the K basis directions e 1 , · · · , e K that generate the data, any two data sources have similar representability in all of them but two directions, especially when the dimension d of the data is large.
Furthermore, when solving the global objective (2.2), the learned representation is also uniform, and its representability differs at most one direction from that of each local data source.
Note that the results hold with highly heterogeneous data across data sources.
In other words, Dec-SSL is not affected significantly by the non-IIDness of the data, justifying the empirical observations in §3.1.
Illustration of the results can also be found in Figure 3.
Intuition & implication.
The main intuition behind Theorem 3.2 is that, the objective of SSL is not biased by the heterogeneous distribution of labels at each local dataset, and tends to learn uniform representations.
Related arguments have also been made in the recent works on the theoretical understanding contrastive learning/SSL (Wang & Isola, 2020; Liu et al., 2021).
In the decentralized setting, this insensitivity to data heterogeneity becomes even more relevant, as it potentially allows each local data source to perform much more local updates, without drifting the iterates significantly.
This enables more communication-efficient decentralized learning schemes, in contrast to most existing ones that are vulnerable to data non-IIDness.
We validate these points next.
D EC -SSL C AN BE F AVORABLE E VEN W HEN L ABELS ARE A VAILABLE
We here seek to address question (ii) in §2.1 – how does the unique property of Dec-SSL, such as the robustness to data heterogeneity, benefit decentralized learning?
While lack of labels seems a limitation, we show that this might not be the case in decentralized learning with heterogeneous data.
First, it is known that decentralized SL in general performs poorly when the data is highly heterogeneous (Zhao et al., 2018; Hsieh et al., 2020).
Further, even in the decentralized representation learning setting when labels are available, Dec-SSL still stands out in face of highly non-IID data.
To make a fair comparison, we mainly compare Dec-SSL with Dec-SLRep (recall the definition in §2.1), which are both decentralized representation learning approaches.
We defer the comparison with Dec-SL to Appendix §B.
We conduct experiments on both ImageNet and CIFAR-10 datasets, and evaluate the performance of the learned representations in terms of the variations of two commonly used metrics in decentralized learning – the number of local updates epochs E , and the participation ratio of data sources ρ .
We observe consistently that Dec-SSL indeed outperforms
Dec-SLRep in learning representations in terms of communication efficiency and participation ratio, especially with highly non-IID data.
We remark that such observations are also consistent with those on object detection and semantic segmentation given in Table 1.
4. E XPERIMENTAL OBSERVATIONS
In this experiment, we train and evaluate the feature backbone on ImageNet-100 in a decentralized setting.
We create non-IIDness across the local datasets based on label skewness and use β = 0 .(each
data source has only 10% of its data coming from the uniform class distributions).
Communication efficiency under high non-IIDness.
In Figure 4, we show that under the nonIID scenario, averaging weights with an infrequent communication schedule causes less trouble to
Dec-SSL than to Dec-SLRep.
In FedAvg , the idea of averaging weights after multiple epochs might sound sub-optimal, but we notice that decentralized SSL is very robust with respect to this parameter.
Intuitively, the robustness of Dec-SSL allows each local model to drift longer, leading to a lower communication frequency for decentralized learning.
[Figure] [Figure]
Participation ratio under high non-IIDness.
In this experiment, we split ImageNet-100 intodata sources and use local update E = 5 epochs.
We measure the performance of decentralized learning algorithms with respect to the participation ratio of data sources at each round.
For instance, when ρ = 1 , at each round, all data sources update their local weights and upload to the server, while ρ = 0 .
05 means that each round a single random data source is selected for update.
On the right of Figure 4, we show that with non-IID data, the convergence of Dec-SSL is more stable to less participants compared to Dec-SLRep.
This allows more efficient decentralized learning, especially when deployed with extremely large number of data sources and unstable communication channels.
4. T HEORETICAL INSIGHTS
To shed light on the above observations, we provide analysis for the feature spaces learned by the local objective of Dec-SLRep, under the same setup as in §3.2.
For Dec-SLRep and each data source k , we consider learning a two-layer linear network g u k ,v k ( x ) :
= v k u k x as classifier, where u k ∈ R m × d and v k ∈ R c × m , and use u k x as the learned representation for downstream tasks.
The network is learned by minimizing ∥ ( u k )
⊤ u k ∥ 2 F + ∥ ( v k ) ⊤ v k ∥ 2 F subject to the margin constraint that [ g u k ,v k ( x )]
y ≥ [ g u k ,v k ( x )]
y ′
+ 1 for all data ( x, y ) in the local dataset k with all y ′ ̸
= y . .
We now have the following proposition on the representations learned by Dec-SLRep across data sources.
Proposition 4.1 (Representations learned by Dec-SLRep across heterogeneous data sources) .
With high probability, the features u k =
[ u k, 1 , · · · , u k,m ]
⊤ ∈
R m × d learned from the local dataset D k satisfies that (cid:80) mi =1 ⟨ u k,i , e j ⟩ 2 ≤ O ( d − 110 ) , for j ∈ [ K ] \{ k } ; while (cid:80) mi =1 ⟨ u k,i , e k ⟩ 2 ≥ 1 − O
( d − 120 ) .
In other words, the correlation between the learned features in w k and e j is small for all j ∈ [ K ] \ { k } , while the correlation between the features and e k is large.
The proposition suggests that the feature spaces learned by Dec-SLRep differ significantly across local data sources, given the highly heterogeneous data.
More specifically, we show that most of the unit bases in { e 1 , · · · , e K } have small correlations with the features learned at each local data source, while these feature spaces themselves vary significantly across data sources.
The unit bases that are not learned might be significant for various other downstream tasks, making the learned representations less favorable.
This heterogeneity among local solutions is not in favor of local updates , as too many local updates would drift the iterates towards its local solution, and the iterates would become too far away from each other, hurting the convergence of decentralized learning.
Hence, compared with the Dec-SSL case and Theorem 3.2, Dec-SLRep can be less robust to data heterogeneity and less communication-efficient.
We note that the advantage of Dec-SSL does not come from using more data , since we use exactly the same data for training Dec-SLRep and DecSSL.
The intuition is also illustrated in Figure 3.
Finally, we remark that the uniformity of features, which is believed to be the key to better transfer performance in SSL (Wang & Isola, 2020; Caron et al., 2020), is not always preferred given specific learning tasks (Burgess et al., 2018).
O UR A LGORITHM – FeatARC (F EATURE A LIGNMENT AND C LUSTERING )
Although Dec-SSL tends to learn relatively uniform features that are robust across datasets, the uniformity itself might not imply the alignment of features across datasets: the representation network from different local data sources can still map the same data point to different regions in the feature space.
This misalignment becomes more significant when the data is highly non-IID and can have an adverse effect on the model aggregation process in decentralized learning (Zhang et al., 2020a).
To mitigate this issue and address question (iii) in §2.1, we propose to use the same feature distance loss as an auxiliary local objective to align the local models with the global model.
The alignment between two features is defined as the negative cosine distance metric D ( z 1 , z 2 ) =
− z 1 · z 2 ||
z
||||
z
|| .
To further improve the Dec-SSL algorithm, we propose to learn multiple models using clusteringbased approach.
In particular, instead of learning a single global model as in (2.2), we learn C models and separate the K data sources into C clusters.
The update of C models and the assignment of data sources to C clusters are conducted alternatively.
When C = K , the algorithm reduces to learning K local models; when C = 1 , it reduces to learning a single global one.
The clustering approach intuitively learns multiple models to interpolate the performance between learning a single global model and K local models, thus achieving a good bias-variance tradeoff when testing on each local dataset (Mansour et al., 2020; Ghosh et al., 2020).
However, unlike the supervised learning case, we do not use the loss of the decentralized learning (i.e., (2.1)) as the metric for clustering.
This is because for contrastive learning, it has been observed that the SSL loss might not be indicative enough for the performance of the representation on downstream tasks (Robinson et al., 2021).
Hence, we here again use the feature alignment distance D ( · , · ) as the metric for clustering.
We adopt the alignment regularization and clustering techniques, and developed a new Dec-SSL algorithm FeatARC , summarized in Algorithm 1 and Algorithm 2 in Appendix.
We show the performance of FeatARC in Figure 5, in comparison with different baselines including FedAvg , under different levels of data heterogeneity and communication frequency.
It is shown that FeatARC outperforms the baselines consistently, including the variants that only uses alignment (“Align Only”) or clustering (“Cluster Only”).
Moreover, on the top of Table 2, we show that FeatARC also outperforms other recent decentralized self-supervised learning algorithms on CIFAR-10 dataset.
E XTENSIONS
In this section, we discuss a few extended experiments of our framework.
Please see Appendix §B for a thorough set of experiments and ablation studies with visualizations.
6. F ULLY DECENTRALIZED CASE AND DIFFERENT NETWORK TOPOLOGY
We conduct experiments on the fully decentralized learning in Appendix §B.5, where the local data sources are only allowed to communicate with their neighbors over a peer-to-peer network, without a centralized server.
In short, most observations we had regarding Dec-SSL in the setting with a centralized server still hold, even under several different network topologies.
This aligns with our theoretical insight provided in Section 3, which came from the benign properties of the solution to the Dec-SSL objective , instead of the properties of specific algorithms (averaging the iterates via a star or other network topologies) that achieves the solution.
6. E XTREMELY HETEROGENEOUS CASE FOR DECENTRALIZED LEARNING
In Figure 13, we show that even in the extremely heterogeneous case where each local source only owns one class, the Dec-SSL framework is still robust to the non-IIDness of the data.
This also holds true when we scale to more clients, as shown in Figure 15.
The Dec-SSL objective would not be biased by the highly heterogeneous class labels at each local dataset, while the Dec-SL objective could be biased by it.
This is also consistent with our theoretical insights in Section §3.2 and the key reason for the success of Dec-SSL is that, despite only having one single class, the information of features obtained from local datasets may still be useful for the jointly classifying of all the classes.
6. C OMPARISON OF F EAT ARC WITH OTHER ALGORITHMS
We also compare our algorithm with the Dec-SSL algorithms that are combined with other federated learning algorithms, including Li et al.
(2020a)
(FedProx) and Li et al.
(2020b) (FedBN).
In Figure 16 (Left), we show that our proposed FeatARC can outperform these two baselines.
C ONCLUSION
We propose the framework of decentralized SSL that learns representations from non-IID unlabeled data and conduct an empirical study on the robustness of Dec-SSL to different types of heterogeneity, communication constraints, and participation rates of data sources.
We also provide findings and theoretical analyses of Dec-SSL compared to its supervised learning counterpart, as well as developing a new algorithm to further address the high heterogeneity in decentralized datasets.
Acknowledgement.
This work is supported in part by Amazon.com Services LLC, PO2D06310236 and Defense Science & Technology Agency, DST00OECI20300823.
L
W. was supported by the MIT EECS Xianhong Wu Graduate Fellowship.
K.Z. also acknowledges support from Simons-Berkeley Research Fellowship.
We thank MIT Supercloud for providing compute resources.
The authors would like to thank many helpful discussions from Phillip Isola at MIT and
Andrew Marchese at Amazon