[논문 리뷰] Specificity-Preserving Federated Learning for MR Image Reconstruction

Chun-Mei Feng, Yunlu Yan외 4명|arXiv (Cornell University)|2021. 12. 09.

MRI in cancer diagnosis인용 수 11

한 줄 요약

이 논문은 k-공간과 영상 공간에서 글로벌로 공유되는 인코더와 클라이언트별 특화된 디코더를 분리하는 특성 유지형 연합 학습 프레임워크인 FedMRI를 제안한다. 국소 기울기 이탈을 보정하기 위해 가중치 기반의 대비 정규화를 도입함으로써, FedMRI는 다양한 기관에서의 데이터에서 최신 기술 수준의 재구성 성능를 달성하면서도 스캐너의 도메인 특수성까지 유지한다.

ABSTRACT

Federated learning (FL) can be used to improve data privacy and efficiency in magnetic resonance (MR) image reconstruction by enabling multiple institutions to collaborate without needing to aggregate local data. However, the domain shift caused by different MR imaging protocols can substantially degrade the performance of FL models. Recent FL techniques tend to solve this by enhancing the generalization of the global model, but they ignore the domain-specific features, which may contain important information about the device properties and be useful for local reconstruction. In this paper, we propose a specificity-preserving FL algorithm for MR image reconstruction (FedMRI). The core idea is to divide the MR reconstruction model into two parts: a globally shared encoder to obtain a generalized representation at the global level, and a client-specific decoder to preserve the domain-specific properties of each client, which is important for collaborative reconstruction when the clients have unique distribution. Such scheme is then executed in the frequency space and the image space respectively, allowing exploration of generalized representation and client-specific properties simultaneously in different spaces. Moreover, to further boost the convergence of the globally shared encoder when a domain shift is present, a weighted contrastive regularization is introduced to directly correct any deviation between the client and server during optimization. Extensive experiments demonstrate that our FedMRI's reconstructed results are the closest to the ground-truth for multi-institutional data, and that it outperforms state-of-the-art FL methods.

연구 동기 및 목표

• 이질적인 스캔 프로토콜과 하드웨어로 인해 발생하는 다기관 MR 영상 재구성에서의 도메인 이동 문제를 해결한다.
• 기존 연합 학습에서 모델 평균화로 인해 손실되는 클라이언트별 영상 특성(예: 스캐너 특화 성질)을 유지한다.
• k-공간과 영상 공간에서의 동시 최적화를 통해 국소 재구성 정밀도를 훼손하지 않으면서도 글로벌 모델의 일반화 능력을 향상시킨다.
• 반복적인 클라이언트 간 정렬을 피하고 통신 오버헤드를 줄이는 경량이며 효율적인 학습 메커니즘을 개발한다.

제안 방법

• MR 재구성 모델을 글로벌로 공유되는 인코더와 클라이언트별 특화된 디코더로 분해하여 일반화된 특징과 도메인 특화된 특징을 동시에 학습할 수 있도록 한다.
• 주파수 영역인 k-공간과 영상 영역에서 동시에 모델 학습을 수행하여 상호 보완적인 표현을 포착한다.
• 최적화 과정에서 국소 기울기 편차를 보정하기 위해 클라이언트 간 모델 가중치(특징이 아닌)를 비교하는 가중치 기반 대비 정규화 항을 도입한다.
• 특징 기반 대비 학습에서 흔히 요구되는 큰 배치 크기를 피하기 위해, 가중치 기반 대비 손실에 L1 노름을 사용하여 안정성과 수렴 성능을 향상시킨다.
• 대비 정규화를 통합한 클라이언트 평균 업데이트 규칙을 사용해 공유 인코더를 학습함으로써 국소 업데이트를 공통의 글로벌 표현으로 정렬한다.
• 모델 공유를 인코더와 첫 번째 다운샘플링 레이어로만 제한하여 디코더 내의 클라이언트별 정보를 유지한다.

Figure 1: ​ ${}_{\!}$ Classical FL algorithm for MR image reconstruction: (a) average all the local client models to obtain a general global model [ 9 ] , or (b) repeatedly align the latent features between the source and target clients [ 13 ] . In contrast, we propose a specificity-preserving mecha

실험 결과

연구 질문

• RQ1 도메인 이동 상황에서 클라이언트별 영상 특성을 유지하면 연합 MR 영상 재구성의 정확도가 향상되는가?
• RQ2 k-공간과 영상 공간에서의 동시 최적화는 단일 영역 학습보다 더 나은 일반화와 정밀도를 이끌어내는가?
• RQ3 특징 수준의 정렬이나 큰 배치 크기를 요구하지 않으면서도 가중치 기반 대비 정규화 메커니즘이 국소 기울기 이탈을 효과적으로 보정할 수 있는가?
• RQ4 공유 모델 파rameter의 비율이 연합 환경에서 재구성 성능와 기밀 보호에 어떤 영향을 미치는가?

주요 결과

• FedMRI는 다기관 MR 데이터셋에서 최고의 PSNR와 SSIM 점수를 기록하며, FedAvg, FedProx, FL-MRCM과 같은 최신 기술 수준의 FL 방법들을 능가한다.
• 절단 실험 결과 전체 로컬 모델을 공유하는 경우(Ours₁) 성능 저하가 발생함을 확인하여, 클라이언트별 디코더를 유지하는 것이 필수적임을 입증한다.
• L1-노름 기반 가중치 기반 대비 정규화는 가장 높은 재구성 정확도를 보이며, L2 손실보다 1.53 dB 높고 NT-Xent 손실보다 1.78 dB 높은 성능를 기록하여 도메인 이동 문제에 대한 우수한 대처 능력을 입증한다.
• FedMRI의 학습 시간은 FedAvg와 FedBN과 유사하며, 대비 정규화로 인한 증가가 미미하여 실세계 구현에 있어 계산적으로 실현 가능하다.
• FL-MRCM에서 발생하는 소스-타겟 클라이언트 간 반복적인 적대적 훈련을 피함으로써 통신 비용을 감소시켰다.
• 결과적으로 k-공간과 영상 공간 양쪽에서 도메인 특화된 특성을 유지하는 것이 단지 영상 공간에서 최적화하는 방법보다 더 정확하고 강건한 재구성 결과를 이끌어낸다.

Figure 2: Overview of the FedMRI framework. Instead of averaging all the local client models, a “globally shared encoder” is used to obtain a generalized representation of both frequency and image spaces, and a “client-specific decoder” is used to explore unique domain-specific information of the tw