[論文レビュー] FedDAR: Federated Domain-Aware Representation Learning

Aoxiao Zhong, Hao He+3 more|arXiv (Cornell University)|Sep 8, 2022

Machine Learning in Healthcare被引用数 4

ひとこと要約

FedDAR は、各クライアントのデータが事前に定義されたドメインの混合であるドメイン混合非 i.i.d. データを対象とする、フェデレーテッドラーニングの新しいフレームワークを提案する。表現学習と予測を分離し、ドメインに特化したヘッドを用いながら全ドメインでエンコーダーを共有することで、交互最適化と2次元集約を用いたドメイン認識型パーソナライゼーションを可能にし、線形収束性と優れた性能を達成した。合成データおよび実世界の医療データセットにおいて優れた性能を示した。

ABSTRACT

Cross-silo Federated learning (FL) has become a promising tool in machine learning applications for healthcare. It allows hospitals/institutions to train models with sufficient data while the data is kept private. To make sure the FL model is robust when facing heterogeneous data among FL clients, most efforts focus on personalizing models for clients. However, the latent relationships between clients' data are ignored. In this work, we focus on a special non-iid FL problem, called Domain-mixed FL, where each client's data distribution is assumed to be a mixture of several predefined domains. Recognizing the diversity of domains and the similarity within domains, we propose a novel method, FedDAR, which learns a domain shared representation and domain-wise personalized prediction heads in a decoupled manner. For simplified linear regression settings, we have theoretically proved that FedDAR enjoys a linear convergence rate. For general settings, we have performed intensive empirical studies on both synthetic and real-world medical datasets which demonstrate its superiority over prior FL methods.

研究の動機と目的

• クライアントのデータを事前に定義されたドメインの混合としてモデル化することで、i.i.d. またはクライアント固有の分布を仮定するのではなく、クロスシロフェデレーテッドラーニングにおけるデータの非同一性を解消すること。
• 既存の FL メソッドがクライアントごとのパーソナライゼーションを行うが、特に人種やスキャナープロトコルなどのドメイン構造を無視するという限界を克服すること。医療分野では、これらのドメイン要因が分布シフトを引き起こす。
• ドメイン間で共有表現を学習しつつ、ドメイン固有の予測ヘッドを可能にすることで、モデルの頑健性と公平性を向上させること。
• 理論的および実験的に、ドメイン混合 FL 環境下でドメイン別パーソナライゼーションがクライアント別パーソナライゼーションを上回ることを検証すること。

提案手法

• FedDAR はモデルを共有エンコーダーとドメイン固有の予測ヘッドに分解し、表現学習と予測を分離することで、一般化性能の向上を図る。
• 局所的ヘッドの更新（複数のローカルエポックを用いて）と、重み付き平均によるエンコーダー重みの集約を交互に実行し、ドメイン公平性を確保するためのサンプル再重み付けを適用する。
• ヘッドの集約には、収束性と最適性の向上を図るための新規な2次最適化戦略を採用する。
• 2段階のトレーニングプロセスを採用：まず FedAvg のウォームアップを行い、その後エンコーダーとヘッドのローカル更新とグローバル集約を交互に実行する。
• トレーニング中に、代表数が少ないドメインからの寄与をバランスさせるために、ドメイン認識型損失再重み付けを適用する。
• 線形および非線形モデルをサポートし、実用的な医療応用において CNN および ResNet-34 を用いて評価した。

実験結果

リサーチクエスチョン

• RQ1 ドメイン混合を明示的にモデル化するフェデレーテッドラーニング手法が、非 i.i.d. 医療データ環境下で、標準的およびパーソナライズド FL ベースラインを上回る性能を発揮できるか？
• RQ2 ドメイン固有の予測と表現学習を分離することで、共同最適化と比較して収束性と性能が向上するか？
• RQ3 予測ヘッドのドメイン認識型集約が、多様なクライアントデータ構成におけるモデルの公平性と一般化性能に与える影響は？
• RQ4 線形設定下で、提案手法に理論的保証（収束速度など）を確立できるか？
• RQ5 実世界の医療データセットにおいて、FedDAR は代表数が少ないドメインのバイアスをどれほど低減し、性能を向上できるか？

主な発見

• 線形回帰設定において FedDAR は線形収束性を達成し、強力な理論的収束特性を示した。
• 合成データにおいて、FedDAR はすべてのドメインにおいて、FedAvg、FedProx、FedPer、FedRep を上回るテスト精度を達成した。
• FairFace データセットにおいて、FedDAR は最先端の手法と比較して、全人種グループの平均精度が高く、公平性が向上しバイアスが低減した。
• 実世界の EXAM データセットにおいて、FedDAR はローカル検証セットで平均 AUC 0.89 を達成し、FedProx、FedRep、FedMinMax を顕著に上回った。
• アブレーションスタディにより、2次元ヘッド集約とドメイン認識型再重み付けが、特に不均衡なドメインにおいて性能に不可欠であることが確認された。
• 1クライアントが1ドメインあたり100件のサンプルしか持たない状況でも、FedDAR は高い性能を維持した。これは、低データドメインに対しても頑健であることを示している。