Deep Domain Isolation and Sample Clustered Federated Learning for Semantic Segmentation

Abstract

Empirical studies show that federated learning exhibits convergence issues in Non Independent and Identically Distributed (IID) setups. However, these studies only focus on label distribution shifts, or concept shifts (e.g. ambiguous tasks). In this paper, we explore for the first time the effect of covariate shifts between participants' data in 2D segmentation tasks, showing an impact way less serious than label shifts but still present on convergence. Moreover, current Personalized (PFL) and Clustered (CFL) Federated Learning methods intrinsically assume the homogeneity of the dataset of each participant and its consistency with future test samples by operating at the client level. We introduce a more general and realistic framework where each participant owns a mixture of multiple underlying feature domain distributions. To diagnose such pathological feature distributions affecting a model being trained in a federated fashion, we develop Deep Domain Isolation (DDI) to isolate image domains directly in the gradient space of the model. A federated Gaussian Mixture Model is fit to the sample gradients of each class, while the results are combined with spectral clustering on the server side to isolate decentralized sample-level domains. We leverage this clustering algorithm through a Sample Clustered Federated Learning (SCFL) framework, performing standard federated learning of several independent models, one for each decentralized image domain. Finally, we train a classifier enabling to associate a test sample to its corresponding domain cluster at inference time, offering a final set of models that are agnostic to any assumptions on the test distribution of each participant. We validate our approach on a toy segmentation dataset as well as different partitionings of a combination of Cityscapes and GTA5 datasets using an EfficientVIT-B0 model, showing a significant performance gain compared to other approaches. Our code is available at https://github.com/MatthisManthe/DDI_SCFL .

Figures (6)

• Figure 1: Overview of the Sample Clustered Federated Learning (SCFL) framework. It is composed of four sequential blocks. 1. Federated pretraining providing a model with informative gradients and a good initialization for further domain specialization. 2. Deep Domain Isolation clustering image domains in a federated fashion. 3. Sample Clustered Federated Refinement performing federated learning on each decentralized image domain. 4. Test-time cluster assignment on a test sample to select its matching cluster model.
• Figure 2: TMNIST-Inv example
• Figure 3: Two samples (image and target categories) from the two image domains of Cityscapes+GTA5 dataset.
• Figure 4: Federated repartition of samples per client
• Figure 5: PaCMAP representations of sample gradients of a model trained with FedAvg on TMNIST-Inv for 30 rounds on the Dirichlet non-IID split. Colors represent the image domains and markers the assigned clusters by Fed-GMM. Plain gradient clustering tends to be more difficult than per-class clustering. Note that PaCMAPs seem to nicely capture the image domains in both cases, while they are not applicable in a federated setup due to privacy constraints.