[論文レビュー] TCT: Convexifying Federated Learning using Bootstrapped Neural Tangent Kernels

Yaodong Yu, Alexander Wei+3 more|arXiv (Cornell University)|Jul 13, 2022

Privacy-Preserving Technologies in Data被引用数 5

ひとこと要約

本稿では、まずフェデレーテッド・アベレージ（FedAvg）を用いてフェデレーテッド学習により特徴を抽出し、その後、経験的ニューラルタングエント・カーネル（eNTK）を用いて凸化されたモデルに対して勾配補正最適化を適用する、TCT（Train-Convexify-Train）フレームワークを提案する。この手法により、極端なデータ非イ.i.d.状態下でもCIFAR10で最大37%の精度向上が達成され、中央集権学習とフェデレーテッド学習の間の性能ギャップが完全に埋まる。

ABSTRACT

State-of-the-art federated learning methods can perform far worse than their centralized counterparts when clients have dissimilar data distributions. For neural networks, even when centralized SGD easily finds a solution that is simultaneously performant for all clients, current federated optimization methods fail to converge to a comparable solution. We show that this performance disparity can largely be attributed to optimization challenges presented by nonconvexity. Specifically, we find that the early layers of the network do learn useful features, but the final layers fail to make use of them. That is, federated optimization applied to this non-convex problem distorts the learning of the final layers. Leveraging this observation, we propose a Train-Convexify-Train (TCT) procedure to sidestep this issue: first, learn features using off-the-shelf methods (e.g., FedAvg); then, optimize a convexified problem obtained from the network's empirical neural tangent kernel approximation. Our technique yields accuracy improvements of up to +36% on FMNIST and +37% on CIFAR10 when clients have dissimilar data.

研究の動機と目的

• データの非イ.i.d.状態下におけるフェデレーテッド学習の性能低下を解消すること、特に非凸な深層ニューラルネットワークに対して。
• 既存のフェデレーテッド最適化手法が、有用であるにもかかわらず、ネットワークの初期層で学習された特徴を活用できない理由を調査すること。
• 非凸性下でも学習済み特徴を保持しつつ、最終層での安定的かつ高精度な学習を可能にする手法を開発すること。
• eNTKによる凸化が、非凸なフェデレーテッド学習環境下で中央集権学習に近い性能を回復できることを示すこと。
• 提案されたTCTフレームワークが、ハイパーパrameterや通信効率に対し、極端なデータ非イ.i.d.状態下でも頑健であることを示すこと。

提案手法

• 最初に、フェデレーテッド学習を用いて深層ニューラルネットワークを訓練し、クライアントデータから有用な特徴を抽出する。
• 2番目に、学習済みモデルの経験的ニューラルタングエント・カーネル（eNTK）近似を計算し、最終層を線形化する。
• 3番目に、eNTK近似特徴に基づく凸問題を解くために勾配補正最適化（例：SCAFFOLD）を適用する。
• 2段階目の最適化における収束速度と通信効率を向上させるために、eNTK特徴を正規化する。
• 1段階目のモデルチェックポイントを用いてeNTKベースの凸最適化を初期化し、特徴の一貫性を保証する。
• 最終層を固定されたeNTK特徴上ですべての層をエンドツーエンドで訓練し、事前に学習済み表現上に線形ヘッドを学習する。

実験結果

リサーチクエスチョン

• RQ1 なぜ現在のフェデレーテッド最適化手法は、データが非イ.i.d.であっても中央集権学習と同等の性能を達成できないのか？
• RQ2 ニューラルネットワークの初期層がフェデレーテッド環境下でどれほど有用な特徴を学習するのか、そしてなぜ最終層がそれらを活用できないのか？
• RQ3 経験的ニューラルタングエント・カーネル（eNTK）による凸化が、データ非イ.i.d.状態下の非凸フェデレーテッド学習で性能を回復できるのか？
• RQ4 1段階目の事前学習期間の選択が、TCTフレームワークにおける最終モデルの精度にどのように影響するのか？
• RQ5 eNTK特徴の正規化が、2段階目の凸最適化における収束速度と通信効率を向上させるのか？

主な発見

• クライアントのデータ分布が著しく異なる状況下で、FMNISTでは最大36%、CIFAR10では最大37%の精度向上を達成した。
• 1段階目に事前学習済み特徴抽出器を用いることで、中央集権学習とフェデレーテッド学習の間の性能ギャップを完全に埋めた。
• FedAvgで訓練されたモデルから抽出したeNTK特徴を用いたTCTは、標準的なFedAvgに比べてテスト精度で約20%高い性能を示した。
• eNTK特徴の正規化は収束速度を顕著に向上させ、TCTが約40回の通信ラウンドでほぼ100%の訓練精度に達するのを可能にした。
• TCTの性能は、1段階目の通信ラウンド数（T₁ ≥ 60）にかかわらず安定しており、ハイパーパrameter選択に対する頑健性を示した。
• FedAvgの性能は、ローカルステップ数の増加（M=1000）に伴い劣化する一方で、TCTで用いられるSCAFFOLDは収束を維持または向上させることから、勾配補正の重要性が示された。