Brain MRI Tumor Segmentation: 3D U-Net to Attention + EfficientNet

Brain tumor segmentation from MRI is one of medical AI's most impactful—and challenging—problems. This project wasn't about chasing SOTA; it was about understanding why architectural choices matter in volumetric medical imaging.

The Challenge

BraTS (Brain Tumor Segmentation) Dataset:

• Multi-modal MRI: T1, T1ce, T2, FLAIR (4 channels)
• Expert annotations: Enhancing tumor, tumor core, whole tumor
• Volumetric: ~240×240×155 voxels per scan
• Class imbalance: Tumor ~2-5% of brain volume
• GPU memory: 3D convolutions explode VRAM usage

Iterative Experimental Methodology

Each experiment changed one architectural component while keeping everything else identical:

Dataset Prep → Baseline → Encoder Swap → Attention → Hyperparams → Evaluation

This incremental approach isolated each improvement's contribution.

Experiment 1: Baseline 3D U-Net

Architecture: Standard encoder-decoder with skip connections

• 4 encoder blocks (32→64→128→256 channels)
• 3D convolutions, InstanceNorm, ReLU
• Dice + CrossEntropy loss
• Mixed precision (AMP), patch size 128³

Result: Validation Dice 0.8057 (whole tumor)

• Learned coarse localization
• Struggled with complex boundaries and small structures
• Stable convergence but plateaued early

Experiment 2: EfficientNet Encoder

Change: Replace baseline encoder with EfficientNet-B0 (3D adapted)

• Compound scaling: depth, width, resolution
• MBConv blocks with squeeze-and-excitation
• Pre-trained on ImageNet (2D → 3D weight inflation)

Result: Validation Dice 0.8157 (whole tumor)

• Better hierarchical feature extraction
• More stable training dynamics
• Justified further architectural exploration

Experiment 3: Attention + EfficientNet

Change: Add attention gates to decoder skip connections on top of EfficientNet encoder

• Attention coefficients modulate encoder features
• Suppresses irrelevant background, highlights tumor regions
• Minimal parameter overhead (~10%)

Result: Validation Dice 0.8256 (whole tumor)

• +1.99 pp improvement over baseline
• Best qualitative results: Cleaner boundaries, fewer false positives
• Improved localization consistency across validation cases
• Became final architecture for the project

Training Optimizations (Non-Architectural)

These improvements applied to all experiments:

• Mixed precision (AMP): 2× memory reduction, 1.5× speedup
• Patch-based training: 128³ patches with overlap, sliding window inference
• Gradient checkpointing: Trade compute for memory
• DiceCELoss: Combined Dice + CE for class imbalance
• Cosine annealing + warmup: Better convergence
• Early stopping + checkpointing: Prevent overfitting, preserve best

Engineering Challenges & Solutions

Challenge	Solution
3D volumes too large for GPU	Patch-based training + sliding window inference
VRAM OOM on 3D conv	Mixed precision + gradient checkpointing + smaller patches
NIfTI/MAT/HDF5 format hell	MONAI transforms + nibabel + custom preprocessing pipeline
Experiment tracking chaos	MLflow + structured config (Hydra) + Git commits per experiment
Class imbalance (tumor < background)	Dice loss + weighted sampling + patch mining near tumor
Long training times (hrs/epoch)	AMP + multi-GPU DDP + early stopping

Key Lessons

1. Experimental methodology > Architectural complexity: One change at a time beats kitchen-sink approaches
2. Preprocessing pipeline = half the battle: NIfTI loading, normalization, orientation correction, patch extraction
3. 3D models are memory-bound: Training strategy (patches, AMP, checkpointing) matters more than layers
4. Dice-based objectives essential: Cross-entropy alone fails on imbalanced segmentation
5. Attention helps where it should: Gates suppress background, highlight tumor—interpretable and effective

Results Summary (Final Validated Configuration)

Model	Validation Dice (Whole Tumor)	Params	VRAM (128³)
Baseline 3D U-Net	0.8057	19M	8.2 GB
EfficientNet Encoder	0.8157	22M	9.1 GB
Attention + EfficientNet	0.8256	21M	8.5 GB

Note: The results above reflect the final validated experimental configuration. Earlier exploratory runs with different preprocessing/evaluation settings produced lower Dice scores (~0.62–0.67); those are superseded by the final configuration reported here.

Future Directions

• Transformer-based: Swin UNETR, nnUNet v2 for volumetric attention
• Self-supervised: Masked autoencoders on unlabeled MRI
• Multi-modal fusion: Explicit cross-modal attention (T1ce + FLAIR)
• Semi-supervised: Consistency regularization with limited labels
• Federated learning: Privacy-preserving multi-institution training
• Explainable AI: Attention visualization for clinical trust
• Foundation models: MedSAM, SAM-Med3D adaptation
• Distributed training: Scale with the Distributed ML Training Framework from this portfolio

Related Project

This blog post accompanies the Brain MRI Tumor Segmentation project case study.

View Project Case Study →

Resources

• BraTS Challenge
• MONAI Framework
• 3D U-Net Paper
• Attention U-Net Paper
• EfficientNet Paper