SegReg: Latent Space Regularization for Improved Medical Image Segmentation

TL;DR

SegReg introduces a latent space regularization method for U-Net based medical image segmentation, improving domain generalization and continual learning without extra parameters.

Contribution

It proposes a novel latent-space regularization framework compatible with standard segmentation losses, enhancing model generalization and continual learning in medical imaging.

Findings

Consistent improvements in domain generalization across multiple datasets.

Latent regularization reduces task drift in continual learning.

Enhances forward transfer without additional memory or parameters.

Abstract

Medical image segmentation models are typically optimised with voxel-wise losses that constrain predictions only in the output space. This leaves latent feature representations largely unconstrained, potentially limiting generalisation. We propose {SegReg}, a latent-space regularisation framework that operates on feature maps of U-Net models to encourage structured embeddings while remaining fully compatible with standard segmentation losses. Integrated with the nnU-Net framework, we evaluate SegReg on prostate, cardiac, and hippocampus segmentation and demonstrate consistent improvements in domain generalisation. Furthermore, we show that explicit latent regularisation improves continual learning by reducing task drift and enhancing forward transfer across sequential tasks without adding memory or any extra parameters. These results highlight latent-space regularisation as a practical approach for building more generalisable and continual-learning-ready models.