# Joint Domain-Class Modeling for Federated Learning Under Feature Skew

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 29 of its 39 references that are in the map.
- First posted on arXiv 19 Sep 2026.

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

- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [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)
- [Federated Learning with Non-IID Data](https://mlsanity.com/paper/1806.00582/) (2018)
- [Model-Contrastive Federated Learning](https://mlsanity.com/paper/2103.16257/) (2021)
- [Ensemble Distillation for Robust Model Fusion in Federated Learning](https://mlsanity.com/paper/2006.07242/) (2020)
- [Federated Multi-Task Learning](https://mlsanity.com/paper/1705.10467/) (2017)
- [Federated Learning for Mobile Keyboard Prediction](https://mlsanity.com/paper/1811.03604/) (2018)
- [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)
- [An Efficient Framework for Clustered Federated Learning](https://mlsanity.com/paper/2006.04088/) (2020)
- [FedMD: Heterogenous Federated Learning via Model Distillation](https://mlsanity.com/paper/1910.03581/) (2019)
- [Data-Free Knowledge Distillation for Heterogeneous Federated Learning](https://mlsanity.com/paper/2105.10056/) (2021)
- [Personalized Federated Learning: A Meta-Learning Approach](https://mlsanity.com/paper/2002.07948/) (2020)
- [Improving Federated Learning Personalization via Model Agnostic Meta Learning](https://mlsanity.com/paper/1909.12488/) (2019)
- [Maxout Networks](https://mlsanity.com/paper/1302.4389/) (2013)
- [FedBABU: Towards Enhanced Representation for Federated Image Classification](https://mlsanity.com/paper/2106.06042/) (2021)
- [FedDC: Federated Learning with Non-IID Data via Local Drift Decoupling and Correction](https://mlsanity.com/paper/2203.11751/) (2022)
- [Federated Meta-Learning with Fast Convergence and Efficient Communication](https://mlsanity.com/paper/1802.07876/) (2018)
- [Multi-Center Federated Learning: Clients Clustering for Better Personalization](https://mlsanity.com/paper/2108.08647/) (2021)
- [Robust Federated Learning in a Heterogeneous Environment](https://mlsanity.com/paper/1906.06629/) (2019)
- [FedProc: Prototypical Contrastive Federated Learning on Non-IID data](https://mlsanity.com/paper/2109.12273/) (2021)
- [Towards Fair Federated Learning with Zero-Shot Data Augmentation](https://mlsanity.com/paper/2104.13417/) (2021)
- [Collaborative Federated Learning For Healthcare: Multi-Modal COVID-19 Diagnosis at the Edge](https://mlsanity.com/paper/2101.07511/) (2021)
- [Federated Learning for Malware Detection in IoT Devices](https://mlsanity.com/paper/2104.09994/) (2021)
- [Federated User Representation Learning](https://mlsanity.com/paper/1909.12535/) (2019)
- [Federated Learning for Ranking Browser History Suggestions](https://mlsanity.com/paper/1911.11807/) (2019)
- [From Challenges and Pitfalls to Recommendations and Opportunities: Implementing Federated Learning in Healthcare](https://mlsanity.com/paper/2409.09727/) (2024)

## Abstract (arXiv)

Federated learning (FL) enables collaborative model training without centralizing private data, but performance often degrades under feature skew: clients share labels while the conditional input distributions pi(x!mid!y) vary due to latent, client-specific appearance factors. We propose Joint Domain-Class Federated Learning (JDFL), a lightweight, optimizer-agnostic extension that makes this latent domain variation usable without sharing raw data. JDFL first infers domain clusters called pseudo-domains from brief local update signals. It then expands the classifier head to output M× C, joint (domain-class) logits. This allows the model to represent domain-conditioned appearance while keeping a shared backbone. To train the expanded head we introduce two complementary supervision strategies based on simple intuitions: a similarity-aware soft-labeling that transfers evidence between nearby inferred domains while allowing domain-specific specialization, and a per-sample randomized target assignment that perturbs supervision across the joint outputs and serves as a low-cost training-time regularizer. JDFL integrates with existing standard FL methods (e.g., FedAvg, SCAFFOLD) with minimal changes. Empirically, both supervision modes consistently improve global test accuracy on standard domain-shifted image benchmarks; ablations and sensitivity studies show the gains stem from the proposed supervision and parametrization rather than mere capacity increase.