[논문 리뷰] FedDis: Disentangled Federated Learning for Unsupervised Brain Pathology Segmentation

Cosmin I. Bercea, Benedikt Wiestler외 2명|arXiv (Cornell University)|2021. 03. 05.

Medical Imaging and Analysis인용 수 25

한 줄 요약

FedDis는 모델 매개변수를 형태(shape)와 외관(appearance)으로 구성하고 연합학습에서 형태만 공유하여 비동일한 도메인 시프트를 완화하며, 여러 데이터셋에서 비지도 뇌 병리 분할 성능을 개선한다.

ABSTRACT

In recent years, data-driven machine learning (ML) methods have revolutionized the computer vision community by providing novel efficient solutions to many unsolved (medical) image analysis problems. However, due to the increasing privacy concerns and data fragmentation on many different sites, existing medical data are not fully utilized, thus limiting the potential of ML. Federated learning (FL) enables multiple parties to collaboratively train a ML model without exchanging local data. However, data heterogeneity (non-IID) among the distributed clients is yet a challenge. To this end, we propose a novel federated method, denoted Federated Disentanglement (FedDis), to disentangle the parameter space into shape and appearance, and only share the shape parameter with the clients. FedDis is based on the assumption that the anatomical structure in brain MRI images is similar across multiple institutions, and sharing the shape knowledge would be beneficial in anomaly detection. In this paper, we leverage healthy brain scans of 623 subjects from multiple sites with real data (OASIS, ADNI) in a privacy-preserving fashion to learn a model of normal anatomy, that allows to segment abnormal structures. We demonstrate a superior performance of FedDis on real pathological databases containing 109 subjects; two publicly available MS Lesions (MSLUB, MSISBI), and an in-house database with MS and Glioblastoma (MSI and GBI). FedDis achieved an average dice performance of 0.38, outperforming the state-of-the-art (SOTA) auto-encoder by 42% and the SOTA federated method by 11%. Further, we illustrate that FedDis learns a shape embedding that is orthogonal to the appearance and consistent under different intensity augmentations.

연구 동기 및 목표

• 여러 기관에 걸친 뇌 병리 분할을 위한 프라이버시 보존 학습의 필요성과 동기를 제시한다.
• shape와 appearance를 해리하고 모양 매개변수만 공유하도록 FedDis를 제안한다.
• 실제 병리 데이터셋에서 이상(segmentation) 분할 성능의 향상을 보여준다.
• FedDis가 appearance와 직교하는 shape 임베딩을 학습하고 강도 변화에 강건함을 보임을 보여준다.

제안 방법

• shape와 appearance 매개변수를 분리하는 해리된 오토인코더 아키텍처.
• 클라이언트 간 글로벌 집계가 모양 매개변수만 평균하는 연합 학습.
• 잠재 공간 대비 손실 LCL을 통해 shape 일관성(SCL)과 shape와 appearance 간의 잠재 직교성(LOL)을 KL-발산 항으로 강제한다.
• 재구성 기반의 비지도 이상 분할에서 잔차 x - x_rec를 이상 신호로 활용한다.
• 하이퍼파라미터: 재구성 대 LCL 균형을 위한 alpha; LCL 내에서 SCL 대 LOL 가중치를 위한 beta; 모양 일관성을 위한 gamma-시프트 증강.

실험 결과

연구 질문

• RQ1 shape와 appearance를 해리함으로써 FedDis가 연합 비지도 뇌 이상 분할에서 사이트 간 도메인 시프트를 완화할 수 있는가?
• RQ2 표준 연합 학습 방법과 비교하여 모양 매개변수만 공유하는 것이 보지 못한 스캐너와 데이터셋에 대한 일반화를 향상시키는가?
• RQ3 잠재 표현이 appearance에 직교하는 shape 임베딩을 나타내고 강도 변화에 대한 강건성을 보이는가?
• RQ4 잠재 대비 손실 구성요소(SCL 및 LOL)가 분할 성능과 해리에 미치는 영향은 무엇인가?

주요 결과

• FedDis는 주요 이상 분할 태스크에서 평균 Dice 점수 0.38을 달성하여 SOTA 오토인코더보다 42%, SOTA 연합 방법보다 11% 더 우수하다.
• FedDis는 실재 병리 데이터셋에서 FedAvg, FedVC, SiloBN 및 FedGN 등 여러 기준선보다 우수하다.
• 잠재 분석과 제거 실험에서 나타나듯 FedDis는 appearance에 직교하는 shape 임베딩을 학습하고 강도 변화하에서 안정적이다.
• 변삭 실험은 잠재 대비 손실이 효과적인 해리 및 성능 향상에 필수적임을 보여준다.
• 데이터 중앙집중 학습과 비교할 때, FedDis는 프라이버시를 보존하는 분산 데이터로 강력한 분할 성능을 제공한다.