HFedATM: Hierarchical Federated Domain Generalization

• HFedATM is a hierarchical federated domain generalization method that enhances unseen-domain robustness by structurally aligning station models.
• It employs filter-wise optimal transport alignment for convolutional layers and shrinkage-aware regularized mean aggregation for linear layers.
• Evaluations on benchmarks show HFedATM improves model accuracy with modest overhead, making it ideal for privacy-sensitive deployments.

Searching arXiv for HFedATM and closely related papers to ground the article. arXiv search: query="HFedATM Hierarchical Federated Domain Generalization via Optimal Transport and Regularized Mean Aggregation" HFedATM, short for Hierarchical Federated Domain Generalization via Optimal Transport and Regularized Mean Aggregation, is a hierarchical, data-free aggregation method for domain generalization in hierarchical federated learning (HFL). It was introduced together with the formal problem setting of Hierarchical Federated Domain Generalization (HFedDG), in which multiple stations aggregate client models locally and a server subsequently merges station models under domain shift, with the target being robust performance on an unseen domain. HFedATM operates specifically at the station-to-server aggregation layer: it aligns convolutional filters across station models through Filter-wise Optimal Transport (FOT) Alignment, then merges aligned models with Shrinkage-aware Regularized Mean (RegMean) Aggregation so that cross-station averaging is semantics-aware for convolutional layers and activation-geometry-aware for linear layers (Nguyen et al., 7 Aug 2025).

1. HFedDG setting and the problem HFedATM addresses

In standard federated learning, a single central server aggregates client updates. HFL introduces an intermediate layer of stations, yielding a three-tier structure

$$\text{server} \rightarrow \text{stations} \rightarrow \text{clients}.$$

This architecture is intended to improve scalability, bandwidth efficiency, and robustness. HFedATM is defined for this hierarchical regime rather than for conventional single-server federated learning (Nguyen et al., 7 Aug 2025).

The motivating difficulty is domain shift. In HFedDG, different clients and stations observe different source-domain distributions, while deployment is evaluated on an unseen target-domain distribution. The formalization uses $E=\{1,\dots,N_E\}$ for stations, $C_e=\{1,\dots,N_e\}$ for clients under station $e$, and client datasets

$$S_{e,i}=\{(x_{e,i}^{(j)},y_{e,i}^{(j)})\}_{j=1}^{n_{e,i}\sim P_{XY}^{e,i}},$$

with an unseen target distribution $P_{XY}^\star$ satisfying $P_{XY}^\star \neq P_{XY}^{e,i}$ for all $(e,i)$. The objective is to learn a hypothesis $h$ minimizing unseen-domain risk,

$$D_{\mathrm{target}}(h)=\mathbb{E}_{(x,y)\sim P^\star_{XY}}[\ell(h(x),y)].$$

Stations do not access raw client data; they aggregate only client models (Nguyen et al., 7 Aug 2025).

The paper’s generalization analysis identifies two residual error sources once local client-side learning has already attempted to produce domain-robust models: intra-station and inter-station distribution mismatch. These are summarized by station-level inner divergence $E=\{1,\dots,N_E\}$0, station-level breadth $E=\{1,\dots,N_E\}$1, server-level inner divergence $E=\{1,\dots,N_E\}$2, and server-level breadth $E=\{1,\dots,N_E\}$3. The hierarchical target-risk bound contains these terms explicitly, which is the principal theoretical motivation for a station-server merger that is more structure-aware than plain hierarchical averaging. In that sense, HFedATM is not a replacement for client-side FedDG training; it is a response to the claim that a substantial remaining performance gap arises from how station models are merged at the server (Nguyen et al., 7 Aug 2025).

2. Placement in the hierarchical training pipeline

HFedATM is a station-to-server aggregation rule compatible with multiple client-side methods, including FedAvg, FedProx, FedSR, and FedIIR. The full pipeline begins with server broadcast of the previous global model $E=\{1,\dots,N_E\}$4 to all stations. Within each station, over $E=\{1,\dots,N_E\}$5 station rounds, an active client subset $E=\{1,\dots,N_E\}$6 is selected, each client trains locally for $E=\{1,\dots,N_E\}$7 epochs, and client parameters $E=\{1,\dots,N_E\}$8 are uploaded to the station (Nguyen et al., 7 Aug 2025).

At the last forward pass of the final local epoch, clients additionally collect dense-layer activation statistics. For each linear layer $E=\{1,\dots,N_E\}$9, client $C_e=\{1,\dots,N_e\}$0 forms

$C_e=\{1,\dots,N_e\}$1

where $C_e=\{1,\dots,N_e\}$2 is the activation matrix for a mini-batch. The station averages client Grams,

$C_e=\{1,\dots,N_e\}$3

and applies diagonal shrinkage,

$C_e=\{1,\dots,N_e\}$4

with default $C_e=\{1,\dots,N_e\}$5. Each station then sends to the server both the station model $C_e=\{1,\dots,N_e\}$6 and the shrunk Gram matrices $C_e=\{1,\dots,N_e\}$7 (Nguyen et al., 7 Aug 2025).

The server chooses station 1 as the reference. It first performs FOT Alignment for every convolutional layer across stations. After alignment, convolutional layers are merged by weighted arithmetic mean, whereas linear layers are merged via RegMean in closed form. The merged model $C_e=\{1,\dots,N_e\}$8 is then broadcast back to all stations (Nguyen et al., 7 Aug 2025).

This design is important because HFedATM leaves local optimization untouched and modifies only the hierarchical merging mechanism. The method is therefore best understood as a top-level model-merging rule for HFL rather than as a full end-to-end replacement for federated training.

3. Filter-wise Optimal Transport Alignment

The first core component addresses a structural weakness of naive parameter averaging in convolutional networks: permutation symmetry of filters. Filter index $C_e=\{1,\dots,N_e\}$9 in one station model need not correspond semantically to filter index $e$0 in another station model. Index-wise averaging can therefore blend unrelated features and degrade the merged representation. HFedATM handles this by solving a one-to-one filter matching problem before cross-station convolutional averaging (Nguyen et al., 7 Aug 2025).

For station $e$1 and convolutional layer $e$2, the filter bank is written as

$e$3

Each filter is flattened and $e$4-normalized,

$e$5

For stations $e$6 and $e$7, the squared Euclidean cost matrix is

$e$8

Using station $e$9 as reference, HFedATM solves

$$S_{e,i}=\{(x_{e,i}^{(j)},y_{e,i}^{(j)})\}_{j=1}^{n_{e,i}\sim P_{XY}^{e,i}},$$0

over the Birkhoff polytope

$$S_{e,i}=\{(x_{e,i}^{(j)},y_{e,i}^{(j)})\}_{j=1}^{n_{e,i}\sim P_{XY}^{e,i}},$$1

The paper states that this discrete OT assignment problem is solved approximately by entropic Sinkhorn (Nguyen et al., 7 Aug 2025).

The resulting transport plan acts as an alignment permutation:

$$S_{e,i}=\{(x_{e,i}^{(j)},y_{e,i}^{(j)})\}_{j=1}^{n_{e,i}\sim P_{XY}^{e,i}},$$2

If filter sizes differ, they are resized to the reference size $$S_{e,i}=\{(x_{e,i}^{(j)},y_{e,i}^{(j)})\}_{j=1}^{n_{e,i}\sim P_{XY}^{e,i}},$$3,

$$S_{e,i}=\{(x_{e,i}^{(j)},y_{e,i}^{(j)})\}_{j=1}^{n_{e,i}\sim P_{XY}^{e,i}},$$4

FOT is therefore applied before any cross-station convolutional averaging, and its purpose is to create a common semantic channel order across station models (Nguyen et al., 7 Aug 2025).

The paper also proves a permutation-invariance lemma for the OT cost,

$$S_{e,i}=\{(x_{e,i}^{(j)},y_{e,i}^{(j)})\}_{j=1}^{n_{e,i}\sim P_{XY}^{e,i}},$$5

for any permutation matrix $$S_{e,i}=\{(x_{e,i}^{(j)},y_{e,i}^{(j)})\}_{j=1}^{n_{e,i}\sim P_{XY}^{e,i}},$$6. This is used to support the interpretation of filter reindexing as semantic alignment rather than an arbitrary transformation. The method is explicitly data-free at this stage: no raw inputs or activation tensors are required for filter alignment, and computational cost depends mainly on the number of kernels $$S_{e,i}=\{(x_{e,i}^{(j)},y_{e,i}^{(j)})\}_{j=1}^{n_{e,i}\sim P_{XY}^{e,i}},$$7, not on dataset size (Nguyen et al., 7 Aug 2025).

4. Shrinkage-aware Regularized Mean Aggregation

The second core component addresses linear layers. Even after convolutional filters have been aligned, direct coordinate-wise averaging of dense-layer weights remains problematic because those weights are coupled to the feature covariance geometry induced by each station’s data. HFedATM therefore uses a RegMean merger driven by activation Gram matrices rather than by naive tensor averaging (Nguyen et al., 7 Aug 2025).

For a dense layer $$S_{e,i}=\{(x_{e,i}^{(j)},y_{e,i}^{(j)})\}_{j=1}^{n_{e,i}\sim P_{XY}^{e,i}},$$8, each client records activations $$S_{e,i}=\{(x_{e,i}^{(j)},y_{e,i}^{(j)})\}_{j=1}^{n_{e,i}\sim P_{XY}^{e,i}},$$9 and computes

$P_{XY}^\star$0

The paper allows two optional privacy mechanisms before transmission to the station: clients may clip $P_{XY}^\star$1, and they may add Gaussian noise for differential privacy. Station-level averaging and shrinkage then yield

$P_{XY}^\star$2

and

$P_{XY}^\star$3

The shrinkage step is described as stabilizing the covariance estimate and reducing sensitivity to noise and sampling variability (Nguyen et al., 7 Aug 2025).

Given aligned station linear weights $P_{XY}^\star$4, the server seeks a merged weight $P_{XY}^\star$5 minimizing

$P_{XY}^\star$6

Replacing $P_{XY}^\star$7 by $P_{XY}^\star$8 gives the reported closed-form aggregation rule

$P_{XY}^\star$9

This is the defining RegMean step in HFedATM (Nguyen et al., 7 Aug 2025).

The contrast between convolutional and linear aggregation is explicit. After FOT, convolutional filters are merged by weighted arithmetic mean,

$P_{XY}^\star \neq P_{XY}^{e,i}$0

typically with $P_{XY}^\star \neq P_{XY}^{e,i}$1, while dense layers are merged via the RegMean solution above. The final global model therefore combines FOT-aligned convolutional averaging with Gram-weighted linear-layer merging (Nguyen et al., 7 Aug 2025).

5. Theoretical analysis and target-risk contraction

The theoretical treatment extends domain-generalization analysis to the hierarchical setting. The appendix defines the $P_{XY}^\star \neq P_{XY}^{e,i}$2-divergence as

$P_{XY}^\star \neq P_{XY}^{e,i}$3

and assumes bounded or Lipschitz loss together with a Hölder continuity condition,

$P_{XY}^\star \neq P_{XY}^{e,i}$4

and

$P_{XY}^\star \neq P_{XY}^{e,i}$5

A lemma then yields

$P_{XY}^\star \neq P_{XY}^{e,i}$6

These ingredients support the hierarchical target-risk analysis for HFedATM (Nguyen et al., 7 Aug 2025).

The central bound states that if client-level training already achieves

$P_{XY}^\star \neq P_{XY}^{e,i}$7

then the HFedATM output satisfies

$P_{XY}^\star \neq P_{XY}^{e,i}$8

The interpretation given in the paper is that HFedATM geometrically contracts divergence and breadth terms over rounds and therefore achieves a tighter bound than plain hierarchical averaging (Nguyen et al., 7 Aug 2025).

That contraction is decomposed into two mechanisms. FOT reduces breadth terms through

$P_{XY}^\star \neq P_{XY}^{e,i}$9

while RegMean reduces divergence terms through

$(e,i)$0

Combining them with

$(e,i)$1

gives

$(e,i)$2

The theoretical message is therefore conditional: HFedATM is most beneficial when client-side training has already made local risks small enough that hierarchical aggregation itself becomes the dominant source of remaining generalization error (Nguyen et al., 7 Aug 2025).

6. Empirical evaluation and comparative behavior

HFedATM is evaluated on four benchmark families: PACS, Office-Home, TerraInc, and Amazon Reviews. The simulated HFL topology uses 10 stations, 100 clients total, and 10 clients per station, with all clients participating each round. Heterogeneity is controlled by $(e,i)$3, where $(e,i)$4 is relatively IID and $(e,i)$5 is maximally heterogeneous. The main backbones are LeNet-5 for vision and RoBERTa-base for NLP, with additional robustness studies on ResNet-18, VGG-11, and DeBERTa-base. Appendix hyperparameters include local epochs $(e,i)$6, station rounds per server round $(e,i)$7, batch size $(e,i)$8, total global rounds $(e,i)$9, SGD with learning rate $h$0 for vision, AdamW with learning rate $h$1 for NLP, cosine scheduler, Sinkhorn regularizer $h$2, shrinkage $h$3, and Sinkhorn iterations $h$4. Experiments use NVIDIA RTX 3090 GPUs and three random seeds $h$5 (Nguyen et al., 7 Aug 2025).

Across all four benchmark families, attaching HFedATM to existing baselines improves unseen-domain accuracy. Representative examples reported in Table 1 include PACS with $h$6, where FedSR + Avg gives $h$7 and FedSR + HFedATM gives $h$8; and Office-Home with $h$9, where FedIIR + Avg gives $$D_{\mathrm{target}}(h)=\mathbb{E}_{(x,y)\sim P^\star_{XY}}[\ell(h(x),y)].$$0 and FedIIR + HFedATM gives $$D_{\mathrm{target}}(h)=\mathbb{E}_{(x,y)\sim P^\star_{XY}}[\ell(h(x),y)].$$1. For Amazon Reviews, the paper states that gains are particularly large and consistent, with FedIIR + Avg around $$D_{\mathrm{target}}(h)=\mathbb{E}_{(x,y)\sim P^\star_{XY}}[\ell(h(x),y)].$$2 and FedIIR + HFedATM around $$D_{\mathrm{target}}(h)=\mathbb{E}_{(x,y)\sim P^\star_{XY}}[\ell(h(x),y)].$$3 (Nguyen et al., 7 Aug 2025).

The paper emphasizes that HFedATM performs best when paired with stronger local DG methods such as FedSR and FedIIR. This is presented as empirical support for the theory: when initial station models are already reasonably domain-generalizable, station-server aggregation quality becomes especially consequential. The architecture-robustness study reports improvements for all tested backbones; for example, FedIIR + Avg vs +HFedATM changes from $$D_{\mathrm{target}}(h)=\mathbb{E}_{(x,y)\sim P^\star_{XY}}[\ell(h(x),y)].$$4 on ResNet-18 PACS, $$D_{\mathrm{target}}(h)=\mathbb{E}_{(x,y)\sim P^\star_{XY}}[\ell(h(x),y)].$$5 on Office-Home, $$D_{\mathrm{target}}(h)=\mathbb{E}_{(x,y)\sim P^\star_{XY}}[\ell(h(x),y)].$$6 on TerraInc, and $$D_{\mathrm{target}}(h)=\mathbb{E}_{(x,y)\sim P^\star_{XY}}[\ell(h(x),y)].$$7 on DeBERTa Amazon Reviews (Nguyen et al., 7 Aug 2025).

Ablation results show that removing either major component lowers accuracy: w/o FOT degrades results, w/o RegMean also degrades results, and the full method is best. The intended interpretation is that the two mechanisms are complementary: FOT corrects semantic filter mismatch, whereas RegMean corrects feature-correlation mismatch in linear layers (Nguyen et al., 7 Aug 2025).

The overhead is reported as modest. Per-round latency increases by about $$D_{\mathrm{target}}(h)=\mathbb{E}_{(x,y)\sim P^\star_{XY}}[\ell(h(x),y)].$$8 on PACS, $$D_{\mathrm{target}}(h)=\mathbb{E}_{(x,y)\sim P^\star_{XY}}[\ell(h(x),y)].$$9 on Office-Home and TerraInc, and $E=\{1,\dots,N_E\}$00 on Amazon Reviews. The paper attributes this to efficient Sinkhorn OT and Gram computation in one final forward pass. Under Gaussian noise added to Gram matrices for $E=\{1,\dots,N_E\}$01-DP, performance remains stable for moderate budgets $E=\{1,\dots,N_E\}$02, usually with less than $E=\{1,\dots,N_E\}$03 accuracy drop, and even at $E=\{1,\dots,N_E\}$04 the degradation is described as graceful (Nguyen et al., 7 Aug 2025).

7. Privacy stance, limitations, and relevance beyond the benchmark suite

HFedATM is described as data-free because no raw examples are transmitted. That description is exact for FOT, which depends only on weights, but the full method also uses Gram matrices of activations for RegMean. The paper argues that these statistics remain relatively privacy-preserving because clients may clip and noise them for DP, and because Gram matrices are many-to-one and not uniquely invertible. The appendix notes that for $E=\{1,\dots,N_E\}$05, $E=\{1,\dots,N_E\}$06 cannot be uniquely recovered; if $E=\{1,\dots,N_E\}$07 for orthogonal $E=\{1,\dots,N_E\}$08, then $E=\{1,\dots,N_E\}$09, and repeated eigenvalues or rank deficiency permit infinitely many compatible activations (Nguyen et al., 7 Aug 2025).

The principal limitations are also explicit. HFedATM assumes homogeneous model architectures across clients and stations. Its benefits depend on the quality of local client-side training: if client models are highly heterogeneous and not yet domain-generalizable, the gains shrink. Gram sharing introduces some privacy surface even if it is weaker than raw-data sharing. The FOT design aligns all stations to station 1; the paper does not deeply investigate multi-reference or barycentric alternatives. These constraints delimit the method’s scope and distinguish it from a fully architecture-agnostic or strictly weights-only merger (Nguyen et al., 7 Aug 2025).

The paper positions HFedATM for large-scale HFL systems with many devices and intermediate edge aggregators, especially in settings with multiple source domains, unseen target domains, and restrictions on raw-data sharing. The intended application categories include healthcare, surveillance, mobile systems, and other multi-tier federated deployments (Nguyen et al., 7 Aug 2025).

A common source of ambiguity is the acronym “HF” itself. In HFedATM, “HF” denotes the hierarchical federated setting, not heart failure. Even so, a plausible implication is that the method is germane to privacy-sensitive heart-failure sensing pipelines. A Chinese heart-failure speech study releases high-level features / Full data by request, explicitly to preserve patient privacy, and identifies individual difference as a major cause of inaccuracy, which suggests a natural motivation for federated or personalized aggregation mechanisms rather than centralized raw-audio pooling (Pan et al., 12 Aug 2025). A further plausible implication is that multimodal wearable heart-failure risk models built from short ECG and sampled long-term HRV could serve as local models inside an HFL stack, since such models already target device heterogeneity, sparse sensing, and privacy-sensitive health data, although that work itself does not study federated learning (González et al., 2024).

Within federated learning research, HFedATM’s distinctive contribution is therefore twofold. It formalizes HFedDG as a hierarchical analogue of federated domain generalization, and it proposes a station-server merger that is explicitly designed for the two structural failures of naive hierarchical averaging: permutation mismatch in convolutional filters and incompatible second-order activation structure in dense layers.