Bridging Memory Gaps: Scaling Federated Learning for Heterogeneous Clients

TL;DR

This paper introduces heir, a scalable federated learning framework that uses sequential block-wise training with curriculum-aware objectives and parameter co-adaptation to enable resource-constrained clients to participate effectively.

Contribution

It proposes a novel block-wise training approach with curriculum mentoring and parameter co-adaptation to address memory limitations in federated learning.

Findings

Improves model performance by up to 84.2%

Reduces peak memory usage by up to 50.4%

Speeds up training by up to 1.9 times

Abstract

Federated Learning (FL) enables multiple clients to collaboratively train a shared model while preserving data privacy. However, the high memory demand during model training severely limits the deployment of FL on resource-constrained clients. To this end, we propose \our, a scalable and inclusive FL framework designed to overcome memory limitations through sequential block-wise training. The core idea of \our is to partition the global model into blocks and train them sequentially, thereby reducing training memory requirements. To mitigate information loss during block-wise training, \our introduces a Curriculum Mentor that crafts curriculum-aware training objectives for each block to steer their learning process. Moreover, \our incorporates a Training Harmonizer that designs a parameter co-adaptation training scheme to coordinate block updates, effectively breaking inter-block information isolation. Extensive experiments on both simulation and hardware testbeds demonstrate that \our significantly improves model performance by up to 84.2\%, reduces peak memory usage by up to 50.4\%, and accelerates training by up to 1.9$\times$.