[論文レビュー] Compare Where It Matters: Using Layer-Wise Regularization To Improve Federated Learning on Heterogeneous Data

Ha Min Son, Moon Hyun Kim+1 more|arXiv (Cornell University)|Dec 1, 2021

Privacy-Preserving Technologies in Data参考文献 20被引用数 3

ひとこと要約

FedCKA は、非独立同分布（non-IID）データにおける性能を向上させるために、層ごとの正則化フレームワークを提案する。この手法は、代表表現の類似度を測るために中心化カーネル整合性（CKA）を用い、自然に類似しているとされる最初の2層にのみ、選択的に正則化を行う。このアプローチにより、ResNet-50 などの深層モデルでも FedAvg と同等の訓練効率を維持しながら、CIFAR-10 で最高 97.58% の精度を達成する。

ABSTRACT

Federated Learning is a widely adopted method to train neural networks over distributed data. One main limitation is the performance degradation that occurs when data is heterogeneously distributed. While many works have attempted to address this problem, these methods under-perform because they are founded on a limited understanding of neural networks. In this work, we verify that only certain important layers in a neural network require regularization for effective training. We additionally verify that Centered Kernel Alignment (CKA) most accurately calculates similarity between layers of neural networks trained on different data. By applying CKA-based regularization to important layers during training, we significantly improve performance in heterogeneous settings. We present FedCKA: a simple framework that out-performs previous state-of-the-art methods on various deep learning tasks while also improving efficiency and scalability.

研究の動機と目的

• 非独立同分布（non-IID）データ分布下でのフェデレーテッドラーニングにおける性能低下を是正すること。
• 異種環境下において正則化が特に重要となるニューラルネットワークの層を特定すること。
• すべての層に正則化を施さないことで、訓練効率とスケーラビリティを向上させること。
• 非独立同分布データ上で学習されたモデル間の表現類似度を測るための CKA の有効性を評価すること。
• 既存の最先端手法を上回る、スケーラブルで高性能な正則化フレームワークの開発

提案手法

• 性能に最も影響を与えるとされる自然に類似している最初の2層にのみ、層ごとの正則化を適用する。
• クライアント間のローカルモデルの活性化表現間の類似度を測るために、中心化カーネル整合性（CKA）を用いる。
• CKA をもとにした正則化項 ℓcka を導入し、ローカル学習中に最初の2層の表現を一致させる。
• 層の数を増やさずに、クライント、グローバル、ターゲットの3モデルによるフォワードパスを用いて CKA 類似度を効率的に計算する。
• 深層アーキテクチャ全体にわたる各層ごとの演算を避けることで、計算コストを最適化する。
• 既存のフェデレーテッドラーニングパイプラインへの最小限の変更で利用可能な、プラグイン型フレームワークとして FedCKA を設計する。

実験結果

リサーチクエスチョン

• RQ1 ニューラルネットワークのどの層が再初期化に対して最も感受性が高く、フェデレーテッドラーニングにおける正則化の対象として特に重要となるか？
• RQ2 非独立同分布データ上で学習されたモデル間の表現類似度を測る際、CKA は l2 やコサイン類似度といった他の類似度指標と比べてどのように優れているか？
• RQ3 自然に類似している層にのみ正則化を適用することで、性能は向上するが、訓練効率は劣化しないか？
• RQ4 特にすべての層に正則化を施す手法と比較して、モデルの深さが増した場合に FedCKA はどのようにスケーリングするか？
• RQ5 CKA を用いた正則化は、CIFAR-10 や Tiny ImageNet といった標準ベンチマークで、収束性と精度の両面で優れた結果をもたらすか？

主な発見

• α = 5.0 の条件下で、CIFAR-10 において FedCKA は 97.58% の精度を達成し、FedAvg（54.82%）や他の最先端手法を大きく上回る。
• CKA を用いた正則化が最も高い性能を示し、カーネル CKA は線形 CKA 程度にわずかに優れるが、計算コストが高くなる。
• Tiny ImageNet における ResNet-50 でも、FedCKA は FedAvg と同等の訓練時間（750.97s）を維持するが、FedProx や SCAFFOLD、MOON は指数関数的に増加する。
• 最初の2層にのみ正則化を施すことで、訓練のオーバーヘッドが削減され、特に層数の多い深層モデルにおいてスケーラビリティが向上する。
• 表現類似度の測定に CKA を用いることで、表現の一致がより正確になり、真正に類似しない更新のみにペナルティが課される。
• 自然に類似している層にのみ正則化を集中させることで、すべての層に均等に正則化を施す手法よりも効果的であることが示された。