Source-linked AI summary
HeteroFL: Computation and Communication Efficient Federated Learning for Heterogeneous Clients
Enmao Diao, Jie Ding, Vahid Tarokh
TL;DR
Federated learning must accommodate clients with widely differing computation and communication capabilities while preserving private local data and producing one global model. HeteroFL trains and aggregates heterogeneous local subnetworks, adds strategies for non-IID data and changing capabilities, and reports better results with fewer communication rounds across multiple architectures.
Problem
Conventional FL assumes local models share the global model’s architecture, forcing global complexity to accommodate the least capable client despite heterogeneous and dynamic capabilities.
Method
HeteroFL trains heterogeneous local models by shrinking complexity within the same model class, then aggregates them into a single global inference model.
Results
The experiments report better results with fewer communication rounds across CNN, PreResNet18, and Transformer architectures, with robustness to statistical heterogeneity and dynamically varying capabilities.
Takeaways & Limitations
HeteroFL makes FL more practical by exploiting clients’ capabilities while maintaining a single global inference model.
Takeaways & Limitations
The approach considers local models with similar architecture that shrink complexity within the same model class, leaving distinct model classes as future work.
Abstract
from arXiv · showhide
Federated Learning (FL) is a method of training machine learning models on private data distributed over a large number of possibly heterogeneous clients such as mobile phones and IoT devices. In this work, we propose a new federated learning framework named HeteroFL to address heterogeneous clients equipped with very different computation and communication capabilities. Our solution can enable the training of heterogeneous local models with varying computation complexities and still produce a single global inference model. For the first time, our method challenges the underlying assumption of existing work that local models have to share the same architecture as the global model. We demonstrate several strategies to enhance FL training and conduct extensive empirical evaluations, including five computation complexity levels of three model architecture on three datasets. We show that adaptively distributing subnetworks according to clients' capabilities is both computation and communication efficient.
1 INTRODUCTION
Heterogeneous client capabilities challenge FL because conventional methods require local models to match the global architecture. HeteroFL trains heterogeneous local models, aggregates them into one global model, and targets efficient, stable training across changing settings.
- Motivation: Conventional FL limits global model complexity to what the least capable client can train.Client computation and communication capabilities may vary significantly and dynamically.
- HeteroFL: HeteroFL trains local models with varying computation complexities while producing a single global inference model.The framework relaxes the assumption that local and global models must share the same architecture.
- HeteroFL: HeteroFL aggregates heterogeneous local models stably and effectively into a single global inference model without additional computation overhead.The authors report performance exceeding state-of-the-art results.
- Evaluation scope: HeteroFL addresses settings where different proportions of clients have distinct capabilities, including dynamically changing model heterogeneity.The reported learning results remain stable and effective in these settings.
- Training strategies: The framework introduces training strategies that are robust to balanced non-IID statistical heterogeneity and can reduce communication rounds to reach state-of-the-art results.The paper evaluates the proposed approach experimentally.
2 RELATED WORK
Prior FL work addresses communication, system, statistical, and privacy challenges, but commonly retains a shared-model architecture assumption. HeteroFL instead permits local architectures to differ from the global model while adding mechanisms for heterogeneous and non-IID settings.
- Existing challenges: FL research addresses communication efficiency, system heterogeneity, statistical heterogeneity, and privacy through compression, split learning, asynchronous communication, active sampling, and personalization.FedAvg is described as the most widely adopted baseline and reduces communication cost through multiple local training iterations.
- HeteroFL’s distinction: HeteroFL is presented as the first work allowing local models to have different architectures from the global model.This heterogeneity lets clients adaptively contribute while optimizing lower-complexity models and communicating fewer parameters.
- HeteroFL’s distinction: The framework uses a Masking Trick for balanced non-IID classification data and modifies Batch Normalization because running estimates raise privacy concerns.These mechanisms target statistical heterogeneity and privacy-related use of advanced deep learning models.
3 HETEROGENEOUS FEDERATED LEARNING
HeteroFL assigns nested subnetworks to clients with different capabilities, aggregates them into one global model, and supports heterogeneous local training. Its design combines width-based model scaling, capability-aware aggregation, static normalization, and representation scaling.
- 3.1 HETEROGENEOUS MODELS: HeteroFL relaxes the shared-architecture assumption by allocating local models with different computation complexities while producing a single global inference model.Local models remain within the same model class by varying hidden-layer width.
- 3.1 HETEROGENEOUS MODELS: Clients receive nested subsets of global parameters according to capability, with multiple computation complexity levels defined by a channel shrinkage ratio.For a layer with global dimensions d_g × k_g, local dimensions are reduced according to the selected level.
- 3.1 HETEROGENEOUS MODELS: Global aggregation averages each parameter over clients whose allocated matrices contain it, fully averaging shared parameters and partially averaging parameters assigned to smaller models.The global model is reconstructed from disjoint parameter partitions, while smaller models use upper-left submatrices of larger models.
- 3.1 HETEROGENEOUS MODELS: Inclusive parameter subsets and a fixed subnetwork stabilize aggregation, and empirically outperform uniformly sampled subnetworks.The fixed subnetwork ensures aggregation across all local models on at least one shared subnetwork.
- 3.3 SCALER: Static Batch Normalization avoids tracking local running statistics, making each communication round independent despite dynamically varying local model sizes.The method normalizes batch data during training and queries representation statistics from clients after training.
- 3.3 SCALER: HeteroFL determines each client's complexity level from capability information, selects active clients, and updates their corresponding parameter submatrices each round.The algorithm can optionally update learning rates and capability information when client conditions change dynamically.
4 EXPERIMENTAL RESULTS
Experiments evaluate HeteroFL across three architectures, three datasets, five computation-complexity levels, and fixed or dynamically varying client assignments. The results show effective heterogeneous training, including strong CIFAR10 performance, stable dynamic aggregation, and reduced communication rounds.
- Experimental setup: Experiments cover MNIST with CNN, CIFAR10 with PreResNet18, and WikiText2 with a Transformer, including balanced non-IID settings.The study trained over 600 individual models and used masked language modeling for WikiText2.
- Experimental setup: Five computation-complexity levels, labeled a–e, represent models ranging from the full model to progressively smaller subnetworks.The levels use hidden-channel shrinkage, with ratios including 0.5, 0.25, 0.125, and 0.0625 relative to the most complex model.
- Experimental setup: Fixed assignments allocate each client a persistent complexity level, whereas Dynamic assignments uniformly resample levels at every communication round.Figures use fixed assignments, while experiments also evaluate dynamically varying client capabilities.
- Results: 90.29% is achieved by the a-e CIFAR10 combination, compared with 77.09% for the standalone e model and 91.99% for the all-a model.The mixed a-e configuration uses larger models for half the clients and smaller models for the other half, while approaching the all-strong-client result.
- Results: 800 communication rounds achieve better results for HeteroFL than the 1800 rounds reported for LG-FedAvg.Gradient clipping stabilizes heterogeneous optimization and supports a universal learning rate across local model sizes.
- Results: Dynamic aggregation remains stable despite varying computation complexities, while aggregation with non-IID data is less stable.The learning curves evaluate both 50% fixed and dynamic assignments.
5 CONCLUSIONS AND FUTURE WORK
HeteroFL coordinates local models smaller than the global model to produce one global inference model. With sBN and Masked Cross-Entropy Loss, experiments report improved results with fewer communication rounds across architectures and heterogeneous settings.
- Conclusions: HeteroFL coordinates local models much smaller than the global model while producing a single global inference model.The conclusion frames this as a way to make federated learning more practical.
- Conclusions: sBN and Masked Cross-Entropy Loss support HeteroFL results with fewer communication rounds across CNN, PreResNet18, and Transformer architectures.The method is reported as robust to statistical heterogeneity and dynamically varying local capabilities.
- Future work: Future work includes distinguishing model classes and applying the proposed methods to heterogeneous few-shot, multi-modal, and multi-task learning.These directions extend beyond the evaluated model-class heterogeneity setting.
A APPENDIX
The appendix reports supplementary experiments, including ablations, hyperparameters, fixed complexity assignments, and complete dynamic-assignment results across MNIST, CIFAR10, and WikiText2.
- Supplementary experiments: Appendix figures show fixed-complexity assignments for MNIST and WikiText2 with IID partitions and for MNIST and CIFAR10 with balanced Non-IID partitions.Figures 3–8 cover these interpolation experiments between global and smaller model complexities.
- Supplementary experiments: Tables 7–9 provide complete results for dynamic complexity assignments on MNIST, CIFAR10, and WikiText2.
- Ablation study: Table 5 reports a Non-IID ablation in which Masked CrossEntropy significantly improves local performance and moderately improves global performance.
- Ablation study: The same ablation finds that static Batch Normalization significantly outperforms InstanceNorm, GroupNorm, and LayerNorm.GroupNorm uses G=4 groups in this comparison.
- Ablation study: The Scaler enables HeteroFL to train models of different sizes and moderately improves the results.
- Experimental configuration: Table 6 lists the hyperparameters and model architectures used in the experiments.