[Paper Review] Partial Variable Training for Efficient On-Device Federated Learning

Tien-Ju Yang, Dhruv Guliani+2 more|arXiv (Cornell University)|Oct 11, 2021

Speech Recognition and Synthesis4 citations

TL;DR

This paper proposes Partial Variable Training (PVT), a method for efficient on-device Federated Learning that trains only a subset of network variables per client, significantly reducing memory usage and communication cost. PVT achieves up to 593× lower communication cost and 1.9× lower memory usage while maintaining accuracy comparable to full model training by leveraging more local steps and clients.

ABSTRACT

This paper aims to address the major challenges of Federated Learning (FL) on edge devices: limited memory and expensive communication. We propose a novel method, called Partial Variable Training (PVT), that only trains a small subset of variables on edge devices to reduce memory usage and communication cost. With PVT, we show that network accuracy can be maintained by utilizing more local training steps and devices, which is favorable for FL involving a large population of devices. According to our experiments on two state-of-the-art neural networks for speech recognition and two different datasets, PVT can reduce memory usage by up to 1.9$ imes$ and communication cost by up to 593$ imes$ while attaining comparable accuracy when compared with full network training.

Motivation & Objective

• Address the challenges of limited memory and high communication cost in on-device Federated Learning (FL).
• Enable efficient FL on edge devices with constrained resources such as smartphones and IoT devices.
• Maintain high model accuracy despite reducing the number of trainable variables per client.
• Develop a method that does not require architectural modifications or network-specific knowledge.
• Support large-scale FL deployments by compensating for variable freezing through increased local steps and client participation.

Proposed method

• PVT selects a subset of trainable variables (weights, biases, scaling factors) and freezes the rest on each client per federated round.
• Only gradients and updates for the trained variables are communicated back to the server, reducing communication overhead.
• Freezing variables eliminates the need to buffer activations for backpropagation, significantly reducing memory usage.
• The method categorizes variables into additive vectors (e.g., biases), multiplicative vectors (e.g., normalization scales), and multiplicative matrices (e.g., convolutional weights), with additive vectors excluded from freezing due to their low cost.
• PVT uses a per-client-per-round (PCPR) scheme to ensure balanced and efficient network update distribution across clients.
• The approach is architecture-agnostic and compatible with standard neural network structures without modification.

Fig. 1 : The illustration of the proposed partial variable training. In this example, $Client\_1$ only trains $Layer\_1$ and sends back its change. Similarly, $Client\_2$ only trains $Layer\_2$ and sends back its change. The server updates both layers after receiving the changes from both clients.

Experimental results

Research questions

• RQ1 Can training only a subset of variables in FL reduce memory usage and communication cost without degrading model accuracy?
• RQ2 How does the number of local training steps affect convergence and accuracy when variables are partially trained?
• RQ3 Can increasing the number of clients compensate for the accuracy loss caused by freezing a large number of variables?
• RQ4 Does the choice of which variables to freeze (e.g., additive vs. multiplicative) impact efficiency and performance?
• RQ5 Can PVT be applied universally across different neural network architectures without architectural changes?

Key findings

• PVT reduces communication cost by up to 593× and memory usage by up to 1.9× compared to full model training on speech recognition models.
• On the non-streaming Conformer model with IID Librispeech, PVT with 1024 clients and 5 local steps achieves WER of 2.1/5.0/2.3/4.9, close to AVT’s 2.0/4.8/2.2/4.6.
• Using 128 local training steps with PVT still achieves converged results, demonstrating robustness to high local update counts.
• The per-client-per-round (PCPR) scheme enables faster convergence than per-round (PR) schemes, which struggle to update the full network efficiently.
• Increasing the number of clients from 128 to 4096 reduces required rounds from 350K to 60K while maintaining stable WER performance.
• Freezing only multiplicative matrices and multiplicative vectors (excluding additive vectors) yields optimal efficiency with minimal accuracy loss.

Fig. 2 : The WER and convergence speed improvement while we sequentially apply PCPR scheme, not freezing additive vectors (NFAV), 5 local steps (5LS), 1024 clients (1024CL) to train non-streaming Conformer on IID Librispeech.