[論文レビュー] 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

ひとこと要約

本稿では、異なるビット幅制約を持つ多様なデバイスにおける効率的なオンデバイス推論を可能にする、量子化に強いフェデレーテッドラーニング手法を提案する。フェデレーテッドアベレージングに、カートシス正則化（KURE）、加法的擬似量子化ノイズ（APQN）、および新規のマルチビット量子化に配慮した訓練（mQAT）フレームワークを統合することで、全精度学習後でさえも、複数の量子化ビット幅で高い精度を達成し、顕著な性能低下を伴わない。

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は、訓練時に使用されたビット幅以外の環境に展開すると性能が低下する。
• 再訓練を必要とせずに、複数の量子化ビット幅に一般化可能なフェデレーテッドトレーニングフレームワークを設計する。
• 非凸なフェデレーテッド設定におけるQATバリアントの収束行動についての理論的分析を含む。

提案手法

• 重みの量子化に対して耐性を持つよう、ローカルクライアントの損失関数に正則化項を追加することで、カートシス正則化（KURE）を導入する。
• ランダムノイズを訓練中に注入することで量子化効果をシミュレートする、加法的擬似量子化ノイズ（APQN）を提案する。これはランダム化スムージングにインspiredされている。
• マルチビット量子化に配慮した訓練（mQAT）を構築し、各ローカルアップデートでクライアントごとにランダムにビット幅をサンプリングすることで、一度のモデル訓練で複数のビット幅に耐性を持つモデルを学習する。
• ストレートスルーエスティメーター（STE）近似を用いた標準的なQATを適応させるが、ビット幅間での一般化を向上させるために訓練プロセスを変更する。
• フェデレーテッドラーニングにおけるQATバリアントの収束性についての理論的分析を提供し、QAT特有の誤差フロアを伴いながらも、FedAvgと同等の収束速度を示す。
• 標準的なFedAvgをベース最適化フレームワークとし、量子化に強い性能を実現するためにクライアントレベルのトレーニング段階に変更を加える。

実験結果

リサーチクエスチョン

• RQ1 フェデレーテッドラーニングに量子化に強い性能を効果的に統合することで、多様なビット幅を持つ異種デバイス上でもモデル精度を維持できるか？
• RQ2 正則化ベースの手法（例：KURE, APQN）は、分布外の量子化に対してQATベースの手法と比較してどれほど耐性があるか？
• RQ3 mQATを用いてフェデレーテッドラーニングで訓練された1つのモデルが、再訓練を伴わず複数の量子化ビット幅に一般化できるか？
• RQ4 非凸な設定におけるQATベースのフェデレーテッドトレーニングの理論的収束挙動はいかなるものか？
• RQ5 量子化に強いメカニズムを統合することで、フェデレーテッド設定における全精度モデルの精度が低下するか？

主な発見

• CIFAR-10でResNet-20を用いた8ビット活性化量子化後、mQATはグローバルバリデーション精度76.42%を達成し、ベースラインQAT（73.72%）および他のバージョンを上回った。
• CIFAR-100でResNet-20を用いた6ビット活性化量子化後、mQATは76.8%の精度を達成し、QAT（71.38%）およびベースライン（64.08%）を著しく上回った。
• TinyImageNetでResNet-18を用いた4ビット重み量子化後、mQATは66.02%の精度を維持したが、QATは66.74%、ベースラインは68.24%に低下した。
• CIFAR-10で8ビット活性化量子化の設定下、APQNは76.38%の精度を達成し、ベースライン（70.06%）およびQAT（73.98%）を上回った。
• 2ビット環境下、LeNet-5を用いたCIFAR-10では、mQATは67.78%の精度を達成したのに対し、QATは58.2%、KUREは26.2%に低下し、強力な耐性を示した。
• 理論的分析により、mQATの収束速度がFedAvgと同等であることが確認され、QAT特有のトレーニングヒューリスティクスに起因するわずかな誤差フロアが存在することが示された。