[論文レビュー] FedZKT: Zero-Shot Knowledge Transfer towards Resource-Constrained Federated Learning with Heterogeneous On-Device Models

Lan Zhang, Dapeng Wu+1 more|arXiv (Cornell University)|Sep 8, 2021

Privacy-Preserving Technologies in Data被引用数 8

ひとこと要約

FedZKTは、リソース制約のある異種のデバイスが、生データを共有せずに協調的にグローバルモデルを学習できる、データフリーでゼロショットの知識伝達フレームワークを提案する。計算負荷の高い知識蒸留処理をサーバーに移管し、アンサンブルされたデバイスモデルから中央の知識を学ぶジェネレータを用いることで、モデルの非均一性、非IIDデータ、遅延デバイスの影響に対しても強固な性能を達成する。デバイス側の計算負荷は最小限であり、プライベートデータの前提条件も不要である。

ABSTRACT

Federated learning enables multiple distributed devices to collaboratively learn a shared prediction model without centralizing their on-device data. Most of the current algorithms require comparable individual efforts for local training with the same structure and size of on-device models, which, however, impedes participation from resource-constrained devices. Given the widespread yet heterogeneous devices nowadays, in this paper, we propose an innovative federated learning framework with heterogeneous on-device models through Zero-shot Knowledge Transfer, named by FedZKT. Specifically, FedZKT allows devices to independently determine the on-device models upon their local resources. To achieve knowledge transfer across these heterogeneous on-device models, a zero-shot distillation approach is designed without any prerequisites for private on-device data, which is contrary to certain prior research based on a public dataset or a pre-trained data generator. Moreover, this compute-intensive distillation task is assigned to the server to allow the participation of resource-constrained devices, where a generator is adversarially learned with the ensemble of collected on-device models. The distilled central knowledge is then sent back in the form of the corresponding on-device model parameters, which can be easily absorbed on the device side. Extensive experimental studies demonstrate the effectiveness and robustness of FedZKT towards on-device knowledge agnostic, on-device model heterogeneity, and other challenging federated learning scenarios, such as heterogeneous on-device data and straggler effects.

研究の動機と目的

• リソース制約があり、オンデバイスのモデル構造が異なるデバイスをフェデレーテッドラーニングに統合する課題に対処すること。古典的手法では、均一なモデル構造を必要とする。
• 知識伝達に、オンデバイスのプライベートデータやパブリックデータセットへの依存を排除することで、データフリーの協調学習を可能とすること。
• 知識蒸留処理をサーバーに移管することで、デバイス側の計算負荷を低減すること。
• 実世界のフェデレーテッド環境におけるモデルの非均一性、非IIDデータ分布、遅延デバイスの影響に対して、耐性を確保すること。
• グローバルモデルの収束性と性能を維持しつつ、デバイス側のモデル設計を独立して行えるようにすること。

提案手法

• デバイスは、自らのローカルハードウェア制約に応じて、コンパクトでリソース効率の良いモデルを独立して学習する。
• サーバーで学習されたジェネレータは、収集されたオンデバイスモデルのアンサンブルと敵対的に訓練され、プライベートデータを必要とせずに中央の知識を蒸留する。
• 構造的に異種のモデル間でゼロショットの知識蒸留を可能にする、新しいSL（ソフトラベル）損失関数を導入する。
• 蒸留された知識は、オンデバイスモデルのパラメータとして戻され、デバイスが容易にグローバル知識を吸収し、ファインチューニングできる。
• フレームワークは、計算負荷の高い蒸留処理をサーバーに割り当てることで、デバイスのリソース効率を保つ。
• 非IIDデータ分布下での耐性を高めるために、オンデバイス学習中にL2正則化を適用する。

Figure 1 : Overview of FedZKT: zero-shot knowledge transfer for federated learning with heterogeneous on-device models. Most compute-intensive workloads of FedZKT are done at the server. Participating devices only need to update the on-device models based on their local data.

実験結果

リサーチクエスチョン

• RQ1 共有データやモデル構造の整合性がなくても、高度に異種のオンデバイスモデルを搭載したデバイス間でフェデレーテッドラーニングを効果的に実現できるか？
• RQ2 パブリックデータセットや事前学習済みデータジェネレータに依存せずに、非均質なモデル間で知識をデータフリーに伝達する方法は何か？
• RQ3 蒸留処理をサーバーに移管した場合、リソース制約のあるデバイスがフェデレーテッドラーニングにどの程度参加できるか？
• RQ4 実世界の展開において、非IIDデータ分布や遅延デバイスの影響に対して、このフレームワークはどの程度耐性を示すか？
• RQ5 オンデバイスモデルのアーキテクチャとデバイス数が、グローバルモデルの収束性と精度に与える影響は何か？

主な発見

• FedZKTは多様なオンデバイスモデルにおいて強力な性能を発揮する：CIFAR-10（IID）でDevice 3（Model C）は88.63%の精度を達成し、Device 8（Model C）は88.34%を記録し、モデルの非均一性への高い適応性を示す。
• 1ラウンドあたり20%のデバイス（p=0.2）しか参加しない状況でも、安定した学習が維持され、遅延デバイスの影響に対する耐性が確認された。
• L2正則化により、非IIDデータ下での性能が向上する：C=5の不均衡下で56.58%から63.89%へ、β=0.5の分布不均衡下で66.17%から69.39%へと向上。
• デバイス数（K ∈ {5,10,15,20}）の変化が性能に与える影響は小さく、異なるK値間で平均精度の差は±2%以内にとどまる。
• 蒸留処理がサーバー最適化されており、知識伝達が軽量なモデルパラメータの更新で行われるため、リソースが限られたデバイスでも効果的に貢献可能である。
• プライベートデータやパブリックデータに依存せず、知識伝達が実現可能であることを確認し、非均質なフェデレーテッド環境下におけるデータフリーでゼロショットの蒸留の実現可能性を裏付けた。

Figure 2 : Norm of gradients w.r.t input data (MNIST, IID). The gradients for the KL-divergence loss tend to vanish, while the gradients for the $\ell_{1}$ norm loss are much larger and unstable during the learning process. The proposed SL loss overcomes both problems in the federated learning.