[Paper Review] AdaptiveFL: Adaptive Heterogeneous Federated Learning for Resource-Constrained AIoT Systems

Chentao Jia, Ming Hu+5 more|arXiv (Cornell University)|Nov 22, 2023

Privacy-Preserving Technologies in Data19 references4 citations

TL;DR

AdaptiveFL proposes a novel federated learning framework for resource-constrained AIoT systems that uses fine-grained width-wise model pruning to generate heterogeneous local models tailored to device-specific resources. By combining reinforcement learning-based dynamic device selection with adaptive model dispatching, AdaptiveFL achieves up to 16.83% higher inference accuracy than state-of-the-art methods across both IID and non-IID data distributions.

ABSTRACT

Although Federated Learning (FL) is promising to enable collaborative learning among Artificial Intelligence of Things (AIoT) devices, it suffers from the problem of low classification performance due to various heterogeneity factors (e.g., computing capacity, memory size) of devices and uncertain operating environments. To address these issues, this paper introduces an effective FL approach named AdaptiveFL based on a novel fine-grained width-wise model pruning strategy, which can generate various heterogeneous local models for heterogeneous AIoT devices. By using our proposed reinforcement learning-based device selection mechanism, AdaptiveFL can adaptively dispatch suitable heterogeneous models to corresponding AIoT devices on the fly based on their available resources for local training. Experimental results show that, compared to state-of-the-art methods, AdaptiveFL can achieve up to 16.83% inference improvements for both IID and non-IID scenarios.

Motivation & Objective

• Address the performance degradation in federated learning caused by hardware heterogeneity and non-IID data in AIoT systems.
• Overcome limitations of existing heterogeneous FL methods that rely on coarse-grained pruning or require external datasets, violating data privacy.
• Enable dynamic, on-the-fly model dispatching based on real-time device resource availability to maximize global model performance.
• Reduce communication overhead by minimizing mismatched model transfers through intelligent device selection.
• Ensure strong generalization and robustness across diverse AIoT device configurations under uncertain, dynamic operating conditions.

Proposed method

• Introduces a fine-grained width-wise model pruning strategy that selectively prunes channels within layers rather than entire layers, preserving model capacity while reducing size.
• Employs a reinforcement learning (RL)-based client selection mechanism that selects devices based on historical model size and resource availability, minimizing communication waste.
• Uses a dual-reward signal in RL—resource reward and curiosity reward—to balance efficiency and exploration in device selection.
• Supports direct training of pruned models on devices without additional adapters or parameter reinitialization, enabling plug-and-play deployment.
• Applies knowledge distillation principles implicitly through parameter transfer, where pruned models inherit knowledge from larger global models.
• Enables model aggregation via FedAvg with heterogeneous local models generated through adaptive pruning, maintaining compatibility with standard FL pipelines.

Experimental results

Research questions

• RQ1 How can federated learning be effectively adapted to AIoT systems with highly heterogeneous devices and dynamically changing resource availability?
• RQ2 To what extent does fine-grained width-wise pruning improve model performance compared to coarse-grained or layer-level pruning in heterogeneous FL?
• RQ3 Can a reinforcement learning-based device selection strategy reduce communication overhead and improve training efficiency in resource-constrained FL settings?
• RQ4 How does AdaptiveFL perform under non-IID data distributions and varying proportions of strong, medium, and weak devices?
• RQ5 What is the trade-off between model size, inference accuracy, and communication efficiency in adaptive, heterogeneous federated learning?

Key findings

• AdaptiveFL achieves up to 16.83% higher test accuracy than state-of-the-art methods in both IID and non-IID scenarios, demonstrating superior generalization.
• The fine-grained width-wise pruning strategy improves inference accuracy by up to 9.38% compared to coarse-grained pruning, especially benefiting deeper models like ResNet18 on CIFAR-100.
• The RL-based client selection mechanism reduces communication waste by 30–40% compared to greedy or random selection strategies, as measured by the communication waste rate.
• On a real test-bed with 17 heterogeneous devices (Raspberry Pi, Jetson Nano, and Jetson Xavier), AdaptiveFL outperformed all baselines in learning curves and final accuracy.
• In ablation studies, the full AdaptiveFL model (with both resource and curiosity rewards) achieved the highest accuracy and lowest communication waste, proving the effectiveness of the dual-reward RL design.
• The method maintains high performance even when the proportion of strong devices is low (e.g., 1:8:1), showing resilience to imbalanced device capabilities.