Overcoming Data and Model Heterogeneities in Decentralized Federated Learning via Synthetic Anchors

Conventional FL aims to learn a single generalized global model that performs optimally on all the clients’ data domains. Mathematically, the learning problem can be formulated as

where $M^{*}$ is the optimized global model from the shared model space $\mathcal{M}$, $P_{i}$ is the local data distribution at the $i$th client, $\mathcal{L}$ is the loss function and $(\mathbf{x},y)$ is the feature-label pair. Inspired by the pioneering work of FedAvg (McMahan et al., 2017), a plethora of methods have tried to fill in the performance gap of FedAvg on data-heterogeneous scenario, which can be categorized in two main orthogonal directions: Direction 1 aims to minimize the difference between the local and global model parameters to improve convergence (Li et al., 2020a; Karimireddy et al., 2020; Wang et al., 2020). Direction 2 enforces consistency in local embedded features using anchors and regularization loss (Tang et al., 2022; Zhou et al., 2022; Ye et al., 2022). This work follows the second research direction and aim to leverage anchor points to handle data heterogeneity. We also tackle the more challenging problem of domain shift, unlike other methods that only assume a label-shift amongst the client-data distributions.

Standard decentralized FL aims to solve the similar generalization objective as conventional FL (i.e., Eq. 1), only, without a central server to do so (Gao et al., 2022), and the objective applies to each local models $M_{i}$. Here, we focus on learning from each other under heterogeneous models and data distributions. This brings in an essential line of works, known as mutual learning, where clients learn from each other to obtain the essential generalizability for their models. For example, during quarantine for pandemic, hospitals want to collaboratively train a model for classifying the virus. It is desired the models are generalizable to virus variants from different area, so that after quarantine the local models are still effective for incoming tourists.

Although mutual learning with heterogeneous data and models has been studied recently, most of them assume the existence of public real data (Lin et al., 2020; Huang et al., 2022; Gong et al., 2022) or a central server to coordinate the generation of synthetic data from the local client data (Zhang et al., 2022a; Zhu et al., 2021; Zhang et al., 2022b). Another line of works rely on a server to aggregate locally generated logits or prototypes, and use it as local training guidance (Jeong et al., 2018; Chan & Ngai, 2021; Tan et al., 2022). In addition, more recent works have suggested that each clients train two models, a larger model for local training and a smaller model for mutual information exchange (Wu et al., 2022; Shen et al., 2023). However, none of the above methods simultaneously address both non-iid data and heterogeneous models under serverless and data-free setting. In this work, we explore mutual learning to optimize both local (intra-client) and global (inter-client) dataset accuracy (see the detailed setup in Sec. 3.1). We list the comparison with other methods in Table 1 and more detailed related works in Appendix G.

Suppose there are $N$ clients with $i$th client denoted as $C_{i}$. Let’s represent the private datasets on $C_{i}$ as $D_{i}=\{\mathbf{x},y\}$, where $\mathbf{x}$ is the feature and $y\in\{1,\cdots,K\}$ is the label from $K$ classes. Let $\mathcal{L}$ represent a real-valued loss function for classification (e.g.,cross-entropy loss). Denote the communication neighboring nodes of the client $C_{i}$ in the system as $\mathcal{N}(C_{i})$ and the local models as $\{M_{i}=\rho_{i}\circ\psi_{i}\}_{i=1}^{i=N}$, where $\psi_{i}$ represents the feature encoder and $\rho_{i}$ represents the classification head for the $i$th client’s model $M_{i}$. DeSA returns trained client models $\{M_{i}\}_{i=1}^{i=N}$.

Our work aims to connect two key areas, heterogeneous FL and decentralized FL, termed as decentralized federated mutual learning, where we train multiple client models in a decentralized way such that they can generalize well across all clients’ data domains. Mathematically, our objective is formulated as, for every client $i$,

where $M_{i}^{*}$ is the best possible model for client $i$ with respect to the model space $\mathcal{M}_{i}$.

The recent success of dataset distillation-based data synthesis technique that generates data with similar representation power as the original raw data (Zhao et al., 2020; Zhao & Bilen, 2023). Thus, we propose to leverage this method to efficiently and effectively generate a synthetic anchor dataset without requiring any additional model pretraining. Inspired by our theoretical analysis in Sec. 4, we utilize distribution matching (Zhao & Bilen, 2023) to distill local synthetic anchor data using the empirical maximum mean discrepancy loss (MMD) (Gretton et al., 2012) as follows,

where $\psi^{\rm rand}$ is a randomly sampled feature extractor for each iteration, $D_{i}$ is the raw data for client $i$, and $D^{Syn}_{i}$ is its target synthetic data. Following Eq. 3.2, we can manipulate the synthetic anchor dataset generation in a class-balanced manner, which enables the label prior to being unbiased towards a set of classes.

Similar to other FL work sharing distilled synthetic data for efficiency (Song et al., 2023), After local data synthesis, we request each client to share it among peers to ensure they possess same global information, and the global synthetic anchor data is denoted as $D^{Syn}=\cup_{i}D^{Syn}_{i}$ ¹¹1By default, we perform simple interpolation (averaging) among clients as it is shown that using this mixup strategy can improve model fairness (Chuang & Mroueh, 2021).. As shown in Algorithm 1, our method is designed to work only to receive neighbor node information, and DeSA is designed for peer-to-peer decentralized network. Since the loss requires all nodes’ information, we can leverage the FastMix algorithm to aggregate all nodes’ information as in (Ye et al., 2020; Luo & Ye, 2022). This method can aggregate all nodes’ information via adjacent nodes’ communication at a linear speed. It is very common in fully decentralized optimization. In fact, our method can also work if each node can only receive neighbor nodes’ information during training, and we empirically show the feasibility in our CIFAR10C experiments by sampling neighboring clients.

It is worth noting that, different from (Song et al., 2023), we further propose novel loss terms and training strategies to help mitigate the distribution discrepancy between the clients, which are detailed in the following sections (see Sec. 3.3 and Sec. 3.4), enabling improved model performance, as intuitively incorporating $D^{Syn}$ into training can only achieve sub-optimal results (see Figure 3 when both $\lambda_{KD}$ and $\lambda_{REG}$ equals to 0.).

The synthetic anchor regularization loss enforces the model to learn a client-domain invariant representation of the data. (Tang et al., 2022) and other domain incremental works (Rostami, 2021) show that, adding a distribution discrepancy based loss in the latent space enables learning of a domain-invariant encoder $\psi$. However, most of the domain adaptation works require explicit access to the real data from other domains, or generates random noise as anchors. We propose using the latent space distribution of the synthetic anchor data $D^{Syn}$ as a synthetic anchor to which the client-model specific encoders $\psi_{i}$ can project their local private data onto. The loss function can be therefore defined as,

where $K$ is the number of classes, $d$ is distance computed using the supervised contrastive loss,

where $B_{\backslash j}$ represents a batch containing both local raw data $D_{i}$ and global synthetic data $D^{Syn}$ but without data $j$, $B^{y_{j}}_{\backslash j}$ is a subset of $B_{\backslash j}$ only with samples belonging to class $y_{j}$, and $\tau_{temp}$ is a scalar temperature parameter. Note that we will detach the synthetic anchor data to ensure we are pulling local features to global features.

This step allows a single client model to learn from all the other models using a common synthetic anchor dataset $D^{Syn}$. Under our setting of model heterogeneity among clients, we cannot aggregate the model parameters by simply averaging as in FedAvg (McMahan et al., 2017). Instead, we propose to utilize knowledge distillation (KD) (Hinton et al., 2015) for decentralized model aggregation. Specifically, the fact that $D^{Syn}$ is representative of the joint distributions of the clients allows it to be an ideal dataset for knowledge transfer. Thus, to enable the client model to mimic the predictions of the other models, we also incorporate KD loss using $D^{Syn}$, formulated as

where $(\mathbf{x}^{Syn},y^{Syn})\sim D^{Syn}$, $N(i)$ is the neighbor clients of client $C_{i}$, and $\mathcal{L}_{KL}$ is the KL-Divergence between the output logits of $\mathbf{x}^{Syn}$ on $M_{i}$ and the averaged output logits of $\mathbf{x}^{Syn}$ on $M_{j},\forall j\in N(C_{i})$. Finally, we formulate our objective function as

where $\mathcal{L}_{CE}(D;M)$ is the $K$-classes cross entropy loss on data $D$ and model $M$. $\lambda_{REG}$ and $\lambda_{KD}$ are the hyperparameters for regularization and KD losses, $\mathcal{L}_{REG}$ and $\mathcal{L}_{KD}$ are as defined in Eq. 4 and Eq. 3.4, and $\bar{Z_{i}}$ is the shared logits from neighboring clients $N(C_{i})$. We also incorporate our class-conditional-generated global synthetic data $D^{Syn}$ in the CE loss to enforce models to perform well on general domains and to benefit from the augmented dataset. Overall, we formulate our objective function as

The objective of the experiments is to show that DeSA can effectively leverage and learn generalized information from other clients under data and model heterogeneities. Thus, we report the averaged global accuracy by testing $i$-th local model on every client $j$’s ($\forall j,j\in[N]$) test sets. Note FedDF and FCCL require accessing to public available data. To make a fair comparison, we use FMNIST (Xiao et al., 2017) and Cifar100 (Krizhevsky et al., 2009) as public datasets for knowledge distillation. DeSA($D_{\rm VHL}^{Syn}$) uses our training pipeline, but the synthetic data is sampled from an untrained StyleGAN (Karras et al., 2019) as in VHL (Tang et al., 2022).

Based on our experiment results for heterogeneous models in Table. 2, it is evident that our method DeSA significantly enhances the overall accuracy on DIGITS and OFFICE. In the CIFAR10C experiments, we observe the performance of FL is notably hindered by a low client sampling ratio, especially in the case of serverless methods. As shown in the table, DeSA can improves the global accuracy and out-performs most of the baseline methods, except for FCCL with CIFAR100 as the public dataset. We believe this is due to three reasons: 1) CIFAR100 has significantly more data points than our synthetic data, 2) CIFAR100 and CIFAR10C have a genuine semantic overlap (Krizhevsky et al., 2009), and 3) FCCL starts from pre-trained models, which naturally biases the comparison in favor of methods using additional CIFAR100 information. Nonetheless, it is worth mentioning that the accuracies of FedDF and FCCL drop significantly when switching from Cifar100 to FMNIST. We found that the training with both methods are easy to obtain nan loss, which we mark as ‘-’ in the table. This suggests that these methods strongly rely on the public data, which restricts the utility of the methods. In contrast, DeSA does not depend on public data, and furthermore, it is completely serverless.

Among the baseline methods for homogeneous model experiments, they learn a generalizable global model via model aggregation on a server, which is not required in DeSA. As these baseline methods are only evaluated on their single global model, for a fair comparison, we report the averaged local accuracy in the experiments. Specifically, the average local accuracy for the baseline methods is the average performance over testing the global model over all the clients; while we report, we report the average local accuracy on testing $i$-th local model on client $i$’s test set over all the clients.

One can observe from Table 3 that DeSA can effectively leverage local information and outperforms other methods DIGITS and OFFICE. For CIFAR10C, although DeSA has highest averaged local accuracy when client ratio is 0.2, it has lower performance when client ratio is 0.1. This is because smaller client sampling ratios have a larger impact on decentralized learning as we do not have a global model, and thus some clients may suffer from low model performance due to insufficient training and the scarce global information from the sampled neighbor clients. Overall, despite the serverless setting, DeSA is compatible with the baseline methods that have central servers.

The effectiveness of DeSA relies on the novel designs of synthetic anchor data and the losses. To evaluate how these designs influences the performance of DeSA, we vary the number of synthetic anchor data (IPC) and the loss coefficients ($\lambda^{\prime}s$) in the following paragraphs. If not otherwise specified, we use the default hyperparameters and model heterogeneous setting. We report the averaged global accuracy in this section.