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

TL;DR

This paper introduces HFedATM, a hierarchical federated learning method that tackles domain shift by aligning models across stations using optimal transport and regularized mean aggregation, improving robustness and efficiency.

Contribution

It presents HFedATM, a novel hierarchical aggregation technique combining optimal transport alignment with regularized mean to enhance domain generalization in federated learning.

Findings

HFedATM outperforms baseline methods on multiple datasets.

It maintains computational and communication efficiency.

Theoretical analysis shows faster convergence and better generalization bounds.

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.