[Paper Review] HeteroFL: Computation and Communication Efficient Federated Learning for Heterogeneous Clients

Enmao Diao, Jie Ding+1 more|arXiv (Cornell University)|Oct 3, 2020

Privacy-Preserving Technologies in Data35 references179 citations

TL;DR

HeteroFL enables training of heterogeneous local models with varying computation and communication capabilities and aggregates them into a single global model without requiring identical architectures. It introduces sBN, Masked Cross-Entropy, and a subnetwork distribution scheme to stabilize aggregation and improve efficiency.

ABSTRACT

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.

Motivation & Objective

• Motivate federated learning when client devices have diverse computation and communication capabilities.
• Propose a framework to train heterogeneous local models and aggregate into a single global model.
• Improve training efficiency and robustness to balanced non-IID data.
• Introduce techniques to stabilize learning and preserve privacy in heterogeneous settings.

Proposed method

• Allow local models to have different widths by distributing subnetworks of the global model to clients at multiple computation complexity levels.
• Aggregate parameters across clients by disjoint partial subsets to form the global model (inclusive union of parameter partitions).
• Use static Batch Normalization (sBN) to avoid sharing running statistics; update global BN statistics after training rounds.
• Attach a Scaler module to scale representations during training to account for subnetwork size differences.
• Apply a Masked Cross-Entropy Loss to balance non-IID data by masking outputs for classes not present at a client.

Experimental results

Research questions

• RQ1 Can heterogeneous local models with different architectures or sizes be trained within a federated learning framework to yield a single accurate global model?
• RQ2 How to stable-aggregate heterogeneous local models and ensure training efficiency across varying client capabilities?
• RQ3 What normalization and loss strategies best preserve privacy and performance under balanced/non-balanced non-IID data?
• RQ4 How do dynamic changes in client capabilities affect convergence and final accuracy?

Key findings

• HeteroFL can train heterogeneous local models and aggregate them into a single global model without extra computation overhead.
• Adaptive subnetwork distribution according to client capabilities yields competitive accuracy with reduced computation and communication costs.
• Static BN and Masked Cross-Entropy improve stability and performance under balanced non-IID data scenarios.
• The approach demonstrates robustness to dynamically varying local capabilities and different model architectures (CNN, PreResNet18, Transformer).