CSE 599
FNO-Diffusion for Brain MRI Segmentation
The hybrid failed, and that was the useful result. Inserting Fourier Neural Operator blocks into a diffusion pipeline disrupted the U-shape — the baseline that actually worked.
BraTS 2021 case study testing whether global spectral modeling could improve diffusion-based brain tumor segmentation. A CSE 599 project that produced negative results worth analyzing.
Key Result
The simple U-shape outperformed the hybrid by 12 DICE points.
The Segmentation Problem
Dataset
BraTS 2021 — 1,251 3D MRI volumes
Input
4 modalities (T1, T1ce, T2, FLAIR)
Output
4 segmentation masks (BG, NCR, ED, ET)
Metric
DICE score (multi-class)
Split
70 / 20 / 10
Compute
NVIDIA V100 (32 GB)
4 tumor sub-regions, each with distinct texture and shape
Background
NCR — Necrotic Core
ED — Peritumoral Edema
ET — Enhancing Tumor
Hypothesis
Global spectral features from FNO could complement local refinement from a diffusion U-Net. If the texture of tumor tissue differs from healthy brain in the frequency domain, FNO should help. The reverse turned out to be true — the spectral path degraded local boundary precision.
What Each Model Tested
Diagnostic Model
FNO (spectral only)
Q: Can global spectral modeling detect tumors?
Result: DICE 65.13%. Captures broad structure, misses fine boundaries.
Diagnostic Model
Diffusion U-shape
Q: How well does iterative denoising segment?
Result: DICE 69.44%. Strong local refinement via skip connections.
Diagnostic Model
FNO + Diffusion Hybrid
Q: Does adding FNO to diffusion improve?
Result: DICE 57.03%. Spectral path disrupted U-shape refinement.
Training Details
FNO: 4 spectral layers, modes=16, width=64, LR 1e-4, 200 epochs
Diffusion: MedSegDiff UNet, 1000 timesteps, LR 1e-4, 200 epochs
Hybrid: FNO blocks in UNet bottleneck, same hyperparameters
The Simple U-shape Won
| Model | DICE Score | vs Baseline |
|---|---|---|
| Diffusion (U-shape) | 69.44% | Best |
| FNO (spectral only) | 65.13% | -4.31% |
| Hybrid (FNO + Diffusion) | 57.03% | -12.41% |
Why the Hybrid Likely Failed
FNO is Global
FNO captures low-frequency structure across the entire image. Good for large-scale patterns, shape priors, overall layout — critical for tasks like weather, turbulence, or full-image classification.
Tumor Masks are Local
Segmentation needs pixel-level precision. Tumor boundaries are small, irregular, and texturally specific. Global spectral operations can't resolve sub-pixel edges.
U-shape Preserves Detail
Skip connections carry fine features from encoder to decoder through every resolution level. Each scale's spatial detail is preserved — exactly what boundary-heavy tasks need.
Hybrid Lost Refinement
Inserting FNO in the diffusion path likely displaced the detailed spatial information the U-shape was carrying. The global path clobbered the local path.
The Lesson
FNO should be a complementary global-context branch, not a replacement for local refinement. When you insert spectral layers into a U-shape, you break the skip-connection path that carries the spatial detail segmentation tasks depend on. The architecture matters more than the representation.
Visual Evidence
FNO (spectral only)
FNO captures the general tumor region but gives blurry, over-smoothed boundaries. The spectral representation can't resolve fine edges.
Diffusion (U-shape baseline)
The U-shape produces the cleanest tumor boundaries. Skip connections carry spatial detail across all resolution levels — directly visible in the DICE gap.
Hybrid (FNO + Diffusion)
Adding FNO to the diffusion path visibly degrades predictions — noisy edges, misclassified regions, and artifacts the plain U-shape doesn't produce.
Why FNO Still Matters
Resolution-Invariant Classification
FNO-trained models classify at any input resolution without retraining. Useful when medical images come at varying slice thicknesses.
arXivMulti-Sized Classification
Continuous FNO handles arbitrary input sizes — no need to resize or crop inconsistently-sized scans.
arXivLow-Quality Image Recognition
FNO-based classifiers degrade gracefully on noisy, compressed, or low-res inputs. Robustness matters in clinical deployment.
arXivNext Design I Would Try
Keep the U-shape Local
Retain the diffusion denoising UNet exactly as it works. Don't interrupt the skip-connection path.
Add FNO as a Global Branch
Run FNO in parallel as a separate spectral encoder. Fuse the output with the U-shape bottleneck, not as a replacement for it.
Controlled Ablations
Test fusion location (bottleneck only vs early fusion vs late fusion), spectral bandwidth, and FNO depth. Isolate which knob moves DICE.
Stabilize Before Scaling
Validate on 2D slices first, then 3D volumes. Get the architecture right at small scale before expanding compute.
Implementation Lessons
This project started with MedSegDiff-V2 as the baseline diffuser and added FNO blocks. Some parts worked on the first try; others required significant debugging.
What Worked
- Loading .npy training data with numpy, converting to torch tensors
- Training the diffusion model with a straightforward denoising objective
- Getting the MedSegDiff U-shape to converge on 2D slices
- Binary segmentation (tumor vs background) — worked on first attempt
What Was Difficult
- Multi-class 4-label segmentation — gradients oscillated badly
- FNO spectral layers destabilized the diffusion denoising trajectory
- V100 memory constraints on 3D volumes required aggressive downsampling
- Hyperparameter search space was too large to explore systematically