Federated Learning for Mammography Classification and Conditional Image Synthesis

A PyTorch, MONAI, and NVFlare framework for privacy-preserving federated learning on the EMBED digital mammography dataset, supporting both breast lesion classification and conditional diffusion-based image synthesis.

The repository was developed to investigate whether locally generated synthetic mammograms can improve federated classification performance without requiring institutions to share patient data.

---

Features

• 🏥 Federated learning using NVFlare
• 🩻 Binary mammography classification (benign vs. malignant)
• 🤖 Conditional diffusion model for synthetic mammogram generation
• 🧠 Multiple classification architectures
  • ConvNeXt Tiny
  • ConvNeXt Small
  • ConvNeXt Base
  • Vision Transformer (ViT)
  • Simple CNN baseline
• 🔄 Classifier-Free Guidance (CFG)
• 📈 Composite diffusion loss combining
  • Mean Squared Error (MSE)
  • Structural Similarity (SSIM/MS-SSIM)
  • High-frequency reconstruction loss
• 📊 Automatic evaluation using
  • Balanced Accuracy
  • Sensitivity
  • Specificity
  • ROC AUC
  • Frechet Inception Distance (FID)
  • Structural Similarity Index (SSIM)
• 🧪 Synthetic image generation for local data augmentation

---

Repository Structure

.
├── src/
│   ├── client_classification.py      # Federated classification client
│   ├── client_synthesis.py           # Conditional diffusion training client
│   ├── generate_samples.py           # Generate synthetic mammograms
│   ├── job.py                        # Launch NVFlare experiments
│   ├── model.py                      # Classification and diffusion models
│   ├── losses.py                     # Composite diffusion losses
│   └── utils.py                      # Utility functions
│
├── config.yml
├── client_config_classification.yml
├── client_config_synthesis.yml
├── environment.yml
├── submit_job.sh
└── README.md

---

Project Overview

This repository contains two complementary learning pipelines.

1. Federated Mammography Classification

Each participating institution trains a classification model using only its own mammograms while exchanging model parameters with a central server.

No patient images are shared during training.

The framework currently supports the FedAvg aggregation algorithm through NVFlare.

Workflow

                Global Model
                     │
                     ▼
      ┌─────────────────────────┐
      │      NVFlare Server     │
      └─────────────────────────┘
             │          │
             ▼          ▼
        Client 1    Client 2 ...
             │          │
     Local Training   Local Training
             │          │
             └──────┬───┘
                    ▼
          Federated Averaging
                    │
                    ▼
             Updated Global Model

Each client

• receives the latest global model
• trains locally
• evaluates on its own validation set
• returns updated model parameters

---

2. Conditional Mammography Image Synthesis

The repository also includes a conditional diffusion model capable of generating realistic 2D mammography images conditioned on lesion class.

The diffusion model can be trained independently from the classifier and later used to generate synthetic mammograms for local augmentation.

Features include

• DDPM training
• classifier-free guidance
• conditional generation
• configurable conditioning strength
• automatic image sampling during training

---

3. Synthetic Data Augmentation

Generated mammograms can be incorporated into the federated classification pipeline.

Selected clients automatically load synthetic images from disk and merge them with their local training datasets.

This allows experiments investigating whether locally generated synthetic data improves classification performance without sharing real patient data across institutions.

---

Dataset

The framework is designed for the EMBED digital mammography dataset.

Images are expected to be organized into benign and malignant directories, for example

train/
    benign/
    malignant/

test/
    benign/
    malignant/

Each client receives only the subset of studies belonging to its institution.

Institution assignments are determined automatically using the metadata file specified in config.yml.

---

Implemented Models

Classification

Implemented in

src/model.py

Supported architectures

Model	Description
ConvNeXt Tiny	Default classification model
ConvNeXt Small	Larger ConvNeXt variant
ConvNeXt Base	Highest-capacity ConvNeXt model
Vision Transformer	MONAI Vision Transformer
SimpleNetwork	Small CNN baseline

All ConvNeXt models are automatically adapted from RGB ImageNet weights to accept single-channel mammograms.

---

Diffusion Model

Image synthesis uses

• MONAI DiffusionModelUNet
• DDPM scheduler
• classifier-free guidance

The model learns to predict the diffusion noise added during forward diffusion and reconstruct clean mammograms through iterative denoising.

---

Diffusion Loss

The synthesis model is optimized using a weighted composite objective

[ L = w_{MSE}L_{MSE} + w_{SSIM}L_{SSIM} + w_{HF}L_{HF} ]

where

• MSE supervises predicted diffusion noise
• SSIM/MS-SSIM preserves structural similarity
• High-frequency loss encourages preservation of fine anatomical detail

Loss weights are fully configurable through client_config_synthesis.yml.

---

Configuration

config.yml

Defines the overall federated experiment

• aggregation algorithm
• participating institutions
• communication rounds
• client training script
• model architecture
• dataset locations
• metadata file

---

client_config_classification.yml

Configures

• preprocessing
• augmentation
• optimizer
• loss
• model
• validation frequency
• personalization / fine-tuning

---

client_config_synthesis.yml

Configures

• diffusion scheduler
• optimizer
• composite loss
• classifier-free guidance
• sampling interval
• validation interval
• diffusion architecture

---

Running Classification

Configure

client_script: client_classification.py

and

model:
    task: classification

Then launch

python src/job.py -c config.yml

The launcher

1. reads the experiment configuration,
2. partitions the EMBED dataset according to institution,
3. launches one NVFlare client per site,
4. performs federated training using FedAvg.

---

Running Diffusion Training

Configure

client_script: client_synthesis.py

and

model:
    task: synthesis

Then run

python src/job.py -c config.yml

Although the same NVFlare infrastructure is used, diffusion experiments are typically trained on a single institution for one communication round.

---

Generating Synthetic Mammograms

After training a diffusion model

python src/generate_samples.py \
    --client_site 1 \
    --num_samples 500

The script

• loads a trained diffusion checkpoint
• generates synthetic mammograms
• saves images as NIfTI files
• stores corresponding labels
• computes image quality metrics

Generated images are written to

generated_nifti-for-training/

These images can subsequently be used for federated classification experiments.

---

Evaluation

Classification

During federated training each client reports

• Balanced Accuracy
• Sensitivity
• Specificity
• ROC AUC

---

Image Synthesis

Generated images are evaluated using

Frechet Inception Distance (FID)

Measures the similarity between the distribution of generated and real mammograms.

Lower values indicate better image quality.

Structural Similarity Index (SSIM)

Measures structural similarity between generated and real images.

Higher values indicate better reconstruction quality.

---

Main Source Files

File	Description
src/job.py	Launches federated learning simulations
src/client_classification.py	Federated classification client
src/client_synthesis.py	Conditional diffusion training client
src/generate_samples.py	Generates synthetic mammograms and evaluates image quality
src/model.py	Classification and diffusion architectures
src/losses.py	Composite diffusion loss functions
src/utils.py	Utility functions and preprocessing helpers

---

Installation

Create the conda environment

conda env create -f environment.yml
conda activate <environment_name>

---

Dependencies

Major packages include

• Python
• PyTorch
• MONAI
• NVFlare
• Torchvision
• TorchMetrics
• NumPy
• Pandas
• NiBabel
• Matplotlib

---

Citation

If you use this repository in your research, please cite the associated publication once available.

---

Acknowledgements

This project builds upon

• NVFlare for federated learning
• MONAI for medical imaging and diffusion models
• PyTorch for deep learning
• TorchMetrics for image quality evaluation
• Torchvision for pretrained classification backbones