[論文レビュー] FedBone: Towards Large-Scale Federated Multi-Task Learning

Yiqiang Chen, Teng Zhang+4 more|arXiv (Cornell University)|Jun 30, 2023

Retinal Imaging and Analysis被引用数 4

ひとこと要約

FedBoneは、サーバーに配置されたグローバル汎用モデルとクライントに配置されたタスク固有モデルに分離するフェデレーテッド・マルチタスク学習フレームワークを提案する。このアプローチにより、リソース制約のあるエッジクライント上でも大規模モデルの学習が可能となる。GPAggregationを導入し、異種タスク間の勾配競合を軽減し、ベンチマークおよび実際の眼科画像データセットにおいて最先端の性能を達成するとともに、クライント側の計算負荷を最小限に抑える。

ABSTRACT

Heterogeneous federated multi-task learning (HFMTL) is a federated learning technique that combines heterogeneous tasks of different clients to achieve more accurate, comprehensive predictions. In real-world applications, visual and natural language tasks typically require large-scale models to extract high-level abstract features. However, large-scale models cannot be directly applied to existing federated multi-task learning methods. Existing HFML methods also disregard the impact of gradient conflicts on multi-task optimization during the federated aggregation process. In this work, we propose an innovative framework called FedBone, which enables the construction of large-scale models with better generalization from the perspective of server-client split learning and gradient projection. We split the entire model into two components: a large-scale general model (referred to as the general model) on the cloud server and multiple task-specific models (referred to as the client model) on edge clients, solving the problem of insufficient computing power on edge clients. The conflicting gradient projection technique is used to enhance the generalization of the large-scale general model between different tasks. The proposed framework is evaluated on two benchmark datasets and a real ophthalmic dataset. Comprehensive results demonstrate that FedBone efficiently adapts to heterogeneous local tasks of each client and outperforms existing federated learning algorithms in most dense prediction and classification tasks with off-the-shelf computational resources on the client side.

研究の動機と目的

• エッジクライントの計算リソースが限られているフェデレーテッド学習環境において、大規模モデルを効率的に学習する課題に対処すること。
• 異なるクライントが異なる目的を最適化する異種タスクの状況下で、モデルアグリゲーション段階で生じる勾配競合を処理すること。
• クライント側の計算能力を高価に要しない条件下でも、多様なタスクにわたるグローバルモデルの一般化性能を向上させること。
• 医療画像診断（例：眼科疾患診断）などプライバシーが重要な分野への大規模モデルの実用的導入を可能とすること。

提案手法

• スプリット学習アーキテクチャ：グローバル汎用モデルはクラウドサーバーに配置され、クライントは軽量なタスク固有ヘッドのみを保持・計算する。
• クライント側で局所的なパッチ埋め込みが実行され、特徴量は大規模な汎用モデルを用いてサーバーで抽出される。
• サーバーは抽出された潜在表現をクライントに送信し、クライントはそれらを用いてタスク固有の出力予測を実行する。
• GPAggregation：歴史的勾配注釈を用いてクライント勾配を再スケーリングする新規な勾配アグリゲーション手法で、異種タスク間の競合を低減する。
• 可変畳み込みと自己注意機構を用いたタスク適応モジュールにより、低レベル特徴量間の相互作用を強化する。
• モデル学習はクライントとサーバー間の交互な順方向および逆方向伝搬により進行し、ローカルなパrameterの保存量と計算量を最小限に抑える。

Figure 1: The overview of our proposed framework

実験結果

リサーチクエスチョン

• RQ1 フェデレーテッドマルチタスク学習フレームワークとして、大規模モデルのスケーリングを実現しつつ、エッジクライントにおける計算要求を低く抑えることは可能か？
• RQ2 異種タスクに起因する勾配競合をフェデレーテッドアグリゲーション段階でどのように軽減できるか？その結果、モデルの一般化性能が向上するか？
• RQ3 グローバル汎用モデルとクライント固有ヘッドを用いたスプリット学習アプローチが、既存のフェデレーテッド学習ベースラインに比べて、どの程度優れた性能を発揮するか？
• RQ4 提案フレームワークは、クライント側リソースを最小限に抑えた状態でも、実世界の医療画像認識タスクで最先端の性能を達成できるか？

主な発見

• 実世界のデータセットにおける13の眼科的タスクのうち12つで、FedAvg、FedProx、pFedMe、FedEMを上回り、平均F1スコアが95.90という最高値を記録した。
• Retina-Lesionsデータセットにおいて、セマンティックセグメンテーション（LSタスク）で最高性能を達成し、医療画像セグメンテーション分野への応用可能性が顕著に示された。
• アブレーションスタディの結果、GPAggregationおよびタスク適応モジュールの両方が性能向上に顕著な寄与を示し、フルFedBoneモデルが全設定で最高の精度を達成した。
• CIFAR-100マルチタスクベンチマークでは98.93％の精度を達成し、オフザシェルクライントハードウェアを用いても他のFL手法を上回った。
• フレームワークにより、クライント側のメモリと計算負荷が著しく削減され、モデルパラメータの大部分と特徴抽出処理がサーバーに移行されたため、低リソースデバイスへの展開が可能になった。
• GPAggregation手法は勾配競合を効果的に低減しており、多様なタスクにわたる収束性と一般化性能の向上が裏付けられた。

Figure 2: The workflow of $\mathtt{FedBone}$ framework. Clients perform patch embedding locally and (1) send embeddings to the cloud server for feature extraction using the general model, and the cloud server (2) sends extracted features back to clients. Clients complete the loss computation and (3)