Decentralized Unsupervised Learning of Visual Representations

Collaborative learning is an effective approach for multiple distributed clients to collaboratively learn a shared model from decentralized data. In the learning process, each client updates the local model by using local data, and then a central server aggregates the local models to obtain a shared model. In this way, collaborative learning enables learning from decentralized data McMahan et al. (2017) while keeping data local for privacy. Collaborative learning can be applied to healthcare applications, where many personal devices such as mobile phones collaboratively learn to provide early warnings to cognitive diseases such as Parkinson’s and to assess mental health Chen et al. (2020b). Collaborative learning can also be used for robotics, in which multiple robots learn a shared navigation scheme to adapt to new environments Liu et al. (2019). Compared with local learning, collaborative learning improves navigation accuracy by utilizing knowledge from other robots.

Existing collaborative learning approaches assume local data is fully labeled so that supervised learning can be used for the model update on each client. However, labeling all the data is usually unrealistic due to high labor costs and the requirement of expert knowledge. For example, in medical diagnosis, even if the patients are willing to spend time on labeling all the local data, the deficiency of expert knowledge of these patients will result in large label noise and thus inaccurate learned model. The deficiency of labels makes supervised collaborative learning impractical. Self-supervised learning can address this challenge by pre-training a neural network encoder with unlabeled data, followed by fine-tuning for a downstream task with limited labels. Contrastive learning (CL), an effective self-supervised learning approach Chen et al. (2020a), can learn data representations from unlabeled data to improve the model. By integrating CL into collaborative learning, clients can collaboratively learn data representations by using a large amount of data without labeling.

In collaborative learning, data collected on clients are inherently far from IID Hsu et al. (2020), which results in two unique challenges when integrating collaborative learning with CL as collaborative contrastive learning (CCL) to learn high-quality representations. The first challenge is that each client only has a small amount of unlabeled data with limited diversity, which prevents effective contrastive learning. More specifically, compared with the global data (the concatenation of local data from all clients), each client only has a subset of the global data with a limited number of classes McMahan et al. (2017); Zhao et al. (2018); Wu et al. (2021b). For instance, in real-world datasets Luo et al. (2019), each client only has one or two classes out of seven object classes. Since conventional contrastive learning frameworks He et al. (2020); Chen et al. (2020a) are designed for centralized learning on large-scale datasets with sufficient data diversity, directly applying them to local learning on each client will result in the low quality of learned representations.

The second challenge is that each client focuses on learning its local data without considering the data on the other clients. As a result, the features of data in the same class but from different clients may not be well-clustered even though they could have been clustered for improved representations. Data are decentralized in collaborative learning and even if two clients have data of the same class, they are unaware of this fact and cannot leverage it to collaboratively learn to cluster these data. Besides, even if one client has knowledge of other’s data, since no labels are available, there is no easy way to identify the correct data clusters and perform clustering for better representations.

To address these challenges, we propose a collaborative contrastive learning (CCL) framework to learn visual representations from decentralized unlabeled data on distributed clients. The framework employs contrastive learning He et al. (2020) for local learning on each client and consists of two approaches to learn high-quality representations. The first approach is feature fusion and it provides remote features as accurate contrastive information to each client for better local learning. To protect the privacy of remote features against malicious clients, we employ an encryption method Huang et al. (2020) to encrypt the images before generating their features. The second approach is neighborhood matching and it further aligns each client’s local features to the fused features such that well-clustered features among clients are learned.

In summary, the main contributions of the paper include:

• •

  Collaborative contrastive learning framework. We propose a framework with two approaches to learning visual representations from unlabeled data on distributed clients. The first approach improves the local representation learning on each client with limited data diversity, and the second approach further learns unified global representations among clients.

• •

  Feature fusion for better local representations. We propose a feature fusion approach to leverage remote features for better local learning while avoiding raw data sharing. The remote features serve as negatives in the local contrastive loss to achieve a more accurate contrast with fewer false negatives and more diverse negatives.

• •

  Neighborhood matching for improved global representations. We propose a neighborhood matching approach to cluster decentralized data across clients. During local learning, each client identifies the remote features to cluster local data with and performs clustering. In this way, well-clustered features among clients can be learned.

Contrastive Learning. Contrastive learning is a self-supervised approach to learn an encoder (i.e. a convolutional neural network without the final classifier) for extracting visual representation vectors from the unlabeled input images by performing a proxy task of instance discrimination Chen et al. (2020a); He et al. (2020); Wu et al. (2018, 2021a). For an input image $x$, its representation vector $z$ is obtained by $z=f(x),\ z\in\mathbb{R}^{d}$, where $f(\cdot)$ is the encoder. Let the representation vectors query $q$ and key $k^{+}$ form a positive pair, which are the representation vectors from two transformations (e.g. cropping and flipping) of the same input image. Let $Q$ be the memory bank with $K$ representation vectors stored, serving as negatives. The positive pair query $q$ and key $k^{+}$ will be contrasted with each vector $n\in Q$ (i.e. negatives) by the loss function:

Minimizing the loss will learn an encoder to generate visual representations. Then a classifier can be trained on top of the encoder by using limited labeled data.

However, existing contrastive learning approaches are designed for centralized learning Chen et al. (2020a); He et al. (2020) and require sufficient data diversity for learning. When applied to each client in collaborative learning with limited data diversity, their performance will greatly degrade. Therefore, an approach to increase the local data diversity on each client while protecting the shared information is needed.

Collaborative Learning. The goal of collaborative learning is to learn a shared model by aggregating locally updated models from clients while keeping raw data on local clients McMahan et al. (2017). In collaborative learning, there are $C$ clients indexed by $c$. The training data $D$ is distributed among clients, and each client $c$ has a subset of the training data $D_{c}\subset D$. There are recent works aiming to optimize the aggregation process Reisizadeh et al. (2020); Nguyen et al. (2020). While our work can be combined with these works, for simplicity, we employ a typical collaborative learning algorithm McMahan et al. (2017). The learning is performed round-by-round. In communication round $t$, the server randomly selects $\beta\cdot C$ clients $C^{t}$ and send them the global model with parameters $\theta^{t}$, where $\beta$ is the percentage of active clients per round. Each client $c\in C^{t}$ updates the local parameters $\theta_{c}^{t}$ on local dataset $D_{c}$ for $E$ epochs to get $\theta_{c}^{t+1}$ by minimizing the loss $\ell_{c}(D_{c},\theta^{t})$. Then the local models are aggregated into the global model by averaging the weights $\theta^{t+1}\leftarrow\sum_{c\in C^{t}}\frac{|D_{c}|}{\sum_{i\in C^{t}}|D_{i}|}{\theta_{c}^{t+1}}$. This learning process continues until the global model converges.

To improve the performance of collaborative learning on non-IID data, Zhao et al. (2018); Jeong et al. (2018) share local raw data (e.g. images) among clients. However, sharing raw data among clients will cause privacy concerns Li et al. (2020). Besides, they need fully labeled data to perform collaborative learning, which requires expert knowledge and potentially high labeling costs. Therefore, an approach to performing collaborative learning with limited labels and avoiding sharing raw data is needed.

We propose a collaborative contrastive learning (CCL) framework to learn representations from unlabeled data on distributed clients. These distributed data cannot be combined in a single location to construct a centralized dataset due to privacy and legal constraints Kairouz et al. (2019). The overview of the proposed framework is shown in Figure 1. There is a central server that coordinates multiple clients to learn representations. The local model on each client is based on MoCo He et al. (2020). CCL follows the proposed feature fusion technique to reduce the false negative ratio on each client for better local learning (Sect. 4). Besides, based on fused features, CCL further uses the proposed neighborhood matching to cluster representations of data from different clients to learn unified representations among clients (Sect. 5).

Before introducing the details of feature fusion and neighborhood matching, we present the proposed CCL process. CCL is performed round-by-round and there are four steps in each round as shown in Figure 1. First, each client $c$ uploads its latest model (consisting of the main model $f_{q}^{c}$ and momentum model $f_{k}^{c}$) and latest features $\overline{Q}_{l,c}$ of encrypted images to the server. Second, the server aggregates main models from clients by $\theta_{q}\leftarrow\sum_{c\in C}\frac{|D_{c}|}{|D|}{\theta_{q}^{c}}$ and momentum models by $\theta_{k}\leftarrow\sum_{c\in C}\frac{|D_{c}|}{|D|}{\theta_{k}^{c}}$ to get updated $f_{q}$ and $f_{k}$, where $|D_{c}|$ is the data size of client $c$ and $|D|$ is the total data size of $|C|$ clients. The server also combines features as $\overline{Q}=\{\overline{Q}_{l,c}\}_{c\in C}$. Third, the server downloads the aggregated models $f_{q}$ and $f_{k}$ and combined features $\overline{Q}$ excluding $\overline{Q}_{l,c}$ as $Q_{s,c}=\{\overline{Q}\ \char 92\relax\ \overline{Q}_{c}\}$ to each client $c$. Fourth, each client updates its local models with $f_{q}$ and $f_{k}$, and then performs local contrastive learning for multiple epochs with local features $Q_{l,c}$ and remote features $Q_{s,c}$ by using loss Eq.(12) including contrastive loss with fused features Eq.(6) and neighborhood matching Eq.(11). During local contrastive learning, to generate features $\overline{Q}_{l,c}$ for uploading in the next round, images $x$ are encrypted by InstaHide Huang et al. (2020) as $\tilde{x}$ and fed into momentum model $f_{k}^{c}$. In this way, even if a malicious client can ideally recover $\tilde{x}$ from the features $\overline{Q}_{l,c}$, which is already very unlikely in practice, it still cannot reconstruct $x$ from $\tilde{x}$ since Instahide effectively hides information contained in $x$. Next, we present the details of local contrastive learning, including feature fusion to reduce false negative ratio in Sect. 4 and neighborhood matching for unified representations in Sect. 5.

Next, we focus on how to perform local CL in each round of CCL. We first present the key challenge of CL on each client, which does not exist in conventional centralized CL. Then we propose feature fusion to tackle this challenge and introduce how to perform local CL with fused features.

Key challenge: Limited data diversity causes a high false negative (FN) ratio on each client. A low FN ratio is crucial to achieving accurate CL Kalantidis et al. (2020). For one image sample $q$, FNs are features that we use as negative features but actually correspond to images of the same class as $q$. In centralized CL, the percentage of FNs is inherently low since diverse data are available. The model has access to the whole dataset $D$ with data from all the classes instead of a subset $D_{c}$ as in collaborative learning. Thus, when we randomly sample negatives from $D$, the FN ratio is low. For instance, when dataset $D$ has 1000 balanced classes and the negatives $n$ are randomly sampled, for any image $q$ to be learned, only $\frac{1}{1000}$ of $n$ are from the same class as $q$ and are FNs.

However, in collaborative learning, the FN ratio is inherently high on each client due to the limited data diversity, which significantly degrades the performance of contrastive learning. For instance, in real-world datasets Luo et al. (2019), one client can have only one or two classes out of seven classes. With limited data diversity on each client, when learning image sample $q$, many negatives $n$ to contrast with will be from the same class as $q$ and are FNs. To perform contrastive learning by minimizing the contrastive loss in Eq.(1), the model scatters the FNs away from $q$, which should have been clustered since they are from the same class. As a result, the representations of samples from the same class will be scattered instead of clustered and degrade the learned representations.

In local contrastive learning, each client focuses on learning its local data without considering data on the other clients. As a result, the features of data in the same class but from different clients may not be well-clustered even though they could be clustered for improved representations.

Challenge: To cluster local features to the correct remote features, one has to identify local data and remote features that are in the same class. However, since no labels are available for local data and remote features, there is no easy way to identify the correct clusters to push local features to.

To address this challenge, we propose a neighborhood matching approach to identify the remote features to cluster local data with and define an objective function to perform the clustering. First, during local learning on one client, as shown in Figure 2, for each local sample we find $N$ nearest features from both the remote and local features as neighbors. Then the features of the local sample will be pushed to these neighbors by the proposed entropy-based loss. Since the model is synchronized from the server to clients in each communication round, the remote and local features are encoded by similar models on different clients. Therefore, the neighbors are likely to be in the same class as the local sample being learned, and clustering them will improve the learned representations of global data. In this way, the global model is also improved when aggregating local models.

Identifying neighbors. To push each local sample close to its neighbors, we minimize the entropy of one sample’s matching probability distribution to either a remote feature or a local feature. To improve the robustness, we match one sample to $N$ nearest features at the same time, instead of only one nearest feature. By minimizing the entropy for $N$ nearest neighbors, the sample’s matching probability to each of the nearest neighbors will be individually certain.

For each local sample $q_{i}$, we regard top-N closest features, either from local or remote features, as neighbors. To find the neighbors, we first define the neighbor candidates as:

where $i\sim\mathcal{U}(|Q_{s}|+K,K)$ samples $K$ integer indices from $[|Q_{s}|+K]$ randomly at uniform. $Q_{s+l,i}$ is the element with index i in the union of remote and local features $Q_{s}\cup Q_{l}$. For one local sample $q_{i}$, the neighbors $P(q_{i})$, which are the top-N nearest neighbor candidates $Q^{\prime}$, is given by:

where $S_{i,m}=sim(q_{i},n_{m})={q_{i}^{T}\cdot n_{m}}/{\|q_{i}\|\|n_{m}\|}$ is the cosine similarity between $q_{i}$ and one neighbor candidate $n_{m}\in Q^{\prime}$.

Neighborhood matching loss. To make the probability of $q_{i}$ matching to each $n_{j}\in P(q_{i})$ individually certain, we consider the set:

where $d$ is the dimension of one feature vector. $L_{j}$ contains one of the top-$N$ nearest neighbors $n_{j}$ and neighbor candidates excluding all other top-$N$ nearest neighbors.

Given $L_{j}$, the probability that sample $q_{i}$ is matched to one of the neighbors $n_{a}\in L_{j}$ is:

The temperature $\tau_{nm}$ controls the softness of the probability distribution Hinton et al. (2015). Since $n_{j}$ has the largest cosine similarity with $q_{i}$ for $n\in L_{j}$, $p_{i,j,a}$ will have the largest value when $n_{a}=n_{j}$ for $n_{a}\in L_{j}$. In this way, when minimizing the entropy of the probability distribution $\{p_{i,j,a}\}_{n_{a}\in L_{j}}$, the matching probability of $q_{i}$ and $n_{j}$ will be maximized.

For one mini-batch $B$, to match each sample to its $N$ nearest neighbors, the entropy for all samples in this mini-batch is calculated as:

where $K-N+1$ is the number of features in $L_{j}$, and $N$ is the number of nearest neighbors to match. By minimizing $\mathcal{L}_{neigh}$, each $i\in B$ will be aligned to its top-$N$ nearest neighbors.

Final loss. Based on the contrastive loss with fused features in Eq.(6) and neighborhood matching loss in Eq.(11), the overall objective is formulated as:

where $\lambda$ is a weight parameter.

Datasets, model architecture, and distributed settings. We evaluate the proposed approaches on three datasets, including CIFAR-10 Krizhevsky et al. (2009), CIFAR-100 Krizhevsky et al. (2009), and Fashion-MNIST Xiao et al. (2017). We use ResNet-18 as the base encoder and use a 2-layer MLP to project the representations to 128-dimensional feature space Chen et al. (2020a); He et al. (2020). For each of the three datasets, we consider one IID setting and two non-IID settings. The detailed collaborative learning settings and training details can be found in the Appendix.

Metrics. To evaluate the quality of learned representations, we use standard metrics for centralized self-supervised learning, including linear evaluation and semi-supervised learning Chen et al. (2020a). Besides, we evaluate by collaborative finetuning for realistic collaborative learning. In linear evaluation, a linear classifier is trained on top of the frozen base encoder, and the test accuracy represents the quality of learned representations. We first perform collaborative learning by the proposed approaches without labels to learn representation. Then we fix the encoder and train a linear classifier on the 100% labeled dataset on top of the encoder. The classifier is trained for 100 epochs by the SGD optimizer following the hyper-parameters from He et al. (2020). In semi-supervised learning, we first train the base encoder without labels in collaborative learning. Then we append a linear classifier to the encoder and finetune the whole model on 10% or 1% labeled data for 20 epochs with SGD optimizer following the hyper-parameters from Caron et al. (2020). In collaborative finetuning, the learned encoder by the proposed approaches is used as the initialization for finetuning the whole model by supervised collaborative learning McMahan et al. (2017) with few locally labeled data on clients. Detailed collaborative finetuning settings can be found in the Appendix.

Baselines. We compare the proposed methods with multiple approaches. Predicting Rotation is a self-supervised learning approach by predicting the rotation of images Gidaris et al. (2018). DeepCluster-v2 is the improved version of DeepCluster Caron et al. (2020, 2018) and achieves SOTA performance. SwAV and SimCLR are SOTA approaches for self-supervised learning Caron et al. (2020); Chen et al. (2020a). We combine these approaches with FedAvg as FedRot, FedDC, FedSwAV, and FedSimCLR. FedCA is the SOTA collaborative unsupervised learning approach with a shared dictionary and online knowledge distillation Zhang et al. (2020). Besides, we compare with two methods as upper bounds. MoCo He et al. (2020) is a centralized contrastive learning method assuming all data are combined in a single location. We compare with MoCo since the local model in the proposed methods is based on it. FedAvg McMahan et al. (2017) is a fully supervised collaborative learning method.

We propose a framework for collaborative contrastive representation learning. To improve representation learning on each client, we propose feature fusion to provide remote features as accurate contrastive data to each client. To achieve unified representations among clients, we propose neighborhood matching to align each client’s local features to the remote ones. Experiments show superior accuracy of the proposed framework compared with the state-of-the-art.