MICCAI 2026 Submission -- Federated RADIO + LoRA + Self-Correlating Aggregation with SAM 2-inspired Memory Attention for volumetric context.
Medical image segmentation across institutions faces a fundamental tension: models need diverse multi-site data to generalize, but privacy regulations prevent centralized data pooling. Existing federated learning approaches either train full models (communication-heavy, prone to catastrophic forgetting) or rely on task-specific backbones that lack the rich representations of modern foundation models.
FedRAD bridges this gap by combining three ideas:
- RADIO as a frozen foundation backbone -- NVIDIA's agglomerative vision model distills knowledge from SAM, DINOv2, and CLIP into a single ViT. We freeze RADIO entirely (~105M params) and only train lightweight adapters.
- Federated LoRA with asymmetric splitting -- Following FedSA-LoRA, we inject low-rank adapters (LoRA) into the last N attention blocks. LoRA-A (shared) is aggregated server-side; LoRA-B (local) stays at each client. Only ~49K parameters are communicated per round.
- Self-Correlating Aggregation (SC-Agg) -- Instead of naive FedAvg, we aggregate LoRA-A matrices using pairwise cosine similarity with temperature-scaled softmax, letting similar clients influence each other more.
- Multi-Teacher SC-Agg (MT-SC-Agg) -- Extends SC-Agg by computing client similarity in feature space using RADIO's frozen teacher adaptors (SAM + DINOv2), capturing structural and semantic similarity rather than weight-space similarity.
- SAM 2-inspired Memory Attention -- For volumetric datasets (MRI, CT), a lightweight cross-attention module (~408K params) gives each 2D slice pseudo-3D awareness by attending to features from neighboring slices in the same volume.
Total trainable parameters per client: ~787K (49K LoRA + 330K seg_head + 408K memory_attn) out of 105M backbone parameters.
┌─────────────────────────────┐
│ Frozen RADIO ViT Backbone │
│ (105M params, shared) │
│ │
Input image ─────────►│ [LoRA-A injected, shared] │
(352x352x3) │ [LoRA-B injected, local] │
└──────────┬──────────────────┘
│
patch_tokens (B, 484, 1280)
│
┌──────────────▼──────────────┐
│ Memory Attention Module │ ◄── context_tokens from
│ (cross-attn, gated residual,│ neighboring slices
│ ~408K params, per-client) │ (no_grad, frozen feats)
└──────────────┬──────────────┘
│
refined_tokens (B, 484, 1280)
│
┌──────────────▼──────────────┐
│ Segmentation Head │
│ Conv→BN→ReLU→Up(4x)→... │
│ (~330K params, per-client) │
└──────────────┬──────────────┘
│
logits (B, C, 352, 352)
For volumetric datasets, consecutive axial slices share spatial context. The Memory Attention module provides pseudo-3D awareness without full 3D convolutions:
target_tokens ──► proj_q(1280→96) ──► LayerNorm ──► ┐
├─► CrossAttention(d=96, 4 heads)
context_tokens ─► proj_kv(1280→96) ─► LayerNorm ──► ┘ │
attn_output
│
proj_out(96→1280)
│
× sigmoid(gate) [gate init=0 → 0.5]
│
target + gated_attn = refined
- Context slices are processed with
torch.no_grad()-- no extra VRAM for backprop through 4 additional images. - Gated residual starts at sigmoid(0)=0.5, so the module initially contributes mildly and learns to increase/decrease its influence.
- At validation/test time, no context is provided -- the module acts as pure identity, so evaluation remains standard 2D.
| Method | Similarity Computation | Personalized? |
|---|---|---|
fedavg |
None (uniform weighted avg) | No |
sc_agg |
Weight-space cosine similarity on LoRA-A | No |
sc_agg_personalized |
Same as SC-Agg | Yes (per-client LoRA-A) |
mt_sc_agg |
Feature-space similarity (SAM + DINOv2 + backbone) | No |
mt_sc_agg_personalized |
Same as MT-SC-Agg | Yes (per-client LoRA-A) |
SC-Agg computes pairwise cosine similarity between flattened LoRA-A matrices, suppresses negative correlations (set to -inf), applies temperature-scaled softmax, and produces a weighted aggregate. In the personalized variants, each client gets its own tailored LoRA-A instead of one global average.
MT-SC-Agg replaces weight-space similarity with feature-space similarity: after local training, each client extracts mean feature vectors from RADIO's frozen SAM and DINOv2 teacher adaptors plus the backbone, producing 3 "teacher views" per client. Similarity is computed across all views, giving a richer signal for aggregation.
MedRadSeg/
├── fedrad/ # Core package
│ ├── configs/
│ │ └── default.yaml # All hyperparameters
│ ├── data/
│ │ ├── polyp.py # Polyp benchmark (4 clients, binary)
│ │ ├── prostate.py # Prostate benchmark (6 clients, binary)
│ │ ├── fundus.py # Fundus benchmark (4 clients, 3-class)
│ │ ├── cardiac.py # Cardiac M&Ms benchmark (4 clients, 4-class)
│ │ ├── brats.py # BraTS 2025 glioma (4 clients, 4-class, IID/skew)
│ │ └── volumetric_wrapper.py # Adds neighboring slice context for memory bank
│ ├── models/
│ │ ├── radio_lora.py # RADIO backbone + LoRA injection
│ │ ├── seg_head.py # Lightweight Conv decoder
│ │ └── memory_attention.py # SAM 2-inspired cross-attention module
│ ├── federated/
│ │ ├── simulation.py # Federated training loop orchestrator
│ │ ├── server.py # Aggregation algorithms (FedAvg, SC-Agg, MT-SC-Agg)
│ │ └── client.py # Client-side losses and metrics
│ ├── evaluate.py # Test evaluation + metrics (Dice, IoU, HD95)
│ ├── train_federated.py # CLI entry point for single experiments
│ └── run_experiments.py # CLI entry point for batch comparison experiments
├── scripts/
│ ├── preprocess_cardiac.py # M&Ms NIfTI → 2D PNG slices
│ ├── preprocess_brats.py # BraTS NIfTI → 2D pseudo-RGB PNG slices
│ ├── run_all_polyp.sh # 5 methods × polyp
│ ├── run_all_prostate.sh # 5 methods × prostate
│ ├── run_all_fundus.sh # 5 methods × fundus
│ ├── run_all_cardiac.sh # 5 methods × cardiac
│ └── run_all_brats.sh # 5 methods × 2 partitions (IID + skew)
├── results/
│ └── plot_convergence.py # Generate paper figures
├── pyproject.toml
└── README.md
- Python >= 3.11
- CUDA-capable GPU (tested on RTX 5000 Ada 32GB)
- uv package manager
git clone <repo-url> MedRadSeg
cd MedRadSeg
uv syncThis installs all dependencies (PyTorch, RADIO via torch.hub, albumentations, nibabel, etc.) into a local .venv. The RADIO backbone (~400MB) is auto-downloaded on first run via torch.hub.
| Benchmark | Modality | Clients | Classes | Images | Heterogeneity |
|---|---|---|---|---|---|
polyp |
RGB endoscopy | 4 (Kvasir, CVC-300/612/T) | 1 (polyp) | ~4,376 | Low (all colonoscopy) |
prostate |
Grayscale MRI | 6 (BIDMC, BMC, HK, I2CVB, UCL, RUNMC) | 1 (prostate) | ~1,867 | High (6 scanners) |
fundus |
RGB retinal | 4 | 3 (disc, cup, bg) | ~2,120 | Medium |
cardiac |
Grayscale MRI | 4 (Siemens, Philips, GE, Canon) | 4 (LV, MYO, RV, bg) | ~2,000 | High (4 vendors) |
brats_iid |
Pseudo-RGB MRI | 4 | 4 (NCR, ED, ET, bg) | ~59,000 | Low (IID split) |
brats_skew |
Pseudo-RGB MRI | 4 | 4 (NCR, ED, ET, bg) | ~59,000 | Medium (40/30/20/10%) |
Polyp, Prostate, Fundus: these datasets should be placed in datasets/ following the structure expected by each data module (see fedrad/data/<benchmark>.py docstrings).
Cardiac (M&Ms):
# 1. Download M&Ms dataset to datasets/cardiac_raw/
# 2. Preprocess NIfTI → 2D PNG slices
uv run python scripts/preprocess_cardiac.py \
--input datasets/cardiac_raw \
--output datasets/cardiacBraTS 2025 Glioma (new):
# Preprocess: T1c + T2-FLAIR + T2w → pseudo-RGB PNG (tumor slices only)
# Default input: /eos/project/d/diagbox/BRATS2025/...
uv run python scripts/preprocess_brats.py
# Or specify custom paths:
uv run python scripts/preprocess_brats.py \
--input /path/to/BraTS2025-GLI-PRE-Challenge-TrainingData \
--output datasets/bratsThis produces datasets/brats/all/images/ and datasets/brats/all/masks/ with filenames like BraTS-GLI-00000-000_z047.png. Expected output: ~59,000 tumor-containing slices from 1,251 patients.
The BraTS pseudo-RGB encoding maps three MRI modalities to color channels:
- Red = T1c (contrast-enhanced T1)
- Green = T2-FLAIR
- Blue = T2w
Segmentation labels: 0=background, 1=NCR (necrotic core), 2=ED (peritumoral edema), 3=ET (enhancing tumor).
Partitioning is by patient (all slices from one patient go to the same client, preventing data leakage):
brats_iid: 4 clients, each gets 25% of patients (random equal split)brats_skew: 4 clients with quantity skew: 40% / 30% / 20% / 10% of patients
Within each client, patients are further split 70/15/15 into train/val/test.
# Basic: FedAvg on polyp, 100 rounds
uv run python -m fedrad.run_experiments \
--benchmark polyp \
--rounds 100 \
--methods fedavg
# BraTS with memory bank (volumetric context)
uv run python -m fedrad.run_experiments \
--benchmark brats_iid \
--rounds 100 \
--methods fedavg \
--memory-bank \
--context-k 2| Flag | Description | Default |
|---|---|---|
--benchmark |
Dataset: polyp, prostate, fundus, cardiac, brats_iid, brats_skew |
polyp |
--rounds |
Number of federated communication rounds | 100 |
--methods |
Space-separated aggregation methods | all 5 |
--memory-bank |
Enable SAM 2-style memory attention (volumetric context) | disabled |
--context-k |
Number of neighbor slices per side (total context = 2K) | 2 |
--device |
cuda, mps, or cpu |
auto-detect |
--config |
Path to YAML config file | fedrad/configs/default.yaml |
--output-dir |
Directory for experiment outputs | ./experiments |
This is the primary script for the volumetric/3D experiments in the paper:
# Foreground (see output live)
bash scripts/run_all_brats.sh
# Background (keeps running after logout)
nohup bash scripts/run_all_brats.sh > logs/all_brats.log 2>&1 &
tail -f logs/all_brats.log # monitor progressThis runs 10 experiments sequentially:
- 5 aggregation methods x
brats_iidpartition - 5 aggregation methods x
brats_skewpartition
Each experiment runs 100 rounds with memory bank enabled (--memory-bank --context-k 2). Results are saved to experiments/brats_iid/<method>/ and experiments/brats_skew/<method>/.
# 2D benchmarks (no memory bank needed)
bash scripts/run_all_polyp.sh # 5 experiments, ~4h
bash scripts/run_all_prostate.sh # 5 experiments
bash scripts/run_all_fundus.sh # 5 experiments
bash scripts/run_all_cardiac.sh # 5 experiments
# Volumetric benchmark with memory bank
bash scripts/run_all_brats.sh # 10 experiments (2 partitions)# Verify memory attention module (no GPU needed)
uv run python -c "
from fedrad.models.memory_attention import MemoryAttention
import torch
m = MemoryAttention(1280)
print(f'Params: {sum(p.numel() for p in m.parameters())}') # ~408K
x = torch.randn(2, 484, 1280)
assert torch.equal(m(x, None), x) # Identity without context
print('OK')
"
# 2-round smoke test on BraTS (requires preprocessed data + GPU)
uv run python -m fedrad.run_experiments \
--benchmark brats_iid \
--rounds 2 \
--methods fedavg \
--memory-bank --context-k 2
# 2-round smoke test without memory bank
uv run python -m fedrad.run_experiments \
--benchmark brats_iid \
--rounds 2 \
--methods fedavgAll hyperparameters are in fedrad/configs/default.yaml:
model:
radio_version: "c-radio_v3-b" # RADIO backbone
lora_rank: 4 # LoRA rank (r)
lora_alpha: 8.0 # LoRA scaling (alpha/rank)
num_lora_layers: 4 # Last N ViT blocks with LoRA
num_classes: 1 # Auto-overridden per benchmark
memory_bank:
enabled: false # Set true for volumetric (or use --memory-bank)
context_k: 2 # Neighbor slices per side
d_proj: 96 # Cross-attention bottleneck dim
num_heads: 4 # Attention heads
dropout: 0.1
federated:
num_rounds: 100
local_epochs: 2
aggregation: "sc_agg" # fedavg | sc_agg | sc_agg_personalized | mt_sc_agg | mt_sc_agg_personalized
tau: 10.0 # SC-Agg temperature
training:
lr: 1.0e-4
weight_decay: 1.0e-4
batch_size: 8
optimizer: "adamw"
scheduler: "cosine"
loss: "dice_bce" # Auto-overridden to dice_ce for multi-classDuring training on volumetric datasets (BraTS, cardiac), the VolumetricContextDataset wrapper augments each sample:
- Filename parsing: Filenames like
BraTS-GLI-00000-000_z047.pngare parsed to group slices by patient and sort by z-index. - Context selection: For each target slice at position
z, the wrapper loadsKslices on each side (z-Ktoz-1andz+1toz+K), clamping at volume boundaries. - Context transforms: Context slices use validation transforms only (resize + normalize, no random augmentation) to preserve spatial coherence.
- Batch format: The dataloader yields
(image, mask, context_images, positions)instead of the standard(image, mask).
In the training loop (simulation.py):
1. Target image → RADIO+LoRA → target_tokens (B, 484, 1280) [with grad]
2. Context images → RADIO+LoRA → ctx_tokens (B, 4*484, 1280) [no_grad]
3. refined_tokens = MemoryAttention(target_tokens, ctx_tokens)
4. logits = SegHead(refined_tokens)
5. loss = DiceCE(logits, mask)
6. Backprop through MemoryAttention + LoRA + SegHead
Context slices are processed with torch.no_grad() to avoid storing activations for 4 extra images, keeping VRAM usage manageable (~8GB additional).
Memory attention parameters are per-client local (not aggregated), similar to the segmentation head. Each client's state includes:
lora_B-- local LoRA-B matricesseg_head-- local segmentation head weightsmemory_attn-- local memory attention weights (only when enabled)lora_A-- personalized LoRA-A (only for personalized aggregation methods)
At validation/test time, no volumetric context is provided. The memory attention module receives context_tokens=None and acts as a pure identity function. This means evaluation is standard per-slice 2D inference, making results directly comparable to the non-memory-bank baseline.
Each experiment produces:
experiments/<benchmark>/<method>/
├── checkpoints/
│ ├── best.pt # Best validation Dice checkpoint
│ └── round_*.pt # Per-round checkpoints
├── logs/
└── results/
└── training_history.json # Full metrics per round + test results
The training_history.json contains:
- Per-round:
avg_train_loss,avg_dice,avg_iou,per_client_dice - Final test: per-client
dice_mean,dice_std,iou_mean,hd95_mean
uv run python results/plot_convergence.py
# Outputs: results/convergence_comparison.pdf + .png
# One panel per benchmark (polyp, prostate, fundus, cardiac, brats_iid, brats_skew)# 1. Preprocess volumetric datasets
uv run python scripts/preprocess_cardiac.py --input datasets/cardiac_raw --output datasets/cardiac
uv run python scripts/preprocess_brats.py
# 2. Run all experiments (sequentially, ~24h total on RTX 5000 Ada)
bash scripts/run_all_polyp.sh
bash scripts/run_all_prostate.sh
bash scripts/run_all_fundus.sh
bash scripts/run_all_cardiac.sh
bash scripts/run_all_brats.sh
# 3. Generate figures
uv run python results/plot_convergence.py# 1. Preprocess
uv run python scripts/preprocess_brats.py
# 2. Run 10 experiments (5 methods x 2 partitions)
nohup bash scripts/run_all_brats.sh > logs/all_brats.log 2>&1 &
# 3. Check results
cat experiments/brats_iid/summary.json
cat experiments/brats_skew/summary.json| Decision | Rationale |
|---|---|
| Freeze RADIO entirely | Foundation model features are already powerful; LoRA fine-tunes efficiently |
| LoRA-A shared, LoRA-B local | FedSA-LoRA: share direction (A), keep magnitude (B) local for personalization |
| Memory attn per-client (not aggregated) | Each site's volumetric characteristics differ; local attention learns site-specific context patterns |
| Context with no_grad | 4 extra images per sample would 5x VRAM if backpropped; frozen features suffice for attention targets |
| Gate initialized to 0 | sigmoid(0)=0.5 gives mild starting influence, prevents collapse while allowing the model to learn |
| Partition by patient (BraTS) | Prevents data leakage: all slices from one patient stay in the same client/split |
| Pseudo-RGB for BraTS | RADIO expects 3-channel input; T1c/T2f/T2w encode complementary tissue contrasts |