AI Research Portal

This paper introduces DSFedMed, a dual-scale federated learning framework for medical image segmentation. The core idea is to enable mutual knowledge distillation between a large, centralized foundation model (FM) and smaller, lightweight client models. This approach aims to balance high accuracy with reduced computational and communication costs, making FMs more deployable in resource-constrained federated settings. Empirically, DSFedMed achieves an average 2% improvement in Dice score while significantly reducing communication and inference costs compared to existing federated FM baselines.

The paper addresses several bottlenecks in deploying foundation models (FMs) within federated learning (FL) for medical image segmentation:

High Computational Demands: Large FMs require substantial computational resources and memory, often exceeding the capabilities of typical client devices in clinical environments.
Communication Overhead: Transmitting large models between clients and servers leads to heavy communication overhead and increased synchronization latency.

A dual-scale federated knowledge transfer framework enabling collaboration between a server-side foundation model and client-side lightweight models for medical image segmentation, focusing on efficient client inference and communication.
A medical image generator that produces controllable and modality-adaptive samples, aligned with the global data distribution, using ControlNet.
A learnability-guided mutual knowledge distillation mechanism, utilizing generated samples, to efficiently identify and transfer informative knowledge between heterogeneous models.

DSFedMed is a dual-scale federated framework designed for medical image segmentation. It involves three main components:

Controllable Medical Image Generation: Each client fine-tunes a ControlNet module to generate controllable and modality-adaptive medical images from structural masks. These ControlNet parameters are then uploaded to the server.
Federated Learning and Alignment: Clients train lightweight segmentation models (e.g., TinySAM) locally on their private data. The server aggregates client model parameters using federated averaging. The server also uses the uploaded ControlNets to generate a global dataset of image-mask pairs.
Learnability-Guided Mutual Distillation: The server uses the generated global dataset and the aggregated foundation model (e.g., SAM) to perform mutual distillation with the lightweight client models. A learnability-guided sample selection strategy dynamically identifies informative samples to facilitate knowledge transfer in both directions: from the foundation model to the lightweight model for generalization, and from the lightweight model to the foundation model for domain adaptation.

The central idea is to leverage generated synthetic data for distillation, avoiding the need for real, sensitive data, and to use a learnability score to prioritize samples that maximize mutual knowledge transfer.

Architecture:
- Foundation Model (Server): SAM-B (ViT-B/16) is used as the foundation model.
- Lightweight Model (Client): TinySAM (a lightweight variant of SAM) is used.
- Generator: ControlNet is used for generating synthetic medical images.
Key Equations:
- ControlNet extension: yc = F(x; Φ) + Z(F(x + Z(cf; Φz1); Φc); Φz2) (Equation 1)
- Generated dataset construction: D = {(zk, m) | zk = gctrl(m; Φ(k)), m ~ Mk} for k=1 to K (Equation 2)
- Lightweight model segmentation loss: min Lsup(fTiny(xk), Yk), (xk, Yk) ~ Dreal (Equation 3)
- Federated averaging: θ = Σ (nk/n) * θk (Equation 4)
- Foundation model fine-tuning loss: min Lsup(fSAM(x), m), (x,m) ∈ Ď (Equation 5)
- Learnability score: Si = λ · LGT + (1 − λ) · LKL, λ ∈ [0, 1] (Equation 6)
- Ground truth loss: LGT = l(fSAM(xi), mi) (Equation 7)
- Mutual KL divergence: LKL = DKL(PSAM(xi) || PTinySAM(xi)) + DKL(PTinySAM(xi) || PSAM(xi)) (Equation 7)
- Sampling distribution: α(x) ∝ exp(si) (Equation 8)
Loss/Objective:
- Segmentation loss (Lsup) for local training.
- Mutual knowledge distillation loss (LKL) and ground truth loss (LGT) for learnability-guided selection.
Training:
- Clients fine-tune ControlNet modules offline.
- Clients perform local training of lightweight segmentation models.
- Server aggregates client models and generates a global dataset.
- Server fine-tunes the foundation model on the generated dataset.
- Mutual distillation occurs between server and client models.
- Asynchronous training is employed, with server-side processing (data generation, distillation) occurring without clients needing to be online.
Augmentation: Not explicitly detailed, but image-mask pairs are generated.
Data Handling: Synthetic data is generated using ControlNet and Stable Diffusion to avoid using real private datasets. Input images are resized to 1024 × 1024, and masks to 256 × 256.
Implementation Details: A low VRAM mode is available for 8GB GPUs. Server-side processing can be done on a single GPU (e.g., NVIDIA A800 or A6000).

Datasets: Five Non-IID federated datasets were constructed: Fundus (Optic Disc/Cup segmentation), Prostate Cancer (MRI segmentation), Nuclei (nuclei segmentation), ISIC (skin lesion segmentation), and CHAOS (liver segmentation).
Metrics: Dice score and Intersection over Union (IoU) are used.
Baselines: SAM (centralized), FedSAM, FedMSA (federated foundation models), FedU-Net, FednnU-Net, FedTinySAM (federated lightweight models).
Numeric Results:
- DSFedMed achieved an average Dice improvement of 3.8% over FedTinySAM.
- It achieved a 1.4% gain over centralized SAM and surpassed FedSAM by 1.9%.
- On the Fundus-OD dataset, DSFedMed achieved a Dice score of 0.949, compared to FedTinySAM's 0.917.
- On the Prostate dataset, DSFedMed achieved a Dice score of 0.813, compared to FedTinySAM's 0.756.
- DSFedMed achieved the highest average Dice score of 0.829 across five datasets.
Runtime/Training Cost:
- DSFedMed reduced communication cost by 88% compared to FedSAM.
- Inference time for DSFedMed is 0.015 seconds.
- Asynchronous training design significantly reduces overall training time compared to synchronous methods. For example, DSFedMed training time is 455.5 min, while FedSAM is 1411.7 min (synchronous).

The paper mentions that future work may include accelerating the data generation stage and validating the framework in real clinical settings.
The paper also suggests extending DSFedMed to multimodal scenarios as a future direction.

Method evidence: "We propose DSFedMed, a dual-scale federated segmentation framework that addresses the challenges of deploying large foundation models in resource-constrained medical imaging environments. Specifically, our contributions: • A dual-scale federated knowledge transfer framework enabling collaboration between server-side foundation model and client-side lightweight model for medical image segmentation, with efficient client inference and communication for real-world deployment." (section: Introduction)
Result evidence: "Compared to FedSAM, it reduces communication cost by a substantial 88% while boosting the Dice score by 2%, achieving highest accuracy of 0.829. Additionally, DSFedMed maintains a low inference time of 0.015 seconds, demonstrating both high accuracy and practical viability." (section: Efficiency Analysis)
Limitation evidence: "Nevertheless, several open challenges remain. These include accelerating data generation stage and validating the framework in real clinical settings. Future work may also extend DSFedMed to multimodal scenarios." (section: Conclusion)