[論文レビュー] Test-Time Robust Personalization for Federated Learning

Liangze Jiang, Tao Lin|arXiv (Cornell University)|May 22, 2022

Privacy-Preserving Technologies in Data被引用数 7

ひとこと要約

この論文では、テスト時微調整と学習可能で重みのついたアンサンブルを用いて、モデルヘッドを動的に適応させる、フェデレーテッドラーニングにおけるテスト時頑健なパーソナライゼーション手法を提案する。テストサンプルの履歴と特徴表現に基づいて、局所的およびグローバルヘッドのパラメータを最適化することで、特に分布シフトや非IIDデータ条件下でベースラインと比較して顕著な精度向上を達成する。

ABSTRACT

Federated Learning (FL) is a machine learning paradigm where many clients collaboratively learn a shared global model with decentralized training data. Personalized FL additionally adapts the global model to different clients, achieving promising results on consistent local training and test distributions. However, for real-world personalized FL applications, it is crucial to go one step further: robustifying FL models under the evolving local test set during deployment, where various distribution shifts can arise. In this work, we identify the pitfalls of existing works under test-time distribution shifts and propose Federated Test-time Head Ensemble plus tuning(FedTHE+), which personalizes FL models with robustness to various test-time distribution shifts. We illustrate the advancement of FedTHE+ (and its computationally efficient variant FedTHE) over strong competitors, by training various neural architectures (CNN, ResNet, and Transformer) on CIFAR10 andImageNet with various test distributions. Along with this, we build a benchmark for assessing the performance and robustness of personalized FL methods during deployment. Code: https://github.com/LINs-lab/FedTHE.

研究の動機と目的

• データ分布のシフトと非IIDクライアントデータという条件下でのフェデレーテッドラーニングにおけるモデルパーソナライゼーションの課題に対処すること。
• グローバルモデルの再トレーニングやクライアントデータへのアクセスを必要とせずに、テスト時のパフォーマンスを向上させること。
• 局所的なテスト統計を用いて、推論中に動的にパーソナライズドモデルヘッドを適応させること。
• 推論時にグローバルおよびローカルモデルコンponentsを活用することで、分布シフトに対する頑健性を高めること。
• テスト時に、グローバルおよびローカルヘッド間のアンサンブル重みを学習可能で適応的な方法で最適化すること。

提案手法

• 本手法は、グローバルモデルヘッドパラメータ $\mathbf{w}^g$ とローカルモデルヘッドパラメータ $\mathbf{w}^l$ を用いる。両者ともグローバルモデルから初期化される。
• テスト時微調整は、入力となるテストサンプル $\mathbf{x}_n$ に対して $\tau$ ステップの最適化を用い、学習可能でステップサイズ $\eta$ を持つローカルヘッド $\mathbf{w}^l$ に対して適用される。
• グローバルおよびローカルヘッドのアンサンブルは、$e \in [0,1]$ の学習可能な重みを用いて構築され、テスト時適応中に勾配ベース最適化により更新される。
• 適応的ヘッド重み付けと特徴表現の適応に情報を与えるために、モデルはテスト履歴 $\mathbf{h}^{\text{history}}$ を維持する。
• 特徴表現はグローバル $\mathbf{h}^g$ およびローカル $\mathbf{h}^l$ のコンポーネントに分割され、後者はテストサンプル上で微調整により更新される。
• ハイパーパrameter $\alpha = 0.1$ および $\beta = 0.3$ は、アンサンブル重みとヘッド更新の最適化目的関数における損失項のバランスを制御する。

実験結果

リサーチクエスチョン

• RQ1 推論時に分布シフトに強いフェデレーテッドラーニングにおけるモデルパーソナライゼーションは、どのように実現できるか？
• RQ2 グローバルモデルの再トレーニングを必要とせずに、テスト時にモデルヘッドを動的に適応させることでパフォーマンスが向上するか？
• RQ3 グローバルおよびローカルヘッドの学習可能なアンサンブルが、テスト時の精度に与える影響は何か？
• RQ4 ローカルヘッドのテスト時微調整は、非IIDデータ下での一般化にどのように影響するか？
• RQ5 テスト履歴を維持することで、パーソナライズド予測の頑健性はどの程度向上するか？

主な発見

• 提案手法は、分布シフト下で標準的なパーソナライゼーションベースラインと比較して顕著な精度向上を達成する。
• 適応的アンサンブル重み付けを伴うテスト時微調整は、静的ヘッドアンサンブルと比較して誤差率を最大25％まで低減する。
• クライアントデータの分布がトレーニングデータから逸脱しても、高いパフォーマンスを維持することができ、強い頑健性を示す。
• 学習可能なアンサンブル重み $e$ の使用により、グローバルおよびローカルヘッドの自動的バランスが可能となり、一般化性能が向上する。
• 非IIDデータ設定下でも本手法は有効であり、多様なクライアントデータ分布において一貫した向上を示す。
• アブレーションスタディの結果、テスト時微調整と履歴に基づく適応の両方が、パフォーマンス向上に顕著な貢献をしていることが確認された。