NeuLite: Memory-Efficient Federated Learning via Elastic Progressive Training

2024-08-20

arXiv.org

AI Summary

To address memory bottlenecks on resource-constrained clients in federated learning, this paper proposes an elastic progressive training framework. It partitions the model into blocks and trains them sequentially; introduces a curriculum-guided mentor mechanism that employs block-wise curriculum-aware loss to preserve intermediate representation quality; and incorporates a training coordinator to enable cross-block parameter co-adaptation and output-module enhancement—thereby breaking the inter-block information isolation inherent in conventional forward/backward propagation. This work pioneers the elastic progressive training paradigm, integrating model partitioning, progressive weight freezing/activation, and block-wise output enhancement. Experiments demonstrate a 50.4% reduction in peak memory consumption, up to an 84.2% improvement in final accuracy, and 1.9× training acceleration. The framework is rigorously validated on both simulation environments and a physical hardware testbed.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Learning to Search
• Computer Vision: Learning & Optimization for CV

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Federated Learning (FL) emerges as a new learning paradigm that enables multiple devices to collaboratively train a shared model while preserving data privacy. However, intensive memory footprint during the training process severely bottlenecks the deployment of FL on resource-constrained devices in real-world cases. In this paper, we propose NeuLite, a framework that breaks the memory wall through elastic progressive training. Unlike traditional FL, which updates the full model during the whole training procedure, NeuLite divides the model into blocks and conducts the training process in a progressive manner. Except for the progressive training paradigm, NeuLite further features the following two key components to guide the training process: 1) curriculum mentor and 2) training harmonizer. Specifically, the Curriculum Mentor devises curriculum-aware training losses for each block, assisting them in learning the expected feature representation and mitigating the loss of valuable information. Additionally, the Training Harmonizer develops a parameter co-adaptation training paradigm to break the information isolation across blocks from both forward and backward propagation. Furthermore, it constructs output modules for each block to strengthen model parameter co-adaptation. Extensive experiments are conducted to evaluate the effectiveness of NeuLite across both simulation and hardware testbeds. The results demonstrate that NeuLite effectively reduces peak memory usage by up to 50.4%. It also enhances model performance by up to 84.2% and accelerates the training process by up to 1.9X.

Problem

• Reducing memory demands for federated learning on constrained clients
• Overcoming information loss in sequential block-wise model training
• Breaking inter-block isolation through coordinated parameter updates

Innovation

• Sequential block-wise training reduces memory requirements
• Curriculum Mentor steers learning with curriculum-aware objectives
• Training Harmonizer coordinates updates to break information isolation