#3048·nnUNet

All classes climb to EMA pseudo-Dice ~0.06 by epoch ~55, then collapse back to ~0 by epoch 1000 (extremely sparse multi-class MRI)

Author: andrubrownCreated Jul 2, 2026Updated Jul 29, 2026

Hello!

On a 5-class MRI venous segmentation task, training converges on the loss but the model ends up predicting near-total background. With high foreground oversampling the validation pseudo-Dice climbs for all five classes early, then collapses over the long schedule:

  • epoch ~40–55: all 5 classes nonzero and rising, EMA pseudo-Dice climbs monotonically to ~0.06
  • epoch ~1000: pseudo-Dice back to ~0.01 across classes; final validation Mean Dice = 0.0002
  • train_loss converges normally throughout (to ≈ −0.90 on the Dice+CE loss)

So the model demonstrably can learn foreground early, then un-learns it. I'd like help understanding the collapse and the right configuration for foreground this sparse.

Data

  • 39 training cases, single-channel MRI (bTFE), 5 foreground labels: IVC + 4 iliac veins.
  • Anisotropic, large FOV (pelvis→knees): median spacing 6.0 × 0.69 × 0.69 mm, median shape ~160 × 407 × 512.
  • Structures extremely sparse: one case has LCIV = 2916 voxels / 39.3M (~0.007%); all 5 labels together ~0.07%.
  • Labels verified non-empty and voxel-aligned to images (identical size/spacing/origin; e.g. counts {1:15432, 2:2916, 3:2615, 4:4205, 5:3144}).
  • nnUNetv2_plan_and_preprocess --verify_dataset_integrity passes.

Configurations tried

  • 3d_fullres, default plans (pseudo-Dice flat near 0 for 1000 epochs)
  • 3d_fullres, nnUNetPlannerResEncM (same)
  • 2d, default plans — same.
  • 3d_fullres + ResEnc + custom trainer with oversample_foreground_percent = 0.9. Climbs to EMA ~0.06 by epoch ~55, then collapses to ~0 by epoch 1000 (this run).

Same 5-fold split throughout; fold 0. Example final-epoch pseudo-Dice: [0.023, 0.003, 0.011, 0.033, 0.008].

Environment

  • nnU-Net v2
  • torch 2.11.0+cu128, CUDA 13.2, RTX 5090 (Blackwell / sm_120), Ubuntu 24.04
  • nnUNet_compile=0 (compile disabled due to a Triton build error on this stack)

Questions

  • What causes climb-then-collapse over a long schedule on data this imbalanced. Is this a known LR / loss-stability interaction, and would a lower initial LR or a different loss help?
  • What's the recommended recipe for structures at ~0.007% of a large-FOV volume? Region-based training, tighter ROI cropping to the pelvis before training, still-higher oversampling, or a different patch/target-spacing choice?
  • Is disabling torch.compile on this CUDA 13.2 / Blackwell stack a plausible contributor, or unrelated?

Happy to share the plans JSON, full training log, and the pseudo-Dice trajectory.

Thank you in advance for your time and any guidance you can provide!