FedSAP: Federated Learning with Structured Adaptive Partitioning for Multi-Domain Heterogeneous Edge Devices

Wentao Yue; Tianyou Lai; Hongji Li; Qingyu Mao; Qilei Li

cs.LGOct 1, 2026

FedSAP: Federated Learning with Structured Adaptive Partitioning for Multi-Domain Heterogeneous Edge Devices

Authors: Wentao Yue, Tianyou Lai, Hongji Li, Qingyu Mao, Qilei Li

Abstract

Federated learning (FL) on heterogeneous edge devices must jointly accommodate unequal resource budgets and domain-shifted local data. Existing resource-adaptive methods decide how much of a model each client trains but not where retained capacity should reside or how it should be shared, whereas federated domain-generalization methods usually assume a shared full architecture. Uniform compression can therefore discard high-utility channels, and a single aggregation path can mix transferable features with domain-sensitive updates. We propose FedSAP, a domain-aware heterogeneous FL framework that casts structured pruning as budget-constrained tri-state channel allocation. FedSAP converts each keep ratio into non-uniform layer budgets, assigns stable channels to a Global pool, useful domain-sensitive channels to pseudo-domain-specific Private pools, and low-utility channels to a Dropped state. This partition lets broadly useful features benefit from cross-client pooling while isolating domain-sensitive updates from incompatible clients. Domain-Guided Assignment infers pseudo-domains from shallow-gradient similarity, while Type-Matched Aggregation restricts each channel to its intended sharing scope. Across three random seeds, FedSAP reaches 76.00% and 72.67% mean global accuracy on Digits and Office-Caltech, exceeding the strongest baseline by 1.70 and 4.92 percentage points while supporting client pruning ratios of up to 80% across heterogeneous clients.

Figures & tables

Panel of Figure 1 — Figure 1: Representative domain shifts. Digits spans handwritten, street-view, and synthetic domains; Office-Caltech spans different backgrounds, viewpoints, resolutions, and camera conditions.

Figure 2: Conventional pruning versus FedSAP on ResNet10. Conventional pruning keeps one channel set, whereas FedSAP distributes the client budget across layers and assigns channels to Global or Private type-matched paths.

Figure 3: Overview of FedSAP: the server discovers pseudo-domain groups, allocates client-level budgets across layers, assigns channels to Global , Private , or Dropped states, and aggregates channels through matched sharing scopes.

FL frameworks	System Heter.	Digits					Office-Caltech
FL frameworks	System Heter.	MNIST	USPS	SVHN	SYN	Global Acc.	Caltech	Amazon	Webcam	DSLR	Global Acc.
FedAvg ( McMahan et al. 2017 )	✗	95.89	86.84	78.39	33.63	71.81	66.07	76.84	65.52	56.67	64.54
MOON ( Li et al. 2021 )	✗	93.03	78.38	84.45	25.97	69.44	65.62	75.79	72.41	53.33	61.86
FedSR ( Caldarola et al. 2022 )	✗	96.77	86.15	81.48	31.64	73.89	62.95	78.95	75.86	50.00	65.47
FPL ( Huang et al. 2023 )	✗	95.54	87.69	83.74	34.73	74.17	63.84	82.63	65.52	60.00	65.45
FedSPU ( Niu et al. 2025 )	✗	96.53	85.23	77.99	39.00	73.44	66.07	73.68	53.45	46.67	59.68

Table 1: Mean model accuracy (%) over three random seeds on Digits and Office-Caltech. FedSAP uses

g_{\mathrm{target}}=0.9

\tau_{\mathrm{var}}=0.05

, and

\beta_{\mathrm{EMA}}=0.5

. Per-domain entries are best-so-far accuracies and may be attained at different communication rounds; Global Acc. is the best-so-far unweighted mean of domain accuracies evaluated at the same round, so it need not equal the mean of the per-domain columns.

Figure 4: Component ablation of FedSAP on Digits under the selected adaptive configuration. Removing DGA, SGP, or both degrades accuracy, showing the complementary benefits of pseudo-domain grouping and stable global-channel protection.

Figure 5: One-factor hyperparameter sensitivity of FedSAP on Digits under the fixed-mask, uniform-budget reference setting. The coordinate-wise best values are transferred to the final adaptive configuration.

Gradient source	Acc. (%)	$\Delta$ from Conv1
Conv1	75.54	–
Layer1	75.76	+0.22
Layer2	75.68	+0.14
Layer3	75.11	-0.43

Table 2: Sensitivity to the DGA gradient source under the fixed-mask reference setting on Digits.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Configuration item	Main adaptive	Fixed-mask reference
Layer-budget rule	adaptive_layer_budget =1	adaptive_layer_budget =0
Layer-wise budget pattern	Importance-adaptive	Uniform
$g_{\mathrm{target}}$	0.9	0.9
$\tau_{\mathrm{var}}$	0.05	0.1
$\beta_{\mathrm{EMA}}$	0.5	0.3
Utility weights $(\lambda_{a},\lambda_{g})$	$(0.5,0.5)$	$(0.5,0.5)$

Appendix

Table 3: FedSAP hyperparameter settings used in the main comparison and diagnostic analyses.

Type	Hyperparameter	Value
FL training	Number of clients	10
	Client participation rate	100%
	Communication rounds	100
	Local epochs	10
Optimization	Batch size	64
	Optimizer	SGD

Appendix

Table 4: General training hyperparameters used in our experiments. FedSAP-specific configuration differences are shown in Table 3 .

Aggregation path	Uniform layer budget	Adaptive layer budget
	( adaptive_layer_budget =0)	( adaptive_layer_budget =1)
Single path (w/o Tri-State)	73.5	74.08
Tri-state type-matched (FedSAP)	75.2	76.00

Appendix

Table 5: Global accuracy (%) on Digits under the

2\times 2

combination of aggregation path and layer-budget rule. Entries are best-so-far global accuracies over training rounds, averaged over three seeds.

Domain	Uniform budget + tri-state	Complete FedSAP	Difference pattern
MNIST	96.5	96.83	Small (near saturation)
USPS	89.5	90.23	Small
SVHN	75.0	76.51	Moderate
SYN	42.0	43.85	Largest (most budget-sensitive)

Appendix

Table 6: Domain-level accuracy (%) for uniform-budget tri-state aggregation and the complete adaptive-budget FedSAP configuration. Each entry is the best-so-far accuracy of the corresponding domain, which may be attained at a different communication round; the unweighted mean of the domain entries therefore need not equal the synchronized global accuracy in Table 5 .

Hyperparameter	Value	Global accuracy (%)
$g_{\mathrm{target}}$	0.5	73.71
$g_{\mathrm{target}}$	0.6	73.15
$g_{\mathrm{target}}$	0.7	73.96
$g_{\mathrm{target}}$	0.8	74.91
$g_{\mathrm{target}}$	0.9	75.54
$g_{\mathrm{target}}$	0.95	75.08

Appendix

Table 7: Hyperparameter sensitivity on Digits under the fixed-mask reference setting.

Component	FedAvg	FedSAP overhead
Local training	100%	0% additional
Type-matched aggregation	100%	about $1\%$
Mask update	–	$<0.1\%$
DGA update	–	about $3\%$ – $5\%$ amortized
Evaluation	100%	0% additional
Communication	100%	small channel-level payload

Appendix

Table 8: Qualitative comparison of additional overhead relative to FedAvg.

Figure 6: Round-wise and best-so-far global accuracy of three FedSAP configuration variants on Digits.

Figure 7: Per-domain accuracy curves on Digits and Office-Caltech.

Figure 8: t-SNE visualization on Digits.

Figure 9: t-SNE visualization on Office-Caltech.

Explore similar work

HASA: Subnet Allocation for Compute-Constrained Model-Heterogeneous Federated Learning
May 30, 2026Amir Hossein Shahdadian, Ahmed M. Abdelmoniem, Mahdi Taheri +2Heterogeneous Federated Learning Federated Learning
University of Naples "Federico II", Italy · Queen Mary University of London, United Kingdom · Brandenburg University of Technology Cottbus-Senftenberg, Germany +2
Fed-BAC: Federated Bandit-Guided Additive Clustering in Hierarchical Federated Learning
May 12, 2026Satwat Bashir, Tasos Dagiuklas, Muddesar IqbalHeterogeneous Federated Learning Federated Learning
Computer Science and Informatics · London South Bank University · London, UK +1
FEAST: Federated Shared-Space Training for Resource-Heterogeneous Clients
Aug 10, 2026Bostan Khan, Masoud DaneshtalabBalanced Learning Taco
Mälardalen University, Västerås, Sweden

May 30, 2026, cs.LGJ/K move · Enter open · S save

HASA: Subnet Allocation for Compute-Constrained Model-Heterogeneous Federated Learning

Amir Hossein Shahdadian, Ahmed M. Abdelmoniem, Mahdi Taheri, Samira Nazari+1

University of Naples "Federico II", Italy · Queen Mary University of London, United Kingdom · Brandenburg University of Technology Cottbus-Senftenberg, Germany+2

Heterogeneous Federated Learning Federated Learning Training Data

Open paper AskPDF arXiv

Edge services increasingly use federated learning to personalize on-device models while keeping sensitive data local. In practice, deployments must handle heterogeneity in both client resources and local data distributions. Model-heterogeneous federated learning lowers client cost by allowing each client to train a subnet of a shared supernet, but most subnet-allocation policies are driven by device constraints and do not explicitly account for statistical heterogeneity. This paper proposes Heterogeneity-Aware Subnet Allocation (HASA), a train-only rule that assigns subnet widths based on client heterogeneity scores computed from local training data while enforcing a fixed size-weighted compute budget. This design enables budget-matched comparisons with alternative allocation policies. On an article-title next-word prediction benchmark with seven clients, HASA improves unweighted mean client test accuracy over uniform allocation across 10 matched seeds, increasing mean client test accuracy from 13.82 percent to 14.32 percent, and improves worst-client accuracy on average. In a matched-budget comparison with representative partial-training baselines, HASA achieves the strongest worst-client and tail-client accuracy on this benchmark. A directionality ablation shows that assigning smaller subnets to more heterogeneous clients degrades both mean and tail performance. A cross-domain image-classification study further shows that the effectiveness of heterogeneity-aware allocation depends on how well the heterogeneity score reflects clients' need for additional model width.

Abstract

Explore similar work

HASA: Subnet Allocation for Compute-Constrained Model-Heterogeneous Federated Learning

Fed-BAC: Federated Bandit-Guided Additive Clustering in Hierarchical Federated Learning

FEAST: Federated Shared-Space Training for Resource-Heterogeneous Clients