[论文解读] Online Model Compression for Federated Learning with Large Models

Tien-Ju Yang, Yonghui Xiao+4 more|arXiv (Cornell University)|May 6, 2022

Speech Recognition and Synthesis被引用 4

一句话总结

本文提出在线模型压缩（OMC），一种轻量级联邦学习框架，通过以压缩的量化格式存储和传输大型神经网络参数，仅在运行时计算过程中才进行解压缩。OMC 通过每变量变换、仅权矩阵量化和部分参数量化等技术，在保持全精度准确率和训练速度的同时，将内存使用量和通信成本降低高达 59%。

ABSTRACT

This paper addresses the challenges of training large neural network models under federated learning settings: high on-device memory usage and communication cost. The proposed Online Model Compression (OMC) provides a framework that stores model parameters in a compressed format and decompresses them only when needed. We use quantization as the compression method in this paper and propose three methods, (1) using per-variable transformation, (2) weight matrices only quantization, and (3) partial parameter quantization, to minimize the impact on model accuracy. According to our experiments on two recent neural networks for speech recognition and two different datasets, OMC can reduce memory usage and communication cost of model parameters by up to 59% while attaining comparable accuracy and training speed when compared with full-precision training.

研究动机与目标

• 解决大规模模型在联邦学习中带来的高设备内存占用和通信成本问题。
• 在不损害模型准确率的前提下，减少边缘设备训练过程中模型参数的内存占用。
• 最小化联邦学习环境中频繁压缩/解压缩操作导致的训练速度下降。
• 实现在资源受限边缘设备上部署最先进的自动语音识别（ASR）模型。
• 通过在线按需解压压缩参数，实现联邦学习中高效参数传输与存储。

提出的方法

• 将模型参数以压缩的浮点格式（如 S1E5M10）存储，而非全精度 FP32，以减少内存和通信开销。
• 仅在前向或反向传播需要时按需在线解压参数，并在使用后立即释放内存。
• 应用每变量变换（PVT）以通过量化前的参数分布归一化来稳定训练。
• 采用仅权矩阵量化，将压缩限制在最耗内存的组件（即线性层）上。
• 通过仅量化 90% 的参数并使用低比特宽（如 11 位）的方案实现部分参数量化（PPQ），同时保留 10% 的参数为更高精度。
• 采用可调节指数和尾数位分配的浮点量化（如 S1E3M7、S1E5M7），以在压缩效率与误差控制之间取得平衡。

Figure 1: The illustration of the framework of the proposed online model compression. The cubes with dashed borderlines are transient variables.

实验结果

研究问题

• RQ1 在联邦训练期间对模型参数进行在线压缩，能否显著降低设备内存占用，同时不损害模型准确率？
• RQ2 每变量变换对量化联邦训练的稳定性和收敛性有何影响？
• RQ3 当仅量化权矩阵或部分参数时，比特宽减少与模型性能之间的权衡如何？
• RQ4 与使用更高比特宽的全参数量化相比，使用低比特宽的部分参数量化是否能在准确率和效率方面表现更优？
• RQ5 OMC 在多大程度上能降低联邦学习中的通信成本，同时保持与全精度训练相当的训练速度？

主要发现

• 与全精度训练相比，OMC 将模型参数的内存使用量和通信成本降低了高达 59%。
• 通过每变量变换，OMC 在 Non-Streaming Conformer 模型上实现了稳定的训练收敛，避免了未使用该技术时出现的准确率下降。
• 使用 11 位格式（S1E3M7）的部分参数量化，其词错误率（WER）低于使用 13 位格式（如 S1E5M7、S1E4M8）的全参数量化。
• 11 位部分量化配置的收敛速度更快，且达到的词错误率（WER）低于所有测试的 13 位全量化配置。
• 尽管存在频繁的在线压缩和解压缩操作，OMC 仍保持了与全精度训练相当的训练速度。
• OMC 通过有效管理内存和通信开销，实现了在边缘设备上训练大型 ASR 模型（如 Non-Streaming Conformer）。

Figure 2: The illustration of per-variable transformation. The cubes with dashed borderlines are transient variables.