Federated learning improves accuracy under feature differences
Joint Domain-Class Modeling for Federated Learning Under Feature Skew
Machine Learning
Summary
Federated learning lets many users build a shared AI model without sharing their private data. However, it struggles when the input images or features on different devices vary a lot, even if the kinds of things being recognized are the same. The authors propose a new method called JDFL that figures out groups of clients with similar data styles and adjusts the model to handle these differences better without sharing raw data. Their approach improves accuracy by training the model to adapt to each group’s specific data appearance, while still learning together.
What this means in practice
- •For mobile app developers: Build image recognition apps that maintain accuracy when user photos differ in style across devices without sharing raw images.
- •For healthcare data teams: Improve federated models analyzing medical images that vary across hospitals without exposing patient data.
Authors
Sina Najafi, Mostafa Tavassolipour, Seyed Pooya Shariatpanahi
Abstract
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 $p_i(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\times 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.