[論文レビュー] Multi-Level Branched Regularization for Federated Learning

Jinkyu Kim, Geeho Kim+1 more|arXiv (Cornell University)|Jul 14, 2022

Privacy-Preserving Technologies in Data被引用数 11

ひとこと要約

本稿では、非i.i.d.な設定におけるモデルドリフトを低減するために、局所的およびグローバルなサブネットワークを複数のハイブリッドパスに統合する多段階分岐正則化手法であるFedMLBを提案する。オンライン知識蒸留を用いて局所的メインパスの表現をハイブリッドブランチの表現と一致させることで、追加の通信コストやメモリコストを増やさずに精度と収束性を向上させ、多様なバックボーンとデータ非同一性の水準において最先端の手法を上回る性能を発揮する。

ABSTRACT

A critical challenge of federated learning is data heterogeneity and imbalance across clients, which leads to inconsistency between local networks and unstable convergence of global models. To alleviate the limitations, we propose a novel architectural regularization technique that constructs multiple auxiliary branches in each local model by grafting local and global subnetworks at several different levels and that learns the representations of the main pathway in the local model congruent to the auxiliary hybrid pathways via online knowledge distillation. The proposed technique is effective to robustify the global model even in the non-iid setting and is applicable to various federated learning frameworks conveniently without incurring extra communication costs. We perform comprehensive empirical studies and demonstrate remarkable performance gains in terms of accuracy and efficiency compared to existing methods. The source code is available at our project page.

研究の動機と目的

• データの非同一性および非i.i.d.なクライアントデータ分布に起因するフェデレーテッドラーニングにおけるモデルドリフトの課題に対処すること。
• 通信コストやメモリオーバーヘッドを増加させることなく、トレーニング中に局所モデルとグローバルモデルの乖離を低減すること。
• グローバル表現知識を保持しつつ、クライアント固有の適応を可能にする正則化技術を開発すること。
• さまざまなフェデレーテッドラーニングフレームワークおよびバックボーンアーキテクチャとの互換性と有効性を確保すること。
• 低参加率および高いデータ非同一性を示すクライアント環境下でも収束速度と最終的な精度を向上させること。

提案手法

• 異なるアーキテクチャレベルでグローバルモデルのサブネットワークを統合することで、各局所モデルに複数のハイブリッドブランチを構築する。
• 各ハイブリッドブランチ内のグローバルサブネットワークのパラメータを固定し、局所モデルのメインパスをトレーニングする。
• 分類能力と知識蒸留の両方を維持するため、すべてのブランチに対してクロスエントロピー損失を適用し、メインパスとハイブリッドブランチの出力間のKLダイバージェンスを用いた知識蒸留を実施する。
• 出力分布の滑らかさを向上させ、蒸留の安定性を向上させるために、温度スケーリングを適用したソフトマックスを知識蒸留に用いる。
• バックプロパゲーション中に、ハイブリッドブランチのパラメータを固定された「学生教師」として維持し、唯一メインパスのパラメータのみを更新する。
• 通信や集約メカニズムを変更せずに、FedAvgなどの標準フェデレーテッドラーニングパイプラインに本手法を統合する。

実験結果

リサーチクエスチョン

• RQ1 通信コストやメモリコストを増加させずに、多段階アーキテクチャ正則化が非i.i.d.フェデレーテッドラーニングにおけるモデルドリフトを低減できるか？
• RQ2 メインパスとハイブリッドブランチ間のオンライン知識蒸留は、標準的な蒸留や正則化と比較して、収束性と精度をどのように向上させるか？
• RQ3 FedMLBは、さまざまなニューラルネットワークアーキテクチャとデータ分布設定において、どの程度一般化可能か？
• RQ4 ハイブリッドパスの数とそのネットワーク深さにおける配置が、性能と効率にどのように影響するか？
• RQ5 低クライアント参加率および高いデータ非同一性下でも、FedMLBは性能を維持できるか？

主な発見

• CIFAR-100において、Dirichlet(0.3)非i.i.d.データ下でResNet-18を用いた場合、FedMLBは54.58%のトップ-1精度を達成し、FedAvg（47.83%）、FedLS-NTD（49.29%）、FedGKD（47.96%）を上回った。
• VGG-9においても、同じ条件下で54.54%の精度を達成し、FedAvg（47.04%）およびFedGKD（48.62%）を顕著に上回った。
• 本手法は、軽量モデルを含むすべてのテストバックボーンで一貫して精度を向上させ、MobileNet（48.34% vs. 47.66% for FedLS-NTD）およびShuffleNet（42.29% vs. 39.27% for FedLS-NTD）においても顕著な改善を示した。
• ハイブリッドパスの数を増やすことで性能が向上し、表8ではパス数の増加に伴い明確な精度上昇トレンドが確認された。
• 局所計算量を高めた場合（例：100回の局所エポック）でもFedMLBは強固な性能を維持したのに対し、FedAVGおよびFedLS-NTDは反復回数の増加に伴い一貫性のない挙動を示した。
• アブレーションスタディの結果、ハイブリッドブランチにおけるクロスエントロピー損失とKLダイバージェンス損失の両方が性能向上に寄与しており、特にKL損失がより重要であることが判明した（表9）。