[論文レビュー] Rethinking Architecture Design for Tackling Data Heterogeneity in Federated Learning

Liangqiong Qu, Yuyin Zhou+6 more|arXiv (Cornell University)|Jun 10, 2021

Privacy-Preserving Technologies in Data参考文献 61被引用数 12

ひとこと要約

本稿では、データの非独立同分布性（non-IID）が顕著な状況下において、畳み込みニューラルネットワーク（CNN）の代替としてビジョントランスフォーマー（ViT）を提示し、特に災難的忘却（catastrophic forgetting）の軽減と収束速度の向上を実証した。主な貢献は、複数のベンチマークとデータ分割設定において、ViT-FLがCNNベースのFLを上回ることを実証的裏付けで示したことである。特に、高いデータ非同一性が見られる状況下で顕著な優位性を示した。

ABSTRACT

Federated learning is an emerging research paradigm enabling collaborative training of machine learning models among different organizations while keeping data private at each institution. Despite recent progress, there remain fundamental challenges such as the lack of convergence and the potential for catastrophic forgetting across real-world heterogeneous devices. In this paper, we demonstrate that self-attention-based architectures (e.g., Transformers) are more robust to distribution shifts and hence improve federated learning over heterogeneous data. Concretely, we conduct the first rigorous empirical investigation of different neural architectures across a range of federated algorithms, real-world benchmarks, and heterogeneous data splits. Our experiments show that simply replacing convolutional networks with Transformers can greatly reduce catastrophic forgetting of previous devices, accelerate convergence, and reach a better global model, especially when dealing with heterogeneous data. We release our code and pretrained models at https://github.com/Liangqiong/ViT-FL-main to encourage future exploration in robust architectures as an alternative to current research efforts on the optimization front.

研究の動機と目的

• デバイス間でデータが非同一性を示す状況下における、非収束と災難的忘却という、長年のフェデレーテッドラーニング（FL）の課題に取り組む。
• 特にビジョントランスフォーマー（ViT）といったアーキテクチャ選択が、FLにおける分布シフトに対して本質的に頑健であるかどうかを調査する。
• フェデレーテッドアルゴリズム、実世界のデータセット、非同一性の強いデータ分割をカバーする、ViTとCNNの包括的で体系的な実証的ベンチマークを提供する。
• 最適化ヒューリスティクスやハイパーパrameterチューニングを追加で行わずに、アーキテクチャの改善が即座に性能向上をもたらす可能性を示す。
• ビジョントランスフォーマーを、既存の最適化ベースの手法とは直交する形で、将来的なフェデレーテッドラーニング研究の出発点として位置づける。

提案手法

• フェデレーテッドラーニングアルゴリズム（FedAVG（並列）およびCWT（逐次））を複数用い、ビジョントランスフォーマー（ViT）とCNN（ResNet, EfficientNet）の広範な実証的ベンチマークを実施した。
• CIFAR-10およびRetinaの画像分類タスクにおいて、標準的なViT実装を用い、一貫したデータ拡張および正則化戦略を適用した。
• リアルワールドの非同一性を再現するために、ラベル分布の偏りが増加する3つのデータ分割戦略（Split-1～Split-3）を用いて評価した。
• UMAP可視化を用いて、ViTとCNNモデルが学習した特徴埋め込みを比較し、クラスタ分離度と表現の頑健性を分析した。
• 既存の最適化ベースのFL手法（FedProx, FedAVG-Share）とViTを組み合わせ、相互運用性と性能向上の有効性を検証した。
• 事前学習および学習率、勾配クリッピング、ローカルエポック数などの訓練ハイパーパrameterが、ViT-FLの性能に与える影響を調査した。

実験結果

リサーチクエスチョン

• RQ1 ビジョントランスフォーマー（ViT）は、CNNと比較して、高いデータ非同一性下のフェデレーテッドラーニングにおいて、災難的忘却を軽減できるか？
• RQ2 異なるフェデレーテッドラーニングアルゴリズム（FedAVG 対 CWT）およびデータ分割戦略において、ViT-FLの性能はCNNベースのFLと比べてどのように異なるか？
• RQ3 アーキテクチャ選択（ViT 対 CNN）が、非IID環境下での収束速度と最終的なモデル精度に与える影響はどの程度か？
• RQ4 大規模データセットが入手できない状況下でも、事前学習がViT-FLの性能に顕著な向上をもたらすか？
• RQ5 ViT-FLは、既存の最適化ベースのFL手法と効果的に組み合わせられ、非同一性の高いデータ環境下でさらに性能を向上できるか？

主な発見

• CIFAR-10の極めて非同一性の高いデータ分割（Split-3）において、ビジョントランスフォーマー（ViT）は、ResNet(50)-FedAVGの59.68％と比較して、事前学習を施した場合に96.40％のテスト精度を達成し、顕著に優れた性能を示した。
• ViT-FLは災難的忘却を軽減する：UMAP可視化において、ViT(S)-FedAVGが学習した特徴は、ResNet(50)-FedAVGと比較して明確なクラス分離を示した。
• ViT-FLは収束を加速し、特に深刻なデータ非同一性下でより優れたグローバルモデルに到達する。性能向上は、データの偏りが高くなるほど顕著に増大した。
• 事前学習はViT-FLにとって不可欠である：大規模な事前学習が行われない状況下でも、スウィン（Swin）（T）-FedAVGはSplit-3で64.50％の精度を達成し、ResNet(50)-FedAVG（59.68％）を上回った。
• ViT-FLは最適化ベースの手法と直交的である：ViTとFedProxやFedAVG-Shareを組み合わせることで、非同一性の高いクライアント環境下でさらに性能向上が得られた。
• 最適なパフォーマンスを得るには、極めて非同一性の高い環境下では、小さなローカルエポック数（ViT-FedAVGではE ≤ 5、ViT-CWTではE = 1）を設定すべきであり、同質な環境では通信回数を減らすためにより大きなE値を用いることができる。