Domain-Skewed Federated Learning with Feature Decoupling and Calibration

Federated learning (FL) allows distributed clients to collaboratively train a global model in a privacy-preserving manner. However, one major challenge is domain skew, where clients’ data originating from diverse domains may hinder the aggregated global model from learning a consistent representation space, resulting in poor generalizable ability in multiple domains. In this paper, we argue that the domain skew is reflected in the domain-specific biased features of each client, causing the local model’s representations to collapse into a narrow low-dimensional subspace. We then propose Federated Feature Decoupling and Calibration ($F^2$DC), which liberates valuable class-relevant information by calibrating the domain-specific biased features, enabling more consistent representations across domains. A novel component, Domain Feature Decoupler (DFD), is first introduced in $F^2$DC to determine the robustness of each feature unit, thereby separating the local features into domain-robust features and domain-related features. A Domain Feature Corrector (DFC) is further proposed to calibrate these domain-related features by explicitly linking discriminative signals, capturing additional class-relevant clues that complement the domain-robust features. Finally, a domain-aware aggregation of the local models is performed to promote consensus among clients. Empirical results on three popular multi-domain datasets demonstrate the effectiveness of the proposed $F^2$DC and the contributions of its two modules. Code is available at https://github.com/mala-lab/F2DC.

💡 Research Summary

Federated learning (FL) enables multiple clients to collaboratively train a shared model without exchanging raw data, thereby preserving privacy. In many real‑world scenarios, however, each client’s data originates from a distinct domain (e.g., different weather conditions, sensor types, or artistic styles) while maintaining a similar class distribution. This “domain skew” leads to divergent conditional distributions P(x|y) across clients, causing local models to overfit domain‑specific cues and to produce feature representations that collapse into narrow, low‑dimensional subspaces. Existing FL methods mainly address label heterogeneity or apply regularization/aggregation tricks, often discarding domain‑biased parameters altogether. Such elimination‑based approaches risk losing valuable class‑relevant information that is entangled with domain‑specific features.

The paper proposes Federated Feature Decoupling and Calibration (F²DC), a three‑stage framework designed to correct, rather than erase, domain‑biased features. The first stage introduces a Domain Feature Decoupler (DFD). For each client k, the feature extractor is split into backbone layers (r_B^k) and flattening layers (r_F^k). Given an intermediate feature map f_i, a lightweight two‑layer CNN (the attribution network A_D) computes a robustness score map S_i for each spatial‑channel unit. To obtain a differentiable binary mask, the authors employ a Gumbel‑Concrete distribution, yielding a pseudo‑binary mask M_i whose entries near 1 indicate domain‑robust units and near 0 indicate domain‑related units. The mask partitions f_i into f_i⁺ = M_i ⊙ f_i (domain‑robust) and f_i⁻ = (1−M_i) ⊙ f_i (domain‑related).

DFD is trained with two complementary losses. The separability loss minimizes the cosine similarity between the embeddings l_i⁺ = r_F^k(f_i⁺) and l_i⁻ = r_F^k(f_i⁻), encouraging the two feature groups to occupy orthogonal directions. The discriminability loss applies a single‑layer MLP m to each embedding, producing class logits; a cross‑entropy term penalizes the ground‑truth label y_i and the most confident incorrect label b_{y_i}, weighted by a temperature τ that controls representation concentration. Together, these objectives push A_D to assign high mask values to truly domain‑invariant cues while suppressing domain‑specific noise.

Because perfect separation is impossible—domain‑related features often still contain useful class signals—the second stage introduces a Domain Feature Corrector (DFC). DFC receives f_i⁻ and, through a series of class‑conditioned attention modules and normalization layers, generates a calibrated feature f_i★ that re‑injects the latent class information while attenuating pure domain artifacts. The final client‑side representation is the concatenation (or weighted sum) of the robust part f_i⁺ and the corrected part f_i★, enabling the local classifier to leverage both invariant and complementary cues.

The third stage performs domain‑aware aggregation on the server. Each client’s domain discrepancy p_k (e.g., measured by the variance of its feature covariance matrix) is used as a weighting factor in a modified FedAvg scheme, ensuring that clients with extreme domain shifts do not dominate the global update. This promotes a more balanced and consensus‑driven global model.

Empirical evaluation spans three widely used multi‑domain benchmarks: PACS, Office‑Home, and DomainNet. Across all settings, F²DC consistently outperforms strong baselines such as FedAvg, FedProx, and the recent FDSE method, achieving 2–5 percentage‑point gains in average accuracy. Singular‑value analysis of the feature covariance matrix demonstrates that vanilla FL suffers severe dimensional collapse (many near‑zero singular values), whereas F²DC yields a more uniform spectrum, confirming that the learned representations retain richer information. Qualitative Grad‑CAM visualizations show that F²DC captures holistic object parts across domains, unlike FDSE which misses critical regions (e.g., giraffe antlers in cartoon sketches). Ablation studies reveal that removing either DFD or DFC leads to a marked performance drop, underscoring the necessity of both decoupling and calibration.

In summary, the paper makes four key contributions: (1) it reframes domain bias correction as a calibration problem rather than elimination, preserving entangled class information; (2) it introduces a novel, differentiable feature‑decoupling mechanism based on Gumbel‑Concrete masking; (3) it proposes a corrective module that explicitly links domain‑related features to semantic signals; and (4) it incorporates domain‑aware aggregation to enhance global consensus. Limitations include the extra computational overhead of sampling masks and the current design’s dependence on specific backbone architectures; future work may explore lightweight mask generators and broader architectural compatibility. Overall, F²DC offers a compelling solution to the domain‑skew challenge in federated learning, advancing both theoretical understanding and practical performance.