Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MedFedPure: A Medical Federated Framework with MAE-based Detection and Diffusion Purification for Inference-Time Attacks

Mohammad Karami, Mohammad Reza Nemati, Aidin Kazemi, Ali Mikaeili Barzili, Hamid Azadegan, Behzad Moshiri

TL;DR

MedFedPure tackles inference-time attacks in federated medical imaging by deploying a client-side defense that blends personalization, MAE-based detection, and diffusion purification. The framework operates entirely on-device, integrating a three-stage pipeline that first personalizes predictions through a mixture of experts, then detects suspicious inputs via a Masked Autoencoder, and finally purifies only the flagged inputs with an adaptive diffusion process before classification. Empirical results on Br35H and CIFAR-10 show substantial gains in adversarial robustness (e.g., Br35H robustness from 49.50% to 87.33% with a clean accuracy of 97.67%), outperforming strong baselines like FedAvg and pFedDef while preserving patient privacy. The work demonstrates a practical, privacy-preserving route for deploying secure, trusted AI tools in clinical workflows, with potential extensions to broader modalities and threat models.

Abstract

Artificial intelligence (AI) has shown great potential in medical imaging, particularly for brain tumor detection using Magnetic Resonance Imaging (MRI). However, the models remain vulnerable at inference time when they are trained collaboratively through Federated Learning (FL), an approach adopted to protect patient privacy. Adversarial attacks can subtly alter medical scans in ways invisible to the human eye yet powerful enough to mislead AI models, potentially causing serious misdiagnoses. Existing defenses often assume centralized data and struggle to cope with the decentralized and diverse nature of federated medical settings. In this work, we present MedFedPure, a personalized federated learning defense framework designed to protect diagnostic AI models at inference time without compromising privacy or accuracy. MedFedPure combines three key elements: (1) a personalized FL model that adapts to the unique data distribution of each institution; (2) a Masked Autoencoder (MAE) that detects suspicious inputs by exposing hidden perturbations; and (3) an adaptive diffusion-based purification module that selectively cleans only the flagged scans before classification. Together, these steps offer robust protection while preserving the integrity of normal, benign images. We evaluated MedFedPure on the Br35H brain MRI dataset. The results show a significant gain in adversarial robustness, improving performance from 49.50% to 87.33% under strong attacks, while maintaining a high clean accuracy of 97.67%. By operating locally and in real time during diagnosis, our framework provides a practical path to deploying secure, trustworthy, and privacy-preserving AI tools in clinical workflows. Index Terms: cancer, tumor detection, federated learning, masked autoencoder, diffusion, privacy

MedFedPure: A Medical Federated Framework with MAE-based Detection and Diffusion Purification for Inference-Time Attacks

TL;DR

MedFedPure tackles inference-time attacks in federated medical imaging by deploying a client-side defense that blends personalization, MAE-based detection, and diffusion purification. The framework operates entirely on-device, integrating a three-stage pipeline that first personalizes predictions through a mixture of experts, then detects suspicious inputs via a Masked Autoencoder, and finally purifies only the flagged inputs with an adaptive diffusion process before classification. Empirical results on Br35H and CIFAR-10 show substantial gains in adversarial robustness (e.g., Br35H robustness from 49.50% to 87.33% with a clean accuracy of 97.67%), outperforming strong baselines like FedAvg and pFedDef while preserving patient privacy. The work demonstrates a practical, privacy-preserving route for deploying secure, trusted AI tools in clinical workflows, with potential extensions to broader modalities and threat models.

Abstract

Artificial intelligence (AI) has shown great potential in medical imaging, particularly for brain tumor detection using Magnetic Resonance Imaging (MRI). However, the models remain vulnerable at inference time when they are trained collaboratively through Federated Learning (FL), an approach adopted to protect patient privacy. Adversarial attacks can subtly alter medical scans in ways invisible to the human eye yet powerful enough to mislead AI models, potentially causing serious misdiagnoses. Existing defenses often assume centralized data and struggle to cope with the decentralized and diverse nature of federated medical settings. In this work, we present MedFedPure, a personalized federated learning defense framework designed to protect diagnostic AI models at inference time without compromising privacy or accuracy. MedFedPure combines three key elements: (1) a personalized FL model that adapts to the unique data distribution of each institution; (2) a Masked Autoencoder (MAE) that detects suspicious inputs by exposing hidden perturbations; and (3) an adaptive diffusion-based purification module that selectively cleans only the flagged scans before classification. Together, these steps offer robust protection while preserving the integrity of normal, benign images. We evaluated MedFedPure on the Br35H brain MRI dataset. The results show a significant gain in adversarial robustness, improving performance from 49.50% to 87.33% under strong attacks, while maintaining a high clean accuracy of 97.67%. By operating locally and in real time during diagnosis, our framework provides a practical path to deploying secure, trustworthy, and privacy-preserving AI tools in clinical workflows. Index Terms: cancer, tumor detection, federated learning, masked autoencoder, diffusion, privacy

Paper Structure

This paper contains 30 sections, 25 equations, 3 figures, 2 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Literature Review
  3. Brain Tumor Detection and MRI Imaging
  4. Federated Learning in Medical Imaging
  5. Challenges and Personalization in Federated Learning
  6. Inference-Time Robustness and Security
  7. Research Gap Analysis
  8. Methodology
  9. Problem Formulation
  10. Federated Learning Setup
  11. Adversarial Threat Model
  12. Proposed Defense Framework
  13. Personalized Federated Defense
  14. Diffusion-based Input Purification
  15. Masked Autoencoder Detection
  16. ...and 15 more sections

Figures (3)

  • Figure 1: Federated learning overview in our setting: each site trains locally and keeps local data on-premise; only model updates $\theta$ are sent to the server for aggregation (e.g., FedAvg), and the updated model is returned to clients.
  • Figure 2: Overview of the proposed federated defense framework. (a) Overall architecture with FedAvg; (b) client workflow integrating MAE detection, diffusion purification (if $s_{\mathrm{det}}(x)>\tau$), and classification.
  • Figure 3: Convergence of clean and adversarial accuracy across federated rounds for the PfedDef vs. our method. (a) BR35H; (b) CIFAR-10.