A Federated Learning-Friendly Approach for Parameter-Efficient Fine-Tuning of SAM in 3D Segmentation

  • 2024-07-31 17:48:06
  • Mothilal Asokan, Joseph Geo Benjamin, Mohammad Yaqub, Karthik Nandakumar
  • 0

Abstract

Adapting foundation models for medical image analysis requires finetuningthem on a considerable amount of data because of extreme distribution shiftsbetween natural (source) data used for pretraining and medical (target) data.However, collecting task-specific medical data for such finetuning at a centrallocation raises many privacy concerns. Although Federated learning (FL)provides an effective means for training on private decentralized data,communication costs in federating large foundation models can quickly become asignificant bottleneck, impacting the solution's scalability. In this work, weaddress this problem of efficient communication while ensuring effectivelearning in FL by combining the strengths of Parameter-Efficient Fine-tuning(PEFT) with FL. Specifically, we study plug-and-play Low-Rank Adapters (LoRA)in a federated manner to adapt the Segment Anything Model (SAM) for 3D medicalimage segmentation. Unlike prior works that utilize LoRA and finetune theentire decoder, we critically analyze the contribution of each granularcomponent of SAM on finetuning performance. Thus, we identify specific layersto be federated that are very efficient in terms of communication cost whileproducing on-par accuracy. Our experiments show that retaining the parametersof the SAM model (including most of the decoder) in their original state duringadaptation is beneficial because fine-tuning on small datasets tends to distortthe inherent capabilities of the underlying foundation model. On Fed-KiTS, ourapproach decreases communication cost (~48x) compared to full fine-tuning whileincreasing performance (~6% Dice score) in 3D segmentation tasks. Our approachperforms similar to SAMed while achieving ~2.8x reduction in communication andparameters to be finetuned. We further validate our approach with experimentson Fed-IXI and Prostate MRI datasets.

 

Quick Read (beta)

loading the full paper ...