[論文レビュー] Robust Split Federated Learning for U-shaped Medical Image Networks

Ziyuan Yang, Yingyu Chen+5 more|arXiv (Cornell University)|Dec 13, 2022

Privacy-Preserving Technologies in Data被引用数 9

ひとこと要約

本稿では、分割学習とフェデレーテッドラーニングを組み合わせたハイブリッドフレームワークであるロバストスプリットフェデレーテッドラーニング（RoS-FL）を提案する。このフレームワークにより、U字型の医療画像ネットワークのプライバシー保護型で計算効率の良い学習が可能になる。モデルをヘッド、ボディ、テイルの3つの部分に分割し、クライアント、計算サーバー、集約サーバーに分散配置することで、入力、ラベル、出力、モデルのプライバシーを保護しつつ並列学習を実現する。主な貢献は、データの非独立同分布（non-IID）性に起因する学習の不安定性を緩和するための動的重み補正戦略（DWCS）であり、これによりモデルのずれが軽減され、クライアント間での性能の一貫性が向上する。

ABSTRACT

U-shaped networks are widely used in various medical image tasks, such as segmentation, restoration and reconstruction, but most of them usually rely on centralized learning and thus ignore privacy issues. To address the privacy concerns, federated learning (FL) and split learning (SL) have attracted increasing attention. However, it is hard for both FL and SL to balance the local computational cost, model privacy and parallel training simultaneously. To achieve this goal, in this paper, we propose Robust Split Federated Learning (RoS-FL) for U-shaped medical image networks, which is a novel hybrid learning paradigm of FL and SL. Previous works cannot preserve the data privacy, including the input, model parameters, label and output simultaneously. To effectively deal with all of them, we design a novel splitting method for U-shaped medical image networks, which splits the network into three parts hosted by different parties. Besides, the distributed learning methods usually suffer from a drift between local and global models caused by data heterogeneity. Based on this consideration, we propose a dynamic weight correction strategy ( extbf{DWCS}) to stabilize the training process and avoid model drift. Specifically, a weight correction loss is designed to quantify the drift between the models from two adjacent communication rounds. By minimizing this loss, a correction model is obtained. Then we treat the weighted sum of correction model and final round models as the result. The effectiveness of the proposed RoS-FL is supported by extensive experimental results on different tasks. Related codes will be released at https://github.com/Zi-YuanYang/RoS-FL.

研究の動機と目的

• U字型医療画像ネットワークにおけるプライバシー懸念を解決すること。ここでの生データ、ラベル、出力は機微な情報である。
• フェデレーテッドラーニングにおけるクライアント側の計算負荷を軽減すること。重いモデル部品をサーバーにオフロードすることで実現する。
• 分散学習環境におけるデータの非IID性に起因するモデルのずれ（drift）を緩和すること。
• クライアントとサーバー間でのデータおよびモデルのプライバシーを維持しつつ、並列学習を可能にすること。
• 非IIDデータを持つクライアント間で学習の安定性と性能の一貫性を向上させるため、新規の動的重み補正メカニズムを導入すること。

提案手法

• U字型ネットワークを3つのコンponentsに分割：クライアント側に配置するヘッド、サーバー側に配置するボディ、再びクライアント側に配置するテイル。これにより計算負荷の分散とプライバシー強化を実現する。
• フェデレーテッドラーニングとスプリットラーニングをハイブリッド的に組み合わせたアーキテクチャを採用し、並列学習を可能にするとともにクライアント側の計算負担を軽減する。
• 生データ、ラベル、最終出力の送信を避けるために、クライアントとサーバー間で中間特徴マップのみを送信する。
• 動的重み補正戦略（DWCS）を実装。これは、現在のグローバルモデルと直前のグローバルモデルの重み付き和を計算することで、モデルのずれを補正する。
• 連続する通信ラウンド間のモデルのずれを定量化するための重み補正損失を定義し、この損失を最小化するように補正モデルを最適化する。
• クライアント側の勾配とサーバー側のモデル更新を集中型の集約サーバーで集約し、グローバルモデルを更新する。

実験結果

リサーチクエスチョン

• RQ1 ハイブリッドスプリットフェデレーテッドラーニングフレームワークは、U字型医療画像ネットワークにおいて入力、ラベル、出力、モデルパラメータを効果的に保護できるか？
• RQ2 医療画像モデルの分散学習において、データの非IID性に起因するモデルのずれはどのように緩和できるか？
• RQ3 提案された動的重み補正戦略（DWCS）は、非IIDデータを持つクライアント間で学習の安定性と性能の一貫性をどの程度向上させるか？
• RQ4 提案されたRoS-FLフレームワークは、標準的なフェデレーテッドラーニングと比較して、クライアント側の計算コストを削減しつつも、競争力のある性能を達成できるか？
• RQ5 RoS-FLの性能は、補正重みμのようなハイパーパrameterの選択にどれほど敏感か？

主な発見

• RoS-FLはクライアント間で一貫した性能を達成し、ベースライン手法で見られた顕著な性能差（モデルのずれに起因）を解消した。
• 修復タスクにおいて、1回の通信ラウンドと500回のローカルエポックで最高のPSNR 42.81を達成し、他の設定を上回った。
• セグメンテーションタスクでは、2回の通信ラウンドと150回のローカルエポックで最高のDiceスコア88.21を達成し、データの非IID性に対しても高いロバストネスを示した。
• ハイパーパrameterμに対しては、広い範囲（1×10⁻⁶から100）で安定した性能を示し、特にμ = 1×10⁻⁴で最良の性能を発揮した。
• 画像処理タスクでは、閾値を超えてローカルエポック数を増加させると性能が低下することが判明した。これは、モデルのずれが増大するためであり、DWCSが学習の安定化に不可欠であることを示している。
• 可視化比較により、RoS-FLは他の手法と比較して、より明確なエッジと構造的整合性の高い出力を得ており、アーチファクトも少ないことが確認された。