SegReg: Latent Space Regularization for Improved Medical Image Segmentation
Overview
Medical image segmentation models are typically trained with voxel-wise losses (e.g. Dice, cross-entropy) that only constrain the output space, leaving the model’s internal latent representations unconstrained. This can limit how well the model generalises to new domains or adapts over a sequence of tasks.
Method
The paper proposes SegReg, a latent-space regularisation framework that acts on the feature maps of U-Net models to encourage more structured embeddings, while staying fully compatible with standard segmentation losses. It’s integrated directly into the nnU-Net framework.
Results
Evaluated on prostate, cardiac, and hippocampus segmentation, SegReg consistently improves domain generalisation. It also improves continual learning by reducing task drift and enhancing forward transfer across sequential tasks — without requiring extra memory or additional parameters.