[论文解读] Quantization Robust Federated Learning for Efficient Inference on Heterogeneous Devices

Kartik Gupta, Marios Fournarakis+3 more|arXiv (Cornell University)|Jun 22, 2022

Privacy-Preserving Technologies in Data被引用 13

一句话总结

本文提出了一种抗量化联邦学习方法，以实现跨异构设备在不同位宽约束下的高效设备端推理。通过将峰度正则化（Kurtosis Regularization）、加性伪量化噪声（Additive Pseudo-Quantization Noise）以及一种新型多比特量化感知训练（mQAT）框架整合到联邦平均（Federated Averaging）中，该方法在无需显著性能下降的情况下，实现了在多种量化位宽下的高精度，即使在全精度训练后亦然。

ABSTRACT

Federated Learning (FL) is a machine learning paradigm to distributively learn machine learning models from decentralized data that remains on-device. Despite the success of standard Federated optimization methods, such as Federated Averaging (FedAvg) in FL, the energy demands and hardware induced constraints for on-device learning have not been considered sufficiently in the literature. Specifically, an essential demand for on-device learning is to enable trained models to be quantized to various bit-widths based on the energy needs and heterogeneous hardware designs across the federation. In this work, we introduce multiple variants of federated averaging algorithm that train neural networks robust to quantization. Such networks can be quantized to various bit-widths with only limited reduction in full precision model accuracy. We perform extensive experiments on standard FL benchmarks to evaluate our proposed FedAvg variants for quantization robustness and provide a convergence analysis for our Quantization-Aware variants in FL. Our results demonstrate that integrating quantization robustness results in FL models that are significantly more robust to different bit-widths during quantized on-device inference.

研究动机与目标

• 解决尽管边缘设备硬件异构性日益增长，联邦学习中仍缺乏量化鲁棒性的问题。
• 使训练后的模型在量化至不同位宽（如2、4、8位）时仍能保持高精度，以实现能效更高的推理。
• 克服标准量化感知训练（QAT）在联邦学习中的局限性，即在非训练位宽下部署时性能下降的问题。
• 设计一种联邦训练框架，使其在无需微调的情况下，可泛化至多个量化位宽。
• 包含对非凸联邦设置下QAT变体收敛行为的理论分析。

提出的方法

• 通过在本地客户端损失函数中添加正则化项，引入峰度正则化（KURE），以鼓励对权重量化具有鲁棒性。
• 提出加性伪量化噪声（APQN），在训练过程中注入随机噪声以模拟量化效应，受随机平滑（randomized smoothing）的启发。
• 开发多比特量化感知训练（mQAT），在每次本地更新时为每个客户端随机采样一个位宽，以训练一个对多种位宽均鲁棒的单一模型。
• 采用标准QAT并使用直通估计器（STE）近似，但对训练过程进行修改，以提升在不同位宽间的泛化能力。
• 为联邦学习中的QAT变体提供了理论收敛性分析，表明其收敛速率与FedAvg相近，但存在一个由QAT特有训练启发式方法导致的误差下限。
• 采用标准FedAvg作为基础优化框架，并在客户端训练阶段应用修改，以集成量化鲁棒性。

实验结果

研究问题

• RQ1 能否有效将量化鲁棒性集成到联邦学习中，以在异构设备上多种位宽下保持模型精度？
• RQ2 与基于QAT的方法相比，基于正则化的方法（如KURE、APQN）在应对分布外量化时的鲁棒性如何？
• RQ3 通过联邦学习训练的单一模型是否能无需微调即泛化至多个量化位宽，使用mQAT？
• RQ4 基于QAT的联邦训练在非凸设置下的理论收敛行为如何？
• RQ5 集成量化鲁棒性机制是否会导致联邦设置下全精度模型精度下降？

主要发现

• 在CIFAR-10数据集上，使用ResNet-20模型，mQAT在8位激活量化后达到76.42%的全局验证精度，优于基线QAT（73.72%）及其他变体。
• 在CIFAR-100数据集上，使用ResNet-20模型，mQAT在6位激活量化后达到76.8%的精度，显著优于QAT（71.38%）和基线（64.08%）模型。
• 在TinyImageNet数据集上，使用ResNet-18模型，mQAT在4位权重量化后保持66.02%的精度，而QAT下降至66.74%，基线下降至68.24%。
• 在CIFAR-10数据集上，APQN方法在8位激活量化后达到76.38%的精度，优于基线（70.06%）和QAT（73.98%）。
• 在2位设置下，mQAT在CIFAR-10上使用LeNet-5模型达到67.78%的精度，而QAT为58.2%，KURE仅为26.2%，显示出极强的鲁棒性。
• 理论分析确认，mQAT的收敛速率与FedAvg相近，其误差下限可归因于QAT特有的训练启发式方法。