Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SegReg: Latent Space Regularization for Improved Medical Image Segmentation

Puru Vaish, Amin Ranem, Felix Meister, Tobias Heimann, Christoph Brune, Jelmer M. Wolterink

TL;DR

This work proposes {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 and demonstrates consistent improvements in domain generalisation.

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.

SegReg: Latent Space Regularization for Improved Medical Image Segmentation

TL;DR

This work proposes {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 and demonstrates consistent improvements in domain generalisation.

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.
Paper Structure (16 sections, 1 theorem, 4 equations, 3 figures, 2 tables)

This paper contains 16 sections, 1 theorem, 4 equations, 3 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Methodology
  3. Latent Distribution Drift
  4. Reference Latent Distributions under Moment Constraints.
  5. Regularisation Components
  6. Proposed Method: SegReg
  7. Experiments
  8. Datasets.
  9. Training.
  10. Metrics.
  11. Results
  12. Domain Generalization.
  13. Continual Learning.
  14. Representational Drift Across Tasks.
  15. Discussion and Conclusion
  16. ...and 1 more sections

Key Result

theorem 1

Let $r \in \mathcal{P}_{\Sigma}$ and $g \sim \mathcal{N}(\mathbf{0},\Sigma)$. Then $h(r)\leq h(g)$, where $h(\cdot)$ denotes differential entropy.

Figures (3)

  • Figure 1: SegReg introduces explicit latent-space regularisation that aligns embeddings with a fixed reference distribution, thereby reducing classifier variability. By stabilising the latent space, SegReg improves domain generalisation and furthermore for continual learning, anchoring representations to a fixed reference distribution mitigates bias accumulation across sequential training stages.
  • Figure 2: a) Forward transfer as a function of task index for prostate datasets. b–c) Shows the relationship between $\Bar{\text{DSC}}$ and $\Bar{\text{FWT}}$/$\Bar{\text{BWT}}$ respectively where the improvement towards top-right shows improved transfer and performance.
  • Figure 3: PCA projections of penultimate-layer features before the segmentation head to 2 principal component directions across sequential training stages for Hippocampus, Cardiac and Prostate datasets. Without SegReg, class distributions shift substantially between tasks, indicating representational drift.

Theorems & Definitions (1)

  • theorem 1: Extremal Property of the Gaussian Distribution