[論文レビュー] Auto-FedAvg: Learnable Federated Averaging for Multi-Institutional Medical Image Segmentation

Yingda Xia, Dong Yang+16 more|arXiv (Cornell University)|Apr 20, 2021

Privacy-Preserving Technologies in Data参考文献 50被引用数 36

ひとこと要約

Auto-FedAvg は、フェデレーテッドラーニングのデータ駆動型集約重みを学習し、Dirichlet分布に基づく微分可能なアプローチでクライアントモデルを組み合わせ、CIFAR-10および多機関医療セグメンテーションタスクで FedAvg を改善します。

ABSTRACT

Federated learning (FL) enables collaborative model training while preserving each participant's privacy, which is particularly beneficial to the medical field. FedAvg is a standard algorithm that uses fixed weights, often originating from the dataset sizes at each client, to aggregate the distributed learned models on a server during the FL process. However, non-identical data distribution across clients, known as the non-i.i.d problem in FL, could make this assumption for setting fixed aggregation weights sub-optimal. In this work, we design a new data-driven approach, namely Auto-FedAvg, where aggregation weights are dynamically adjusted, depending on data distributions across data silos and the current training progress of the models. We disentangle the parameter set into two parts, local model parameters and global aggregation parameters, and update them iteratively with a communication-efficient algorithm. We first show the validity of our approach by outperforming state-of-the-art FL methods for image recognition on a heterogeneous data split of CIFAR-10. Furthermore, we demonstrate our algorithm's effectiveness on two multi-institutional medical image analysis tasks, i.e., COVID-19 lesion segmentation in chest CT and pancreas segmentation in abdominal CT.

研究の動機と目的

• クロスシロ医用画像でフェデレーションクライアント間の非独立同分布（非i.i.d.）データ分布に対処する。
• 固定データサイズベースの重みを使うのではなく、集約重みを学習する微分可能なメカニズムを開発する。
• 追加通信を最小限に抑えたプライバシー保護フェデレーションを保証する。
• 医用画像タスクにおいてFedAvgおよび最先端のFL手法より改善を実証する。
• ネットワーク単位およびレイヤー単位の集約変種と Dirichlet 対 Softmax のパラメータ化を検討する。

提案手法

• 局所モデルパラメータとグローバル集約パラメータを反復更新のために分離する。
• 集合重みを可学习パラメータとして表現し（softmaxまたはDirichlet）、単体制約を課す。
• Dirichlet分布を用いて集約重みをモデル化し、勾配ベースの最適化を可能にする。
• 生データを共有せずに集約重みを更新する通信効率の高い LearnAggWeight ステップを導入する。
• ネットワーク単位およびレイヤー単位の集約方式を提案し、層ごとまたはネットワーク全体としてクライアントの寄与を柔軟に重み付けする。
• 追加通信と性能のバランスを取るために集約重み学習間隔 t0 を分析する。

実験結果

リサーチクエスチョン

• RQ1 学習可能でデータ駆動型の集約重みは、クロスシロ FL における非i.i.d. データのもとでフェデレーテッド平均化を改善できるか？
• RQ2 Dirichlet分布の集約重みは、Softmaxベースの重みより性能と汎化能力が向上するか？
• RQ3 ネットワーク単位 vs レイヤー単位の集約が性能と通信コストに与える影響は？
• RQ4 集約重み学習間隔 t0 が、跨機関設定における性能と効率にどう影響するか？

主な発見

• ネットワーク単位 Dirichlet 重みを用いる Auto-FedAvg-N-Dirichlet は CIFAR-10 のバリアントの中で最良の結果を達成し、最終精度 88.98%、FedMA (87.53%) を 1.45% 上回る。
• 多国籍COVID-19病変セグメンテーションでは、Auto-FedAvg-N-Dirichlet バリアントが FedAvg よりグローバルモデル性能、平均ローカルモデル性能、および局所一般化を改善。例として、報告された実行でグローバルテスト平均が FedAvg より約2パーセントポイント向上。
• Auto-FedAvg は CIFAR-10 および2つの医用画像タスク（COVID-19 病変セグメンテーションと膵臓セグメンテーション）で、FedAvg、FedProx、FedMA のベースラインを一貫して上回る。
• Dirichlet 集約重みはタスク全般で Softmax より良い結果をもたらす傾向があり、ネットワーク単位の集約は多くの設定でレイヤー単位の集約より上回る。
• 訓練初期に集約重みを積極的に学習すると収束を加速し、後半に重みを安定化させると一般化をさらに高める。
• 膵臓セグメンテーション実験では、Auto-FedAvg-N-Dirichlet* が最良のグローバルおよびローカル性能を達成し、FedAvg を顕著に上回った。