[论文解读] 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 Synthesis被引用 4

一句话总结

本文提出部分变量训练（PVT），一种用于高效设备端联邦学习的方法，该方法在每个客户端仅训练网络变量的子集，显著降低内存使用和通信成本。PVT 通过增加本地训练步数和客户端参与度，在保持与完整模型训练相当的准确率的同时，将通信成本降低高达 593 倍，内存使用降低 1.9 倍。

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.

研究动机与目标

• 解决设备端联邦学习（FL）中内存受限和通信成本过高的挑战。
• 实现在智能手机和物联网设备等资源受限边缘设备上的高效联邦学习。
• 在减少每个客户端可训练变量数量的同时，保持高模型准确率。
• 开发一种无需架构修改或特定网络知识的方法。
• 通过增加本地训练步数和客户端参与度来补偿变量冻结的影响，支持大规模联邦学习部署。

提出的方法

• PVT 在每个联邦轮次中，为每个客户端选择可训练变量（如权重、偏置、缩放因子）的子集并冻结其余变量。
• 仅将训练变量的梯度和更新结果传回服务器，从而降低通信开销。
• 冻结变量可消除反向传播中激活值的缓冲需求，显著降低内存使用。
• 该方法将变量分为加法向量（如偏置）、乘法向量（如归一化缩放因子）和乘法矩阵（如卷积权重），其中加法向量因成本较低而被排除在冻结之外。
• PVT 采用每客户端每轮（PCPR）机制，确保客户端间网络更新分布均衡且高效。
• 该方法与网络架构无关，无需修改即可兼容标准神经网络结构。

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.

实验结果

研究问题

• RQ1 在联邦学习中仅训练变量子集是否能降低内存使用和通信成本，同时不损害模型准确率？
• RQ2 当变量部分训练时，本地训练步数的多少如何影响收敛性和准确率？
• RQ3 增加客户端数量是否能补偿因冻结大量变量而导致的准确率损失？
• RQ4 冻结变量的选择（如加法变量与乘法变量）是否会影响效率和性能？
• RQ5 PVT 是否能普遍适用于不同神经网络架构，而无需进行架构修改？

主要发现

• 与完整模型训练相比，PVT 在语音识别模型上将通信成本降低高达 593 倍，内存使用降低高达 1.9 倍。
• 在非流式 Conformer 模型与 IID Librispeech 数据集上，PVT 使用 1024 个客户端和 5 个本地训练步数时，WER 达到 2.1/5.0/2.3/4.9，与 AVT 的 2.0/4.8/2.2/4.6 接近。
• 使用 128 个本地训练步数的 PVT 仍能获得收敛结果，表明其对高本地更新次数具有鲁棒性。
• 每客户端每轮（PCPR）机制相比每轮（PR）机制收敛更快，后者难以高效更新完整网络。
• 将客户端数量从 128 增加到 4096，所需轮次数从 350K 降至 60K，同时保持稳定的 WER 表现。
• 仅冻结乘法矩阵和乘法向量（排除加法向量）可实现最优效率，且准确率损失最小。

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.