HFedATM: Hierarchical Federated Domain Generalization via Optimal Transport and Regularized Mean Aggregation

2025-08-07

AI Summary

To address the degradation in generalization performance on unseen target domains caused by client domain shifts in hierarchical federated learning (HFL), this paper proposes HFedDG—the first domain-generalizable hierarchical federated learning framework. Methodologically, it introduces (1) a filter-level optimal transport alignment mechanism that achieves fine-grained cross-domain convolutional kernel distribution alignment in feature space, and (2) a contraction-aware regularized mean aggregation strategy that substantially tightens the generalization error bound and enhances training stability. Extensive experiments demonstrate that HFedDG consistently outperforms existing federated domain generalization (FedDG) methods across multiple benchmark datasets. Moreover, it achieves low communication overhead, rapid convergence, and strong robustness to domain heterogeneity—without requiring access to target-domain data during training.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Computer Vision: Learning & Optimization for CV
• Humans and AI: Human-in-the-loop Machine Learning

Application Category

• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing

Abstract

Federated Learning (FL) is a decentralized approach where multiple clients collaboratively train a shared global model without sharing their raw data. Despite its effectiveness, conventional FL faces scalability challenges due to excessive computational and communication demands placed on a single central server as the number of participating devices grows. Hierarchical Federated Learning (HFL) addresses these issues by distributing model aggregation tasks across intermediate nodes (stations), thereby enhancing system scalability and robustness against single points of failure. However, HFL still suffers from a critical yet often overlooked limitation: domain shift, where data distributions vary significantly across different clients and stations, reducing model performance on unseen target domains. While Federated Domain Generalization (FedDG) methods have emerged to improve robustness to domain shifts, their integration into HFL frameworks remains largely unexplored. In this paper, we formally introduce Hierarchical Federated Domain Generalization (HFedDG), a novel scenario designed to investigate domain shift within hierarchical architectures. Specifically, we propose HFedATM, a hierarchical aggregation method that first aligns the convolutional filters of models from different stations through Filter-wise Optimal Transport Alignment and subsequently merges aligned models using a Shrinkage-aware Regularized Mean Aggregation. Our extensive experimental evaluations demonstrate that HFedATM significantly boosts the performance of existing FedDG baselines across multiple datasets and maintains computational and communication efficiency. Moreover, theoretical analyses indicate that HFedATM achieves tighter generalization error bounds compared to standard hierarchical averaging, resulting in faster convergence and stable training behavior.

Problem

• Addresses domain shift in hierarchical federated learning
• Improves model performance on unseen target domains
• Enhances scalability and robustness in federated systems

Innovation

• Hierarchical aggregation for domain generalization
• Filter-wise Optimal Transport Alignment
• Shrinkage-aware Regularized Mean Aggregation