# Towards Optimal Customized Architecture for Heterogeneous Federated Learning with Contrastive Cloud-Edge Model Decoupling

arXiv: https://arxiv.org/abs/2403.02360 · HTML: https://mlsanity.com/paper/2403.02360/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 38 references that are in the map.
- First posted on arXiv 4 Mar 2024.

## Builds on (most cited first, 26 of 26)

- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Ditto: Fair and Robust Federated Learning Through Personalization](https://mlsanity.com/paper/2012.04221/) (2020)
- [TabNet: Attentive Interpretable Tabular Learning](https://mlsanity.com/paper/1908.07442/) (2019)
- [Personalized Federated Learning with Moreau Envelopes](https://mlsanity.com/paper/2006.08848/) (2020)
- [FedBN: Federated Learning on Non-IID Features via Local Batch Normalization](https://mlsanity.com/paper/2102.07623/) (2021)
- [Federated Learning with Matched Averaging](https://mlsanity.com/paper/2002.06440/) (2020)
- [Federated Learning Based on Dynamic Regularization](https://mlsanity.com/paper/2111.04263/) (2021)
- [Federated Learning with Personalization Layers](https://mlsanity.com/paper/1912.00818/) (2019)
- [Exploiting Shared Representations for Personalized Federated Learning](https://mlsanity.com/paper/2102.07078/) (2021)
- [Clustered Federated Learning: Model-Agnostic Distributed Multi-Task Optimization under Privacy Constraints](https://mlsanity.com/paper/1910.01991/) (2019)
- [Bayesian Nonparametric Federated Learning of Neural Networks](https://mlsanity.com/paper/1905.12022/) (2019)
- [Think Locally, Act Globally: Federated Learning with Local and Global Representations](https://mlsanity.com/paper/2001.01523/) (2020)
- [No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data](https://mlsanity.com/paper/2106.05001/) (2021)
- [Personalized Federated Learning with First Order Model Optimization](https://mlsanity.com/paper/2012.08565/) (2020)
- [FedBABU: Towards Enhanced Representation for Federated Image Classification](https://mlsanity.com/paper/2106.06042/) (2021)
- [Convergence of score-based generative modeling for general data distributions](https://mlsanity.com/paper/2209.12381/) (2022)
- [SOFT: Softmax-free Transformer with Linear Complexity](https://mlsanity.com/paper/2110.11945/) (2021)
- [Multi-Center Federated Learning: Clients Clustering for Better Personalization](https://mlsanity.com/paper/2108.08647/) (2021)
- [Recovering Private Text in Federated Learning of Language Models](https://mlsanity.com/paper/2205.08514/) (2022)
- [Performance Optimization for Federated Person Re-identification via Benchmark Analysis](https://mlsanity.com/paper/2008.11560/) (2020)
- [Fed2: Feature-Aligned Federated Learning](https://mlsanity.com/paper/2111.14248/) (2021)
- [MetaFed: Federated Learning among Federations with Cyclic Knowledge Distillation for Personalized Healthcare](https://mlsanity.com/paper/2206.08516/) (2022)
- [pFedSim: Similarity-Aware Model Aggregation Towards Personalized Federated Learning](https://mlsanity.com/paper/2305.15706/) (2023)
- [Federated Two Stage Decoupling With Adaptive Personalization Layers](https://mlsanity.com/paper/2308.15821/) (2023)

## Built on by (most cited first, 2 of 2)

- [Federated Client-tailored Adapter for Medical Image Segmentation](https://mlsanity.com/paper/2504.18020/) (2025)
- [FedOBP: Federated Optimal Brain Personalization through Cloud-Edge Element-wise Decoupling](https://mlsanity.com/paper/2604.16574/) (2026)

## Abstract (arXiv)

Federated learning, as a promising distributed learning paradigm, enables collaborative training of a global model across multiple network edge clients without the need for central data collecting. However, the heterogeneity of edge data distribution drags the model towards the local minima, which can be distant from the global optimum. Such heterogeneity often leads to slow convergence and substantial communication overhead. To address these issues, we propose a novel federated learning framework called FedCMD, a model decoupling tailored to the Cloud-edge supported federated learning that separates deep neural networks into a body for capturing shared representations in Cloud and a personalized head for migrating data heterogeneity. Our motivation is that, by the deep investigation of the performance of selecting different neural network layers as the personalized head, we found rigidly assigning the last layer as the personalized head in current studies is not always optimal. Instead, it is necessary to dynamically select the personalized layer that maximizes the training performance by taking the representation difference between neighbor layers into account. To find the optimal personalized layer, we utilize the low-dimensional representation of each layer to contrast feature distribution transfer and introduce a Wasserstein-based layer selection method, aimed at identifying the best-match layer for personalization. Additionally, a weighted global aggregation algorithm is proposed based on the selected personalized layer for the practical application of FedCMD. Extensive experiments on ten benchmarks demonstrate the efficiency and superior performance of our solution compared with nine state-of-the-art solutions. All code and results are available at https://github.com/elegy112138/FedCMD.