[논문 리뷰] FedBone: Towards Large-Scale Federated Multi-Task Learning

Yiqiang Chen, Teng Zhang외 4명|arXiv (Cornell University)|2023. 06. 30.

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.

연구 동기 및 목표

• 엣지 클라이언트의 계산 자원이 제한된 환경에서 대규모 모델을 훈련하는 데 도전하는 데 목적이 있다.
• 다른 클라이언트들이 서로 다른 목표를 최적화함에 따라 모델 집합 과정에서 기울기 갈등을 유발하는 작업 이질성 문제를 다루는 데 목적이 있다.
• 높은 클라이언트 측 계산 능력이 필요 없이 다양한 작업 간의 전역 모델 일반화를 향상시키는 데 목적이 있다.
• 의료 영상과 같이 개인정보가 민감한 분야에서 대규모 모델의 실용적 구현을 가능하게 하는 데 목적이 있다. 예를 들어 망막 질환 진단과 같은 분야이다.

제안 방법

• 분할 학습 아키텍처: 전역 일반 모델은 클라우드 서버에 배치되며, 클라이언트는 경량의 작업별 특화 헤드만을 저장하고 계산한다.
• 클라이언트 측에서 패치 임bedding을 현지에서 수행하고, 특징을 서버로 전송하여 대규모 일반 모델을 사용해 추출한다.
• 서버는 추출된 잠재 표현을 클라이언트로 다시 전송하여 작업별 출력 예측을 수행한다.
• 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 데이터셋에서 FedBone은 의미적 세그멘테이션(LS 작업)에서 최고 성능을 기록하여 의료 영상 세그멘테이션 잠재력이 높음을 시사한다.
• 제거 실험을 통해 GPAggregation와 작업 적응 모듈이 모두 성능 향상에 기여하는 것으로 확인되었으며, 전체 FedBone 모델이 모든 설정에서 가장 높은 정확도를 달성했다.
• CIFAR-100 다중 작업 벤치마크에서 FedBone은 오프더쇼프 클라이언트 하드웨어를 사용함에도 불구하고 다른 FL 방법들을 능가하는 98.93%의 정확도를 기록했다.
• 이 프레임워크는 대부분의 모델 파라미터와 특징 추출을 서버로 이관함으로써 클라이언트 측 메모리와 계산을 줄여 자원이 제한된 장치에 배포 가능하게 했다.
• 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)