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

2026-05-30

AI Summary

This work addresses a critical limitation in heterogeneous federated learning, where existing subnet allocation strategies predominantly consider device computational capacity while neglecting the statistical heterogeneity of client data distributions, thereby constraining model performance. To overcome this, the authors propose Heterogeneity-Aware Subnet Allocation (HASA), a novel approach that explicitly incorporates data statistical heterogeneity into subnet assignment decisions for the first time. HASA computes a heterogeneity score from each client’s local data and dynamically adjusts subnet width under a fixed computational budget, effectively balancing device constraints with data distribution characteristics. Experimental results on text prediction tasks demonstrate that HASA significantly improves performance, raising average client accuracy from 13.82% to 14.32%, while also enhancing fairness by outperforming baseline methods on both worst-case and tail clients, thus achieving a superior trade-off between overall accuracy and equitable performance.

Abstract

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.

Problem

• federated learning
• model heterogeneity
• subnet allocation
• statistical heterogeneity
• compute-constrained

Innovation

• model-heterogeneous federated learning
• subnet allocation
• statistical heterogeneity
• compute budget
• heterogeneity-aware