[論文レビュー] Closing the Generalization Gap of Cross-silo Federated Medical Image Segmentation

An Xu, Wenqi Li+8 more|arXiv (Cornell University)|Mar 18, 2022

Privacy-Preserving Technologies in Data被引用数 7

ひとこと要約

本稿では、個人化モデルとモデルセレクタを導入することで、医療画像セグメンテーションにおけるクロスシロ型FLと集中型学習の一般化ギャップを解消する新規のフェデレーテッドラーニングフレームワーク、FedSMを提案する。新しい最適化手法であるSoftPullを用いて、局所的知識とグローバル知識のバランスを図り、クライアントドリフトを回避しながら、集中型学習と同等の性能を達成する。

ABSTRACT

Cross-silo federated learning (FL) has attracted much attention in medical imaging analysis with deep learning in recent years as it can resolve the critical issues of insufficient data, data privacy, and training efficiency. However, there can be a generalization gap between the model trained from FL and the one from centralized training. This important issue comes from the non-iid data distribution of the local data in the participating clients and is well-known as client drift. In this work, we propose a novel training framework FedSM to avoid the client drift issue and successfully close the generalization gap compared with the centralized training for medical image segmentation tasks for the first time. We also propose a novel personalized FL objective formulation and a new method SoftPull to solve it in our proposed framework FedSM. We conduct rigorous theoretical analysis to guarantee its convergence for optimizing the non-convex smooth objective function. Real-world medical image segmentation experiments using deep FL validate the motivations and effectiveness of our proposed method.

研究の動機と目的

• 非IIDデータとクライアントドリフトに起因する、クロスシロ型フェデレーテッドラーニングにおける医療画像セグメンテーションの一般化ギャップを是正すること。
• すべてのクライアントに適合する単一のグローバルモデルを必要としないために、各クライアントごとに個人化されたモデルを学習すること。
• データ分布の類似性に基づいて、任意のテスト入力に対して最適な個人化モデルを動的に選択するモデルセレクタの設計。
• 新しい個人化FL最適化定式化と、局所的およびグローバルモデル更新のバランスを取ることで一般化性能を向上させる手法、SoftPullの開発。
• SoftPullによって最適化される非凸かつ滑らかな目的関数に対する収束の理論的分析と、実験的検証により、FedSMが集中型学習と同等の性能を達成することの裏付け。

提案手法

• 各クライアントごとに個人化モデルの集合を学習し、テスト入力を最も適切な個人化モデルにルーティングする共有モデルセレクタを備えたフレームワークであるFedSMを提案。
• 局所的およびグローバルモデル更新の重み付き補間として定式化された、新規の個人化FL目的関数を導入。この関数は学習可能な係数λでパrameter化されている。
• SoftPullを提案。この手法は、局所データとグローバルモデルの両方に整合するよう促すソフト制約を適用することで、個人化目的関数を最適化する。
• セグメンテーションヘッドの最終層からの特徴埋め込みに基づくモデルセレクタを採用。これにより、テスト入力に対して最も適した個人化モデルを動的に選択可能となる。
• モデルセレクタには閾値γを導入。これにより、どの個人化モデルに対しても十分な信頼性がない場合にグローバルモデルにフォールバックできるようにし、未学習クライアントに対する耐性を向上。
• SoftPullが最適化する非凸かつ滑らかな目的関数に対して、標準的なFLの仮定に基づく厳密な収束解析を提供。

Figure 1 : The proposed FedSM framework with “super model”.

実験結果

リサーチクエスチョン

• RQ1 フェデレーテッドラーニングフレームワークは、クロスシロ型医療画像セグメンテーションにおけるFLと集中型学習の一般化ギャップを解消できるか？
• RQ2 クライアントドリフトを個人化モデルによって回避することで、従来のFL手法（例：FedAvg）と比較して一般化性能が向上するか？
• RQ3 モデルセレクタは、データ分布の類似性に基づいて、与えられたテスト入力に対して最も適した個人化モデルを効果的に同定できるか？
• RQ4 SoftPull最適化手法は、既存の個人化FLベースラインと比較して、より優れた一般化性能を達成するか？
• RQ5 個人化目的関数における調整可能な補間係数λは、小規模な局所データセットにおける過学習を防ぐことで、一般化性能の向上に寄与するか？

主な発見

• FedSMは、集中型学習と同等のテストDiceスコアを達成し、実世界の医療画像セグメンテーションタスクにおいて、一般化ギャップを効果的に解消した。
• 網膜セグメンテーションタスクでは、FedAvgおよびローカル学習を上回り、特に小規模データセットを持つクライアントにおいて顕著な性能向上を示した。最適な閾値γを用いることで、最大3%のスコア向上を達成した。
• SoftPullは、APFL、Per-FedAvg、FTなどの他の個人化手法を上回り、次に良いベースラインと比較して0.62%高いグローバルDice係数を達成した。
• 低データ類似度タスク（例：網膜）では最適なλ値が0.7、高類似度タスク（例：前立腺）では0.3であることが判明。SoftPull定式化の適応性を裏付ける。
• 損失関数の表面解析により、SoftPullはFedAvgと比較してより平坦な損失ランドスケープを達成しており、わずかな収束誤差があるものの、より優れた一般化性能を示していることが示された。
• モデルセレクタは、未学習クライアントのデータ分布に最も類似したクライアントを正しく同定できており、特徴可視化（t-SNE）と実測性能の両方と一致する選択頻度を示した。

Figure 2 : Training curves comparison. The curves are non-decreasing because we record the best result during training.