[論文レビュー] FedMix: Mixed Supervised Federated Learning for Medical Image Segmentation

Jeffry Wicaksana, Zengqiang Yan+5 more|arXiv (Cornell University)|May 4, 2022

Privacy-Preserving Technologies in Data被引用数 10

ひとこと要約

FedMixは、医療画像分類のための新しいフェデレーテッドラーニングフレームワークを提案する。このフレームワークは、ピクセル単位、ボクシングボックス、画像単位の複数のラベルタイプを1つのトレーニングプロセスに統合し、偽ラベルの最適化と動的サンプル選択を用いてモデル品質を向上させる。適応的クライアント重み付けを導入することで、グローバルモデルの収束性が向上し、乳がん腫瘍および皮膚病変分類タスクで最先端の手法を大きく上回る性能を発揮する。

ABSTRACT

The purpose of federated learning is to enable multiple clients to jointly train a machine learning model without sharing data. However, the existing methods for training an image segmentation model have been based on an unrealistic assumption that the training set for each local client is annotated in a similar fashion and thus follows the same image supervision level. To relax this assumption, in this work, we propose a label-agnostic unified federated learning framework, named FedMix, for medical image segmentation based on mixed image labels. In FedMix, each client updates the federated model by integrating and effectively making use of all available labeled data ranging from strong pixel-level labels, weak bounding box labels, to weakest image-level class labels. Based on these local models, we further propose an adaptive weight assignment procedure across local clients, where each client learns an aggregation weight during the global model update. Compared to the existing methods, FedMix not only breaks through the constraint of a single level of image supervision, but also can dynamically adjust the aggregation weight of each local client, achieving rich yet discriminative feature representations. To evaluate its effectiveness, experiments have been carried out on two challenging medical image segmentation tasks, i.e., breast tumor segmentation and skin lesion segmentation. The results validate that our proposed FedMix outperforms the state-of-the-art method by a large margin.

研究の動機と目的

• 既存のフェデレーテッドラーニング手法がクライアント間で均一なラベル監視レベルを仮定しているという制限を解決する。
• 強力（ピクセル単位）、弱い（ボクシングボックス）、最も弱い（画像単位）の異種ラベルタイプを、医療画像分類において効果的に活用できるようにする。
• 生データを共有せずに、容易に入手可能な弱いラベルを活用することで、医療専門家によるアノテーション負荷を軽減する。
• トレーニング中に予測の一貫性に基づいて高一貫性の偽ラベル付きサンプルを動的に選択することで、モデルの頑健性と収束性を向上させる。
• ローカルデータの質と量に基づいてクライアント固有の重みを計算する適応的アグリゲーション機構を開発し、グローバルモデルのパフォーマンスを向上させる。

提案手法

• ピクセル単位、ボクシングボックス、画像単位のラベルを統合した、混合監視をサポートする統一されたFedMixフレームワークを導入する。
• 監視タイプに特化した最適化戦略を適用し、弱いラベル付きデータから高品質な偽ラベルを生成することで、ラベルの正確性を向上させる。
• 2つの独立に訓練されたモデル間の予測一貫性を用いた動的サンプル選択を実装し、低信頼度の偽ラベル付きサンプルをフィルタリングする。
• 予測一致度に基づいて評価サンプルを選択するための整合性閾値εを用い、トレーニングに使用する偽ラベルは信頼性の高いものに限定する。
• ローカルデータの質と量に基づいてクライアント固有の重みを計算する適応的アグリゲーション関数を設計し、グローバルモデルの更新安定性を向上させる。
• ラベル付きデータと偽ラベル付きデータの両方を用いてローカルモデルをトレーニングし、クライアントモデルの重み付きアグリゲーションによりグローバルモデルを更新する。

Figure 1 : Examples of different levels of medical image labels, where the image-level class labels in (b) contain only the lesion category. The bounding box labels in (c) contain not only the lesion category, but also a coarse location. The pixel-level labels in (d) contain both the lesion category

実験結果

リサーチクエスチョン

• RQ1 生データ共有なしに、フェデレーテッドラーニングフレームワークが、ピクセル単位、ボクシングボックス、画像単位の複数の画像監視レベルを医療画像分類で効果的に活用できるか？
• RQ2 2つのモデル間の予測一貫性に基づく動的サンプル選択が、偽ラベルの品質とモデルの収束性に与える影響は何か？
• RQ3 異種フェデレーテッド環境下で、均一なアグリゲーションと比較して、適応的クライアント重み付けがグローバルモデルパフォーマンスをどの程度向上させるか？
• RQ4 FedMixは、FedAvg や FedST といった最先端手法と比較して、完全監視および半監視フェデレーテッドラーニングの両状況でどの程度の性能を発揮するか？
• RQ5 提案されたフレームワークは、多様な医療画像クライアントにおいて、アノテーション負荷を軽減しながら、分類精度を維持または向上させることができるか？

主な発見

• 乳がん腫瘍分類において、FedMixはBUSデータセット（C1）で66.92%のDice係数（DC）、BUSIS（C2）で88.49%、UDIAT（C3）で86.95%を達成し、FedAvg や FedST を大きく上回る。
• 皮膚病変分類では、C1で81.30% DC、C2で78.10%、C3で94.43%、C4で91.11%を記録し、平均86.24%を達成し、クライアント間での強力な一般化性能を示している。
• 動的サンプル選択機構により性能が向上：最適なε = 0.9を用いた場合、BUS（C1）で86.07%、BUSIS（C2）で89.1%、UDIAT（C3）で79.6%を達成し、不正確な偽ラベルのフィルタリングにおけるその有効性が裏付けられた。
• 適応的アグリゲーションにより、乳がん腫瘍分類ではFedAvgより平均DCが1.00%向上し、皮膚病変分類では全体で0.69%向上。特に低品質データを持つクライアントでは1.44%の高い向上を示した。
• アブレーションスタディにより、動的サンプル選択を削除すると性能が低下（例：選択なしで平均DC 37.67% 対 選択ありで80.78%）することが確認され、安定したトレーニングのための必要性が証明された。
• 半監視環境下ではFedSTを上回り、完全監視状況下でもFedAvgを一貫して上回る性能を発揮し、フレームワークの頑健性とスケーラビリティが確認された。

Figure 2 : An illustration of our proposed Mixed Supervised Federated Learning (FedMix) framework. The local client update utilizes every available supervision for training. Based on which, an adaptive weight aggregation procedure is used for the global federated model update. Compared to the existi